技術要約: AdaJEPA: 適応型潜在世界モデル

問題提起

潜在世界モデル、特に結合埋め込み予測アーキテクチャ（JEPA）は、コンパクトな潜在空間で将来の状態を予測することにより、効率的なプランニングを可能にします。これらのモデルは通常、報酬のないオフラインの軌跡を用いて学習され、目標指向タスクのためにモデル予測制御（MPC）と組み合わせて使用されます。しかし、決定的な制限が存在します。それは、世界モデルが学習後に固定されてしまうという点です。

展開時、テスト時の分布シフト（例：ノイズや照明の変化などの視覚的変化、あるいは摩擦や質量などの動的な変化）によってモデルの予測が不正確になった場合、MPCプランナーは誤った目的関数を最適化してしまいます。潜在的な「想像」の中では効果的に見えるアクションも、実際の環境では失敗する可能性があります。わずかな一歩の誤差がプランニングのホライゾン（期間）にわたって蓄積し、大幅な性能低下を招きます。現在のアプローチでは、追加のエキスパートによるデモンストレーション、ターゲットドメインへのファインチューニング用データ、あるいは複雑な外部ループが必要となることが多く、これが実世界への応用を妨げています。

手法: AdaJEPA

著者らは、MPCのクローズドループ内でテスト時適応（test-time adaptation）を行う適応型潜在世界モデルであるAdaJEPAを提案しています。モデルを静的なものとして扱うのではなく、AdaJEPAはエージェント自身の相互作用から得られる自己教師あり信号を用いて、世界モデルを継続的に再校正します。

プラン・実行・適応・再プランニングのループ

コアとなるメカニズムは、各MPCステップにおいて以下のように動作します。

1. プラン（Plan）: エージェントは、現在の（おそらく古くなった）世界モデルを使用して、潜在的な目標到達コストを最小化するように一連のアクションを計画します。
2. 実行（Execute）: エージェントは、環境内で最初のアクションチャンク（$a_t$）を実行します。
3. 観測（Observe）: エージェントは次の観測（$o_{t+1}$）を受け取ります。
4. 適応（Adapt）: 遷移 $\{o_t, a_t, o_{t+1}\}$ がリプレイバッファに追加されます。この新しい遷移における潜在予測誤差を最小化するために、少数の勾配ステップ（多くの場合、わずか1ステップ）を経てモデルパラメータが更新されます。
5. 再プラン（Replan）: 更新されたモデルが、直ちに次のプランニングイテレーションに使用されます。

技術的構成要素

• アーキテクチャ: 基盤となる世界モデルは、感覚エンコーダ（$\mathcal{E}_s$）、アクションエンコーダ（$\mathcal{E}_a$）、および予測器（$f_\theta$）で構成されています。これは、ピクセルを再構成するのではなく、将来の潜在表現（$z_{t+1}$）を予測するようにオフラインで学習されます。
• 適応損失: 適応には、事前学習と同じ自己教師あり潜在予測目的関数が使用されます。
  $$\mathcal{L}_{ada}(\mathcal{B}) = \frac{1}{|\mathcal{B}|} \sum_{(o_i, a_i, o_{i+1}) \in \mathcal{B}} \ell(f_\theta(z_i, \mathcal{E}_a(a_i)), \text{sg}(z_{i+1}))$$
  ここで、$\mathcal{B}$ は最近の遷移のバッファであり、$\ell$ は潜在予測損失（例：MSE）であり、$\text{sg}(\cdot)$ は崩壊を防ぐためのストップグラディエント演算子です。
• パラメータ更新: 効率性を確保するため、適応は軽量に行われます。著者らは、ネットワーク全体ではなく、パラメータの小さなサブセット（例：エンコーダと予測器の最終層）のみを更新します。
• バッファ戦略: リプレイバッファ $\mathcal{B}$ は、「Recent-N」（直近の $N$ 個の遷移を保持する）または「Hard-N」（予測誤差が最大の遷移を保持する）によって管理できます。

