Multi-Turn On-Policy Distillation with Prefix Replay
2026-07-06 • Machine Learning
Machine LearningArtificial IntelligenceComputation and Language
AI summaryⓘ
The authors study a method called on-policy distillation (OPD), where a student AI learns by imitating a teacher AI through multiple-step interactions in an environment. They identify that fully online OPD is slow and costly because it needs new environment runs and teacher feedback for every update. To fix this, they propose Replayed-Prefix On-Policy Distillation (ReOPD), which reuses stored teacher interaction data so the student can learn without fresh environment runs. Their method also carefully selects which parts of the data to train on, reducing unreliable learning. Tests show ReOPD is faster and just as accurate as traditional OPD, making it a scalable way to train AI agents offline.
on-policy distillationlarge language modelsmulti-turn interactionagentic tasksoffline learningteacher-student learningdistribution shifttrajectory replayenvironment interactionreinforcement learning
Authors
Baohao Liao, Hanze Dong, Christof Monz, Xinxing Xu, Li Dong, Furu Wei
Abstract
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 training step than OPD. ReOPD therefore turns expensive agent-environment interaction into a reusable offline resource, enabling scalable distillation across tools, tasks, and environments.