Multi-Turn On-Policy Distillation with Prefix Replay
Baohao Liao, Hanze Dong, Christof Monz, Xinxing Xu, Li Dong, Furu Wei
ReOPD reuses offline teacher trajectories for on-policy distillation, avoiding environment interaction while maintaining dense supervision.
How can we distill agentic LLM policies from a teacher without the high cost of running live environment interactions during training?
Multi-turn on-policy distillation (OPD) is prohibitively expensive because it requires fresh environment rollouts and teacher queries at every training step. Replayed-Prefix On-Policy Distillation (ReOPD) solves this by training off-environment: it replays pre-collected teacher trajectories as prefixes, letting the student act at specific steps while using the teacher’s recorded conditional as a dense supervision target. This method uses a reliability-aware step-decay schedule to prioritize early, low-shift prefixes, effectively balancing student relevance against teacher reliability. Across mathematical reasoning and search tasks, ReOPD matches or exceeds the accuracy of online OPD while eliminating environment interaction and reducing training time by at least 4×.
Paper Primer
The core challenge in multi-turn OPD is the "prefix trap": a two-sided distribution shift where student-generated histories improve relevance but risk querying the teacher on unreliable, out-of-distribution states. ReOPD treats this as a reliability-aware distribution design problem, using a geometric bridge to select prefixes that are both likely under the student and within the teacher's reliable support.
ReOPD achieves accuracy parity or improvement over online OPD while removing the need for live environment interaction.
Empirical validation across mathematical reasoning (Python-tool) and search environments using varying teacher-student model scales. At least 4× faster per rollout than standard online OPD with zero tool calls during student training.
The method implements a step-decay schedule $\omega(t; \kappa) = \kappa^t$, where $\kappa$ is a single interpretable knob derived from the average teacher-student KL divergence. This schedule concentrates supervision on early trajectory steps where the teacher's signal is most reliable, effectively mitigating the compounding errors inherent in multi-turn generation.
Why is "on-policy" distillation still relevant if the method uses pre-collected teacher trajectories?
The student remains on-policy at the specific supervised step $t$ because it generates its own action $A_t$ given the prefix. This ensures the training signal is relevant to the student's own behavior, unlike off-policy supervised fine-tuning which relies entirely on teacher-generated actions.
When does ReOPD outperform standard on-policy distillation?
ReOPD excels when the teacher-student capability gap is wide, as the teacher's signal degrades on student-induced histories. In these regimes, the reliability-aware decay prevents the student from learning from unreliable teacher targets at deep trajectory steps.
ReOPD transforms expensive agent-environment interaction into a reusable offline resource, enabling scalable distillation for complex agentic tasks without the overhead of online rollouts.
Introduction to Agentic Distillation
We expose the live‑environment cost of on‑policy distillation and motivate an offline replay solution.
On‑policy distillation (OPD) lets a student imitate a teacher while staying on the student’s own trajectory, which is essential for agentic tasks that span many interaction turns. The bottleneck is that every training step must launch a fresh environment rollout and query the teacher, making the process prohibitively expensive. ReOPD sidesteps this cost by replaying pre‑recorded teacher prefixes, turning the expensive online loop into an offline, reusable resource.
**Figure 1.** ReOPD keeps the benefits of on-policy distillation while removing environment interaction. ReOPD matches or improves OPD accuracy, but trains much faster per step and eliminates tool calls during training by replaying teacher-recorded prefixes instead of executing fresh environment rollouts. The student/teacher models are Qwen3-4B-Instruct-2507 and Qwen3-8B, respectively.
The dominant cost of agentic distillation is live environment interaction; replaying teacher prefixes removes that cost without sacrificing performance.
Formulating Multi-Turn Distillation
Formulate on‑policy distillation and introduce the weighted‑history trick that makes offline replay viable.
On‑policy distillation (OPD) forces the student to interact with a live environment at every training update, which makes scaling to many tasks or long horizons prohibitively expensive.
**Figure 2.** Comparison between OPD and ReOPD for agentic tasks. Up: OPD with online environment. The environment is always alive during training. And all steps equally contribute to the loss. Down: ReOPD with offline environment. The environment is only alive for collecting the teacher's trajectories, which can happen during the training of the teacher agent by using RL, like GRPO. Afterwards, the environment is not needed for the training of the student agent. The earlier steps contribute more to the loss.
Formally, a trajectory is built step‑by‑step: at step $t$ the student sees history $H_t=(O_1,A_1,\dots,O_t)$, selects $A_t\sim\pi(\cdot\mid x,H_t)$, and the environment returns $O_{t+1}\sim E(\cdot\mid x,H_t,A_t)$. The induced distribution over histories is the occupancy $d_t^{\pi,E}(h_t\mid x)$.
Instead of training on whatever the student happens to encounter, we reshape the sampling distribution so that each step sees histories that are both reachable by the current student and lie inside the teacher’s reliable support.
Step 1: set $\gamma_1=0.8$ (early step, favor student). Compute unnormalized $\rho_1^{*}(h^{(1)})=0.7^{0.8}\cdot0.4^{0.2}\approx0.63$, $\rho_1^{*}(h^{(2)})=0.3^{0.8}\cdot0.6^{0.2}\approx0.27$.
Normalize: $Z_1=0.63+0.27=0.90$, so $\rho_1^{*}(h^{(1)})\approx0.70$, $\rho_1^{*}(h^{(2)})\approx0.30$.
Weight $w_1^{*}= \rho_1^{*}/P_1$: if $P_1$ is uniform (0.5 each), then $w_1^{*}(h^{(1)})\approx1.40$, $w_1^{*}(h^{(2)})\approx0.60$.
Step 2: set $\gamma_2=0.5$ (mid‑trajectory). Re‑compute similarly, obtaining a more balanced $\rho_2^{*}$.
Step 3: set $\gamma_3=0.2$ (late step, favor teacher). The bridge shifts toward the teacher’s distribution, down‑weighting prefixes the student rarely visits.
Early steps keep the student’s own dynamics dominant, preserving on‑policy relevance; later steps gradually hand control to the teacher’s support, mitigating unreliable teacher signals on out‑of‑distribution histories.
How does this “geometric‑bridge” weighting differ from naïve importance sampling that simply reweights by the likelihood‑ratio $r_t$?
Naïve importance sampling uses the exact density ratio $r_t=\prod_{s<t}\pi_{\theta^{\text{old}}}(a_s)/\pi_T(a_s)$, which can have extremely high variance because the two policies may diverge sharply. The bridge instead forms a convex combination of the two occupancies in probability space (raising each to a power) before normalizing, yielding a smoother, lower‑variance weight $w_t^{*}$ that can be tuned per step via $\gamma_t$.
**Figure 3.** Training a shared student agent on multiple heterogeneous environments with OPD and ReOPD. Left: With an increasing number of environments, the operational complexity grows for OPD due to the heavy deployment of the environments. Right: ReOPD doesn't require the deployment of all environments at the same time. The teacher's trajectories can be collected separately for different environments, and then be merged into a unified pool for the later training of the student agent without any online environment.
In summary, the core trick is to replace the raw student‑on‑policy distribution with a step‑dependent bridge $\rho_t^{*}$ that interpolates toward the teacher’s support where the teacher is trustworthy, and to realize this bridge via normalized importance weights $w_t^{*}$.
Prior Work in Distillation
We situate ReOPD among prior distillation, RL, and self‑training approaches.
Related work clusters around three themes: (1) distilling language models, (2) reinforcement‑learning post‑training, and (3) self‑training with rejection sampling.
Early knowledge‑distillation pipelines (Hinton et al., 2015) compress a teacher into a smaller student; sequence‑level extensions (Kim & Rush, 2016) train on teacher‑generated trajectories.
Off‑policy variants let the student imitate teacher prefixes but require the student to recover from its own prefixes at inference.
On‑policy distillation (OPD) queries the teacher on student‑sampled prefixes; recent methods include MiniLLM (reverse KL), GKD (mixed generations), DistiLLM (refined divergence), and speculative distillation (interleaved tokens).
These works frame OPD as dense supervision contrasted with sparse RL, and explore black‑box OPD and on‑policy context distillation, highlighting that a stronger teacher does not guarantee better student performance.
Our contribution extends this line from single‑turn generation to multi‑turn training using a fixed pool of teacher trajectories, treating the replayed prefix distribution as the primary design variable.
Reinforcement‑learning post‑training optimizes scalar feedback rather than teacher conditionals. RLHF pairs preference rewards (Bradley–Terry) with PPO, while contrastive losses such as SLIC‑HF, DPO, IPO, and GPO replace online reward optimization.
Reasoning‑oriented RL variants (GRPO, DAPO) and critic‑free REINFORCE derivatives (ReMax, RLOO, Reinforce‑Rej) further diversify the optimizer landscape; RAFT and rejection sampling treat high‑reward trajectories as imitation data.
ReOPD is orthogonal to this optimizer family: it retains dense teacher supervision while eliminating live environment interaction.
Agentic LLM research augments RL‑style post‑training with tool use, search‑augmented agents, self‑hinting, budget‑aware reasoning, and distillation of teacher‑generated reasoning‑and‑acting traces.
Self‑training cycles repeatedly turn model samples into new data; expert‑iteration, STaR, and ReST alternate generation, filtering, and retraining, while rejection sampling and RAFT instantiate the same principle for LLMs.
ReOPD also reuses offline traces, but rather than merely accepting or rejecting them, it queries the teacher at specific prefixes, making the prefix distribution the core variable.
Distribution‑shift literature describes the “prefix trap”: behavior‑cloning on expert states suffers covariate shift once the learner deviates, a problem addressed by DAgger, scheduled sampling, and Professor Forcing.
Recent reasoning‑distillation work identifies a dual exposure bias—teacher‑forced traces mismatch student inference, while fully student‑generated contexts can render teacher queries unreliable.
The ReOPD Algorithm
ReOPD replaces live roll‑outs with teacher‑forced prefixes and a reliability‑aware step schedule.
Standard OPD stalls because each training step must query the environment, inflating compute cost. ReOPD sidesteps this bottleneck by replaying recorded teacher prefixes and weighting them according to a reliability‑aware schedule, keeping the supervision on‑policy without any live interaction.
ReOPD reuses a static pool of teacher trajectories, forces the early part of each trajectory to be exactly the teacher’s actions, lets the student act only at a single sampled step, and then applies a decay schedule that emphasizes those early, low‑shift steps.
How does prefix replay differ from the “teacher‑forced” roll‑in used in vanilla OPD?
In vanilla OPD the student rolls out its own actions for the entire prefix, requiring live environment interaction. ReOPD instead freezes the prefix to the teacher’s recorded actions, so the roll‑in is purely offline while the student still acts on‑policy at the sampled step.
The exact bridge weight for a replayed prefix is the product of per‑step student‑to‑teacher probability ratios; it quantifies how likely the current student would have followed the teacher’s recorded actions up to step t.
Step 1 ratio = 0.2 / 0.5 = 0.4.
Step 2 ratio = 0.3 / 0.6 = 0.5; cumulative r₂ = 0.4 × 0.5 = 0.20.
Step 3 ratio = 0.4 / 0.7 ≈ 0.57; cumulative r₃ = 0.20 × 0.57 ≈ 0.11.
The product quickly decays, illustrating why later steps receive smaller weights: the student is far less likely to have reproduced the teacher’s exact actions as the prefix grows.
Why is the likelihood‑ratio rₜ usually below 1 for later steps?
Each factor compares the student’s probability of the teacher’s recorded action to the teacher’s own probability. Since the teacher chose the action, its probability is typically higher than the student’s, making each ratio < 1; multiplying many such ratios drives the product downward.
Because the exact likelihood‑ratio weight shrinks monotonically with depth, we can replace the noisy per‑history factor with a simple geometric decay that depends only on the step index.
How should I pick the decay factor $\kappa$ in practice?
Estimate the average per‑step KL divergence \bar{c} between teacher and current student on a held‑out subset of the teacher pool; then set $\kappa$ ≈ exp(−$\gamma$ₜ \bar{c}). Empirically sweep $\kappa$ in [0.2, 1.0] and pick the value that maximizes validation performance.
Initialize $\theta_{old}$ ← $\theta$.
For each OPD iteration, assemble candidate positions (x, hₜ) from the teacher pool $D_T$, where hₜ contains teacher actions $A_{<t}$ and observations $O_{ \leq t}$.
Sample positions with probability pₜ ∝ $\omega$(t; $\kappa$) = $\kappa^{t}$ (the step‑decay schedule).
For each sampled position (x, hₜ):
Generate the student’s action sequence Aₜ ∼ $\pi_{θ_{old}$}(· | x, hₜ) autoregressively.
Compute the token‑level KL loss ∑_{j} KL($\pi$_$\theta$(· | x, hₜ, $a_{<j}$) ‖ $\pi_{T}$(· | x, hₜ, $a_{<j}$)).
Accumulate the loss over all sampled positions and update $\theta$ by gradient descent.
Set $\theta_{old}$ ← $\theta$ and repeat.
**Figure 4.** **Step index is a strong proxy for the likelihood-ratio weight.** For each supervised position we compute the per-position weight $\hat{r}_t = \prod_{s<t} \pi_{\theta_{old}}(a_s)/\pi_T(a_s)$ – the student-to-teacher likelihood ratio along the teacher prefix – and plot it against the step index $t$. The empirical weight decays with depth, because each factor compares the student’s and teacher’s probability of the teacher’s own recorded action and is typically below one, so the product shrinks as the prefix lengthens. This makes step index a cheap proxy for the expensive prefix-level ratio and motivates replacing the explicit ratio with the one-parameter step-decay schedule $\omega(t) = \kappa^t$ introduced below.
**Figure 5.** **Favoring early, low-shift steps improves distillation.** Both panels use ReOPD's sampling implementation and put more training mass on early interaction steps, where the student stays close to the teacher and the two-sided shift is small. (a) *Sampling by chunk*: drawing supervised positions from early chunks yields the best student, and pushing mass to later chunks degrades it. (b) *Sampling with decay $\kappa$*: applying the step-decay $\kappa^t$ as the sampling probability, a moderate decay outperforms uniform sampling ($\kappa=1$) and tracks the same trend. Final performance is governed by how strongly early steps are favored, validating the step-decay of Section 4. The student model is Qwen3-4B-Instruct-2507, and the models in the legends are teacher model.
**Figure 6.** The resources used by different environments. For math, the processes are utilized for the execution of Python code in parallel. For search, the GPU memory is utilized for deploying the embedding model and the storage of the Wikipedia embeddings.
Empirical Evaluation
ReOPD matches OPD accuracy while cutting rollout time by up to 9.1×.
ReOPD eliminates the need for live environment interaction during distillation by replaying teacher prefixes, preserving performance while saving compute.
ReOPD reduces rollout time by up to 9.1× compared to OPD.
Figure 8 shows OPD requires 169 s per step on math tasks versus 14 s for ReOPD + teacher‑prefix resampling on search, a 9.1× speedup.
We evaluate Qwen3‑family models across three teacher scales (4B‑Instruct‑2507, 8B, 30B‑A3B‑Instruct‑2507) distilled into a 4B‑Instruct‑2507 student, and also a 8B student under a 30B teacher. All models undergo a cold‑start SFT phase before the teacher receives GRPO reinforcement learning.
The gap measures how much stronger the teacher’s policy is than the student’s, influencing how much offline replay can faithfully substitute live interaction.
How does the capability gap differ from simply using a larger teacher model?
The gap is not just size; it quantifies distributional shift between teacher and student policies. A larger model can still produce prefixes that lie within the student’s reachable space, keeping the gap small, whereas a mismatched architecture may create out‑of‑distribution trajectories that hurt offline distillation.
**Figure 7.** Prefix source: the teacher's own model gives the best prefixes, not a stronger generator. Teacher Qwen3-8B, student Qwen3-4B-Instruct-2507; we fix the distillation target to the teacher and vary only the prefix generator. Performance peaks when prefixes come from the teacher itself and degrades with larger or stronger generators – a direct test of the teacher-reliability view: the teacher's conditional is trustworthy only on prefixes within its own support.
**Figure 8.** Rollout time when taking the teacher's rollout into account. When the teacher's trajectories from RL are not available, one could use the teacher to sample the trajectories for all prompts at once. We take this time into account for the comparison. For this setting, the environment is only online for the teacher's rollout.
**Table 8.** Teacher's prefix resampling. For some cases, we don't have the teacher's prefix from RL. For example, the teacher is already a strong open-sourced model. And we don't have resource to do RL on it. We can deploy the environment for the teacher to sample all trajectories for all prompts at once, and use these trajectories as prefixes for ReOPD. Table 7 already shows that the performance is similar to the one using prefixes from RL. In Table 8, we take the time used for the teacher's resampling into account when comparing to OPD. Our ReOPD is still efficient, > 2x. Notably, the environment is only online for the teacher's resampling, and deactivated for the training of student, which is resource efficient.