技術サマリー：PriorZero

問題定義

大規模言語モデル（LLM）を強化学習（RL）エージェントに統合することは、豊かな世界知識を活用して、長期視野かつスパース報酬の意思決定タスクを解決するための有望な道筋を提供する。しかし、既存のアプローチは、根本的な「事前知識とダイナミクスの不一致」という課題に阻まれている。

現在のパラダイムは、以下の特定の構造的ボトルネックに直面している：

1. 単純な方策（P1）： LLM をゼロショット方策として使用すると、言語的な尤度が環境の報酬と一致しないため、近視眼的な意思決定を招き、「知っていることと実行することのギャップ」が生じる。
2. 単純な RL 微調整（P2）： RL による LLM のエンドツーエンド微調整は、特にスパース報酬環境において、分散の大きいクレジット割り当てと最適化の不安定性に悩まされる。
3. ダイナミクスモデリング（P3）： LLM を用いて状態遷移を予測しても、粒度の不一致により失敗する。LLM はトークンレベルのテキスト生成で事前学習されているのに対し、潜在ダイナミクスモデルは状態 - 行動レベルの遷移予測を必要とするためである。
4. テキストエンコーダ（P4）： 凍結された LLM 特徴量を静的な表現として使用すると、タスク固有の適応が制限され、性能の天井が課される。

核心的な課題は、LLM の静的で意味的な事前知識と、効果的な計画に必要な動的かつ環境固有の遷移モデルとの間のギャップを埋めることである。

手法：PriorZero フレームワーク

PriorZero は、LLM と世界モデルを分離し、分離されたロールアウト・トレーニング設計を通じて統合する統一フレームワークを提案する。このシステムは、LLM と環境の間に学習可能な世界モデルを配置し、LLM が意味的なガイダンスを提供する一方で、世界モデルが遷移モデリング、価値推定、多ステップ計画を処理できるようにする。

1. ルート事前知識注入（ロールアウト段階）

ロールアウト段階において、PriorZero は潜在世界モデル（UniZero に基づく）を用いて、潜在空間内でモンテカルロ木探索（MCTS）を実行する。

• メカニズム： 新規のルート事前知識注入メカニズムが、LLM の意味的行動事前知識を MCTS のルートノードのみに限定して融合する。
• 実装： LLM は、現在のテキストコンテキストに基づいて状況分析（Chain-of-Thought）と行動分布（$\pi_{LLM}$）を生成する。これをルートにおいて世界モデルの方策（$\pi_{WM}$）と融合する：
  $$P_{root}(a | z_t) = (1 - \alpha) \pi_{WM}(a | z_t) + \alpha \pi_{LLM}(a | C_t)$$
• 制約： 探索木内のルート以外のノードは、完全に世界モデルの潜在ダイナミクスに依存する。これにより、LLM が想像上の潜在状態に対する状態遷移を予測することを強制されないため、パラダイム P3 の「粒度の不一致」が防止される。
• 利点： この手法は、意味的に有望な行動へと探索を誘導し、MCTS が無意味な軌道（例えば、報酬ループ）に陥るのを防ぎつつ、世界モデルの深い先読み能力を維持する。

2. 交互強化学微調整（トレーニング段階）

パラダイム P2 のクレジット割り当ての失敗に対処するため、PriorZero は世界モデルの学習と LLM の適応を分離する交互微調整戦略を採用する。

• スケジュール： トレーニングは非対称なスケジュールに従う（例：LLM 更新 1 回に対して世界モデル更新 10 回）。
• 世界モデル学習： 世界モデルは、相互作用データに基づいて更新され、そのダイナミクス、価値、方策予測を洗練させる。
• LLM 適応： 分散の大きいオンポリシー報酬の代わりに、LLM は世界モデルの価値推定から導出される分散の小さいアドバンテージ信号を用いて微調整される。
  • アドバンテージは、$n$ ステップのブートストラップ TD ターゲットを用いて計算される：$A(z_t, a_t) = Q_n(z_t, a_t) - v_\theta(z_t)$。
  • LLM は、クリップされた PPO 目的関数を通じて更新され、特定の重み付けスキームが適用される。これにより、行動トークンにはフルの重みが、Chain-of-Thought 推論トークンには安定性を維持するために軽減された重み（$w_{cot}=0.1$）が適用される。
