Multi-Turn On-Policy Distillation with Prefix Replay¶

🕒 Published (v1): 2026-07-06 07:56 UTC · Source: Arxiv · link

Why this paper was selected

Li Dong (MSR); multi-turn on-policy distillation with prefix replay; practical for agentic student training

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

TL;DR¶

ReOPD (Replayed-Prefix On-Policy Distillation) replaces expensive live environment interaction in multi-turn on-policy distillation with offline replay of pre-collected teacher trajectories, where the student acts at selected steps and the teacher provides per-step supervision from fixed prefixes. It identifies a "prefix trap" — a two-sided distribution shift between student occupancy and teacher reliability — and addresses it with a step-decaying weight schedule. ReOPD matches or improves OPD accuracy while eliminating environment calls during training and training at least 4× faster per step.

Problem¶

On-policy distillation (OPD) for multi-turn agentic tasks requires fresh student rollouts through the environment and teacher queries at every visited history per training update — incurring both environment-interaction and teacher-inference cost at each step. Standard off-policy SFT avoids this but suffers covariate shift when the student deviates from teacher-generated prefixes. Multi-turn OPD introduces a further tension: making roll-in histories more student-on-policy improves relevance to the student but pushes the teacher's conditional into histories where it is no longer a reliable improvement signal, creating a two-sided distribution shift. Neither fully student-on-policy nor teacher-forced roll-in is uniformly optimal.

Method¶

ReOPD decouples student training from live environment interaction by replaying a fixed pool of pre-collected teacher trajectories as prefixes. At training step \(t\), the student is conditioned on a teacher-recorded prefix \((o_1, a_1, \ldots, o_t)\) and generates its own action \(a_t\); the teacher supplies the per-step target distribution \(\pi_T(\cdot \mid x, H_t)\) without executing any new environment query.

The theoretical core is Proposition 1, which bounds the gap between the ideal interactive objective \(R^\star(\theta; \theta_\text{old})\) (evaluated on student-environment occupancy \(d^t_{\theta_\text{old}}\)) and the replayed objective \(L_\rho(\theta; \theta_\text{old})\) (evaluated on effective prefix distribution \(\rho_t = w_t P_t\)):

\[|R^\star - L_\rho| \leq \mathbb{E}_{x}\!\sum_t \alpha_t \!\left[2B\,\mathrm{TV}\!\left(d^t_{\theta_\text{old}}, \rho_t\right) + \mathbb{E}_{H_t \sim \rho_t}\!\left[\epsilon^{\theta}_{T,t}\right]\right]\]

The first term penalizes student occupancy mismatch; the second penalizes teacher reliability error \(\epsilon^{\theta}_{T,t} = |\ell(\pi_\theta, q^\star_t) - \ell(\pi_\theta, \pi_T)|\). The optimal effective distribution is the geometric bridge:

\[\rho^\star_t(h_t \mid x) \propto \left(d^t_{\theta_\text{old}}\right)^{\gamma_t}\!\!\left(d^t_T\right)^{1-\gamma_t}, \quad \gamma_t = \frac{\lambda_\text{stu,t}}{\lambda_\text{stu,t}+\lambda_\text{tea,t}}\]

Because teacher reliability degrades at later interaction steps (as student-induced histories drift from the teacher's support), the optimal \(\gamma_t\) decreases with depth. Rather than tuning \(\gamma_t\) directly, ReOPD implements a step-decaying weight schedule over the fixed teacher-pool roll-in: earlier steps receive higher weight, shifting training mass away from deep steps where the teacher signal is least trustworthy. This is applied as a sampling/reweighting over positions, not a per-history reweight, keeping the method simple and environment-free.

Key Contributions¶

Identifies the prefix trap in multi-turn OPD and decomposes it into a temporal compounding-error layer and a distributional two-sided-shift layer.
Derives Proposition 1 bounding the objective gap into an explicit student occupancy mismatch term and a teacher reliability term, showing fully student-on-policy distillation is not always optimal.
Frames multi-turn OPD as reliability-aware prefix distribution design with a geometric bridge between student and teacher occupancies.
Proposes ReOPD: offline prefix replay with a step-decay schedule as a practical, training-environment-free realization of the bridge.
Demonstrates that pre-collected teacher RL rollouts serve as a free trajectory pool, enabling reuse across multiple heterogeneous environments with a single unified offline pool.

Results¶

Math with Python (Qwen3-4B student, Qwen3-8B teacher): ReOPD 53.7% vs OPD 45.0% vs SFT 39.1% accuracy — ReOPD improves over OPD when the teacher–student gap is large.
Search/QA: ReOPD accuracy essentially matches OPD (consistent with the two-regime analysis — teacher stays reliable on student histories in this setting).
Training speed: ReOPD is 4.2× faster per step on Math with Python (≈40 s vs 175 s) and 9.1× faster on Search.
Tool calls during training: ReOPD uses 0 tool calls per trajectory vs ~40 (Math+Python) and ~7 (Search) for OPD.
Gains hold across single- and multi-environment settings and varying teacher/student model scales.

Limitations¶

Prefix replay is not fully on-policy along the roll-in prefix; the student experiences on-policy conditioning only at the single evaluated step per training sample, not along the full trajectory.
Requires a pre-existing pool of teacher trajectories; quality and diversity of this pool bound ReOPD's ceiling.
The step-decay schedule is a positional proxy for teacher reliability rather than a per-history estimate; no explicit reliability signal is computed per prefix.
When the teacher is already reliable on student-induced histories (small gap, short horizon), ReOPD yields no accuracy gain over OPD — it only recovers parity in that regime.
Paper text is truncated; additional ablations and limitations may not be fully visible.

Relevance to Agentic AI / LLM Agents¶

ReOPD directly addresses a core bottleneck in training LLM agents for multi-turn tool-use: the need to keep environments live and interactive throughout student training. By converting teacher RL rollouts into a reusable offline resource, it enables scalable distillation pipelines where environment deployment is a one-time cost. The two-regime analysis — on-policy prefixes are best when the teacher is reliable, teacher-supported prefixes are best when reliability degrades — provides a principled framework for thinking about curriculum and supervision quality in sequential agent training. The decoupled offline pool also generalizes naturally to multi-environment agent training, relevant to heterogeneous tool-use settings (code execution, web search, APIs) where maintaining all environments simultaneously is prohibitively complex.