🤖 AI

Step-Level On-Policy Distillation: Interpolating Between On-Policy Distillation and Supervised Fine-Tuning

The paper proposes Step-Level On-Policy Distillation (SOPD), a novel method that interpolates between Supervised Fine-Tuning and On-Policy Distillation to provide complete, step-level corrections on student-generated trajectories, thereby substantially outperforming both conventional approaches in reasoning and agent tasks.

Original authors: Changhui Sun, Lanbo Liu, Hang Lei, Tong Ling, Jiahang Xie, Zhiyong Zheng, Yujia Wang, Hao Liu, Feng Xiao, Lu Liu, Yanlong Du, Zifeng Cheng, Ziwei Jiang, Qing Gu

Published 2026-08-18

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Changhui Sun, Lanbo Liu, Hang Lei, Tong Ling, Jiahang Xie, Zhiyong Zheng, Yujia Wang, Hao Liu, Feng Xiao, Lu Liu, Yanlong Du, Zifeng Cheng, Ziwei Jiang, Qing Gu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the world of artificial intelligence, large language models are often compared to students who must learn from a teacher. For years, the standard way to teach these digital students has been to show them perfect examples written by an expert and ask them to memorize the sequence. This method, known as supervised fine-tuning, is effective but rigid; it assumes the student will always follow the expert's path perfectly, ignoring the mistakes the student is actually prone to making. A newer approach, called on-policy distillation, tries to fix this by letting the student generate its own answers and then correcting it token by token as it goes. While this allows the student to learn from its own errors, it creates a different problem: the corrections are so fragmented that the student never sees a complete, coherent path to the right answer. It is like trying to learn to drive by receiving a single-word instruction every time you touch the steering wheel, rather than being shown how to navigate a whole curve.

Researchers have now developed a new training method that bridges this gap, combining the best of both worlds. By introducing a technique called Step-Level On-Policy Distillation, the team allows the student to complete a full thought or action sequence on its own, and then asks the teacher to provide a single, coherent correction for each major step of that journey. This approach was tested on complex tasks ranging from solving difficult math problems to navigating virtual household environments. The results show that this method helps the student learn significantly faster and more accurately than previous techniques, achieving success rates that were previously out of reach for models of this size. The key discovery is that by breaking the learning process into natural, meaningful chunks rather than individual words or random corrections, the model can understand not just what the right answer is, but how to get there from the specific mistakes it made.

The core of this new method lies in how it handles the interaction between the student model and the teacher. In the traditional approach of on-policy distillation, the teacher looks at the student's work and offers a tiny correction for every single word generated. If the student makes a mistake early on, the teacher tries to fix it word by word, but the student is still stuck in a confused state, trying to follow a broken path. The new method, Step-Level On-Policy Distillation, changes the rhythm of this interaction. First, the student is allowed to generate a complete response or a full sequence of actions without interruption. Once the student has finished, the system breaks that response down into natural steps, such as a complete sentence in a math proof or a full turn of conversation in a game.

At this point, the teacher steps in, but not to rewrite the whole story. Instead, the teacher looks at the student's starting point for each specific step and generates a single, correct version of that step. The student then learns to match this correct step while keeping the rest of its own original context. This means the student learns to correct its own trajectory without the teacher taking over the entire process. The teacher provides a clear, local guide for each part of the journey, ensuring that the student sees a complete, logical path forward from the exact spot where it went wrong. This preserves the student's experience of making mistakes while giving it a structured way to fix them.

The researchers tested this approach in two very different environments to see if it held up under pressure. The first test involved an agent navigating a simulated home, a task known as ALFWorld. In this environment, the model had to figure out how to perform tasks like finding a key or cleaning a room by interacting with a text-based world. The team used a powerful teacher model to guide a smaller student model. When they compared the new method against the old token-by-token correction style, the results were striking. The student trained with the new method improved its success rate by over 18 percentage points on tasks it had seen before and by over 21 percentage points on tasks it had never seen. Furthermore, it completed these tasks in fewer attempts, showing that it was not just guessing better but learning more efficient ways to solve problems.

