VLA-JEPA

これは、Qwen3-VL言語バックボーン、自己教師あり動画ワールドモデル（V-JEPA2）、およびフローマッチングDiTアクションヘッドを組み合わせるビジョン・言語・アクションモデルであるVLA-JEPAのLeRobot移植版です。

アーキテクチャ概要

VLA-JEPAには3つの主要なコンポーネントがあります:

データフロー

学習:

1. num_video_frames フレームの動画クリップは、V-JEPA2 によってフレームごとのパッチトークンにエンコードされます。
2. Qwen3-VLバックボーンは、マルチビュー画像とタスク指示を処理し、世界モデルの条件付け用の特殊なアクショントークンとエンボディドトークンを含むコンテキストトークンのシーケンスを生成します。
3. アクションヘッドは、これらのコンテキストトークンをクロスアテンションのキー/バリューとして受け取り、フローマッチングを通じてノイズ除去されたアクションチャンクを予測します。
4. ワールドモデル予測器は、Qwenから抽出されたアクショントークンを使用して、将来のV-JEPA2フレーム埋め込みを予測します。これらの予測に対する回帰損失が、アクション損失に追加されます。

推論: Qwenとアクションヘッドのみを使用します。推論時にはワールドモデルは不要です。

アクションヘッドの詳細

action_model_typeで利用可能なプリセット:

プリセットはアテンションのジオメトリのみを設定し、各エントリはaction_num_heads / action_attention_head_dimによって上書きできます。これに基づいて2つの幅が導出されます：

• DiTの内部幅はheads x head_dim（DiT-Bの場合は768）であり、設定ではなく導出されます。
• DiTの出力幅、およびアクションデコーダーとステートエンコーダーのMLPの幅は、action_hidden_size（デフォルトは1024）です。

つまり、DiT-Bは幅768のtransformerを実行し、1024に射影します。これらは独立しています。

ワールドモデルの詳細

動画予測器はViTスタイルのTransformer(ActionConditionedVideoPredictor)であり、以下を受け取ります:

• フレームトークン: predictor_embed_dim に射影された V-JEPA2 のパッチ埋め込み
• アクショントークン: predictor_embed_dim に射影された Qwen アクショントークンの埋め込み

ブロック因果注意機構を使用しているため、各時間ステップはそれ以前のすべてのステップに注意を払うことができます。予測器の入力embed_dimはnum_views × video_encoder_hidden_sizeに等しくなります（事前学習済みチェックポイントでは、2 views × 1024 = 2048 など）。

事前学習済みチェックポイント

LeRobot組織のこちらに、ginwind/VLA-JEPAから変換された3つのチェックポイントが直接利用できます: lerobot/VLA-JEPA

すべてのチェックポイントは、言語バックボーンとしてQwen/Qwen3-VL-2B-Instructを使用しています。

設定

VLAJEPAConfigの主なパラメータ:

pre_snap_gripper_action と binarize_gripper_action は starVLA の LIBERO 評価ループの移植であり、LIBERO のアクション規約でのみ正しく動作します。pre_snap は 正規化空間に {0, 1} を書き込み、非正規化処理がそれらを中間値と最大値にマッピングし、binarize がその 物理値を gripper_threshold (0.5) と比較します。グリップが度、mm、または [0, 100] で測定される場合、両方の値が閾値を超え、コマンドされたグリップが一定値になります。そのため、デフォルトでは False になっています。LIBERO 形式のセットアップでのみ有効にし、有効にする場合は gripper_threshold をグリップ自身の単位で設定してください。プロセッサファクトリは、データセットの統計情報から範囲が機能できないことを示す場合に警告を発します。

トレーニング

学習ステップ数は、データセットのサイズや計算リソースの予算によって変動する可能性があります。元の論文では、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 に依存するレイヤーは以下の通りです：

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

期待される結果:

カメラ数が異なるデータセットでのファインチューニング

事前学習済みのワールドモデル予測器は、embed_dim = world_model_num_views × 1024、つまり2つのカメラビューで学習されました。

このビュー数は以前、jepa_tubelet_sizeから読み取られていました。このパラメータは、JEPAエンコーダーの時間軸のチューレットサイズの名前としても使われていました。現在は、world_model_num_viewsがそのためのフィールドです。Noneのままにしておくと、jepa_tubelet_sizeにフォールバックするため、公開されているチェックポイントは変更なしで読み込むことができます。

デフォルトの動作 — パディング/トリミングの表示（対応不要）

VLA-JEPA-Pretrainからのファインチューニングでは、モデルは自動的にワールドモデルに与えるビューの数をworld_model_num_viewsに合わせて調整します:

• 単一視点データセット（例：BridgeV2）: 単一視点の潜在変数を複製して2視点のワールドモデル入力を作成し、重みの不一致なしでJEPAの自己教師ありシグナルを保持します。
• >2ビューのデータセット（例：3ビューのDROID）: すべてのビューがQwenバックボーンに渡されます（より豊かなコンテキストを得るため）が、ワールドモデルには最初のworld_model_num_viewsビューのみ（設定されたビュー順序に従い、1つの手首視点と1つの第三者視点）が使用されます。

オプション 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 — 予測器の入力射影を再初期化する

world_model_num_views を 2 以外の値に変更する場合は、strict=False を指定してチェックポイントをロードし、新しい embed_dim に対して model.video_predictor.predictor_embed を再初期化してください。その他の予測器ブロックの重み（attention、MLP、norm、出力投影）はカメラ数に依存しないため、事前学習済みチェックポイントから再利用できます。

引用

@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 License）に基づいて配布されます。LeRobot統合コードは、Apache 2.0 Licenseに従います。