プレフィックスリプレイを用いたマルチターンオン方策蒸留

要旨

本論文では、エージェントタスクにおけるオン方策蒸留（OPD）について研究する。ここで、LLMエージェントは環境と複数ターンにわたって相互作用し、生徒モデルはこれらのマルチターン相互作用履歴に基づいて教師モデルを模倣する。完全オンラインOPDはコストが高く、更新のたびに生徒が環境を通じて新たなロールアウトを実行し、訪問した履歴において教師にクエリを送る必要がある。我々は、Replayed-Prefix On-Policy Distillation（ReOPD）を提案する。これはオフ環境の代替手法であり、事前収集した教師の軌跡を再生プレフィックスとして再利用する。生徒は選択されたステップで行動し、教師は新たな環境相互作用を実行することなく、ステップごとの密な教師情報を提供する。我々は、マルチターンOPDが「プレフィックストラップ」を導入することを示す。すなわち、履歴をより生徒オン方策にすることで生徒への関連性は向上するが、教師のターゲットが信頼できない履歴に対して教師にクエリを送る可能性がある。これにより、生徒の占有分布と教師の信頼性の間で両側の分布シフトが生じる。ReOPDは、マルチターンOPDを信頼性考慮型のプレフィックス分布設計として捉え、これを初期のシフトが少ないプレフィックスを重視する単純なステップ減衰サンプリングスケジュールで実装することにより、この問題に対処する。Pythonを用いた数学的推論や検索環境において、複数の教師・生徒モデル規模にわたって、ReOPDはOPDレベルの精度を維持または向上させ、生徒のトレーニング中にツール呼び出しを一切行わず、ロールアウトあたりの速度がOPDよりも少なくとも4倍高速である。したがってReOPDは、高コストなエージェント-環境相互作用を再利用可能なオフラインリソースに変換し、ツール、タスク、環境を横断したスケーラブルな蒸留を可能にする。

English

We study on-policy distillation (OPD) for agentic tasks, where an LLM agent interacts with an environment over multiple turns and a student imitates a teacher over these multi-turn interaction histories. Fully online OPD is costly because each update requires fresh student rollouts through the environment and teacher queries at visited histories. We propose Replayed-Prefix On-Policy Distillation (ReOPD), an off-environment alternative that reuses pre-collected teacher trajectories as replayed prefixes: the student acts at selected steps, while the teacher provides dense per-step supervision without executing new environment interactions. We show that multi-turn OPD introduces a prefix trap: making histories more student-on-policy improves relevance to the student, but can query the teacher on histories where its target is unreliable. This creates a two-sided distribution shift between student occupancy and teacher reliability. ReOPD addresses this by treating multi-turn OPD as a reliability-aware prefix distribution design and implements it with a simple step-decaying sampling schedule that emphasizes early, lower-shift prefixes. Across mathematical reasoning with Python and search environments over multiple teacher and student model scales, ReOPD preserves or improves OPD-level accuracy, uses zero tool calls during student training, and is at least 4times faster per rollout than OPD. ReOPD therefore turns expensive agent-environment interaction into a reusable offline resource, enabling scalable distillation across tools, tasks, and environments.