技術要約: FF-JEPA: 潜在プランナーを用いた世界モデルにおける長期的計画策定

問題提起

視覚的世界モデル、特にLeWorldModel (LeWM) のようなJoint Embedding Predictive Architectures (JEPA) の最近の進展は、将来の観測を予測し、Cross-Entropy Method (CEM) のようなサンプリングベースのプランナーを介したモデルベース制御をサポートする能力を示してきました。しかし、著者らは、これらの「フラット」な世界モデルを現実世界の長期的なタスクに展開することを妨げている3つの決定的な限界を特定しています：

1. 計算上の困難: 長期的なホライゾン（期間）にわたって計画を立てるには、繰り返しのロールアウトと、ますます増大するアクションシーケンスが必要となり、CEMによる最適化が計算量的に実行不可能になります。
2. 誤差の累積: 自己回帰的な予測誤差は多くのステップにわたって蓄積され、複雑な軌道に対してプランナーを逸脱させます。
3. 目標への依存性: 既存の手法は通常、計画を駆動するために明示的な目標画像または状態 ($o_g$) を必要とします。多くの現実世界のシナリオでは、このような具体的な目標指定は利用できないか、あるいは非現実的です。

階層的手法や生成的アプローチがこれらの問題に対処しようとしてきましたが、それらは依然としてピクセル空間で動作しているか、あるいは外部で指定された目標画像や軌跡に依存していることが多く、オープンエンドな環境における適用性を制限しています。

手法: Forward-Forward JEPA (FF-JEPA)

本論文では、標準的なアクション条件付きフォワードモデル ($P$) と並行して動作する潜在プランナー ($G$) を導入することで、JEPAスタイルの世界モデルを拡張する階層的フレームワークであるFF-JEPAを提案します。このアプローチは、明示的な目標画像を必要とせずに、長期的な計画策定を、扱いやすい短期的な最適化問題のシーケンスへと分解します。

コアアーキテクチャ

システムは、凍結された学習済み世界モデルエンコーダ ($E$) によって定義される潜在空間内で学習される、主に2つのコンポーネントで構成されています：

1. 世界モデル ($P$): 前のステートの履歴とアクションに基づいて次の潜在状態を予測する、標準的なアクション条件付きフォワードダイナミクスモデル（LeWMに基づく）：$\hat{z}_{t+1} = P(z_{t-W_P+1:t}, a_t)$。
2. 潜在プランナー ($G$): アクションを用いないフォワードダイナミクスモデルであり、過去のサブゴール潜在状態の履歴のみに基づいて、次のサブゴール潜在状態 ($\hat{z}_{sg}$) を予測します。これは、より高い時間解像度（環境ステップ $H$ ごと）で動作します。
   • 決定論的プランナー ($G_{Det}$): 予測された将来のサブゴールと真の将来のサブゴールの間の平均二乗誤差を最小化するように訓練された、トランスフォーマーアーキテクチャ（世界モデルの予測器と同様だが、アクション条件付けはない）。
   • 拡散プランナー ($G_{DM}$): 将来のサブゴールのシーケンスを予測するために、デノイジング・スコアマッチング目的関数を用いて訓練されたDiT (Diffusion Transformer) バックボーン。

計画推論 (Planning Inference)

推論プロセスは階層的です：

1. サブゴールの予測: 潜在プランナー $G$ が、現在のサブゴール履歴に基づいて、次のサブゴール状態 $\hat{z}_{sg, m+1}$ を予測します。
2. アクションの最適化: 世界モデル $P$ は、予測された状態を短いホライゾン $H$ 以内に予測されたサブゴール $\hat{z}_{sg, m+1}$ へと向かわせるためのアクションシーケンス ($a_{1:H}$) を見つけるためにCEMを使用します。
3. 反復: このサイクルを繰り返すことで、短期間の計画を連鎖させ、長期的な目標を達成します。

決定的なのは、この手法が推論時にターゲットとなる目標画像を使用する必要を排除している点です。プランナーは、暗黙的に定義された目標に向かう軌道を生成します。

主な貢献

• 目標フリーの長期計画策定: 本フレームワークは、潜在プランナーを使用して中間的なサブゴールを生成することにより、明示的な目標画像なしで世界モデル内での計画策定を可能にします。
• 階層的分解: 複雑な軌道を、短期間で解決可能なサブ問題へと分解することで、「長期ホライゾン崩壊」の問題に対処し、誤差の累積と計算コストを軽減します。
• ラベルなし学習: 潜在プランナーは、ラベルなしの成功したデモンストレーション（世界モデルの潜在空間からサブサンプリングされたもの）で訓練できるため、アクションラベル付きデータへの依存度を低減できます。
• 統一されたフレームワーク: 予測モデリングと制御を単一の整合した潜在空間内で統合し、世界モデルを、想像された潜在軌道上で動作する逆ダイナミクスモジュールとして再解釈します。

実証評価

著者らは、PushTタスクにおける3つの実験設定（短期ホライゾン：25ステップ、長期ホロンゾン：75ステップ、ランダム初期化：ランダムな開始位置での300ステップ）を用いてFF-JEPAを評価しました。

結果

• 長期ホライゾンのパフォーマンス: ベースラインのフラットなLeWMは、長期タスクにおいて崩壊し（$t=75$ で成功率3.52%）、ランダム初期化下では完全に失敗しました（0.00%）。対照的に、FF-JEPA (Diffusion) は、$t=75$ で 91.80%、ランダム初期化下で 82.42% の成功率を達成しました。
• 目標フリーの優位性: 推論時に目標画像（goal image）を必要とするベースライン（LeWM, DINO-WM）とは異なり、FF-JEPAはそれを必要としません。それにもかかわらず、FF-JEPAは短期（89.0% vs 96.09%）および長期（61.0% vs 91.80%）の両方の設定において、階層的DINO-WMベースラインを上回りました。
• データ効率: アブレーション研究により、200個の高品質なデモンストレーションのみ（40倍の削減）で訓練された拡散プランナーが、フルデータを用いた場合（82.42%）と同等の競争力のあるパフォーマンス（76.17%）を示したことが明らかになり、高いデータ効率が示唆されました。
• 計算オーバーヘッド: 決定論的プランナーによる推論時間のオーバーヘッドは、CEM（926ms）と比較して無視できる程度（2.1ms）であり、拡散プランナーは約242msでした。これらはいずれも、長期的なフラットな計画策定に必要とされる繰り返しのロールアウトよりも大幅に高速です。

失敗分析

失敗事例は、エージェントが予測されたサブゴールの位置から離れすぎた場合、または潜在プランナーが低品質なサブゴール状態（例：エージェントを見落としている）を生成した場合に発生しました。

重要性と主張

本論文は、FF-JEPAをスケーラブルで目標フリーなロボット学習への有望な方向性として位置付けています。アクション条件付きの世界モデルを、ラベルなしデータで訓練された潜在プランナーと組み合わせることで、フラットな計画策定の限界を克服できることを示すことで、著者らは、このアプローチが明示的な目標仕様が利用不可能なオープンエンドな環境においてエージェントを動作させることを可能にすると主張しています。本研究は、計画を短期的なサブゴールへと分解することが、現在の世界モデルのアプローチに固有の計算量および誤差累積の問題を緩和するための有効な戦略であることを強調しています。