技術要約: CAST (Credit Assignment from Solver Teachers)

問題提起

Sokoban、Minesweeper、Rush Hourといった古典的なゲームのような、長期的なホライゾンを持つ環境において、汎用エージェントとして動作するように大規模言語モデル（LLM）を訓練することは、**疎な報酬信号（sparse reward signals）**という決定的なボトルネックに直面している。検証可能な報酬を用いた強化学習（RLVR）はシングルターンのタスクでは成功しているが、最終的な結果の報酬（勝ち／負け）に依存しており、どの特定の中間決定が成功や失敗につながったのかについての情報を提供しない。これは、すべての軌道内のトークンが同じ粗いフィードバックを受け取るという、根本的なクレジット割り当て（credit-assignment）の問題を生み出し、モデルが長期的なプランニング、目標指向の探索、およびエラーリカバリを学習することを妨げている。コストのかかる探索、学習されたプロセス報酬モデル、あるいは軌道間の比較によってこれらの信号を密にする既存の試みは、計算コスト、監督の質、および信号の信頼性の間のトレードオフにしばしば苦しんでいる。

手法: CAST

著者らは、ドメイン固有のゲームソルバーを活用して、LLMエージェントに対してきめ細かなターンレベルの監督を提供するフレームワークである**CAST (Credit Assignment from Solver Teachers)**を提案する。

コアメカニズム

1. ターンレベルの教師としてのソルバー: ソルバーをエキスパートの軌道を生成するためだけに使うのではなく（これはエキスパートが訪問した状態への学習を制限する）、CASTはソルバーを使用してLLMのオンポリシー・アクションをリアルタイムで評価する。
2. ソルバー・アドバンテージの構築:
   • ソルバーは、状態 $s$ から勝利するために必要な最小ステップ数を示す**コスト・トゥ・ゴー（cost-to-go）**値 $N(s)$ を計算する。
   • 状態 $s_t$ から $s_{t+1}$ へ遷移するLLMによるアクション $a_t$ に対して、この手法はコスト・トゥ・ゴーの変化量を計算する。
   • **シフトされたソルバー・アドバンテージ（shifted solver advantage）**は以下のように定義される：
     $$\tilde{A}^{\pi_{Solver}}(s_t, a_t) = N(s_t) - \mathbb{E}[N(s_{t+1})]$$
   • この値は、アクションがゴールへの距離を縮める場合は正、中立の場合はゼロ、距離を広げるか行き止まりに導く場合は負となる。
3. 信号整形（Signal Shaping）: 重い裾を持つ値（例：行き止まりによる大きなペナルティ）による訓練の不安定性や、ゲーム間のスケールの違いを防ぐために、CASTは2つの変換を適用する：
   • Asinh圧縮: $g(x) = \text{asinh}(x)$。これは小さな値を線形に保持しつつ、極端な外れ値を対数的に圧縮する。
   • バッチレベルのRMS正規化: 信号をバッチ内の平方根平均二乗（RMS）マグニチュードで再スケールする。これにより、「進展なし」というゼロの意味を保持するために平均を引くことなく、一貫したスケールを確保する。
4. RLVRとの統合: 整形されたソルバー・アドバンテージは、標準的な軌道レベルの結果アドバンテージ（例：GRPOによるもの）に加算され、結合されたアドバンテージ信号を形成する：
   $$\hat{A}_{i,t} = \hat{A}^{outcome}_i + \alpha \cdot h(\tilde{A}^{\pi_{Solver}}_{i,t})$$
   この結合された信号は、ポリシーを更新するためのクリップされたサロゲート目的関数に使用される。

理論的洞察: ロジットフリーのオンポリシー蒸留

本論文は、ソルバー・アドバンテージを最大化することと、ソフト最適（soft-optimal）な教師からの**オンポリシー蒸留（OPD）**との間の理論的な等価性を確立している。

