GRASP: 長期ホライズンにおけるワールドモデルの勾配ベース計画
BAIR は GRASP (Gradient RelAxed Stochastic Planner) を導入しました。これは、学習されたワールドモデルに対して長期ホライズン計画を実用的にすることを目的とした新しい勾配ベースプランナーです。GRASP は、軌道を仮想状態に持ち上げて並列最適化を行い、状態イテレートに確率的変動を注入して探索を行い、勾配を脆弱な状態入力勾配ではなく安定したアクションヤコビアンに依存させるように再形成することで、長期計画の脆弱性を解決します。
ワールドモデルにおける長期ホライズン計画の課題
大規模なワールドモデルは汎用シミュレータとして機能し得ますが、制御や計画に使用する際はしばしば脆弱です。長期ホライズン計画は主に次の三つの要因で失敗します。
- Ill-conditioned Computation Graphs（条件の悪い計算グラフ）: モデルを自分自身に繰り返し適用して微分（Backpropagation Through Time）すると、勾配が爆発または消失します。ヤコビアンの条件数は時間ホライズン $T$ に対して指数的にスケールします。
- Non-Greedy Landscapes（非貪欲的なランドスケープ）: 長期タスクはしばしば非貪欲的な振る舞い（例: 障害物を回避するために目標から離れる）を必要とします。ホライズンが伸びるにつれて最適化空間が拡大し、局所最小値に陥る可能性が高まります。
- State-Input Gradient Sensitivity（状態入力勾配の感度）: 深層学習ベースのワールドモデルでは、状態入力に対する勾配 ($D_s F_\theta$) が信頼できないことが多いです。ワールドモデルは低次元データ多様体上で訓練されるため、状態のわずかな摂動でモデルを望む結果に騙す「敵対的」サンプルに対し脆弱であり、最適化ランドスケープが「粘着」し信頼性が低くなります。
GRASP の技術的アプローチ
GRASP は、コロケーションベースの計画と特定の勾配再形成・探索戦略を組み合わせることでこれらの失敗を克服します。
1. 動力学制約のリフティング（Collocation）
シリアルロールアウト $s_{t+1}=F_\theta(s_t, a_t)$ の代わりに、GRASP は動力学をソフト制約として扱います。アクションと状態の両方を同時に最適化するペナルティ関数を最小化します。
$$\min_{\mathbf{s},\mathbf{a}} \mathcal{L}(\mathbf{s}, \mathbf{a}) = \sum_{t=0}^{T-1} \big|F_\theta(s_t,a_t) - s_{t+1}\big|_2^2, \quad \text{with } s_0 \text{ fixed and } s_T=g$$
この「リフティング」された定式化により、すべての $T$ 項を時間軸上で並列に計算でき、勾配信号が深い $T$ ステップの合成を逆伝搬する必要がなくなります。
2. Stop-Gradient による勾配リシェイピング
状態ヤコビアンの敵対的感度を回避するため、GRASP は最適化がアクションヤコビアン ($D_a F_\theta$) のみを利用するようにします。これらは通常低次元で、より密に訓練されています。
GRASP は stop-gradient dynamics loss を実装し、勾配がワールドモデルの状態入力に流れ込むのを遮断します。システムが自明な最小値に陥るのを防ぐため、GRASP は dense goal shaping 項を加えます。
$$\mathcal{L}(\mathbf{s},\mathbf{a}) = \mathcal{L}{\text{dyn}}^{\text{sg}}(\mathbf{s},\mathbf{a}) + \gamma \mathcal{L}{\text{goal}}^{\text{sg}}(\mathbf{s},\mathbf{a})$$
この目的関数は、脆弱な状態勾配に依存せずに、実現可能な動力学と目標への一貫したシグナルをバランスさせます。
3. 探索のための状態イテレート確率的変動
長期計画の非凸ランドスケープを探索するため、GRASP は最適化中に仮想状態更新に直接ガウスノイズを注入します。
$$s_t \leftarrow s_t - \eta_s \nabla_{s_t}\mathcal{L} + \sigma_{\text{state}} \xi, \qquad \xi\sim\mathcal{N}(0,I)$$
状態にノイズを加えることで、プランナーはリフティング空間の異なるベースン間を「跳び」ながら、アクション更新はガイドされたままです。
4. 定期的同期
リフティングされた stop‑gradient 目的は近似であるため、GRASP は $K_{\text{sync}}$ イテレーションごとに「同期」フェーズを実行します。$s_0$ から現在のアクションでロールアウトし、元のシリアル損失上で少数の小さな勾配ステップを取ることで、計画が実際の軌道に根ざしたままになることを保証します。
パフォーマンスと結果
GRASP は、Push‑T 環境においてベースラインプランナー（Cross‑Entropy Method (CEM) や標準 Gradient Descent (GD)）に比べ、成功率と計算速度の両面で大幅な改善を示します。ホライズン $H$ が増加しても、GRASP は高い成功率と低い中央値達成時間を維持します。
| ホライズン | CEM | GD | LatCo | GRASP |
|---|---|---|---|---|
| H=40 | 61.4% / 35.3s | 51.0% / 18.0s | 15.0% / 598.0s | 59.0% / 8.5s |
| H=50 | 30.2% / 96.2s | 37.6% / 76.3s | 4.2% / 1114.7s | 43.4% / 15.2s |
| H=60 | 7.2% / 83.1s | 16.4% / 146.5s | 26.2% / 49.1s | |
| H=70 | 7.8% / 156.1s | 12.0% / 103.1s | 0.0% / — | 16.0% / 79.9s |
| H=80 | 2.8% / 132.2s | 6.4% / 161.3s | 0.0% / — | 10.4% / 58.9s |
(成功率 % / 中央達成時間)
今後の方向性
研究者らは、GRASP を拡張するためのいくつかの道筋を提案しています。具体的には、拡散ベースのワールドモデルへの適用、より洗練されたノイズ付加戦略の開発、プランナーを閉ループシステムや適応的長期ホライズン計画のための RL ポリシー学習に統合することなどです。