The second test focused on mathematical reasoning, a domain where logical consistency is paramount. Here, the models were asked to solve competition-level math problems. The researchers found that the new method allowed the student to organize its thinking into clearer, more structured steps. As the training progressed, the student began to generate more distinct reasoning steps, and the teacher's corrections became richer and more detailed. On four different math benchmarks, the student trained with this new method outperformed the previous best approach by a wide margin, improving its average accuracy by nearly 10 percentage points. This demonstrated that the method was not just a technique for simple tasks but a robust way to improve complex reasoning.

One of the most practical advantages of this new technique is that it does not require the teacher to reveal its internal calculations. In many previous methods, the student needed access to the teacher's raw probability scores for every word, which is often impossible if the teacher is a closed, black-box system. The new method only requires the teacher to generate text, making it much easier to apply to real-world scenarios where the best models are not fully open. Additionally, because the teacher only generates one step at a time based on the student's existing work, the process is faster and avoids the delays caused by waiting for the environment to react to every single move.

The study confirms that the way we structure learning matters as much as the quality of the teacher. By moving away from fragmented, word-by-word corrections and toward coherent, step-level guidance, the researchers have created a training dynamic that respects the student's own path while providing a clear map for correction. This approach successfully combines the flexibility of learning from one's own mistakes with the clarity of expert guidance. The findings suggest that for artificial intelligence to become more capable and reliable, especially in complex, interactive tasks, we must design training methods that allow models to see the whole picture of their errors and learn how to repair them in a single, logical motion.

Technical Summary: Step-Level On-Policy Distillation (SOPD)

1. Problem Statement

The paper addresses limitations in current On-Policy Distillation (OPD) and Supervised Fine-Tuning (SFT) paradigms for training Large Language Models (LLMs), particularly in reasoning and agent tasks.

Limitations of Standard OPD: While OPD aligns a student model with a teacher's logit distribution on student-generated trajectories (avoiding the state-distribution mismatch of off-policy methods), it operates at the token level. This results in "fragmented corrections." When a student makes an error, the teacher provides one-token guidance conditioned on the erroneous prefix. Consequently, the teacher cannot unfold a complete, coherent repair path; subsequent targets remain conditioned on the original erroneous student trajectory rather than a corrected one.
Limitations of Standard SFT: SFT provides coherent, long-horizon supervision but is inherently off-policy. It trains on teacher-generated trajectories, creating a mismatch between the training state distribution and the states actually visited by the student during inference.
Limitations of Existing Hybrid Approaches: Methods like DAgger or On-Policy Expert Correction (OEC) attempt to bridge this gap by inserting teacher responses into student trajectories. However, these approaches inherently alter the trajectory the student would have produced, changing the state distribution and weakening the on-policy nature of the training.

2. Methodology: Step-Level On-Policy Distillation (SOPD)

The authors propose Step-Level On-Policy Distillation (SOPD), a method that interpolates between the long-horizon correction of SFT and the on-policy advantages of OPD.

Core Mechanism

SOPD operates on complete student-generated trajectories but partitions them into natural steps (e.g., reasoning paragraphs in math, or environment interaction turns in agents).

Student Rollout: The student model $\pi_\theta$ generates a complete trajectory $\tau_S = (s_1, s_2, \dots, s_K)$ , where each $s_k$ is a natural step. The sequence of prefixes $h_k$ (the state before step $k$ ) is recorded.
Teacher Query: After the student completes the trajectory, the teacher model is queried once per step. For each recorded prefix $h_k$ $h_{k}$ , the teacher generates a single target step $\tilde{s}_k \sim q_T(\cdot | h_k)$ $\tilde{s}_{k} \sim q_{T} (\cdot ∣ h_{k})$ .
- Crucially, the teacher generates this target independently based only on the student's prefix $h_k$ . It does not see previous teacher outputs, nor does its output alter the student's trajectory.
- This preserves the student's state occupancy ( $h_k \sim d_{\pi_\theta}$ ) while providing a coherent, multi-token correction path for that specific step.
Training Objective: The student is trained to predict the teacher's generated step $\tilde{s}_k$ given the prefix $h_k$ . The loss function is a step-balanced cross-entropy:
$L_{SOPD} = \frac{1}{|S_B|} \sum_{(h_k, \tilde{s}_k) \in S_B} \frac{1}{|\tilde{s}_k|} \sum_{j=1}^{|\tilde{s}_k|} -\log \pi_\theta(\tilde{s}_{k,j} | h_k, \tilde{s}_{k,<j})$
This formulation averages the loss within each step and then averages across all steps, preventing long teacher steps from dominating the gradient updates.

