技術的概要：パースimonious 潜在変数へ回帰：視覚的基盤からタスク中心の世界モデルを学習する

問題定義

世界モデルは、エージェントが行動に条件付けられた将来のダイナミクスを予測することを可能にする能力であり、これは計画と制御の中核をなす。しかし、既存のアプローチは、潜在表現において根本的なトレードオフに直面している：

1. ピクセルベースのモデル（例：Dreamer、MuZero）は、ピクセルから直接、または報酬信号によって形成されたコンパクトな潜在変数からダイナミクスを学習する。これらは効果的であるが、しばしば高い計算コストを伴うか、潜在空間を形成するためにオンラインな相互作用や報酬監督を必要とする。
2. 基盤埋め込みモデル（例：DINO-WM、V-JEPA）は、凍結された事前学習済み視覚エンコーダを活用し、高次元の意味空間内で直接予測を行う。これらは豊かな意味的事前知識とゼロショット汎化を提供する一方で、埋め込みには制御に影響を与えない過剰なタスク無関係の詳細（テクスチャ、照明、背景）が含まれている。

報酬なしのオフライン設定、すなわちエージェントが報酬監督なしに固定された軌跡から学習しなければならない状況では、これらの問題がさらに悪化する。生きた基盤埋め込みを状態空間として使用すると、モデルは無関係な変動に容量を浪費せざるを得なくなり、モデル化効率の低下、制御の解釈性の低下、特に高次元の行動空間（例：7 自由度のロボットアーム）における表現の崩壊（representation collapse）のリスクを招く。核心的な課題は、オンラインな相互作用なしに、これらの豊かだがノイズの多い基盤埋め込みから、コンパクトでタスクに十分な潜在表現を抽出することである。

手法：TC-WM

著者らは、**TC-WM（Task-Centric World Model：タスク中心世界モデル）**を提案する。これは、事前学習された基盤埋め込みを最終的な状態空間として扱うのではなく、タスク関連のダイナミクスを蒸留するための「意味的足場（semantic scaffold）」として扱うフレームワークである。

中核アーキテクチャ

TC-WM は、視覚的埋め込み（$x_t$）と固有受容信号（$s^p_t$）のパイプライン上で動作する：

1. 結合埋め込み：凍結された視覚エンコーダ（例：DINOv2）が画像を視覚的特徴にマッピングし、学習可能な線形層が固有受容を符号化する。これらを連結して結合埋め込み $x_t$ を形成する。
2. 潜在エンコーダ（射影）：線形エンコーダが高次元の結合埋め込みをコンパクトな潜在状態 $z_t$ に射影する。重要なのは、$z_t$ が 2 つの部分空間に分解される点である：
   • $z^s_t$（状態関連）：物理状態と整合する。
   • $z^c_t$（補完的）：再構成制約を満たすために、残りの意味構造（外観、背景）を保持する。
3. 潜在ダイナミクス：Transformer ベースの予測器（$f_\theta$）が、コンパクトな潜在空間内での時間的遷移をモデル化する（$\hat{z}_{t+1} = f_\theta(z_{t-H:t}, a_{t-H:t})$）。
4. タスク中心の整合：モデルは、$z^s_t$ 部分空間を固有受容状態 $s^p_t$ と整合させるために、対照的目的関数（InfoNCE）を採用する。これにより、エンコーダはタスク関連の物理的要因をノイズとなる変動から能動的に分離するよう強制される。
5. 埋め込みデコーダ：線形デコーダが $z_t$ から元の視覚的埋め込みを再構成する。これは正則化として機能し、コンパクトな潜在変数が基盤多様体に固定されたまま保たれ、退化した解を防ぐ。

訓練目的関数

総損失は 4 つの成分を組み合わせる：

• 潜在ダイナミクス損失（$L^z_{dyn}$）：潜在空間内での予測誤差を最小化する。
• TC ダイナミクス損失（$L^s_{dyn}$）：タスク中心のダイナミクスを強制するために、過去から将来の固有受容を直接予測する。
• 整合損失（$L_{align}$）：$z^s_t$ と $s^p_t$ を整合させる対照的損失。
• 再構成損失（$L_{rec}$）：視覚構造を保持するために、$z_t$ から視覚的埋め込みを再構成する。

