世界モデル（World Model）は、学習されたダイナミクスシミュレーターとしてますます高性能になっています。高次元の視覚空間で将来の観測の長いシーケンスを予測し、数年前には想像もできなかった方法でタスクを越えて一般化することができます。しかし、強力な予測モデルを持っていることは、それを制御や学習、計画に効果的に使用できることと同じではありません。実際には、現代の世界モデルを用いた長い時間スケールでの計画は依然として脆弱です。最適化が ill-conditioned になり、非貪欲な構造が悪い局所最小値を生み出し、高次元の潜在空間が微妙な失敗モードを導入します。
これらの問題に対処するため、カリフォルニア大学バークレー校のAI研究チーム（BAIR）のMichael Psenka、Mike Rabbat、Aditi Krishnapriyan、Yann LeCun、Amir Barは、GRASP（Gradient RelAxed Stochastic Planner）を提案しました。GRASPは、軌道を仮想状態に持ち上げる、状態反復に確率性を追加する、勾配を再形成するという3つの革新により、長い時間スケールでの計画を実用的にします。
世界モデルと計画の課題
世界モデルは、現在の状態と将来の行動のシーケンスが与えられたときに、次に何が起こるかを予測する学習モデルです。計画では、最も単純な方法は、モデルを通じて行動シーケンスを展開し、終端誤差を最小化することです。しかし、長い時間スケールの展開は深く ill-conditioned な計算グラフを作成し、勾配の爆発や消失を引き起こします。さらに、長い時間スケールのタスクは通常、非貪欲な行動（例えば障害物を回避する）を必要とし、損失ランドスケープは局所的な罠で満たされます。
自然な解決策の1つは、動的制約をソフト制約として扱い、行動と状態の両方を同時に最適化するコロケーション（collocation）アプローチです。コロケーションはすべてのタイムステップを並列計算可能にし、勾配が深い展開を介して逆伝播されるのを防ぎ、長期依存の問題を回避します。しかし、深層学習ベースの世界モデルでは、状態を直接最適化すると敵対的ロバスト性の問題に直面します。状態入力勾配は非常に脆弱です。というのも、訓練プロセスはデータ多様体の接線方向のみを正則化し、法線方向の振る舞いは制御されないため、モデルは未観測の状態に対して非常に敏感になるからです。
GRASPの核心的革新
GRASPはコロケーションフレームワークに基づいていますが、深層学習モデルの勾配特性に特化して設計されています。重要な洞察は、状態勾配は信頼できないが、行動勾配は通常良好に振る舞うということです。行動空間は低次元で十分に訓練されているためです。したがって、GRASPは計画に行動勾配のみを使用します。
GRASPは2つの主要な要素から構成されます：
- 状態ノイズによる探索：最適化中に仮想状態更新にガウスノイズを注入し、異なる損失の谷間へのジャンプを促進します。行動は勾配によって導かれます。
- 勾配の再形成：状態から世界モデルへの勾配伝播を停止し、敵対的脆弱性を回避します。同時に、密な目標項（各ステップで目標に近づくよう促す）を追加し、停止勾配のダイナミクス損失とバランスをとることで、自明な解に陥るのを防ぎます。
さらに、GRASPはK反復ごとに短い「同期」フェーズを実行します。現在の行動を使用して初期状態から展開し、元の系列損失に対して少数の勾配ステップを実行することで、状態と行動を実際の軌道と一致させます。
実験結果
Push-Tタスクにおいて、GRASPはCEM、標準勾配降下法（GD）、LatCoなどの手法と比較されました。結果は、計画の時間スケールが長くなるにつれて、GRASPが成功率と計画速度の両方で他の手法を大幅に上回ることを示しています。例えば、時間スケールH=60では、GRASPの成功率は26.2%であり、CEMの7.2%、GDの16.4%を上回りました。また、GRASPの計画時間の中央値は49.1秒であり、他の手法よりもはるかに高速でした。GRASPは長い時間スケールでより高い成功率を達成し、より高速に計画を実行します。
将来の展望
研究者らは、現在は世界モデル計画に取り組むのにエキサイティングな時期であると考えています。古典的な計画と制御の文献は非常に成熟していますが、現代の大規模世界モデルにおける純粋な計画最適化はまだ十分に探索されていません。将来の方向性としては、拡散世界モデルへの拡張、より洗練された最適化器とノイズ戦略の採用、GRASPの閉ループシステムや強化学習ポリシーへの統合などが挙げられます。