Multi-Turn On-Policy Distillation with Prefix Replay

The Core Problem and Solution

On-policy distillation (OPD) for agentic tasks is expensive. When an LLM agent learns from a teacher over multiple turns—solving a math problem step-by-step, or navigating a search environment—the standard approach requires constant fresh interaction: the student acts in the environment, the teacher responds to those exact histories, and both get queried repeatedly. This is slow and costly.

The authors propose Replayed-Prefix On-Policy Distillation (ReOPD), which sidesteps the environment entirely during training. Instead, they collect teacher trajectories once, then replay them. The student acts at selected checkpoints within these pre-recorded paths while the teacher provides supervision at every step—but without new environment calls. The result: at least 4× faster per rollout than OPD, zero tool calls during training, and comparable or better accuracy.

The Prefix Trap: Why Naive Replay Fails

The key insight is subtle and crucial. In multi-turn tasks, there's a tension between two goals:

1. Relevance: Make the histories the student sees as close as possible to what the student would naturally produce (increase student on-policyness).
2. Reliability: Ensure the teacher's supervision is actually trustworthy on those histories.

These pull in opposite directions. When you force histories to be more student-on-policy—by having the student take more steps—you drift further from the teacher's training distribution. The teacher may have been trained on expert demonstrations, so asking it to supervise paths the student naturally produces can put it in low-confidence, unreliable territory. The authors call this the prefix trap: optimizing for one goal worsens the other, creating a two-sided distribution shift.

How ReOPD Works

ReOPD treats this as a reliability-aware prefix distribution design problem. The solution is elegantly simple: use a step-decaying sampling schedule that heavily emphasizes early prefixes (close to the start of trajectories) and downweights later ones.

Early prefixes are shorter and closer to the teacher's original training distribution—the teacher is most reliable there. Later prefixes drift further, introducing shift. By sampling early prefixes more often, ReOPD maximizes the overlap between where the student acts and where the teacher is confident.

The method is straightforward to implement: shuffle the pre-collected trajectories, then sample where the student should begin acting according to a decay schedule. The teacher always supervises densely (every step) from that point onward, but reuses the same data across multiple training epochs.

Results and Impact

Tested on mathematical reasoning (with Python execution) and search-based environments across multiple teacher-student model pairs, ReOPD matches or exceeds OPD accuracy while eliminating environment interaction during training. The 4× speedup per rollout and zero tool calls translate directly to cost savings at scale.

The broader contribution is methodological: the paper reframes multi-turn OPD as a distribution alignment problem and shows that a simple, theoretically motivated fix (step decay) outperforms naive approaches. This turns expensive, one-time teacher rollouts into a reusable offline asset—enabling cheaper distillation across different tools, tasks, and environments without re-querying the teacher or environment.