• 利点： これにより、世界モデルが LLM に対して安定した微細な監督を提供する閉ループが形成され、LLM の事前知識は、エンドツーエンド最適化の不安定性なしに環境のダイナミクスと共進化できるようになる。

主要な貢献

1. パラダイムの分析： 本論文は、LLM の知識を長期視野の RL に転送する際の中心的な障壁として事前知識とダイナミクスの不一致を特定し、既存のアプローチ（P1–P4）とその特定の失敗モードを分類する。
2. PriorZero フレームワーク： LLM の事前知識と世界モデルの計画を 2 つの中核メカニズムを通じて相乗的に統合する新規アーキテクチャ：
   • ルート事前知識注入： 潜在ダイナミクスを汚染することなく探索を促進するため、MCTS ルートに意味的ガイダンスを注入する。
   • 交互強化学微調整： 世界モデルの価値推定を用いて、LLM 適応のための安定した分散の小さいクレジット割り当て信号を構築する。
3. 実証的検証： Jericho（テキストベースのアドベンチャーゲーム）およびBabyAI（グリッドワールドの指示追従）における包括的な実験により、PriorZero がベースラインと比較して、探索効率と漸近的性能を一貫して向上させることが示された。

実験結果

• Jericho ベンチマーク： PriorZero は、4 つの異なるゲーム（Detective, Acorncourt, Zork1, Omniquest）において、LLM の事前知識を持たない強力な世界モデルベースラインである UniZero を上回った。
  • 長期視野かつスパース報酬のZork1において、PriorZero はより高い漸近報酬とより速い収束を達成した。
  • Detectiveにおいて、PriorZero は初期の利益の後に UniZero を上回り、タスク固有のダイナミクスへの事前知識の適応の価値を示した。
• BabyAI ベンチマーク： BabyAI グリッドワールドの 18 レベルにおいて、PriorZero はより高い報酬プラトー（UniZero の 0.79 に対して 0.82）に到達し、LLM の事前知識が目標をサブタスクに分解するのに役立つ構成的タスク（例：SynthLoc）において顕著な利益を示した。
• アブレーション研究：
  • 交互スケジュール： 安定性に不可欠であり、非交互のバリアントはノイズの多い価値信号に起因する「知識の汚染」に苦しむ。
  • ルートのみ注入： 先読み中に LLM が世界モデルのダイナミクスをオーバーライドするのを防ぐために不可欠である。
  • CoT 損失重み： バランスの取れた重み（$w_{cot}=0.1$）が必要である。CoT を除去すると数値的不安定性が生じ、フルの重みは不安定な学習をもたらす。
  • MCTS： MCTS を無効化すると性能が崩壊し、計画によって改善されたターゲットの必要性が確認された。

意義と主張

本論文は、PriorZero が事前転送の構造的ボトルネックを解決することで、LLM 支援の意思決定のための有望なフレームワークを確立すると主張する。その意義は以下の点にある：

• 分離： LLM の役割（意味的事前知識の提供）と世界モデルの役割（ダイナミクスの学習と計画）を成功裏に分離し、これにより従来の手法の粒度の不一致とクレジット割り当ての失敗を回避する。
• 安定性： 世界モデルを分散の小さいクリティカルとして活用することで、長期視野かつスパース報酬の環境における LLM の安定した微調整を可能にする。
• スケーラビリティ： 複雑なタスクにおいて 7B モデルで性能が向上するなど、フレームワークが LLM の容量に応じてスケーリングすることを示しつつ、交互トレーニングスケジュールを通じて堅牢性を維持する。

著者らは、PriorZero を世界モデルや LLM の代替としてではなく、LLM が探索に焦点を当て、世界モデルが継続的な相互作用を通じて LLM の事前知識を洗練させる、相乗的な統合として位置づけている。