VLA-JEPA

Esta es la versión de LeRobot de VLA-JEPA, un modelo de Visión-Lenguaje-Acción que combina un backbone de lenguaje Qwen3-VL con un modelo de mundo de video auto-supervisado (V-JEPA2) y una cabeza de acción DiT de flow-matching.

Resumen de la arquitectura

VLA-JEPA tiene tres componentes principales:

Flujo de datos

Entrenamiento:

1. Un clip de video de num_video_frames fotogramas se codifica mediante V-JEPA2 en tokens de parche por fotograma.
2. El backbone Qwen3-VL procesa imágenes multivista + la instrucción de la tarea y produce una secuencia de tokens de contexto que incluye tokens de acción especiales (para la condicionamiento del modelo de mundo) y tokens de encarnación.
3. La cabeza de acción recibe esos tokens de contexto como claves/valores de atención cruzada y predice un trozo de acción desruidado mediante flow matching.
4. El predictor del modelo de mundo utiliza los tokens de acción extraídos de Qwen para predecir los embeddings de los cuadros futuros de V-JEPA2; una pérdida de regresión sobre esas predicciones se añade a la pérdida de acción.

Inferencia: Solo se utilizan Qwen y la cabeza de acción. El modelo del mundo no es necesario en el momento de la inferencia.

Detalles de la cabeza de acción

Presets disponibles mediante action_model_type:

El preset solo establece la geometría de atención, y cada entrada puede sobrescribirse mediante action_num_heads / action_attention_head_dim. De ella se derivan dos anchos:

• el ancho interno del DiT es heads x head_dim (768 para DiT-B), derivado y no configurado;
• el ancho de la salida del DiT, y el ancho de los MLPs del decodificador de acciones y del codificador de estado, es action_hidden_size (predeterminado 1024).

Por lo tanto, DiT-B ejecuta un transformer de 768 de ancho que proyecta a 1024. Ambos son independientes.

Detalles del modelo de mundo

El predictor de video es un transformer de estilo ViT (ActionConditionedVideoPredictor) que toma:

• Fichas de fotogramas: incrustaciones de parches de V-JEPA2 proyectadas a predictor_embed_dim
• Tokens de acción: incrustaciones de tokens de acción de Qwen proyectadas a predictor_embed_dim

Utiliza atención causal por bloques, de modo que cada paso temporal puede atender a todos los pasos anteriores. La entrada del predictor embed_dim es igual a num_views × video_encoder_hidden_size (p. ej., 2 vistas × 1024 = 2048 para los puntos de control preentrenados).

Puntos de control preentrenados

Hay tres puntos de control disponibles directamente dentro de la organización LeRobot aquí: lerobot/VLA-JEPA, convertidos desde ginwind/VLA-JEPA:

Todos los puntos de control utilizan Qwen/Qwen3-VL-2B-Instruct como columna vertebral del lenguaje.

Configuración

Parámetros clave en VLAJEPAConfig:

pre_snap_gripper_action y binarize_gripper_action son una adaptación del bucle de evaluación de LIBERO de starVLA y solo son correctos para la convención de acciones de LIBERO. pre_snap escribe {0, 1} en el espacio normalizado, el desnormalizador mapea esos valores al punto medio y al máximo, y binarize luego compara ese valor físico con gripper_threshold (0.5). Para un actuador de agarre medido en grados, mm o [0, 100], ambos valores quedan por encima del umbral y el actuador de agarre comandado se convierte en una constante. Por esa razón, su valor predeterminado es False; actívalos solo para configuraciones de estilo LIBERO y, si lo haces, establece gripper_threshold en las unidades propias del actuador de agarre. La fábrica de procesadores emite una advertencia cuando las estadísticas del conjunto de datos indican que el rango no puede funcionar.

Entrenamiento

El número de pasos de entrenamiento puede variar según el tamaño del conjunto de datos y el presupuesto de cómputo. El artículo original realizó el preentrenamiento durante 50k pasos en ssv2 + droid de forma conjunta, y luego 30k pasos adicionales para LIBERO, pero es posible que con menos pasos se obtenga un buen rendimiento al realizar el ajuste fino a partir de los puntos de control preentrenados proporcionados.

Entrenamiento completo desde cero

lerobot-train \
  --policy.type=vla_jepa \
  --policy.repo_id=your_org/your_repo \
  --dataset.repo_id=your_org/your_dataset

Ajuste fino desde un punto de control preentrenado

lerobot-train \
  --policy.path=lerobot/VLA-JEPA-Pretrain \
  --policy.repo_id=your_org/your_repo \
  --dataset.repo_id=your_org/your_dataset

Si desea congelar el backbone de Qwen y entrenar solo la cabeza de acción, establezca policy.freeze_qwen=True:

lerobot-train \
  --policy.path=lerobot/VLA-JEPA-Pretrain \
  --policy.repo_id=your_org/your_repo \
  --policy.freeze_qwen=true \
  --dataset.repo_id=your_org/your_dataset

Ajuste fino en una encarnación diferente

