技術要約：長期的予測のための潜在世界モデルにおけるロールアウト・デコード再構成（Rollout-Decoded Reconstruction）

1. 問題提起

潜在世界モデル（Latent world models）は、観測を潜在状態へと圧縮し、潜在空間内で遷移モデルを前方向に反復実行（フリーランニング）し、得られた潜在変数を再び観測空間へとデコードすることで動作する。ここで、学習とデプロイメントの間に決定的な乖離が存在する：

• 学習時： デコーダは、「事後（posterior）」の潜在変数、すなわち真の観測値や教師強制（teacher-forcing）による予測（観測から1ステップ離れた状態）に紐付けられた状態に対してのみ学習される。
• デプロイメント時： モデルは何百ステップものフリーランニングを行い、事後分布から逸脱した潜在変数の軌跡を生成する。
• 乖離（The Gap）： 標準的な目的関数は、これらのドリフトしたフリーランニング中の潜在変数に対するデコーダの誤差を罰しない。その結果、観測データに対しては優れた性能を示すデコーダであっても、要求される長期間のロールアウトをデコードする際に壊滅的に失敗することがある。Plaٕnetは以前、この問題（「観測オーバーシューティング」と呼ばれる）を特定していたが、画像ドメインにおける計算コストの懸念から棚上げされ、その後の潜在世界モデル（例：Dreamer）においても優先順位は付けられてこなかった。

2. 手法：ロールアウト・デコード再構成（RDR）

著者らは、デコーダの学習分布とデプロイメント分布の間のギャップを埋めるために設計された学習目的関数として、**ロールアウト・デコード再構成（Rollout-Decoded Reconstruction: RDR）**を提案する。

コアメカニズム

RDRは、評価時と全く同じように、モデル自身のフリーランニング中のロールアウトをデコードさせることを強制する単一の損失項を導入する。

1. フリーランニング・ロールアウト： 学習中、モデルは初期の観測をエンコードし、外部からの観測入力なしに、遷移モデルを$K$ステップ前方に反復実行する。
2. デコーディング： このフリーランニングの軌跡におけるすべての潜在状態$\hat{z}_{t+k}$が、デコーダ$D_\psi$を通過する。
3. 再構成誤差： 出力は真の観測値$u_{t+k}$と比較され、その再構成誤差が罰せられる。

目的関数

全損失$L$は、標準的な項と新しいRDR項を組み合わせたものである：
$$L = L_{TF} + \alpha_e L_R + L_{OBS} + L_{recon} + \lambda L_{RDR}$$
ここで：

• $L_{TF}$: 教師強制による1ステップの潜在変数予測。
• $L_R$: 多段階の潜在変数ロールアウトの一貫性（潜在対潜在）。
• $L_{OBS}$: 教師強制による予測のデコーディング。
• $L_{recon}$: オンラインエンコーダに対する再構成アンカー。
• $L_{RDR} = \frac{1}{K} \sum_{k=1}^K \| D_\psi(\hat{z}_{t+k}) - u_{t+k} \|^2$: フリーランニング中のロールアウトにおける誤差を罰する新しい項。

実装の詳細

• 新規パラメータなし： RDRは新たなアーキテクチャ・パラメータを追加しない。既存のデコーダおよび遷移コンポーネントを利用する。
• 計算コスト： 増加するコストは学習時の計算量に限定される。ベースラインおよびRDRの両方の構成において、すでに$L_R$のためにフリーランニング・ロールアウトを計算しているため、RDRはこれらのロールアウトステップにおけるデコーダの評価とバックプロパゲーションを追加するだけであり、推論時には影響しない。
• カリキュラム： ロールアウト・ブランチ（およびRDR項）は、2エポックのウォームアップ後に有効化される。重み$\lambda$は0.3に設定されるが、結果は$\lambda \in \{0.1, 0.3, 0.6, 1.0\}$において頑健である。

3. 実験設定

• ドメイン： 正確な真値を持つ偏微分方程式（PDE）である、カオス的な蔵本・シヴァスキー（Kuramoto–Sivashinsky: KS）方程式（$L=22$）。
• 指標： 正規化された誤差閾値0.5における有効予測時間（Valid Prediction Time: VPT）。これは、予測誤差が気候学的標準偏差の半分を超えた時点までの時間を測定する。
• ベースライン：
  • 事後のみ（Posterior-Only）： RDRを含まない標準的な目的関数。
  • 観測空間プッシュフォワード（Observation-Space Pushforward）： 自身の展開された出力に対して直接学習を行う予測器（潜在ボトルネックなし）。
  • 分割デコーダ・バリアント（Split-Decoder Variants）： 事後潜在変数とロールアウト潜在変数のために別々のヘッドを持つモデル。これは容量が制限要因であるかどうかをテストするためのものである。
