VLA-JEPA

Il s’agit de la portage LeRobot de VLA-JEPA, un modèle Vision-Language-Action qui combine un backbone linguistique Qwen3-VL avec un modèle de monde vidéo auto-supervisé (V-JEPA2) et une tête d’action DiT par flow-matching.

Vue d'ensemble de l'architecture

VLA-JEPA comporte trois composants principaux :

Flux de données

Entraînement :

1. Un extrait vidéo de num_video_frames images est encodé par V-JEPA2 en jetons de patch par image.
2. Le backbone Qwen3-VL traite les images multi-vues + l'instruction de la tâche et génère une séquence de jetons de contexte qui inclut des jetons d'action spéciaux (pour le conditionnement du modèle du monde) et des jetons incarnés.
3. La tête d'action reçoit ces jetons de contexte comme clés/valeurs d'attention croisée et prédit un segment d'action débruité via le flow matching.
4. Le prédicteur de modèle du monde utilise les jetons d’action extraits de Qwen pour prédire les embeddings de futures images de V-JEPA2 ; une perte de régression sur ces prédictions est ajoutée à la perte d’action.

Inférence : Seuls Qwen et la tête d’action sont utilisés. Le modèle du monde n’est pas nécessaire au moment de l’inférence.

Détails de la tête d'action

Préréglages disponibles via action_model_type :

Le préréglage ne définit que la géométrie de l’attention, et chaque entrée peut être remplacée par action_num_heads / action_attention_head_dim. Deux largeurs en découlent :

• la largeur interne du DiT est heads x head_dim (768 pour DiT-B), dérivée plutôt que configurée ;
• La largeur de la sortie du DiT, ainsi que la largeur des MLP du décodeur d’action et de l’encodeur d’état, est action_hidden_size (1024 par défaut).

Ainsi, DiT-B exécute un transformateur de largeur 768 qui projette vers 1024. Les deux sont indépendants.

Détails du modèle du monde

Le prédicteur vidéo est un transformateur de type ViT (ActionConditionedVideoPredictor) qui prend :

• Jetons de trame : embeddings de patch de V-JEPA2 projetés sur predictor_embed_dim
• Jetons d'action : embeddings des jetons d'action de Qwen projetés sur predictor_embed_dim

Elle utilise une attention causale par blocs, de sorte que chaque étape temporelle peut prendre en compte toutes les étapes précédentes. L’entrée du prédicteur embed_dim est égale à num_views × video_encoder_hidden_size (par exemple, 2 vues × 1024 = 2048 pour les points de contrôle pré-entraînés).

Points de contrôle pré-entraînés

Trois points de contrôle sont disponibles directement dans l'organisation LeRobot ici : lerobot/VLA-JEPA, convertis depuis ginwind/VLA-JEPA :

Tous les points de contrôle utilisent Qwen/Qwen3-VL-2B-Instruct comme colonne vertébrale linguistique.

Configuration

Paramètres clés de VLAJEPAConfig :

pre_snap_gripper_action et binarize_gripper_action sont une adaptation de la boucle d’évaluation LIBERO de starVLA et ne sont correctes que pour la convention d’action de LIBERO. pre_snap écrit {0, 1} dans l’espace normalisé, le dénormaliseur mappe ces valeurs vers le point médian et le maximum, puis binarize compare cette valeur physique à gripper_threshold (0,5). Pour une pince mesurée en degrés, en mm ou dans l’intervalle [0, 100], les deux valeurs se situent au-dessus du seuil et la pince commandée devient constante. C’est pourquoi leur valeur par défaut est False ; activez-les uniquement pour les configurations de type LIBERO, et définissez gripper_threshold dans les unités propres à la pince si vous le faites. L’usine de processeurs émet un avertissement lorsque les statistiques du jeu de données indiquent que l’intervalle ne peut pas fonctionner.

Entraînement