推論時の計画

学習されたタスク中心の潜在空間は、多様な計画戦略をサポートする：

• 低次元制御（ナビゲーション）：潜在空間内で行動シーケンスを最適化するために、クロスエントロピー法（CEM）を用いたモデル予測制御（MPC）。
• 高次元操作：潜在軌跡から実行可能な行動を生成するために、逆ダイナミクスモデル（IDM）と組み合わせた潜在拡散計画（LDP）。

主要な貢献

1. タスク中心の分解：本論文は、物理状態のダイナミクスを意味的なノイズ変数から明示的に分離することにより、高次元の基盤埋め込みをコンパクトでタスクに十分な潜在空間に分解する手法を導入する。
2. 能動的識別可能性：識別可能性のためにメカニズムの多様性を受動的に仮定する先行研究とは異なり、TC-WM は固有受容との対照的整合を用いて、異なる物理状態に対して能動的に異なる条件付き分布を誘発する。これにより、アフィン変換までタスク中心の要因を回復することが理論的に保証される。
3. 報酬なしのオフライン学習：このフレームワークは、報酬信号なしに、視覚的観測と固有受容のサイド情報のみに依存して、オフライン設定で高品質な世界モデル化と計画を達成する。
4. 理論的保証：著者らは、穏やかな仮定（文脈的ダイナミクス、潜在 - 観測変動、十分な整合）の下で、学習されたコンパクトな潜在変数が識別可能であり、再パラメータ化まで真の物理状態を回復することを証明する定理を提供する。

実験結果

著者らは、TD-MPC2、DreamerV3、MuZero、DINO-WM などのベースラインに対して、9 つのオフラインベンチマーク（Maze、Wall、Push-T、Lift、Can、Square、Reacher、Cheetah、Hopper）で TC-WM を評価した。

• 世界モデルの品質：TC-WM は、特に長期的なダイナミクスが正確であることが重要な操作環境において、ほぼすべてのタスクで最も低い潜在ロールアウト予測誤差（MSE）を達成する。画像空間の再構成は競争力があるが、主な利点は潜在ダイナミクスの安定性と精度にある。
• 計画性能：
  • ナビゲーション（低次元の行動）において、TC-WM は TD-MPC2 などの強力なベースラインと同等の性能を達成する。
  • 操作（高次元で接触の多い行動）において、TC-WM はすべてのベースラインを大幅に上回る。例えば、Robomimic の「Lift」タスクでは、TC-WM は 54% の成功率を達成し、DINO-WM の 29% や TD-MPC2 の 31% を上回った。
• 表現分析：
  • 線形プロービング：TC-WM の潜在空間で訓練された線形回帰モデルは、ベースラインよりも高い $R^2$ スコアで固有受容と物体の状態を回復し、潜在変数が物理的に根ざしていることを確認する。
  • 崩壊の防止：TC-WM は、生埋め込みを直接使用した際に表現崩壊に陥る DINO-WM と比較して、潜在空間のより良い利用（より高い実効ランク）を示す。
  • ロバスト性：TC-WM は、DINO-WM よりも視覚的摂動（ガウスノイズ、カラージッター）に対して性能を維持し、SSIM の低下が小さい。

意義と主張

本論文は、基盤埋め込みは直接の状態表現ではなく、意味的足場として活用するのが最適であると主張する。これらの埋め込みを物理信号と整合したタスク中心の部分空間に射影することで、エージェントは以下のことを達成できる：

• パースimonious な潜在変数：制御に必要なダイナミクスを保持しつつ、無関係な視覚的詳細をフィルタリングする次元削減。
• 改善された制御：ノイズや無関係な変動のために標準的な基盤モデルが失敗する複雑で高次元の環境における、向上した計画性能。
• 理論的妥当性：凍結された視覚エンコーダからタスク関連の要因を回復するために、線形射影と対照的整合が十分であることの形式的な正当化。

著者らは、TC-WM を基盤モデルの汎化能力と物理的エージェントに必要な制御可能性の間の架け橋として位置づけ、特に報酬なしのオフライン学習シナリオにおけるギャップに対処している。