VLA-JEPA

This is the LeRobot port of VLA-JEPA, a Vision-Language-Action model that combines a Qwen3-VL language backbone with a self-supervised video world model (V-JEPA2) and a flow-matching DiT action head.

Architecture Overview

VLA-JEPA has three main components:

Data flow

Training:

1. A video clip of num_video_frames frames is encoded by V-JEPA2 into per-frame patch tokens.
2. The Qwen3-VL backbone processes multi-view images + the task instruction and produces a sequence of context tokens that includes special action tokens (for world model conditioning) and embodied tokens.
3. The action head receives those context tokens as cross-attention keys/values and predicts a denoised action chunk via flow matching.
4. The world model predictor uses the action tokens extracted from Qwen to predict future V-JEPA2 frame embeddings; a regression loss on those predictions is added to the action loss.

Inference: Only Qwen + the action head are used. The world model is not needed at inference time.

Action head details

Available presets via action_model_type:

The preset only sets the attention geometry, and each entry can be overridden by action_num_heads / action_attention_head_dim. Two widths follow from it:

• the DiT’s internal width is heads x head_dim (768 for DiT-B), derived rather than configured;
• the DiT’s output width, and the width of the action-decoder and state-encoder MLPs, is action_hidden_size (default 1024).

So DiT-B runs a 768-wide transformer that projects to 1024. The two are independent.

World model details

The video predictor is a ViT-style transformer (ActionConditionedVideoPredictor) that takes:

• Frame tokens: V-JEPA2 patch embeddings projected to predictor_embed_dim
• Action tokens: Qwen action token embeddings projected to predictor_embed_dim

It uses block-causal attention so each temporal step can attend to all previous steps. The predictor’s input embed_dim equals num_views × video_encoder_hidden_size (e.g. 2 views × 1024 = 2048 for the pretrained checkpoints).

Pretrained Checkpoints

Three checkpoints are available directly inside the LeRobot org here: lerobot/VLA-JEPA, converted from ginwind/VLA-JEPA:

All checkpoints use Qwen/Qwen3-VL-2B-Instruct as the language backbone.

Configuration

Key parameters in VLAJEPAConfig:

pre_snap_gripper_action and binarize_gripper_action are a port of the starVLA LIBERO eval loop and are only correct for LIBERO’s action convention. pre_snap writes {0, 1} into normalized space, the unnormalizer maps those to the midpoint and the max, and binarize then compares that physical value against gripper_threshold (0.5). For a gripper measured in degrees, mm or [0, 100], both values land above the threshold and the commanded gripper becomes a constant. They default to False for that reason; enable them only for LIBERO-style setups, and set gripper_threshold in the gripper’s own units if you do. The processor factory warns when the dataset stats show the range cannot work.

Training

Number of training steps may vary based on dataset size and compute budget. The original paper pretrained for 50k on ssv2 + droid jointly, then additional 30k steps for LIBERO, but fewer steps may still yield good performance when fine-tuning from the provided pretrained checkpoints.

Full training from scratch

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

Fine-tuning from a pretrained checkpoint

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

If you want to freeze the Qwen backbone and only train the action head, set 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 on a different embodiment

When the target robot has a different action or state dimensionality than the pretrained checkpoint, the input/output projection layers of the action head will have mismatched shapes and cannot be loaded directly. reinit_modules lets you list the key prefixes that are allowed to mismatch — those layers are randomly re-initialised while every other weight is reused from the checkpoint. Any shape mismatch outside the listed prefixes raises an error.

The layers that depend on action_dim and state_dim are:

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

If your robot has no proprioceptive state, omit model.action_model.state_encoder from the list.

Reproducing the LIBERO results

Training on LIBERO: starts the training from the Pretrain checkpoint, trains for 30k steps on the LIBERO dataset. Original paper mentions training across 8 GPUs with a batch size of 32, meaning global batch size of 256.

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

Evaluating the pretrained LIBERO-10 checkpoint:

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

To evaluate a subset of tasks only:

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

Expected results:

Fine-tuning on datasets with a different number of cameras

The pretrained world model predictor was trained with embed_dim = world_model_num_views × 1024, i.e. two camera views.

This view count used to be read from jepa_tubelet_size, which also names the JEPA encoder’s temporal tubelet size. world_model_num_views is the field for it now; leaving it at None falls back to jepa_tubelet_size so the published checkpoints keep loading unchanged.

Default behaviour — view padding / trimming (no action required)

When fine-tuning from VLA-JEPA-Pretrain the model automatically adjusts the number of views fed to the world model to match world_model_num_views:

• Single-view datasets (e.g. BridgeV2): the single-view latent is duplicated to produce a two-view world-model input, preserving the JEPA self-supervised signal without any weight mismatch.
• >2-view datasets (e.g. DROID with 3 views): all views are passed to the Qwen backbone (for richer context), but only the first world_model_num_views views (one wrist + one third-person, following the configured view order) are used for the world model.

Option 1 — Disable the world model

Set enable_world_model=False to skip the JEPA loss entirely. Only the Qwen backbone and action head are loaded and trained. This is sufficient for good action performance.

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 — Reinitialize the predictor input projection

If you want to change world_model_num_views to a value other than 2, load the checkpoint with strict=False and reinitialize model.video_predictor.predictor_embed for the new embed_dim. All other predictor block weights (attention, MLP, norm, output projection) are camera-count-agnostic and can be reused from the pretrained checkpoint.

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

License

Weights are distributed under the license terms of the original ginwind/VLA-JEPA repository (Apache 2.0 License). The LeRobot integration code follows the Apache 2.0 License.