Implementation Details

Packed Attention Mask: To ensure efficiency, the input prompt, the full student trajectory, and all teacher target steps are packed into a single sequence. The attention mask ensures that each teacher target can only attend to the prompt, its corresponding student prefix, and its own causal tokens. Future student steps and other teacher targets are masked.
Black-Box Compatibility: SOPD requires only teacher-generated text, not teacher logits. This makes it applicable to black-box distillation scenarios.
Parallel Generation: Since teacher queries are conditionally independent (based on fixed student prefixes), they can be generated in parallel after the student rollout is complete, avoiding the latency of sequential teacher-environment interactions.

Theoretical Limits

The authors demonstrate that SOPD connects two granularity limits:

One-Step Limit: If the entire response is treated as a single step, SOPD reduces to standard Sequence SFT.
One-Token Limit: If every step is a single token, SOPD approximates a sampled Forward-KL OPD objective.

3. Key Contributions

Proposal of SOPD: A new distillation paradigm that combines the on-policy state coverage of OPD with the coherent, multi-token supervision of SFT.
Black-Box Applicability: Unlike standard OPD which requires teacher logits, SOPD only requires teacher text generation, making it suitable for black-box teachers.
Theoretical Unification: The paper establishes that SOPD mathematically reduces to SFT or approximates Forward-KL OPD depending on the step length, offering a unified perspective on these distillation paradigms.
Empirical Validation: Demonstrated superiority over both SFT and Vanilla OPD across reasoning and agent tasks.

4. Experimental Results

The authors evaluated SOPD on two distinct domains: ALFWorld (interactive agent tasks) and Mathematical Reasoning.

ALFWorld (Agent Interaction)

Setup: Distilling a 3B student (Qwen2.5-3B) from a 7B RL teacher (Qwen2.5-7B-RL).
Results: SOPD achieved a 84.29% success rate on "Valid Seen" tasks and 82.09% on "Valid Unseen" tasks.
- This represents a +18.57 and +21.64 point improvement over Vanilla OPD, respectively.
- SOPD also reduced the mean number of environment interaction rounds by ~3.5 to 4.3 compared to OPD, indicating more efficient policies.
- On "Hard" tasks, SOPD maintained the shortest interaction trajectory among student methods while achieving the highest success rate.

Mathematical Reasoning

Setup: Distilling a 4B non-thinking student from a 4B RL teacher (Qwen3-4B-Non-Thinking-RL-Math) using the DeepMath dataset.
Benchmarks: AIME24, AIME25, HMMT25 (Feb), and HMMT25 (Nov).
Results: SOPD significantly outperformed Vanilla OPD across all benchmarks.
- Average Accuracy: SOPD achieved 57.7% average accuracy compared to 47.7% for OPD (a ~10 percentage point gain).
- Specific gains included +9.9 points on AIME24 and +13.1 points on HMMT25-Nov.
Training Dynamics: Analysis showed that as training progressed, the student model generated more natural reasoning steps, and the teacher provided richer supervision (more tokens) per step, indicating the model learned to organize more structured reasoning processes.

5. Significance and Claims

The paper claims that SOPD offers a new perspective on distillation methods by resolving the trade-off between state-distribution alignment (OPD) and coherent supervision (SFT).

Coherent Correction: Unlike token-level OPD, SOPD allows the teacher to provide a complete, executable correction path for a specific segment of a student's trajectory without altering the trajectory itself.
Efficiency: By decoupling the teacher generation from the student's environment interaction (in agent tasks) and allowing parallel queries, SOPD reduces the practical cost of black-box distillation to be comparable to standard SFT.
General Effectiveness: The results suggest that step-level supervision is a robust mechanism for improving both reasoning capabilities and agent interaction policies, outperforming existing state-of-the-art distillation and fine-tuning baselines.

The authors conclude that SOPD effectively turns student-visited intermediate states into black-box teacher queries, providing dense, local supervision while preserving the integrity of the student's state distribution.