技術要約: IMWM (Intuition Models Complement World Models for Latent Planning)

1. 問題提起

本論文は、生のピクセルからの制御におけるモデルベース強化学習（RL）における根本的な限界に対処している。それは、学習された潜在世界モデル（World Model: WM）と、有限の予算を持つサンプルベースのプランナー（例：Cross-Entropy Method, CEM）への依存である。世界モデルの精度を高めることが、プランニング性能を向上させるための標準的なアプローチであるが、著者らはその直感は不十分であると主張している。

特定された核心的な問題は、「予測ではなく、探索が制約となる可能性がある」ということである。たとえプランナーが「完璧な」世界モデル（真の環境ダイナミクスを用いてロールアウトを行うオラクルとして定義される）を保有していたとしても、有限の探索予算によって成功するアクションシーケンスを提案できなかった場合、目標に到達できない可能性がある。論文は、ボトルネックはフォワード予測器の正確さではなく、「提案のボリューム（提案量）」、すなわち、目標達成に繋がる候補を見つけ出す能力にあると断じている。

2. 診断と理論的根拠

解決策を提案する前に、著者らは探索が失敗の要因であることを示す実験的および理論的な証拠を提示している。

• 実験的診断: 標準的な潜在プランナーにおいて、学習されたフォワード予測器を理想化された真の環境ダイナミクスのロールアウト（Oracle Dynamics）に置き換えた。完璧なダイナミクスを備えていても、プランナーは OGBenc-Cube や Two-Room といったタスクにおいて依然として大幅に失敗した。失敗したエピソードの分析により、ほとんどのケース（例：Two-Roomでは98.9%）において、再プランニングのどのステップにおいても、CEMの集団内に目標達成となる候補がゼロであったことが明らかになった。
• 理論的境界: 著者らは、任意の潜在コスト表面に対するブラックボックス・クエリの有限な予算の下では、成功確率はクエリによってカバーされるアクション空間の体積によって制限されることを証明した（定理 A.1）。この境界は、予測器がどのように訓練されたかには依存せず、提案分布が成功のマニホールドをカバーしていない場合、より優れた予測器であっても探索のボトルネックを克服できないことを裏付けている。

3. 手法: IMWM

世界モデルを改善することなく探索のボトルネックに対処するため、著者らはIMWM (Intuition Model + World Model) を提案する。このフレームワークは、凍結された潜在世界モデルと、デモンストレーションから学習された独立した**直感モデル（Intuition Model）**を組み合わせるものである。直感モデルは、世界モデルを置き換えるのではなく、有望なアクションを認識するためのガイドとして機能する。

3.1 直感モデル

直感モデル ($D_{\psi}^{inv}$) は、InfoNCE型の目的関数を用いてデモンストレーションのウィンドウを用いて訓練された、コントラスティブなスカラー・スコアラーである。これは、（開始潜在変数、目標潜在変数、アクションチャンク）のタプルを受け取り、適合度スコアを出力する。

• 逆条件付け (Inverse Conditioning): 世界モデル（$s_t, a_t$ から $s_{t+1}$ を予測する）とは異なり、直感モデルは開始と目標のエビデンスを与えられた状態でアクションをスコアリングする。
• 分離された潜在空間: これは、世界モデルのフォワード予測空間 ($z_L$) から切り離された、独自の潜在空間 ($z_I$) で動作し、ロールアウト誤差の累積を防ぐ。
• 訓練: 同一のタスクコンテキスト内において、デモンストレーションのアクションチャンクを負のサンプルから区別するように訓練される。

3.2 3つの軽量な協調コンポーネント

直感モデルは、以下の3つの具体的なメカニズムを通じてプランナーと相互作用する。

1. リトリーバル初期化 (Retrieval Initialization):

   • コスト関数が候補を評価する前に、プランナーは固定されたデモンストレーション・バンクからアクションチャンクをリトリーブする。
   • リトリーバルは、直感エンコーダの潜在空間における、現在の（開始, 目標）の潜在変数と保存されたデモンストレーション・キーとの間のコサイン類似度に基づいている。
   • CEMの提案分布は、このリトリーブされたアクションチャンクを中心とした等方的なガウス分布として初期化され、デモンストレーションによって支持されている領域へと探索をバイアスさせる。

2. ハイブリッド・コスト (Hybrid Cost):

   • プランナーは、直感スコアと世界モデルのロールアウト誤差を組み合わせた複合コスト関数 $J$ を最小化する。
   • $J = \alpha_{inv} \cdot zS(-D_{\psi}^{inv}) + \beta_{eff} \cdot zS(\text{MSE}(\text{rollout}, \text{goal}))$。
   • 両方の項は、尺度不一致を扱うために候補セット全体でzスコア化される。重み $\alpha$ と $\beta$ は、直感と世界モデルへの依存度を決定する。