• プロトコル： 最終報告のために、新鮮なシード（選択には使用されていない）を用いた厳格な事前登録（preregistration）を実施。すべての比較は、固定されたデータ、シード、予算、およびパラメータ数を用いたA/Bテストである。

4. 主な結果

主要な予測結果（KS方程式）

• 性能向上： RDRは、有効予測時間（VPT）を3.87 ± 0.23から6.97 ± 0.42へと増加させた。これは、予測ホライゾンの1.80倍の改善を意味する。
• 一貫性： この改善は、10個中10個の事前登録された構成（勝者の新鮮なシードによる再実行を含む）において保持されており、改善比率は1.71倍から2.50倍の範囲にある。
• パラメータ効率： この利得は、事後のみのベースラインと比較して、同一のパラメータ数（学習可能なパラメータ数193,568個）で達成されている。
• 潜在幅への感度： 潜在次元（16から48）のスイープにより、潜在空間が広がるにつれて、事後のみのベースラインの性能は低下する一方で、RDRの性能は向上または維持されることが示された。RDRとベースラインの性能比は、潜在幅が増大するにつれて1.41倍から2.50倍へと上昇しており、RDRが高次元でドリフトした潜在空間に対するデコーダの学習に特に効果的であることを示唆している。

アブレーションおよび制御実験

• 分割デコーダ： 事後潜在変数とロールアウト潜在変数のために別々のデコーダヘッドを持つバリアント（約25%のパラメータを追加）は、10個中7個の構成において、共有デコーダを用いたRDRよりも劣った性能を示した。これにより、「追加された容量」が改善の理由であることを排除し、単一の堅牢なデコーダが十分であることを示唆している。
• 観測空間との比較： 同等の予算で学習された観測空間プッシュフォワード予測器は、VPT 7.00 ± 0.26を達成し、RDR潜在モデルと一致した。これは、この完全観測システムにおいては、潜在ボトルネック自体が、一致した予算における直接的な観測空間予測に対して固有の優位性を持たないことを示している。しかし、潜在内での比較（RDR vs. 事後のみ）は、依然として本論文の主要な貢献である。
• 制御タスク（Pendulum/Cartpole）： 制御タスクに関する予備的な結果では、RDRはプランナーの学習におけるロールアウトの不一致（例：状態のリセット）に対してより頑健であることが示されている。ただし、オプティマイザのステップ数を一致させた場合（学習データ量を制御した場合）、性能の差はほぼ消失しており、RDRは生のデータ効率よりも、オプティマイザのステップ数に関するサンプル効率を主に改善することを示している。

5. 意義と主張

本論文は、潜在世界モデルにおいて、デコーダの学習分布をデプロイメント分布に適合させることの価値を、分離して測定できることを主張している。

• 主要な貢献： RDRは、標準的な目的関数が、長期間のフリーランニング中にデコードしなければならない特定の潜在変数に対して、デコーダを十分に学習させていないことを実証した。単一の損失項を追加して、これらのフリーランニング中のロールアウトに対する再構成誤差を罰することで、モデルはアーキテクチャの変更やパラメータの膨張なしに、有意な（1.80倍の）有効予測時間の延長を実現した。
• メカニズム： 改善は、以下の2つの経路から生じている：
  1. デコーダの堅牢性： デコーダが、分布外（ドリフトした）潜在変数をより効果的に扱えるようになる。
  2. 潜在ダイナミクス： 勾配信号がロールアウトの連鎖を通じてバックプロパゲーションされ、潜在変数を真の状態分布に近づけるための遷移モデルの能力をわずかに向上させる。
• 限界：
  • 結果は現在、KS方程式および低次元の制御タスクに特化している。
  • 他のシステム、バックボーン、または観測モダリティ（例：ピクセル、部分観測性）への転移については未検証である。
  • 本手法は、学習時の計算量オーバーヘッド（デコーダ評価の約40%増、ウォーククロック時間の5〜10%増）を伴う。
  • 潜在ボトルネックは、この完全観測システムにおいて、直接的な観測空間予測を上回ることは示されなかった。

要約すると、本論文は、ロールアウト・デコード再構成が、潜在世界モデルにおいて、デコーダの学習分布とデプロイメント分布の間の分布的ギャップを埋め、長期間の予測能力を大幅に拡張するために必要かつ効果的な学習目的関数であることを確立している。