Source-linked AI summary
Step-Level On-Policy Distillation: Interpolating Between On-Policy Distillation and Supervised Fine-Tuning
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
TL;DR
Token-level OPD can give fragmented corrections that do not form a complete repair path. SOPD provides step-level supervision on student trajectories, outperforming SFT and OPD across reasoning and agent tasks, including a 10-point average gain over OPD in mathematical reasoning.
Problem
Token-level OPD provides fragmented corrections along erroneous student trajectories rather than complete, coherent repair paths.
Method
SOPD asks the teacher to generate one step from each student prefix, combining on-policy trajectories with long-horizon supervision and recovering SFT or forward-KL OPD at limiting granularities.
Results
Across reasoning and agent tasks, SOPD outperforms SFT and OPD; in mathematical reasoning, it improves over OPD by an average of 10 percentage points.
Takeaways & Limitations
SOPD offers a unified perspective connecting SFT and OPD while providing coherent local supervision without letting the teacher control later student states.
Abstract
from arXiv · showhide
On-policy distillation (OPD) aligns a student model with a teacher's logit distribution on student-generated trajectories. This approach has achieved strong empirical gains and can often surpass conventional off-policy distillation with substantially less data. However, standard token-level OPD can provide only fragmented corrections along an erroneous student trajectory and cannot unfold a complete and correct repair path. Motivated by this limitation, we propose \emph{Step-Level On-Policy Distillation} (SOPD), which combines the long-horizon correction of supervised fine-tuning (SFT) with the on-policy advantage of OPD to provide step-level supervision over complete student-generated trajectories. We show that, at different limits of step length, SOPD reduces to SFT or approximates OPD. Compared with SFT, the teacher responses in SOPD are conditioned on student trajectories and therefore align more closely with student-visited states; compared with OPD, SOPD provides longer-horizon corrections rather than fragmented token-level guidance. Across both reasoning and agent tasks, SOPD substantially outperforms conventional SFT and OPD. For example, on ALFWorld, SOPD improves the average success rate by 13.4 points over Vanilla OPD. We hope this work offers a new perspective for future research on distillation methods.
1 INTRODUCTION
Step-Level On-Policy Distillation (SOPD) combines OPD’s student-generated supervision with SFT’s long-horizon correction by teaching complete trajectories through teacher-generated steps. It interpolates between SFT and approximate forward-KL OPD while outperforming both across reasoning and agent tasks.
- Motivation and Practicality: Unlike standard OPD, which fragments correction into one-token guidance, SOPD supplies continuous segment-level correction and requires only teacher-generated responses.This makes SOPD more directly applicable to black-box distillation, with teacher-generation costs close to one complete response per sample.
- Method: SOPD combines OPD’s on-policy supervision with SFT’s long-horizon correction, providing step-level teacher supervision over complete student-generated trajectories.After each student trajectory, the teacher generates one step from the student prefix at every student step.
- Method: SOPD reduces to SFT when the step count approaches one and approximates forward-KL OPD when step length approaches one token.
- Results: 13.4 points: On ALFWorld, SOPD improves average success rate over Vanilla OPD.
- Results: SOPD outperforms both SFT and OPD across reasoning and agent tasks, improving mathematical reasoning over OPD by an average of 10 percentage points.For agent tasks, steps correspond to environment-interaction turns; for mathematical reasoning, they follow natural reasoning boundaries.
2 RELATED WORK
Prior work spans sequence-level distillation, student-sampled on-policy distillation, and intervention-based guidance for interactive agents. SOPD addresses fragmented token-wise corrections and intervention-heavy training by supervising coherent local corrections across complete student trajectories.
- Knowledge distillation: Knowledge-distillation methods establish sequence generation and student sampling as distinct design choices, including sequence-level distillation, reverse-KL student sampling, and teacher-logit supervision at student-generated states.The cited approaches include Hinton et al. (2015), Kim & Rush (2016), Gu et al. (2024), and Agarwal et al. (2024).
- On-policy language-model distillation: Token-level OPD can provide only fragmented corrections along an erroneous student trajectory, rather than unfolding a complete, correct repair path (Jiang et al., 2026).Trajectory-Refined Distillation identifies this limitation even when the teacher is perfect.
- On-policy language-model distillation: SOPD preserves the complete student trajectory while generating one coherent local teacher correction within each natural step.This differs from TRD, which rewrites the complete student trajectory to construct a repair path.
- Distillation and guidance for multi-turn agents: Existing agent-guidance methods use expert control, external feedback, or temporal curricula to modify which parts of multi-turn trajectories the student generates.Examples include OEC, Agent-RLVR, TCOD, and Guided-OPD, which differ in state generation, teacher intervention, and whether teacher actions are executed.
- Distillation and guidance for multi-turn agents: SOPD remains more faithfully on-policy by conditioning one independent teacher fragment on each corresponding student prefix across the complete student trajectory.Unlike intervention-based methods, teacher guidance does not replace control of the generated trajectory at environment turns.
3 PRELIMINARIES: SFT AND OPD
SFT provides coherent long-range supervision along complete teacher sequences but uses teacher-derived contexts, whereas OPD trains on student-visited prefixes with dense token-level feedback. OPD’s position-wise corrections may fail to form an executable repair path after the student enters an erroneous trajectory.
- Supervised fine-tuning (SFT): SFT teacher-forces complete teacher or expert sequences, yielding coherent long-range targets along the teacher’s preceding tokens.Its contexts come from teacher data rather than the current student, creating an off-policy state-distribution mismatch.
- On-policy distillation (OPD): OPD samples student trajectories and queries the teacher distribution at every student prefix, aligning training states with inference states while providing dense token-level feedback.A common implementation uses a reverse-KL objective and requires teacher logits.
- On-policy distillation (OPD): After an erroneous student prefix, OPD’s next-position teacher targets remain conditioned on increasingly erroneous histories and need not form a coherent, executable correction path.This limitation is attributed to position-wise targets in standard OPD (Jiang et al., 2026).
4 STEP-LEVEL ON-POLICY DISTILLATION
SOPD partitions student-generated trajectories into natural steps and generates one teacher target per student-visited prefix, combining OPD’s state coverage with locally coherent multi-token supervision. Its one-step and one-token limits recover sequence SFT and sampled forward-KL OPD, respectively.
- Method: SOPD preserves student-visited states while providing coherent multi-token teacher corrections within each natural step.Each target begins from a student prefix, then unfolds autoregressively only within that step, so teacher text does not redefine later student states.
- Method: SOPD partitions complete student trajectories into domain-specific natural steps and records every prefix before independently querying the teacher.Mathematical steps follow paragraph or explicit boundaries, while ALFWorld steps pair an assistant response with its parsed action and subsequent observation.
- Teacher generation: SOPD uses generated teacher text rather than logits or gradients, making it compatible with black-box teachers and enabling parallel, execution-free target generation.Teacher queries exclude earlier teacher targets, so they are conditionally independent and require no new environment observations.
- Objective and implementation: The training loss averages token cross-entropy within each teacher step and then averages across steps, preventing longer steps from dominating shorter decisions.The packed implementation shares repeated prefixes while masking each target to its student prefix and own causal target prefix.
- Relation to existing methods: SOPD occupies the granularity between sequence SFT and token-level OPD by retaining student-state coverage while allowing executable multi-token corrections.The method replaces OPD’s logit access with local teacher generation while avoiding teacher-written context for subsequent student steps.
- Limits: With one response as the sole step, SOPD becomes sequence SFT; with one-token steps, it becomes sampled forward-KL OPD on student-visited prefixes.Caching the generated target yields ordinary teacher-data SFT, while a single sampled teacher token estimates the full-distribution forward-KL objective; greedy decoding instead performs hard imitation.
5 EXPERIMENTS
Across interactive and mathematical reasoning tasks, SOPD consistently outperforms OPD, improving success or accuracy while producing more efficient or structured solution trajectories. The experiments evaluate long student-induced trajectories using natural environment turns or reasoning steps.
- Experimental design: The evaluations use natural environment turns for ALFWorld and natural reasoning steps for mathematics, matching each task’s long student-induced trajectory structure.SOPD partitions mathematical responses at natural reasoning boundaries, whereas OPD optimizes sampled-token reverse KL.
- ALFWorld: On ALFWorld, SOPD raises Seen success from 65.72 to 84.29 and Unseen success from 60.45 to 82.09 versus Vanilla OPD, while reducing mean rounds by 3.53 and 4.33.SOPD achieves the highest success and fewest rounds among student methods on both splits.
- ALFWorld: On Hard ALFWorld tasks, SOPD reaches 10.74% success with 28.15 mean rounds, the shortest interaction trajectory among student methods.Together with Seen and Unseen gains, this indicates improved learning of common and unseen interaction states while maintaining stable difficult-task behavior.
- ALFWorld: SOPD improves Seen and Unseen success during ALFWorld training as interaction lengths decrease, while Hard performance remains comparatively stable.The full-split trajectory shows improving task completion alongside more efficient interaction policies.
- Mathematical reasoning: SOPD leads OPD across all four mathematical benchmarks, improving accuracy by 9.9, 10.4, 6.4, and 13.1 points and raising average accuracy from 47.7 to 57.7.These results aggregate 32 decoding samples for each of 30 problems at the final checkpoint.
6 CONCLUSION
SOPD uses student-visited intermediate states for black-box teacher queries, producing step-level targets along complete trajectories. Its granularity connects sequence-level SFT with sampled forward-KL OPD, while results show coherent, dense local supervision in mathematical reasoning and ALFWorld.
- 6 CONCLUSION: SOPD queries a black-box teacher at student-visited intermediate states while preserving the complete student trajectory.It generates one teacher target at every natural step.
- 6 CONCLUSION: At response-level granularity, SOPD recovers sequence SFT; at token-level granularity, it corresponds to sampled forward-KL OPD.The method optimizes a step-balanced cross-entropy loss across its natural steps.
- 6 CONCLUSION: Results in mathematical reasoning and ALFWorld show that SOPD provides coherent, dense local supervision.
REPRODUCIBILITY STATEMENT
The paper supports reproducibility by documenting experimental configurations, providing detailed implementation settings, and releasing scripts and aggregated records for reported results.
- REPRODUCIBILITY STATEMENT: Section 5 and Appendix A specify model pairs, data, rollout and optimization protocols, decoding and evaluation settings, random seeds, and ALFWorld interaction limits.The submission source also includes figure-generation scripts and aggregated records used for the reported tables and figures.
A TRAINING AND EVALUATION DETAILS · A.1 MATHEMATICAL REASONING
The mathematical experiments train a Qwen3-4B teacher with filtered DeepMath data, then compare SOPD and OPD under matched prompt and rollout budgets. SOPD uses independently generated teacher steps aligned to retained student prefixes, with evaluation based on repeated sampled solutions and symbolic answer verification.
- A.1 MATHEMATICAL REASONING: The mathematical teacher is trained from Qwen3-4B non-thinking using 57,000 DeepMath examples of difficulty at least 6, binary final-answer rewards, and 500 optimizer updates.The teacher follows the mathematical GRPO configuration of G-OPD (Yang et al., 2026).
- A.1 MATHEMATICAL REASONING: SOPD and OPD use the same fixed prompt order, 1,024 prompts per update, one student rollout per prompt, and ten updates, totaling 10,240 prompt instances and rollouts each.Prompts are capped at 2,048 tokens and responses at 32,768 tokens; rollouts use temperature 1.0 and top-p 1.0.
- A.1 MATHEMATICAL REASONING: SOPD uses a 2 × 10^-6 learning rate and generates each teacher step greedily at temperature 0.These settings are applied within the matched ten-update training procedure.
- A.1 MATHEMATICAL REASONING: SOPD terminates natural steps at blank-line paragraph boundaries or explicit headings such as “Step” or “Solution,” using a 1,024-token safety cap otherwise.This defines how teacher-response spans are segmented during mathematical training.
- A.1 MATHEMATICAL REASONING: Each retained student prefix receives one independently generated teacher step, and SOPD averages loss first across target tokens and then across retained steps.The aggregation explicitly balances supervision within each target and across retained steps.
- A.1 MATHEMATICAL REASONING: Terminal evaluation samples 32 independent generations for each of 30 problems in AIME24, AIME25, HMMT25-Feb, and HMMT25-Nov, yielding 960 responses per benchmark.Generation uses temperature 1.0, top-p 1.0, a 32,768-token cap, and seeds 42 + 32i + j; final boxed answers are checked with Math-Verify.
A.2 ALFWORLD
The ALFWorld experiments use a Qwen2.5-3B-Instruct student paired with a 7B teacher, asynchronous exploration and training, and evaluation across seen, unseen, and hard task splits.
- Implementation: The setup pairs a Qwen2.5-3B-Instruct student with the GiGPO-Qwen2.5-7B-Instruct-ALFWorld teacher and independently advancing Explorer and Trainer counters.The Explorer samples 16 tasks with one rollout each, while the Trainer consumes 64 experiences per batch.
- Implementation: Training uses learning rate 10^-6, temperature 1.0, staleness two, 512-token responses, 10,240-token prompts, and at most 30 environment turns.These settings are applied within the released asynchronous launcher.
- Evaluation: Evaluation covers Valid Seen-140, Valid Unseen-134, and Hard-121 using temperature 0.4, a 4,096-token response cap, task-level success mean, and mean environment rounds.The evaluation therefore spans both familiar and unfamiliar tasks, including a hard subset.
B LLM USAGE
Generative AI tools supported language editing, translation, figure-layout iteration, and LaTeX preparation, while the authors retained responsibility for the research and manually verified the final manuscript and results.
- B LLM USAGE: Generative AI assisted language editing, translation, figure-layout iteration, and LaTeX preparation, while the authors determined the research claims, methods, experiments, and conclusions.The authors also manually verified the final manuscript and reported results.