Multi-Turn On-Policy Distillation with Prefix Replay

2026-07-06

AI Summary

This work addresses the high cost of multi-turn online policy distillation (OPD), which suffers from frequent environment interactions, repeated teacher queries, and bidirectional distribution shift between student and teacher policies. To overcome these limitations, we propose ReOPD, an efficient offline distillation framework that replays pre-collected teacher trajectories as prefixes and lets the student act at critical steps under stepwise teacher supervision. We identify a “prefix trap” inherent in OPD and reformulate distillation as a reliability-aware prefix distribution design problem. Our method employs a step-decay sampling strategy that prioritizes early-stage prefixes with lower distributional shift. Across diverse model scales and tasks—including mathematical reasoning and search—ReOPD matches or improves distillation accuracy without requiring tool calls, while accelerating training by at least 4× and substantially reducing agent training costs.

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.

Problem

• on-policy distillation
• multi-turn interaction
• distribution shift
• prefix trap
• LLM agent

Innovation

• on-policy distillation
• prefix replay
• distribution shift
• multi-turn agent learning
• offline distillation