Source-linked AI summary
Multi-Turn On-Policy Distillation with Prefix Replay
Baohao Liao, Hanze Dong, Christof Monz, Xinxing Xu, Li Dong, Furu Wei
TL;DR
Multi-turn on-policy distillation is expensive and can query teachers on unreliable histories as student prefixes drift. ReOPD reuses teacher trajectories with reliability-aware prefix sampling, preserving or improving OPD-level accuracy while avoiding environment interaction during training.
Problem
Multi-turn OPD requires costly student rollouts and faces a two-sided shift between student relevance and teacher reliability.
Method
ReOPD reuses offline teacher trajectories and samples replayed prefixes with a step-decaying schedule that balances student relevance against teacher reliability.
Results
ReOPD improves AIME24 from 28.3 to 36.7 under a Qwen3-8B teacher while improving the 4B-student math average across teacher settings.
Takeaways & Limitations
ReOPD supports efficient off-environment multi-turn distillation by balancing student occupancy against teacher reliability.
Takeaways & Limitations
ReOPD assumes a pre-collected teacher-trajectory pool, whose coverage and quality bound learning, while step depth is only a coarse reliability surrogate.
Abstract
from arXiv · showhide
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 4$\times$ 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.
1 Introduction
The introduction frames multi-turn on-policy distillation as a costly sequential problem with a two-sided prefix-distribution trap. It proposes ReOPD, which replays teacher prefixes and uses reliability-aware sampling to retain efficient, effective supervision.
- Motivation: Multi-turn OPD queries the teacher at each student-visited history while incurring environment-interaction and teacher-inference costs at every step.Prefix replay instead uses teacher-recorded prefixes and per-step teacher targets without executing the recorded actions.
- Replayed-Prefix On-Policy Distillation: ReOPD rolls the student into teacher-recorded prefixes, generates the evaluated-step action, and reuses recorded observations instead of executing environment actions.The teacher supplies dense per-step supervision on these replayed prefixes.
- Prefix Trap: The prefix trap combines compounding temporal errors with two-sided distribution shift between student occupancy and teacher reliability.Fully student-on-policy prefixes improve relevance but may place the teacher on histories where its target is unreliable.
- Empirical Validation: ReOPD is validated on mathematical reasoning and search environments across teacher and student scales, improving over SFT where reported and matching or outperforming OPD across regimes.It improves on mathematical reasoning when the teacher–student gap is wide and essentially matches OPD when the teacher remains reliable on student-induced histories.
- Reliability-Aware Design: A simple step-decay schedule implements the effective prefix distribution by sampling earlier, lower-shift prefixes more heavily.The introduction presents this schedule as the practical realization of the reliability-aware distribution design.
- Reliability-Aware Design: The analysis decomposes the objective gap into student occupancy mismatch and teacher reliability terms, showing that fully student-on-policy distillation is not automatically optimal.This motivates treating multi-turn OPD as reliability-aware prefix distribution design rather than maximizing student on-policy-ness alone.
2 Related Work
Prior work spans teacher-conditioned knowledge distillation, scalar-feedback reinforcement learning, agentic reasoning, and self-training. ReOPD differs by retaining dense teacher supervision while focusing on reliable replayed-prefix selection under distribution shift.
- Knowledge distillation for language models: Knowledge distillation transfers teacher behavior to smaller students, with sequence-level variants training autoregressive models on teacher trajectories.This paradigm supports instruction and reasoning models trained from teacher-written solutions or rationales.
- RL algorithms for LLM post-training: RL post-training optimizes scalar feedback rather than teacher conditionals, using preference rewards, policy optimization, or direct preference objectives.The passage contrasts RLHF and PPO with simplified contrastive or implicit-reward losses.
- RL algorithms for LLM post-training: ReOPD is orthogonal to these optimizer families because it preserves dense teacher supervision while selecting replayed prefixes where that supervision is reliable.This distinguishes ReOPD from binary-reward approaches such as RAFT and rejection sampling, which select and imitate successful trajectories.
- Reasoning and agentic LLMs: Agentic post-training extends RL-style methods to tool- and environment-interacting systems, including search-augmented agents and tool-integrated mathematical reasoning.Related approaches also improve reasoning through self-hinting, elastic reasoning budgets, online experiential learning, and teacher distillation.
- Self-training and rejection sampling: Self-training repeatedly converts model samples into training data through expert iteration, self-generated rationale bootstrapping, scalar-feedback filtering, or rejection sampling.RAFT and rejection sampling instantiate this broader self-training principle in LLM post-training.
- Distribution shift and compounding errors: The prefix trap relates to covariate shift and exposure bias, motivating methods that reduce mismatch between expert-state or teacher-forced training and learner-induced or free-running inference.DAgger addresses learner-state mismatch, while scheduled sampling, sequence-level training, and Professor Forcing target autoregressive exposure bias.
3 Problem Formulation and Analysis
Multi-turn OPD is formulated as learning over environment-generated interaction histories, but offline replay changes the effective prefix distribution and creates two distinct mismatches: student occupancy shift and teacher reliability shift. The analysis therefore frames prefix selection as balancing student relevance against teacher reliability, with deeper steps favoring teacher-supported histories.
- Problem formulation: Multi-turn OPD models each decision step as a policy action followed by an environment response over an interaction history.The current student policy collects histories, while the optimized student is trained using those histories and teacher targets.
- Ideal and replayed objectives: The ideal objective evaluates the updated student on histories encountered by the current student, whereas offline replay trains on a fixed teacher-trajectory pool.Replayed prefixes avoid repeatedly rolling the student through the environment and querying the teacher at every visited history.
- Effective history distribution: The central design variable is the effective prefix distribution ρ_t, obtained by reweighting collected histories rather than optimizing the raw weights alone.The weights reshape the collection distribution P_t into ρ_t, which determines the histories used by the replayed objective.
- Two-sided mismatch: Two mismatches govern replay quality: ρ_t may differ from student occupancy, and teacher targets may be unreliable on histories far from teacher occupancy.The decomposition identifies these as occupancy shift and teacher-reliability mismatch, respectively.
- Two regimes: The occupancy term favors fully student-on-policy prefixes, while the reliability term favors teacher-supported roll-ins; their balance should vary by interaction step.Teacher reliability worsens with depth, so the preferred mixture shifts from student-relevant early prefixes toward teacher-supported deep prefixes.
4 Replayed-Prefix On-Policy Distillation
ReOPD performs on-policy distillation off-environment by replaying teacher-recorded prefixes, letting the student act only at the supervised step while receiving teacher supervision. A reliability-aware step-decay schedule emphasizes early, low-shift prefixes to balance student relevance against teacher reliability.
- Algorithm: ReOPD reuses a fixed pool of teacher trajectories and chooses an effective prefix distribution that balances student relevance with teacher reliability.Teacher training rollouts can serve as the reusable trajectory pool without dedicated collection.
- Off-environment prefix construction: At each supervised step, the prefix is replayed entirely from the teacher trajectory, while the student generates its own action and the teacher supplies the distillation target.The student’s generated action and token contexts make the evaluated step on-policy despite the teacher-forced prefix.
- Reliability-aware schedule: ReOPD emphasizes early, low-shift positions and downweights late, high-shift positions because replayed teacher prefixes increasingly diverge from histories the student would reach.The schedule controls how teacher-supported histories stand in for student occupancy.
- Reliability-aware schedule: The step-decay surrogate uses wt = ω(t; κ) = κt, with κ ∈(0, 1] controlling steepness and normalization performed across trajectory positions.The overall scale cancels under normalization, so the surrogate has no separate prefactor.
- Setting the decay: κ ≈ exp(−γt¯c) is estimated from the average DKL(πT ∥πθold) on the teacher pool; κ = 1 is uniform and smaller κ concentrates supervision earlier.The weight always decreases with depth, while the teacher–student gap determines how quickly it falls off.
- Interpretation: ReOPD is distillation with an on-policy supervised step and a reliability-aware prefix distribution, not off-policy reinforcement learning or fresh environment rollout.It uses a fixed teacher distribution, does not optimize rewards or estimate advantages, and avoids environment interaction during training.
5 Experiments
Experiments across mathematical reasoning, search, and joint multi-environment training show that ReOPD preserves OPD-level accuracy while improving mathematical reasoning when teacher reliability limits student-on-policy roll-ins. ReOPD also removes student-training tool calls and is at least 4× faster per rollout than OPD.
- Single environment: ReOPD consistently improves mathematical-reasoning student averages over OPD, with gains from 55.1 to 57.2, 51.0 to 53.7, 51.1 to 52.5, and 56.5 to 56.8 across teacher–student settings.The gains are largest with wider teacher–student capability gaps, where student roll-ins can reach histories with less reliable teacher targets.
- Single environment: On search, OPD and ReOPD obtain nearly identical average accuracies: 40.6 versus 40.5 with a Qwen3-4B teacher and 39.1 versus 39.0 with Qwen3-8B.This reflects a reliable-teacher regime in which teacher-anchored replay provides smaller benefits.
- Multiple environments: A single Qwen3-4B student trained jointly on math and search remains on par with OPD in both domains while avoiding online environment interaction during student training.The replay-based pipeline combines heterogeneous environment data from domain-specific teachers into a shared offline training pool.
- Efficiency: 0 tool calls during student training and at least 4× faster per rollout distinguish ReOPD from OPD’s fresh environment rollouts and tool calls at every student update.ReOPD instead replays teacher-recorded prefixes, removing OPD’s main training-time bottleneck.
- Two regimes in practice: Teacher-reliability shift dominates mathematical reasoning as the teacher–student gap widens, whereas student-occupancy shift dominates search and multi-hop QA where teacher reliability remains high on student-induced histories.ReOPD’s teacher-anchored prefixes therefore improve math performance while essentially matching student-on-policy OPD in reliable-teacher settings.
6 Conclusion
ReOPD reuses fixed teacher trajectories for off-environment multi-turn distillation while addressing temporal and distributional prefix traps. Its effectiveness depends on replay-pool coverage and quality, and its step-decaying reliability surrogate remains coarse.
- Conclusion: ReOPD replays teacher-forced prefixes from pre-collected trajectories to supervise student actions without new environment interaction.The method preserves multi-turn temporal structure despite removing the environment.
- Conclusion: The prefix trap has temporal and distributional components: compounding errors and a two-sided distribution shift.The analysis also bounds the gap to an ideal interactive implementation, though the supplied passage truncates the bound.
- Limitations and future work: ReOPD requires a pre-collected pool of teacher trajectories with recorded observations, whose coverage and quality bound student learning.In the experiments, these trajectories come from the teacher’s own reinforcement-learning rollouts at no additional cost.
- Limitations and future work: The reliability surrogate reduces to a coarse step-decaying weight that reflects prefix depth rather than directly measuring each prefix’s shift.The schedule is effective but does not directly assess how far an individual prefix lies from the desired distribution.