• 強力なソルバーが、アクション確率が $\exp(Q/\tau)$ に比例するソフト最適ポリシーとして機能するという仮定の下で、ソルバー・アドバンテージは数学的に、そのアクションのソルバーのポリシーにおける対数確率に比例する。
• 主要な貢献: これは、スカラー値のソルバー・アドバンテージが、本質的に教師の選好分布をエンコードしていることを意味する。したがって、CASTは教師のロジット（logits）を必要とせずに（完全なアクション分布ではなく、スカラー値のみに依存して）蒸留を実現する。これにより、学生となるLLMは、ソルバーの価値関数から効率的に学習することができる。

主要な貢献

1. ソルバー由来のクレジット: ゲームソルバーの状態価値の変化を、LLM訓練のためのきめ細かなターンレベルのアドバンテージに変換する手法の導入。
2. ロジットフリーの蒸留: ソルバー・アドバンテージを最大化することが、完全なアクション分布ではなく、スカラー値のみを使用して効率的な訓練を可能にするオンポリシー蒸留と等価であるという理論的証明。
3. 安定した信号整形: 極端な値を扱い、クロスドメインのスケール変動を処理するために、asinh変換とバッチレベルのRMS正規化を使用し、RLVRへの安定した統合を保証する。
4. 実用性: 正確なソルバーが加えるオーバーヘッドは無視できる程度であり、正確なソルバーが利用できない場合には学習された価値ネットワーク（例：DQN）が効果的な近似として機能することのデモンストレーション。

実験結果

本手法は、Sokoban、Minesweeper、Rush Hourで評価され、ALFWorldおよびWebShopへのゼロショット転移も行われた。

• インドメイン・パフォーマンス: CASTは、インドメインおよび未経験の難易度の設定の両方において、3つのゲームすべてで、すべての訓練済みベースライン（GRPO, GSPO, DAPO）および訓練フリーのプロンプティング手法（Claude-Opus-4.6のような強力なクローズドソースモデルを含む）を上回った。
  • 平均して、CASTはインドメインのタスクで**62.1%の成功率（Avg@4）を達成し、未経験の難易度では28.4%**を達成した。これは、バックボーンであるDAPO（ID 44.7%、Unseen 18.7%）を大幅に上回る数値である。
• サンプル効率: CASTは、3つのゲームにおいて、DAPOのピーク検証パフォーマンスに1.7倍から2.0倍少ない訓練ステップで到達した。
• ゼロショット汎化: さらなるファインチューニングなしで、CASTエージェントはALFWorldおよびWebShopに転移し、全手法の中で最高の平均ゼロショット・パフォーマンス（全体で30.3%）を達成し、次に優れたベースラインを5.6ポイント上回った。
• アブレーション研究:
  • ソルバー・アドバンテージの重み（$\alpha$）は極めて重要であり、$\alpha=0.1$ がプロセス・ガイダンスと結果目的関数の間で最良のバランスを提供した。
  • asinh圧縮とRMS正規化の両方が不可欠であった。どちらか一方を削除すると性能が低下し、特にasinhの削除は、極端な値による勾配の不安定性により最も顕著な低下を引き起こした。
• 学習されたソルバー: 正確なソルバーを学習されたDQN価値ネットワークに置き換えても、ほとんどの性能上の利点が維持されたことは、本手法が正確なソルバーが存在しない場合にも適用可能であることを示している。

意義と主張

本論文は、ドメイン・ソルバーから派生したスケーラブルで詳細な信号を提供することにより、CASTが長期的なRLVRにおける根本的なクレジット割り当ての課題に対処することを主張している。ソルバー・アドバンテージをロジットフリーのオンポリシー蒸留と等価であると証明することで、本手法は、完全なアクション分布を生成する計算負荷なしに、LLMをエキスパートの価値関数に適合させるための実用的な経路を提供する。実験結果は、信頼できる状態評価（正確なソルバーまたは学習された価値関数のいずれかによる）が、疎な結果から密で実行可能なフィードバックへと至る堅牢なルートを提供し、LLMが結果のみのアプローチよりも効率的に複雑な多段階の意思決定タスクを習得できることを示唆している。