主な貢献

1. プランニング中の初のJEPA適応: 著者の知る限り、これはJEPA世界モデルを、事前学習やオフラインのファインチューニングだけに頼るのではなく、プランニングのフェーズ中に適応させる初めての研究です。
2. クローズドループ自己教師あり学習: この手法は、実行されたアクションの直接的な結果を自己教師あり信号として利用するため、報酬ラベル、エキスパートのデモンストレーション、または個別のデータ収集フェーズを必要としません。
3. サンプルおよび計算効率: このフレームワークは、最小限の計算オーバーヘッドで大幅な性能向上を実現します。多くの実験において、単一の勾配ステップが、MPCの再プランニングステップあたりのステップとして十分であることが示されています。
4. 分布シフトへの堅牢性: このアプローチは、視覚的シフト（ノイズ、ぼかし、色の変化）、動的シフト（質量、減衰）、および構造的シフト（未知の迷路レイアウト）に対して堅牢性を示します。

実験結果

著者らは、様々な訓練・テスト間の分布シフト条件下で、PushT（ブロック押し）およびPointMaze（ナビゲーション）のベンチマークを用いてAdaJEPAを評価しました。

• イン・ディストリビューション（分布内）性能: 学習データに近いタスクにおいても、AdaJEPAは成功率を向上させ（例：特定のPushObj形状で20%以上の向上）、あるいは性能を低下させることなくベースラインの性能を維持しており、イン・ディストリビューションにおいて適用することが安全であることを証明しています。
• アウト・オブ・ディストリビューション（分布外）性能:
  • 形状の変化: 未知のブロック形状において、AdaJEPAは凍結されたモデルと比較してプランニングの成功率をほぼ倍増させました。
  • 視覚的変化: ガウスぼかし、ごま塩ノイズ、暗い照明の下でも明確な利点を示しました。
  • ダイナミクスおよびレイアウトの変化: PointMazeにおいて、AdaJEPAは低質量や高減衰によって引き起こされる失敗から回復し、未知の迷路レイアウト（例：特定のレイアウトシフトで25.3%の向上）における成功率を大幅に改善しました。
• レイテンシ: 適応ステップは、パラメータの小さなサブセットのみを更新するため、再プランニングあたりのレイテンシは無視できる程度です（約0.01〜0.03秒）。さらに、予測誤差からより速く回復することで、エージェントは総再プラン回数が少なくなることが多く、全体的な実行時間を短縮します。
• データ不足への汎化: 低データ環境において、テスト時適応は限定的なトレーニングカバレッジを補完しました。例えば、トレーニングデータが限られている場合でも、AdaJEDAは16倍の軌跡を用いて学習された凍結モデルを上回りました。

意義と主張

本論文は、潜在世界モデルは学習後に固定されるべきではなく、展開中に遭遇した経験から継続的に改善されるべきであると主張しています。

• 生物学的インスピレーション: このアプローチは、変化する入力やダイナミクスに基づいて行動を調整するシステム（小脳など）のような、生物学的な感覚運動適応に基づいています。
• データスケーリングの補完: より大規模で多様なトレーニングデータは汎化性能を向上させますが、AdaJEPAは展開中のモデルを洗練するためのサンプル効率の高い経路を提供し、データスケーリングを補完するものです。
• 今後の方向性: 著者らは、次のステップとして、この軽量なテスト時適応を継続学習および能動学習と組み合わせ、時間の経過とともに世界モデルのカバレッジを拡大し、変化する環境におけるより強靭な知覚とプランニングを可能にすることを示唆しています。

著者らは、彼らの手法が特定の世界モデルの実装に依存しない（異なるアーキテクチャ（ResNet、Transformer、DINOベース）やプランナー（勾配降下法、CEM）にわたって一貫した改善を示している）ことを強調しています。