技術的概要：プリズマティック・ワールド・モデル（PRISM-WM）

1. 問題定義

ロボティクス領域におけるモデルベース強化学習（MBRL）と計画は、物理ダイナミクスのハイブリッド性によって本質的に制限されています。現実世界のシステム、特に脚式移動や操作においては、連続的な運動が離散イベント（接触、衝撃、固定対滑りなど）によって断絶されます。これらの相互作用は、基礎的なベクトル場における急激な不連続性を伴う、非滑らかで高度に非線形なダイナミクスを生み出します。

従来の潜在ワールド・モデルは、通常、大域的な連続性を強制する単一構造のニューラルネットワークを採用しています。これらは理論的には万能近似が可能ですが、実際には異なるダイナミクスモードを平均化した表現を学習し、それにより過度に平滑化されてしまいます。軌道最適化プランナーにとって、この過度の平滑化は、特に物理的な境界において、長視野の先読み中に誤差が蓄積する原因となります。その結果、以下が生じます：

• 接触境界全体にわたるスペクトルのぼやけにより、顕著な勾配の不一致が生じます。
• 高次元の構成空間における誤差の蓄積的ドリフトが発生し、長期的な計画を信頼できなくします。
• 標準的なモデルが安定領域と不安定領域を区別できない接触の多い環境において、方策の崩壊が発生します。

2. 手法：プリズマティック・ワールド・モデル（PRISM-WM）

これらの制限に対処するため、著者は複雑なハイブリッド・ダイナミクスを構成可能なプリミティブに分解するように設計された構造化アーキテクチャであるPRISM-WMを導入します。このモデルは、新しい潜在直交化目的関数を備えた、**文脈認識型エキスパート混合（MoE）**フレームワークに基づいています。

中核アーキテクチャ

• 文脈認識型ゲーティング: ゲーティングネットワーク（$G$）は、状態 $z_t$、動作 $a_t$、および任意のタスク埋め込み $e_{task}$ からなる文脈ベクトル $c_t$ を受け取り、ソフトマックス正規化された重み $w_t$ を生成します。このメカニズムは、明示的な教師なしで、飛行対接地など、アクティブな物理的レジームを暗黙的に識別します。
• 専門化されたエキスパート: 状態は、$K$ 個のエキスパートネットワーク $\{E_k\}$ のセットへルーティングされます。各エキスパートは、特定の局所的なダイナミクス・レジーム内での連続的な流れを学習することに特化しており、競合するモードからの干渉を回避します。
• 残差予測: 次の潜在状態は、現在の状態に加えられたエキスパート出力の重み付き組み合わせとして予測されます：
  $$\hat{z}_{t+1} = z_t + \sum_{k=1}^{K} w_{t,k} \cdot v_k(x_t)$$
  ここで、$v_k$ は $k$ 番目のエキスパートによって予測される残差ダイナミクスを表します。

主要な革新：潜在直交化

モード崩壊（エキスパートが冗長な特徴を学習すること）を防ぎ、多様性を確保するために、PRISM-WM は微分可能なグラム・シュミット直交化層を導入します。

• 各エキスパートが最終層手前のベクトル $u_k$ を出力した後、モデルは安定化されたグラム・シュミット過程を適用して、直交基底ベクトル $v_k$ を生成します。
• これにより、エキスパート表現がシュティーフェル多様体上に存在することが強制され、残差ダイナミクスに対して最大ランクで冗長性のない部分空間を張ることが保証されます。
• 物理的な接触モードは順序付けられていませんが、この直交化により、エキスパートは残差多様体の異なる成分を捉えるようになり、ゲーティングネットワークが重み付き和を通じて順序付けられていない物理的モードを再構成できるようになります。

計画との統合

PRISM-WM は、TD-MPC計画フレームワークに統合されています。複雑なダイナミクスを非冗長なエキスパートに因数分解することで、モデルは 1 ステップあたりのモデル残差を $\epsilon$ から約 $\epsilon/K\alpha$ に削減します。この構造的分解は誤差の蓄積的ドリフトを軽減し、高忠実度の長視野ロールアウトと、方策学習のための安定した勾配伝播を可能にします。

