# VLA-JEPA これは、Qwen3-VL言語バックボーン、自己教師あり動画ワールドモデル（V-JEPA2）、およびフローマッチングDiTアクションヘッドを組み合わせるビジョン・言語・アクションモデルである**VLA-JEPA**のLeRobot移植版です。 --- ## アーキテクチャ概要 VLA-JEPAには3つの主要なコンポーネントがあります: | コンポーネント | モジュール | 役割 | | ----------------------- | --------------------------------- | ------------------------------------------------------- | | **Qwen3-VL バックボーン** | `Qwen3VLInterface` | 画像と言語指示をコンテキストトークンに融合します | | **DiT-B アクションヘッド** | `VLAJEPAActionHead` | アクションチャンクに対するフローマッチング拡散 | | **V-JEPA2 ワールドモデル** | `ActionConditionedVideoPredictor` | 自己教師あり動画予測損失（トレーニング専用） | ### データフロー **学習:** 1. `num_video_frames` フレームの動画クリップは、V-JEPA2 によってフレームごとのパッチトークンにエンコードされます。 2. Qwen3-VLバックボーンは、マルチビュー画像とタスク指示を処理し、世界モデルの条件付け用の特殊なアクショントークンとエンボディドトークンを含むコンテキストトークンのシーケンスを生成します。 3. アクションヘッドは、これらのコンテキストトークンをクロスアテンションのキー/バリューとして受け取り、フローマッチングを通じてノイズ除去されたアクションチャンクを予測します。 4. ワールドモデル予測器は、Qwenから抽出されたアクショントークンを使用して、将来のV-JEPA2フレーム埋め込みを予測します。これらの予測に対する回帰損失が、アクション損失に追加されます。 **推論:** Qwenとアクションヘッドのみを使用します。推論時にはワールドモデルは不要です。 ### アクションヘッドの詳細 `action_model_type`で利用可能なプリセット: | プリセット | ヘッド | ヘッド次元 | | ------- | ----- | -------- | | `DiT-B` | 12 | 64 | | `DiT-L` | 32 | 48 | プリセットはアテンションのジオメトリのみを設定し、各エントリは`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](https://huggingface.co/ginwind/VLA-JEPA)から変換された3つのチェックポイントが直接利用できます: [`lerobot/VLA-JEPA`](https://huggingface.co/collections/lerobot/vla-jepa) | チェックポイント | データセット | カメラ | ワールドモデル | アクション次元 | | ----------------------------- | ----------------- | ----------------------- | ----------- | ---------- | | `lerobot/VLA-JEPA-LIBERO` | LIBERO-10 | 2 (agentview + wrist) | 有効 | 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 | 1回の推論呼び出しごとに予測されるアクション数 | | `n_action_steps` | 7 | 再計画の前に予測チャンクから実行されるステップ | | `num_video_frames` | 8 | ワールドモデルに渡される動画クリップの長さ | | `enable_world_model` | `True` | V-JEPA2のpredictorを読み込み、学習するかどうか | | `world_model_loss_weight` | 0.1 | アクション損失に対するJEPA予測損失の重み | | `causal_world_model_context` | `False` | 世界モデルのコンテキストを因果的にエンコードし（コンテキストの各位置ごとにプレフィックス専用の V-JEPA2 パスを 1 回実行）ことで、双方向アテンションによって未来のフレームが予測器の入力に漏れるのを防ぎます。これにより、`t_enc_ctx` 回の追加エンコーダー呼び出しが発生します | | `num_inference_timesteps` | 4 | アクションのデノイジングのためのEuler積分ステップ | | `freeze_qwen` | `False` | Qwen3-VLのバックボーンをフリーズし、アクションヘッドのみを学習します | | `reinit_modules` | `None` | 読み込み時にランダムに再初期化が許可されるキーのプレフィックス（クロスエンボディメント転送については、[異なるエンボディメントでのファインチューニング](#fine-tuning-on-a-different-embodiment)を参照） | | `resize_images_to` | `None` | `(height, width)` は、Qwen3-VL のビジョンタワーに渡される前に、すべてのカメラフレームがリサイズされるサイズです。`None` を指定するとネイティブ解像度が維持され、Qwen3-VL のパッチ数がそれに伴って増加するため、720x1280 のカメラでは GPU メモリを使い切ってしまう可能性があります。公開されているチェックポイントは `[224, 224]` を使用しています。 | | `gripper_dim` | 6 | アクションベクトルにおけるグリップラ次元のインデックス。`gripper_joint_names`がデータセットのアクション名と一致する場合、無視されます | | `gripper_joint_names` | `["gripper"]` | グリップを識別するアクション次元の名前。一致したインデックスが`gripper_dim`よりも優先されます | | `gripper_threshold` | 0.5 | `pre_snap_gripper_action` および `binarize_gripper_action` で使用される閾値です。なお、`binarize` は正規化解除の_後_に実行されるため、この値はグリッパーの物理的な値と比較されます。 | | `pre_snap_gripper_action` | `False` | unnormalizationの前に、gripper dimを{0, 1}にスナップします。LIBERO固有の処理です。詳細は下記を参照してください | | `binarize_gripper_action` | `False` | unnormalization後、gripper dimを{-1, 1}に二値化します。LIBERO固有の処理であり、詳細は下記を参照してください | | `clip_normalized_actions` | `True` | `ACTION`が`MIN_MAX`を使用する場合にのみ適用されます。`MEAN_STD`では無視され（警告が表示されます）、1シグマで切り捨てられてしまうためです。非正規化の前に、正規化されたアクションを[-1, 1]にクリップします。 | | `world_model_num_views` | `None` | Camera world-model predictorが構築されるための視点。チェックポイントの形状に組み込まれています。`None`の場合、公開されているチェックポイントがエンコードしている`jepa_tubelet_size`にフォールバックします | `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ステップの学習を行いましたが、提供されている事前学習済みチェックポイントからファインチューニングを行う場合、より少ないステップ数でも良好なパフォーマンスが得られる可能性があります。 ### ゼロからフルトレーニング ```bash lerobot-train \ --policy.type=vla_jepa \ --policy.repo_id=your_org/your_repo \ --dataset.repo_id=your_org/your_dataset ``` ### 事前学習済みチェックポイントからのファインチューニング ```bash 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`を設定してください。 ```bash 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` | ```bash 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で学習したと記載されています。 ```bash lerobot-train \ --policy.path=lerobot/VLA-JEPA-Pretrain \ --policy.repo_id=your_org/your_repo \ --dataset.repo_id=HuggingFaceVLA/libero \ --steps=30000 ``` **事前学習済みLIBERO-10チェックポイントの評価:** ```bash 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 ``` タスクのサブセットのみを評価するには: ```bash 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 = 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 バックボーンとアクションヘッドのみが読み込まれて学習されます。これでも十分なアクション性能が得られます。 ```bash 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、出力投影）はカメラ数に依存しないため、事前学習済みチェックポイントから再利用できます。 --- ## 引用 ```bibtex @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](https://huggingface.co/ginwind/VLA-JEPA)リポジトリのライセンス条項（**Apache 2.0 License**）に基づいて配布されます。LeRobot統合コードは、**Apache 2.0 License**に従います。