Le nombre d'étapes d'entraînement peut varier en fonction de la taille du jeu de données et du budget de calcul. L'article original a effectué un pré-entraînement pendant 50k étapes sur ssv2 + droid conjointement, puis 30k étapes supplémentaires pour LIBERO, mais un nombre inférieur d'étapes peut toujours offrir de bonnes performances lors de l'ajustement fin à partir des points de contrôle pré-entraînés fournis.

Entraînement complet à partir de zéro

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

Ajustement fin à partir d'un point de contrôle pré-entraîné

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

Si vous souhaitez geler le backbone Qwen et n'entraîner que la tête d'action, définissez 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

Fine-tuning sur une autre incarnation

Lorsque le robot cible possède une dimensionnalité d’action ou d’état différente de celle du point de contrôle pré-entraîné, les couches de projection d’entrée/sortie de la tête d’action présentent des formes incohérentes et ne peuvent pas être chargées directement. reinit_modules permet de lister les préfixes de clés autorisés à présenter une incohérence de forme — ces couches sont réinitialisées aléatoirement tandis que tous les autres poids sont réutilisés depuis le point de contrôle. Toute incohérence de forme en dehors des préfixes listés provoque une erreur.

Les couches qui dépendent de action_dim et state_dim sont :

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 votre robot n'a pas d'état proprioceptif, omettez model.action_model.state_encoder de la liste.

Reproduction des résultats de LIBERO

Entraînement sur LIBERO : lance l’entraînement à partir du checkpoint Pretrain et entraîne le modèle pendant 30 k étapes sur le jeu de données LIBERO. L’article original mentionne un entraînement sur 8 GPU avec une taille de lot de 32, soit une taille de lot globale de 256.

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

Évaluation du point de contrôle pré-entraîné 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

Pour évaluer uniquement un sous-ensemble de tâches :

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

Résultats attendus :

Fine-tuning sur des jeux de données avec un nombre de caméras différent

Le prédicteur de modèle du monde préentraîné a été entraîné avec embed_dim = world_model_num_views × 1024, c'est-à-dire deux vues de caméra.

Ce nombre de vues était auparavant lu depuis jepa_tubelet_size, qui désigne également la taille du tubelet temporel de l’encodeur JEPA. world_model_num_views est désormais le champ dédié ; si sa valeur est laissée à None, le système retombe sur jepa_tubelet_size afin que les points de contrôle publiés continuent de se charger sans modification.

Comportement par défaut — affichage du bourrage / de la suppression (aucune action requise)

Lors de l’ajustement fin à partir de VLA-JEPA-Pretrain, le modèle ajuste automatiquement le nombre de vues transmises au modèle du monde pour correspondre à world_model_num_views :

• Jeux de données à vue unique (par ex. BridgeV2) : le latent à vue unique est dupliqué pour produire une entrée de modèle du monde à deux vues, préservant le signal auto-supervisé de JEPA sans aucun désaccord de poids.
• >Jeux de données à 2 vues (par ex. DROID avec 3 vues) : toutes les vues sont transmises au backbone Qwen (pour un contexte plus riche), mais seules les world_model_num_views premières vues (un poignet + une vue tierce personne, selon l’ordre des vues configuré) sont utilisées pour le modèle du monde.

Option 1 — Désactiver le modèle du monde

Définissez enable_world_model=False pour ignorer entièrement la perte JEPA. Seuls le backbone Qwen et la tête d'action sont chargés et entraînés. Cela suffit pour de bonnes performances en matière d'action.

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

Option 2 — Réinitialiser la projection d'entrée du prédicteur

Si vous souhaitez modifier world_model_num_views pour une valeur autre que 2, chargez le checkpoint avec strict=False et réinitialisez model.video_predictor.predictor_embed pour la nouvelle embed_dim. Les poids de tous les autres blocs du prédicteur (attention, MLP, norm, projection de sortie) sont indépendants du nombre de caméras et peuvent être réutilisés depuis le checkpoint pré-entraîné.

Citation

@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},
}

Licence

Les poids sont distribués selon les conditions de la licence du dépôt original ginwind/VLA-JEPA (Apache 2.0 License). Le code d'intégration LeRobot suit la Apache 2.0 License.