3. 主要な貢献

1. アーキテクチャ的分解: 単一構造の大域的近似から脱却し、文脈認識型 MoE アーキテクチャと潜在直交化を組み合わせてハイブリッド・ダイナミクスを分解する PRISM-WM の導入。
2. 直交化目的関数: 微分可能なグラム・シュミット直交化を用いた新しい学習目的関数により、エキスパートの多様性を強制しモード崩壊を防ぎ、離散的な物理的遷移と潜在空間を構造的に整合させること。
3. 理論的および実証的検証: この分解が高次元空間における誤差の蓄積を削減し、計画（TD-MPC）および直接方策学習（PWM）の両方に対して損失ランドスケープの幾何学を改善することを示実。
4. ゼロショット転移能力: 直交基底が、特定のダイナミクスモード（摩擦対垂直支持など）を分離することで、異なる接触物理を持つシミュレーター間（例：ソフト接触の IsaacGym から剛性接触の MuJoCo へ）でモデルが一般化できることを示す証拠。

4. 実験結果

著者は、高次元のヒューマノイド、マルチタスク設定、およびゼロショットのシム・トゥ・シム転移を含む連続制御ベンチマークにおいて PRISM-WM を評価しました。

• 高次元移動（DiffRL および Humanoid-Bench）: PRISM-WM は、単一構造のベースライン（TD-MPC2、DreamerV3、PWM、SAC）と比較して、高いサンプル効率と漸近的性能を達成しました。特に、単一構造モデルは非滑らかな接触イベントを平均化することにより二足歩行タスク（Humanoid、SNU-Humanoid）で方策の崩壊を経験したのに対し、PRISM-WM は安定した最適化を維持しました。
• マルチタスク一般化: 30 タスクの DeepMind Control Suite ベンチマークにおいて、PRISM-WM は平均正規化スコア0.531を達成し、単一構造の TD-MPC2 ベースライン（0.430）に対して23.5% の改善を示しました。異質なダイナミクス間のネガティブ転移を効果的に軽減し、30 タスク中 26 タスクでベースラインを上回りました。
• ゼロショット・シム・トゥ・シム転移: IsaacGym から MuJoCo への転移において、単一構造の MLP ベースラインは47.8% の性能低下と 12% の成功率を被りました。PRISM-WM はこの低下を**11.8%**に抑え、88% の成功率を達成しました。直交化により、モデルはグローバルなダイナミクス多様体を損なうことなく、特定のエキスパート（例：水平摩擦）の重みを調整することができました。
• 安定性と効率性: 消融実験により、直交化が制約のない MoE モデルで見られる報酬推定の劣化を防ぐことが確認されました。NVIDIA RTX 4090 上での推論効率テストでは、PRISM-WM（K=4）は7,712 FPSを達成し、最小限のメモリオーバーヘッドで単一構造のベースライン（7,499 FPS）をわずかに上回りました。

5. 意義と主張

本論文は、PRISM-WM がモデルベースエージェントのための基盤モデルとして機能する、軌道最適化アルゴリズムのための高忠実度基盤を提供すると主張しています。複雑な物理を直交プリミティブに分解することで、このモデルは：

• 計画視野における長期的ドリフトを軽減し、接触の多い環境での信頼性の高い長視野推論を可能にします。
• 方策学習のための勾配伝播を安定化させ、分断された損失ランドスケープを機能的に滑らかな凸結合に変換します。
• 過度の平滑化により標準的なワールド・モデルが失敗するハイブリッドシステムにおいて、堅牢な意思決定を可能にします。

著者は、複雑な物理を直交プリミティブに分解することが「接触の多い制御を習得するために本質的に不可欠である」と結論付けています。現在のフレームワークはエキスパート数 $K$ の手動調整を必要としますが、この研究は、将来の具身エージェントのための汎用的で物理的に根ざしたワールド・モデルを構築するための構造的基盤を確立しています。