SOD: Step-wise On-policy Distillation for Small Language Model Agents

2026/05/08 by Qiyong Zhong, Mao Zheng, Mingyang Song +6
Computer Science · #Explainable Artificial Intelligence (XAI) #Multimodal Machine Learning Applications #Topic Modeling #cs.AI #cs.CL

paper · pdf · doi:10.48550/arxiv.2605.07725

openalex publication_date 2026/05/08 · openalex created_date 2026/05/12 · openalex updated_date 2026/07/28 · arxiv created 2026/08/03 · arxiv updated 2026/08/04

Abstract

Tool-integrated reasoning (TIR) is difficult to scale to small language models due to instability in long-horizon tool interactions and limited model capacity. While reinforcement learning methods like group relative policy optimization provide only sparse outcome-level rewards. Recently, on-policy distillation (OPD) has gained popularity by supplying dense token-level supervision from a teacher on student-generated trajectories. However, our experiments indicate that applying OPD to TIR leads to a critical failure mode: erroneous tool calls tend to cascade across subsequent reasoning steps, progressively amplifying student-teacher divergence and rendering the teacher's token-level supervision increasingly unreliable. To address this, we propose SOD, a step-wise on-policy distillation framework for small language model agents, which adaptively reweights distillation strength at each step based on step-level divergence. Therefore, SOD can attenuate potentially misleading teacher signals in high-divergence regions while preserving dense guidance in well-aligned states. Experiments on challenging math, science, and code benchmarks show that SOD achieves up to 20.86% improvement over the second-best baseline. Notably, our 0.6B student achieves 26.13% on AIME 2025, demonstrating effective transfer of agentic reasoning to lightweight models. Our code is available at https://github.com/YoungZ365/SOD.

Citations

• SODA: Semi On-Policy Black-Box Distillation for Large Language Models
• Learning beyond Teacher: Generalized On-Policy Distillation with Reward Extrapolation
• Reinforcement Learning via Self-Distillation
• Self-Distillation Enables Continual Learning
• Black-Box On-Policy Distillation of Large Language Models
• VOLD: Reasoning Transfer from LLMs to Vision-Language Models via On-Policy Distillation
• Demystifying Reinforcement Learning in Agentic Reasoning
• SimpleTIR: End-to-End Reinforcement Learning for Multi-Turn Tool-Integrated Reasoning
• Chain-of-Agents: End-to-End Agent Foundation Models via Multi-Agent Distillation and Agentic RL
• MegaScience: Pushing the Frontiers of Post-Training Datasets for Science Reasoning
• Replacing thinking with tool usage enables reasoning in small language models
• ReasonFlux-PRM: Trajectory-Aware PRMs for Long Chain-of-Thought Reasoning in LLMs
• AgentDistill: Training-Free Agent Distillation with Generalizable MCP Boxes
• Co-Evolving LLM Coder and Unit Tester via Reinforcement Learning
• Skywork Open Reasoner 1 Technical Report
• Distilling LLM Agent into Small Models with Retrieval and Code Tools
• RLKD: Distilling LLMs' Reasoning via Reinforcement Learning
• Structured Agent Distillation for Large Language Model
• Qwen3 Technical Report
• Agentic Reasoning and Tool Integration for LLMs via Reinforcement Learning
• LeetCodeDataset: A Temporal Dataset for Robust Evaluation and Efficient Training of Code LLMs
• ToolRL: Reward is All Tool Learning Needs
• ReTool: Reinforcement Learning for Strategic Tool Use in LLMs
• ToRL: Scaling Tool-Integrated RL
• DAPO: An Open-Source LLM Reinforcement Learning System at Scale
• Search-R1: Training LLMs to Reason and Leverage Search Engines with Reinforcement Learning
• R1-Searcher: Incentivizing the Search Capability in LLMs via Reinforcement Learning
• Reinforcement Learning for Long-Horizon Interactive LLM Agents
• s1: Simple test-time scaling
• WebRL: Training LLM Web Agents via Self-Evolving Online Curriculum Reinforcement Learning
• On-Device Language Models: A Comprehensive Review
• DigiRL: Training In-The-Wild Device-Control Agents with Autonomous Reinforcement Learning
• SWE-agent: Agent-Computer Interfaces Enable Automated Software Engineering
• LiveCodeBench: Holistic and Contamination Free Evaluation of Large Language Models for Code
• ArCHer: Training Language Model Agents via Hierarchical Multi-Turn RL
• A Survey on Knowledge Distillation of Large Language Models
• BioXP-0.5B: Explainable Medical-AI via RL-GRPO
• GPQA: A Graduate-Level Google-Proof Q&A Benchmark
• FireAct: Toward Language Agent Fine-tuning
• The Rise and Potential of Large Language Model Based Agents: A Survey
• On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes
• The False Promise of Imitating Proprietary LLMs
• Toolformer: Language Models Can Teach Themselves to Use Tools
• Distributing Accountability, Not Capability: Phase Separation and the LLM Workflow Quadrant in Autonomous AI Agent Architectures
• Fine-Tuning Language Models from Human Preferences
• Proximal Policy Optimization Algorithms
• Deep reinforcement learning from human preferences
• Policy Distillation
• A Reduction of Imitation Learning and Structured Prediction to No-Regret Online Learning