技術要約: SAGE – 潜在世界モデルのプランニングにおけるサブゴール条件付きアクション生成

問題提起

潜在世界モデル（Latent world models）は、ピクセルレベルの再構成を行うことなく、コンパクトな表現空間内でアクション条件付きの予測ダイナミクスを学習することで、効率的なプランニング・パラダイムとして確立されている。しかし、プランニングのホライゾン（期間）が延びるにつれ、その有効性は提案されるアクションの質によってますます制約を受けるようになる。固定された候補予算の下では、一般的なランダム分布でアクション探索を初期化すると、指数関数的に拡大するアクション空間を十分にカバーすることができない。その結果、目標指向のシーケンスがサンプリングされることは稀であり、世界モデルが高品質な候補となる未来を評価するためにさらされることがなくなる。この「プロポーザル・ギャップ（提案の乖離）」は、たとえ基礎となる予測ダイナミクスが正確であっても、長期間のプランニング性能を著しく制限する。

手法: SAGEフレームワーク

このプロポーザル・ギャップに対処するため、著者らはSAGE（Subgoal-Conditioned Action Generation：サブゴール条件付きアクション生成）という、事前分布による条件付けを用いたプランニング・フレームワークを導入する。SAGEは、既存の潜在世界モデル（具体的にはLeWM）を凍結したまま、長期間のプランニングを、密接に結合された2つの学習された決定へと分解することで機能する。

1. 目標条件付きサブゴール生成器 (Goal-Conditioned Subgoal Generator)

第一のコンポーネントは、可変長の潜在サブゴール生成器である。現在の観測履歴、遠方の目標、および要求された将来の継続時間（$\tau$）が与えられると、このモジュールは、時刻 $t+\tau$ における到達可能な中間潜在状態（サブゴール）を予測する。

• アーキテクチャ: 履歴トークン、現在の状態、遠方の目標トークン、および時間埋め込み（$\delta, \tau$）を入力とする4層のTransformerデコーダ。
• メカニズム: 遠方の目標の潜在表現に対する残差補正を出力し、特定の期間 $\tau$ 内で到達可能なローカルターゲット $\hat{z}_{t+\tau}$ を生成する。
• 機能: この生成器は、遠方の目標を様々な時間スケールにおける到達可能な中間状態へと分解し、微細なローカルダイナミクス（短い $\tau$）と高次のタスク進行（長い $\tau$）の両方を捉える。

2. サブゴール条件付きアクション生成器 (Subgoal-Conditioned Action Generator)

第二のコンポーネントは、現在の履歴、遠方の目標、および予測されたローカルサブゴールを、アクションの選択肢の分布へとマッピングする。

• アーキテクチャ: 学習されたアクションクエリを用いて、条件付けメモリ（履歴、状態、サブゴール、遠方の目標）にアテンションを払う3層のTransformerデコーダ。
• 出力: ネットワークは、軌道レベルのガウス混合モデル（GMM）である $q_\phi(a_{t:t+\tau-1} | \dots) = \sum \pi_m \mathcal{N}(a; \mu_m, \Sigma_m)$ をパラメータ化する。
• 機能: 一般的な分布からサンプリングする代わりに、プランナーは特定のローカルサブゴールに条件付けられたアクションシーケンスをサンプリングする。これにより、提案されるアクションが、即時的かつ到達可能なターゲットに対して構造的に整合していることが保証される。

3. プランニングと洗練のループ (Planning and Refinement Loop)

プランニングプロセスは、これらの学習されたコンポーネントを凍結された世界モデルと統合する：

1. 分解 (Decomposition): 長期間のクエリ $H$ は、期間のスケジュール $(\tau_1, \dots, \tau_R)$ へと分解される。
2. サブゴール予測 (Subgoal Prediction): 各ステージにおいて、サブゴール生成器は残りのホライゾンに基づき、ローカルターゲット $\hat{z}_{t+\tau}$ を予測する。
3. プロポーザル生成 (Proposal Generation): アクション生成器は、このローカルターゲットに条件付けされた $K$ 個の候補アクションシーケンスを生成する。
4. 評価と洗練 (Evaluation & Refinement): 凍結されたLeWMがこれらの候補の潜在的な未来を予測し、ローカルサブゴールに対するコストを評価する。その後、交差エントロピー法（CEM）が30ラウンドにわたってエリート候補セットを洗練させる。
5. 実行 (Execution): エージェントは選択されたシーケンスの先頭セグメントを実行し、新しい状態を観測して、再プランニングを行う。