Cuando el robot objetivo tiene una dimensionalidad de acción o estado diferente a la del punto de control preentrenado, las capas de proyección de entrada/salida de la cabeza de acción tendrán formas no coincidentes y no se pueden cargar directamente. reinit_modules permite enumerar los prefijos de clave que se permiten que no coincidan: esas capas se reinicializan aleatoriamente, mientras que todos los demás pesos se reutilizan desde el punto de control. Cualquier discrepancia de forma fuera de los prefijos enumerados genera un error.

Las capas que dependen de action_dim y state_dim son:

lerobot-train \
  --policy.path=lerobot/VLA-JEPA-Pretrain \
  --policy.repo_id=your_org/your_repo \
  --policy.freeze_qwen=true \
  --policy.reinit_modules='["model.action_model.action_encoder", "model.action_model.action_decoder", "model.action_model.state_encoder"]' \
  --dataset.repo_id=your_org/your_dataset

Si su robot no tiene un estado propioceptivo, omita model.action_model.state_encoder de la lista.

Reproduciendo los resultados de LIBERO

Entrenamiento en LIBERO: inicia el entrenamiento desde el checkpoint de Pretrain y entrena durante 30k pasos en el conjunto de datos LIBERO. El artículo original menciona el entrenamiento en 8 GPUs con un tamaño de lote de 32, lo que significa un tamaño de lote global de 256.

lerobot-train \
  --policy.path=lerobot/VLA-JEPA-Pretrain \
  --policy.repo_id=your_org/your_repo \
  --dataset.repo_id=HuggingFaceVLA/libero \
  --steps=30000

Evaluación del punto de control preentrenado de LIBERO-10:

lerobot-eval \
  --policy.path=lerobot/VLA-JEPA-LIBERO \
  --env.type=libero \
  --env.task=libero_spatial,libero_object,libero_goal,libero_10 \
  --eval.n_episodes=10 \
  --eval.batch_size=5

Para evaluar solo un subconjunto de tareas:

lerobot-eval \
  --policy.path=lerobot/VLA-JEPA-LIBERO \
  --env.type=libero \
  --env.task=libero_10 \
  --env.task_ids='[0,1,2]' \
  --eval.n_episodes=10 \
  --eval.batch_size=5

Resultados esperados:

Ajuste fino en conjuntos de datos con un número diferente de cámaras

El predictor del modelo de mundo preentrenado se entrenó con embed_dim = world_model_num_views × 1024, es decir, dos vistas de cámara.

Este recuento de vistas solía leerse de jepa_tubelet_size, que también denomina el tamaño de tubelet temporal del codificador JEPA. world_model_num_views es el campo para ello ahora; dejarlo en None recurre a jepa_tubelet_size para que los puntos de control publicados sigan cargando sin cambios.

Comportamiento predeterminado: visualización del relleno/recorte (no requiere acción)

Al realizar el ajuste fino desde VLA-JEPA-Pretrain, el modelo ajusta automáticamente el número de vistas que se envían al modelo del mundo para que coincida con world_model_num_views:

• Conjuntos de datos de una sola vista (p. ej., BridgeV2): el latente de una sola vista se duplica para producir una entrada de modelo de mundo de dos vistas, preservando la señal autorregulada de JEPA sin ningún desajuste de pesos.
• >Conjuntos de datos de más de 2 vistas (p. ej., DROID con 3 vistas): todas las vistas se pasan al backbone de Qwen (para un contexto más rico), pero solo las primeras world_model_num_views vistas (una muñeca + una tercera persona, siguiendo el orden de vistas configurado) se utilizan para el modelo de mundo.

Opción 1: deshabilitar el modelo de mundo

Establezca enable_world_model=False para omitir por completo la pérdida de JEPA. Solo se cargan y entrenan el backbone de Qwen y la cabeza de acción. Esto es suficiente para un buen rendimiento de la acción.

lerobot-train \
  --policy.path=lerobot/VLA-JEPA-Pretrain \
  --policy.enable_world_model=false \
  --policy.repo_id=your_org/your_repo \
  --dataset.repo_id=your_org/single_camera_dataset

Opción 2: Reinicializar la proyección de entrada del predictor

Si desea cambiar world_model_num_views a un valor distinto de 2, cargue el checkpoint con strict=False y reinicialice model.video_predictor.predictor_embed para el nuevo embed_dim. Los pesos de todos los demás bloques del predictor (atención, MLP, norm, proyección de salida) no dependen del número de cámaras y pueden reutilizarse desde el checkpoint preentrenado.

Cita

@misc{sun2026vlajepaenhancingvisionlanguageactionmodel,
  title         = {VLA-JEPA: Enhancing Vision-Language-Action Model with Latent World Model},
  author        = {Jingwen Sun and Wenyao Zhang and Zekun Qi and Shaojie Ren and Zezhi Liu and Hanxin Zhu and Guangzhong Sun and Xin Jin and Zhibo Chen},
  year          = {2026},
  eprint        = {2602.10098},
  archivePrefix = {arXiv},
  primaryClass  = {cs.RO},
  url           = {https://arxiv.org/abs/2602.10098},
}

Licencia

Los pesos se distribuyen bajo los términos de la licencia del repositorio original ginwind/VLA-JEPA (Apache 2.0 License). El código de integración de LeRobot sigue la Apache 2.0 License.