VLA-JEPA

这是 VLA-JEPA 的 LeRobot 移植版，这是一个视觉-语言-动作模型，将 Qwen3-VL 语言骨干网络与自监督视频世界模型（V-JEPA2）以及基于流匹配的 DiT 动作头相结合。

架构概述

VLA-JEPA 包含三个主要组件：

组件	模块	作用
Qwen3-VL 骨干网络	`Qwen3VLInterface`	将图像与语言指令融合为上下文令牌
DiT-B 动作头	`VLAJEPAActionHead`	对动作块进行基于流匹配的扩散
V-JEPA2 世界模型	`ActionConditionedVideoPredictor`	自监督视频预测损失（仅训练阶段使用）

数据流

训练：

num_video_frames 帧的视频片段由 V-JEPA2 编码为逐帧块令牌。
Qwen3-VL 骨干网络处理多视角图像和任务指令，生成一系列上下文令牌，其中包含用于世界模型条件化的特殊动作令牌和具身令牌。
动作头将这些上下文令牌作为交叉注意力的键/值，通过流匹配预测去噪后的动作块。
世界模型预测器利用从 Qwen 中提取的动作令牌来预测未来的 V-JEPA2 帧嵌入；这些预测上的回归损失被加入动作损失中。

推理： 仅使用 Qwen 和动作头。推理时不需要世界模型。

动作头详情

通过 action_model_type 可选的预设：

预设	隐藏维度	注意力头数	头维度
`DiT-B`	768	12	64
`DiT-L`	1536	32	48

世界模型详情

视频预测器是一个类 ViT 的 Transformer（ActionConditionedVideoPredictor），输入为：

帧令牌：投影到 predictor_embed_dim 的 V-JEPA2 块嵌入
动作令牌：投影到 predictor_embed_dim 的 Qwen 动作令牌嵌入

它使用块因果注意力，使每个时间步都能关注所有之前的步骤。预测器的输入 embed_dim 等于 num_views × video_encoder_hidden_size（例如，预训练检查点中 2 个视角 × 1024 = 2048）。

预训练检查点

LeRobot 组织中直接提供三个检查点：lerobot/VLA-JEPA，转换自 ginwind/VLA-JEPA：

检查点	数据集	摄像头	世界模型	动作维度
`lerobot/VLA-JEPA-LIBERO`	LIBERO-10	2（agentview + 腕部）	启用	7
`lerobot/VLA-JEPA-Pretrain`	DROID 1.0.1	2（外部左侧视图）	启用	7
`lerobot/VLA-JEPA-SimplerEnv`	OXE Bridge / RT-1	1（视图复制 ×2）	启用	7

所有检查点均使用 Qwen/Qwen3-VL-2B-Instruct 作为语言骨干网络。

配置

VLAJEPAConfig 中的关键参数：

参数	默认值	描述
`chunk_size`	7	每次推理调用预测的动作数量
`n_action_steps`	7	重新规划前从预测块中执行的步骤数
`num_video_frames`	8	输入世界模型的视频片段长度
`enable_world_model`	`True`	是否加载和训练 V-JEPA2 预测器
`world_model_loss_weight`	0.1	JEPA 预测损失相对于动作损失的权重
`num_inference_timesteps`	4	动作去噪的欧拉积分步数
`freeze_qwen`	`False`	冻结 Qwen3-VL 骨干网络，仅训练动作头
`reinit_modules`	`None`	加载时允许随机重新初始化的键前缀（用于跨具身迁移，参见在不同具身上微调）
`gripper_dim`	6	动作向量中夹爪维度的索引（例如，7 自由度手臂夹爪为最后一个关节时为 6）
`gripper_threshold`	0.5	`pre_snap_gripper_action` 和 `binarize_gripper_action` 用于二值化夹爪维度的阈值
`pre_snap_gripper_action`	`True`	在反归一化之前将夹爪维度对齐到 {0, 1}。对于没有二值夹爪的机器人，设为 `False`
`binarize_gripper_action`	`True`	在反归一化之后将夹爪维度二值化为 {-1, 1}。对于没有二值夹爪的机器人，设为 `False`

训练

训练步骤数取决于数据集大小和计算预算。原始论文在 ssv2 + droid 联合预训练 50k 步，然后在 LIBERO 上额外训练 30k 步，但从提供的预训练检查点进行微调时，更少的步骤仍可取得良好性能。

从零开始完整训练

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

从预训练检查点微调

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

若要冻结 Qwen 骨干网络并仅训练动作头，设置 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

在不同具身上微调

当目标机器人的动作或状态维度与预训练检查点不同时，动作头的输入/输出投影层形状不匹配，无法直接加载。reinit_modules 允许你列出允许不匹配的键前缀——这些层将被随机重新初始化，而检查点中的其他所有权重都会被复用。列出前缀之外的任何形状不匹配都会引发错误。

依赖 action_dim 和 state_dim 的层包括：

层	键前缀
动作编码器（action_dim → inner_dim）	`model.action_model.action_encoder`
动作解码器（hidden_size → action_dim）	`model.action_model.action_decoder`
状态编码器（state_dim → inner_dim）	`model.action_model.state_encoder`

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

若你的机器人没有本体感知状态，可从列表中省略 model.action_model.state_encoder。

复现 LIBERO 结果

在 LIBERO 上训练： 从 Pretrain 检查点开始，在 LIBERO 数据集上训练 30k 步。原始论文提到在 8 张 GPU 上以批大小 32 训练，即全局批大小 256。

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

评估预训练 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

若只评估部分任务：

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

预期结果：

套件	片段数	成功数	成功率
libero_spatial	100	93	95.0%
libero_object	100	100	100.0%
libero_goal	100	98	98.0%
libero_10	100	96	93.0%
总计	400	387	96.5%

在不同摄像头数量的数据集上微调

预训练世界模型预测器以 embed_dim = jepa_tubelet_size × 1024（默认 jepa_tubelet_size=2）训练。

默认行为 — 视角填充/裁剪（无需操作）

从 VLA-JEPA-Pretrain 微调时，模型会自动调整输入世界模型的视角数量以匹配 jepa_tubelet_size：

单视角数据集（如 BridgeV2）： 单视角潜变量被复制以产生双视角世界模型输入，在不造成权重不匹配的情况下保留 JEPA 自监督信号。
超过 2 个视角的数据集（如含 3 个视角的 DROID）： 所有视角都传递给 Qwen 骨干网络（以获得更丰富的上下文），但世界模型只使用前 jepa_tubelet_size 个视角（按配置的视角顺序，一个腕部 + 一个第三人称视角）。

选项 1 — 禁用世界模型

设置 enable_world_model=False 以完全跳过 JEPA 损失。仅加载和训练 Qwen 骨干网络和动作头，对于良好的动作性能已经足够。

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

选项 2 — 重新初始化预测器输入投影

若要将 jepa_tubelet_size 更改为 2 以外的值，以 strict=False 加载检查点，并为新的 embed_dim 重新初始化 model.video_predictor.predictor_embed。所有其他预测器块权重（注意力、MLP、归一化、输出投影）与摄像头数量无关，可从预训练检查点复用。

引用

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

许可证

权重依据原始 ginwind/VLA-JEPA 仓库的许可证条款（Apache 2.0 许可证）分发。LeRobot 集成代码遵循 Apache 2.0 许可证。