主な貢献

1. マルチホライゾン潜在サブゴール生成器: 遠方の目標を、異なる時間スケールにおける到達可能な中間状態へと分解する新しいモジュールであり、現在状態と遠方の目標との間のギャップを効果的に埋める。
2. サブゴール条件付きアクション生成器: 予測されたサブゴールに条件付けられたアクション候補を生成するメカニズムであり、一般的なランダムな初期化を、構造化されたタスク認識型の事前分布へと置き換える。
3. 構造化されたプロポーザルの実証的検証: 可変長のサブゴール生成と条件付きアクション生成を組み合わせることで、凍結された世界モデルの予測ダイナミクスを変更することなく、長期間のプランニング性能を大幅に向上させつつ、短期間の性能を維持できることを示した。

実験結果

本手法は、凍結されたLeWMチェックポイント、同一の候補予算（300）、およびCEMによる洗練（30ラウンド）を用いた制御された設定を用いて、PushT（持続的な平面接触を必要とする）およびOGBench Cube（腕の動き、把握、配置を必要とする）で評価された。

• 長期間の性能 (Long-Horizon Performance): プランニングホライゾンが長くなるにつれ、SAGEはベースライン（標準的なLeWMおよびPRISM）に対して大幅な向上を示した。
  • ターゲットオフセット H=150 において：
    • PushT: 成功率は 12.7% (LeWM) から 64.7% (SAGE) へと向上した。
    • OGBench Cube: 成功率は 26.7% (LeWM) から 67.3% (SAGE) へと向上した。
• 短期間の性能 (Short-Horizon Performance): SAGEは、目標が近い場合（例：H=25でのPushTの94.0%）においても強力な性能を維持しており、構造化された事前分布が性能を低下させないことを示している。
• アブレーション研究 (Ablation Studies):
  • 生成器のみ (Generator-Only): 生成されたローカルターゲットを使用し、標準的なガウス分布によるプロポーザル（学習されたアクション生成器なし）を用いた場合でも、依然として大幅な改善が見られた（例：H=150のPushTにおいて58.7% vs 12.7%）。これは、ローカルターゲットの分解の価値を裏付けている。
  • SAGE Prior Top: LeWMによるランキングやCEMによる洗練を行わず、学習された事前分布の最高重みオプションを実行した場合、長期間の性能は低かった（H=150で16.0%）。これは、世界モデルによる誘導付きの探索と洗練の必要性を強調している。
• 時間スケジュール (Temporal Schedules): 異なる期間スケジュール（例：15、20、25ステップのコミットメントの混合比率の変化）を用いた実験により、再プランニングの頻度と期間の順序の両方が成功率に大きく影響することが明らかになった。

意義と主張

本論文は、SAGEが、予測ダイナミスの再学習という計算コストをかけることなく、既存の潜在世界モデルのプランニング能力を拡張するための実用的な経路を提供すると主張している。アクションのプロポーザル（提案）の構造を改善すること、具体的には、遠方の目標を到達可能なローカルターゲットへと分解し、それらのターゲットに条件付けられたアクション生成を行うことにより、フレームワークは長期間のプランニングに固有のプロポーザルの質というボトルネックを効果的に緩和する。

著者らは、彼らのアプローチが「提案」メカニズムと「予測」メカニズムを分離していることを強調している。世界モデルは未来を評価するための固定されたシミュレータとして残る一方で、SAGEは、シミュレータに提示される候補が妥当であり、タスクの構造に沿ったものであることを保証する。このモジュール的な改善により、基礎となる世界モデルの効率性と堅牢性を維持したまま、複雑で長期間のロボットタスクにおいて大幅な利得を得ることが可能となる。