3. 信頼性ゲート (Reliability Gate):

   • 直感は普遍的に有用ではないため（例：デモンストレーション・バンクに関連する例が存在しない場合）、ゲートがタスクやセルごとに、直感モデルをどの程度信頼するかを決定する。
   • ゲートは、リトリーブされたチャンクから2つの診断指標（スコア・マージン $r_{inv}$ およびラグ1自己相関 $r_{lag}$）を計算する。
   • 凍結された閾値に基づき、ゲートは以下の3つの事前指定された「レシピ」のいずれかを選択する：
     • Trusted Hybrid: 直感と世界モデルの両方に高い依存度を示す。
     • Intuition-Dominant: 直感への依存度が高く、世界モデルへの依存度が低い。
     • Forward-Only Fallback: 直感を無効化し、プランナーは標準的な世界モデルのみの設定に戻る。
   • 重要なのは、この選択が手動チューニングではなく、診断に基づく自動的なものである点である。

4. 主な貢献

1. 有限CEM潜在プランニングの診断: 本論文は、実験的および理論的に、たとえオラクル・ダイナミクスを持つ世界モデルであっても、有限予算のプランナーは予測エラーではなく、探索の制限（提案のボリューム）によって失敗することを実証した。
2. IMWM手法: リトリーバル初期化、ハイブリッド・コスト、および信頼性ゲートを介して、コントラスティブな直感モデルを世界モデルと統合する新しいアーキテクチャを提案した。
3. 実証的な利得: パラメータのタスクごとのチューニングなしに、4つのピクセルベースの目標到達タスクにおいて、世界モデルのみのベースラインに対して大幅な改善を実現した。

5. 実験結果

本手法は、stable-worldmodel プラットフォームを用い、48セルのグリッド（4タスク $\times$ 12シード）で評価された。タスクは Two-Room, Reacher, Push-T, OGBench-Cube である。

• Two-Room: IMWMは99.2%の成功率を達成し、ベースライン（87.7%）に対し11.5パーセントポイント (pp) の向上を実現した。
• OGBenc-Cube: IMWMは94.7%の成功率を達成し、ベースライン（66.2%）に対し28.5 pp という劇的な向上を実現した。
• Push-T: IMWMは**92.7%**を達成し、2.8 pp の向上（ベースライン 89.8%）となった。
• Reacher: IMWMは**83.8%**を達成し、ベースライン（83.2%）とほぼ同等（+0.7 pp）であった。このケースでは、信頼性ゲートが正しく「forward-only fallback」レシピへとルーティングしており、直感信号が有益ではなかったことを示している。

アブレーション研究:

• リトリーバル初期化: リトリーバルのみをオンにすると、探索制限のあるタスク（例：OGBench-Cubeでは66.2%から88.3%へ）において性能が大幅に向上したが、世界モデルが既に強力なタスク（Reacher）では性能を低下させた。これは、ゲートの必要性を裏付けている。
• ハイブリッド・コスト: 完全な構成（直感と世界モデルの両方のコスト）は、どちらか一方の項のみを使用する場合よりも優れた性能を示した。
• 信頼性ゲート: 単一の固定レシピが全タスクで最適となることはなかった。適応型ゲートは、事前に計算された診断に基づき、各タスクに対して最適なレシピを自動的に選択することに成功し、手動でチューニングされたベストケースの性能に一致した。

6. 意義と主張

本論文は、潜在世界モデルによるプランニングにおける主要なボトルネックは、多くの場合、予測の正確さではなく探索戦略であると主張している。直感モデルを導入することで、IMWMはデモンストレーションから導かれたメカニズムを用いて、探索予算を有望な領域に効果的に集中させ、盲目的なサンプリングによるボリュームの限界を回避できる。

著者らは以下を強調している：

• 改善は、より良い世界モデルやタスクごとのハイパーパラメータ調整によるものではない。アルゴリズムのファミリーと訓練されたアーティファクトは固定されており、ゲートのレシピ選択のみが変化する。
• このアプローチは補完的である。直感モデルは世界モデルを置き換えるのではなく、それをガイドするものである。
• 限界: 信頼性ゲートは（解釈可能ではあるものの）固定された閾値を用いて設計されている。また、結果は単一のベンチマーク・ファミリー（ピクセルベースのタスク）に限定されており、さらなる研究なしに実ロボット制御やより広範なベンチマークへ転用できるとは主張していない。

結論として、IMWMは、世界モデルをリトリーバルベースの直感モデルと組み合わせることで、純粋な世界モデル・プランニングが苦戦する特定の目標到達タスクにおいて、「探索の失敗」モードを解決し、最先端の性能を達成できることを示している。