Closing the Train-Test Gap in World Models for Gradient-Based Planning

基本情報

• arXiv ID は 2512.09929v1 である（https://arxiv.org/abs/2512.09929 ）
• 著者は Arjun Parthasarathy、Nimit Kalra、Rohun Agrawal らである
• 所属機関は Columbia University および New York University である
• 投稿日は2024年12月13日である
• カテゴリは cs.AI、cs.LG である

簡単に説明すると

この論文は、機械学習における世界モデルを用いた勾配ベース計画法の性能を具体的に改善する手法を提案しています。
世界モデルとは環境の動的システムを学習したモデルのことです。
従来、世界モデルは次状態予測の目的で訓練されます。
しかし実際の使用時には行動シーケンスの最適化に利用されるため、訓練時と使用時でのギャップが問題となっていました。

本研究では「Online World Modeling」と「Adversarial World Modeling」という2つの新しい手法を提案しています。
これらの手法によってギャップを解消し、勾配ベース計画の信頼性と性能を向上させることに成功しました。
実験では従来のサンプリングベース手法（Cross-Entropy Method）と同等以上の性能を実現しています。
さらに計算時間を10分の1に短縮しました。
GitHubリポジトリは https://github.com/nimitkalra/robust-world-model-planning で公開されています。

1. 研究概要

1.1 背景と動機

ロボットタスクにおいて、エージェントの行動が環境状態にどのような影響を与えるかを予測することは、予測と計画の両方において基本的な要素です。
古典的な手法では環境進化モデルを第一原理から解析的に導出しますが、これは環境やエージェントの事前知識に依存します。
一方、学習ベースの手法は複雑な動的システムを捉えることができ、汎化性と不確実性への頑健性を改善します。

世界モデルは現在の状態と行動が与えられたときの次状態を予測する強力なパラダイムとして登場しました。
これらのモデルは正確な状態情報または画像などの高次元的な感覚入力から学習します。
後者の設定では、事前学習済み視覚表現を活用して生画像から直接、知覚・予測・制御を可能にし、実際には困難な精密な環境状態測定の必要性を排除します。

世界モデルは多くのモデルベース計画アルゴリズムと互換性があります。
従来の手法（DDP、iLQR）は線形・二次近似から導出された最適化問題を反復的に解きますが、大規模世界モデルでは計算が困難になります。
CEMやMPPIなどのサーチベース手法は候補解を反復的にサンプリングして評価する必要があり、高次元空間ではスケーラビリティに問題があります。

勾配ベース手法は世界モデルの微分可能性を直接活用して行動をエンドツーエンドで最適化し、コストの高いロールアウトを排除します。
しかし、この手法は根本的な訓練・テスト間のギャップに悩まされています。
世界モデルは通常、専門家軌道データセットでの次状態予測の目的で訓練されますが、テスト時には行動シーケンスに対する計画目的の最適化に使用されます。

1.2 主要な貢献

この研究では、上記の課題への対処として2つのアルゴリズムを提案している。

• Online World Modeling：勾配ベース計画により生成された軌道を環境シミュレータで修正し、修正された軌道で世界モデルをファインチューニングする
• Adversarial World Modeling：専門家軌道に敵対的摂動を加えてファインチューニングし、計画損失面を平滑化する
• 両手法により勾配ベース計画の性能が18-30%改善され、CEMと同等以上の性能を10分の1の計算時間で達成した
• Adversarial World Modelingが計画損失面を平滑化し、両手法が世界モデル誤差の訓練・テスト間ギャップを逆転させることを実証しました

2. 提案手法

2.1 手法の概要

世界モデルは状態空間S×行動空間Aから状態空間Sへの動的関数h: S×A→Sを近似します。
実際には、観測空間Oから低次元の潜在空間Zへの埋め込み関数Φμ: O→Zを用いて、潜在世界モデルfθ: Z×A→Zを学習する。

訓練では、軌道データセットTから（ot, at, ot+1）の3つ組をサンプリングし、真の次潜在状態と予測の次潜在状態間のL2距離を最小化する。
計画時には、学習済み世界モデルを用いて目標状態に到達する行動シーケンスを最適化し、勾配ベース計画は勾配降下により計画目的を解きます。

2.2 技術的詳細

Online World Modelingは、勾配ベース計画中に最適化される行動シーケンスが訓練時の専門家の行動分布内に制約されないという問題に対処する。
これにより世界モデルが分布外状態に駆動され、予測誤差が蓄積します。

この手法は以下のステップを実行する。

1. 専門家軌道の初期・目標潜在状態を用いて勾配ベースで計画する
2. 予測行動シーケンスを真の動的シミュレータで実行して修正軌道を取得する
3. 修正軌道をデータセットへ追加して世界モデルを再訓練する

Adversarial World Modelingは、世界モデルの入力勾配の振る舞いを改善することを目的とします。
敵対的訓練は入力勾配の振る舞いを改善し、結果として入力損失面を平滑化することが示されています。
この手法は、世界モデルが性能不良となることが予想される状態・行動空間の領域を明示的にターゲットとする敵対的な訓練目的を使用する。

敵対的サンプルは、Fast Gradient Sign Method（FGSM）を用いて生成され、予測誤差を最大化する最悪ケース摂動を効率的に近似します。

2.3 新規性

既存の世界モデル手法と比較して、本研究の新規性は以下の点にある。

従来の世界モデルは次状態予測の精度改善に焦点を当てていたが、本研究は勾配ベース計画の収束性と信頼性向上を特に目標としている。
Online World Modelingは、計画アルゴリズム自体が生成する軌道を用いてデータセットを集約する点で、低次元の状態空間での従来手法とは異なる。
Adversarial World Modelingは、計画時の最適化の安定性向上のため、世界モデル訓練に敵対的摂動を直接適用する新しいアプローチである。

両手法とも高次元の潜在世界モデルのコンテキストで動作し、事前学習済み視覚エンコーダ（DINOv2）との組み合わせで効果を発揮する。

3. 実験結果

3.1 実験設定

実験は3つのタスク（PushT、PointMaze、Wall）で実施され、初期構成から目標構成への到達成功率で評価されました。
オープンループとMPC（Model Predictive Control）の両方の設定で計画結果を報告しています。

ベースライン世界モデルとしてDINO-WMを使用し、埋め込み関数にはDINOv2エンコーダを採用しました。
遷移モデルはViTアーキテクチャで実装され、埋め込み関数は固定したまま遷移モデルのみをファインチューニングしました。

行動シーケンスの初期化には標準正規分布からのランダムサンプリングと初期化ネットワークの両方を評価し、ランダム初期化の方が優れた性能を示しました。
勾配ベース計画では、最後の状態からだけでなく各予測状態から勾配を得るために重み付き目標損失を使用しました。

3.2 主要な結果

全3タスクにおいて、提案手法はDINO-WMを勾配降下で上回り、より高価なCEMと同等または上回る性能を達成しました。

オープンループ設定では以下の成功率向上を達成しました：

• PushT：+18%向上
• PointMaze：+20%向上
• Wall：+30%向上

MPC設定では、AdamによるAdversarial World Modelingが特に優秀で、PointMazeとWallでCEMを上回り、PushTでCEMと同等の性能を達成しました。

計算効率の観点では、勾配ベース計画はCEMより桁違いに高速であることが示されました。
これは実世界の計画アプリケーションにおける本手法の実用性を示しています。

3.3 既存手法との比較

Online World ModelingとAdversarial World Modelingは異なる分布を誘導します。
Online World Modelingは計画時に見られる分布を予測してカバーするのに対し、Adversarial World Modelingは世界モデルの現在の損失面を利用して専門家軌道近傍での局所平滑性を促進します。

全環境において、AdamでGBPを実行する場合、Adversarial World ModelingがOnline World Modelingを上回りました。
これは、Adversarial World Modelingによる損失面の平滑化効果が、複数の初期化からのより安定した最適化を可能にするためと考えられます。

DINO-WMの元論文で報告されたWall環境でのCEM結果を再現できなかったため、提供されたチェックポイントとコードリポジトリでは32%の成功率を得ました（論文記載は74%）。

4. 実用性評価

4.1 実装の容易性

提案手法は既存の世界モデルアーキテクチャ（DINO-WM、IRIS）に容易に適用可能です。
Online World Modelingは環境シミュレータへのアクセスが必要ですが、多くのロボティクス研究環境では利用可能です。
Adversarial World Modelingはシミュレータ不要で実装でき、計算コストもFGSMの使用により管理可能です。

両手法とも標準的な深層学習フレームワークで実装可能で、既存のコードベースへの統合が容易です。
ハイパーパラメータ調整も比較的簡単で、摂動の大きさ（εa、εz）が主要な調整対象となります。

4.2 計算効率

勾配ベース計画の主要な利点は計算効率であり、本研究でもこれが実証されています。
CEMと比較して10倍の高速化を達成しており、リアルタイム制御アプリケーションに適しています。

Adversarial World Modelingのファインチューニングコストは追加的ですが、敵対的サンプル生成にFGSMを使用することで効率的です。
PGDなどのより強力な攻撃手法と比較して、FGSMは同等の改善を提供しながら大幅に計算効率が高いです。

4.3 応用可能性

本手法はロボット操作タスクや自律ナビゲーションシステムに広く適用可能です。
特に、リアルタイム性が要求される環境や、高次元状態・行動空間を持つタスクでの有効性が期待されます。

多時間スケールや階層的世界モデルとの組み合わせにも適しており、異なる抽象化レベルでの計画安定性向上により長期計画が可能となります。
実世界のノイズや環境摂動に対する頑健性向上も期待されるため、実際のロボットシステムでの展開に有望です。

5. まとめと所感

5.1 論文の意義

この論文は、世界モデルベースの勾配計画における根本的な問題である訓練・テスト間ギャップに対処する実用的な解決策を提供しています。
提案された2つの手法は理論的に妥当で、実験的に有効性が実証されており、実際のロボティクスアプリケーションでの採用が期待されます。

特に、計算効率の大幅な改善（10倍高速化）は実用的な観点から非常に重要です。
これにより勾配ベース計画がサンプリングベース手法の実用的な代替案となる可能性が示されました。

著者陣には著名な研究者（Yann LeCunなど）が含まれており、研究の信頼性と影響力を高めています。
実装コードの公開も研究の再現性と発展に貢献しています。

5.2 今後の展望

今後の研究方向として、実世界システムでの評価が重要です。
シミュレーション環境での成功を実際のロボットシステムに転移できるかが鍵となります。

敵対的訓練の副次効果として、環境の敵対者や確率的要素に対する世界モデルの頑健性向上も期待されます。
より長期的には、マルチエージェント環境や部分観測環境での適用も興味深い研究方向です。

計算効率のさらなる改善や、より複雑な敵対的攻撃手法との組み合わせも検討価値があります。
また、異なる世界モデルアーキテクチャでの汎用性検証も重要な課題です。