A Good Self-Teacher Meets the Student Where They Are: Joint On-Policy Learning and Teaching
cs.LG, cs.AI
Submitted: 2026-10-07
Updated: 2026-10-07
Code: https://github.com/hamishivi/tmax
License: http://creativecommons.org/licenses/by/4.0/
The gist: Reinforcement Learning (RL) from outcome rewards suffers from sparse supervision, particularly on difficult, long-horizon tasks where successful trajectories are rare and costly to generate.
Abstract
Reinforcement Learning (RL) from outcome rewards suffers from sparse supervision, particularly on difficult, long-horizon tasks where successful trajectories are rare and costly to generate. On-Policy Distillation (OPD) offers an attractive alternative by providing dense token-level supervision from a stronger teacher along the student's own generations. Self-distillation methods further remove the need for a separate teacher model by conditioning the same policy on privileged information to serve as its own teacher. However, privileged conditioning alone does not guarantee that the resulting distillation update improves the student. Indeed, privileged information can lead the teacher to solve tasks through shortcuts unavailable to the student, producing supervision poorly matched to the student's current behavior. Consequently, even a higher-performing teacher can provide guidance that degrades student performance. To address this, we analyze how the choice of privileged teacher affects the student's update. We derive a necessary and sufficient condition for the teacher's local distillation update to be a positive multiple of the student's reward gradient. Our analysis suggests that the teacher should not only perform well on the task, but also provide guidance suited to the student's current capabilities. This characterization motivates a practical teacher-training surrogate that combines outcome rewards with token-level Kullback-Leibler (KL) regularization toward the student. Based on this result, we propose Joint On-Policy Learning and Teaching (JOLT), which jointly trains a single policy in two roles: a privileged teacher using a KL-regularized objective, and an unprivileged student using dense on-policy distillation. Across mathematical reasoning, coding, tool use, and terminal use, JOLT improves training efficiency and performance, with further gains from student rewards.
Sources
- Maximum a Posteriori Policy Optimisation
- Deep Reinforcement Learning at the Edge of the Statistical Precipice
- Reinforcement Learning from Rich Feedback with Distributional DAgger
- Unmasking On-Policy Distillation: Where It Helps, Where It Hurts, and Why
- Reinforcement Learning Teachers of Test Time Scaling
- Evaluating Large Language Models Trained on Code
- Training Verifiers to Solve Math Word Problems
- DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model
- DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning
- HDPO: Hybrid Distillation Policy Optimization via Privileged Self-Distillation
- Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes
- Measuring Mathematical Problem Solving With the MATH Dataset
- Distilling the Knowledge in a Neural Network
- Reinforcement Learning via Self-Distillation
- LiveCodeBench: Holistic and Contamination Free Evaluation of Large Language Models for Code
- Entropy-Aware On-Policy Distillation of Language Models
- Rethinking On-Policy Self-Distillation for Thinking Models
- The Art of Scaling Reinforcement Learning Compute for LLMs
- Why Does Self-Distillation (Sometimes) Degrade the Reasoning Capability of LLMs?
- Long-Horizon-Terminal-Bench: Testing the Limits of Agents on Long-Horizon Terminal Tasks with Dense Reward-Based Grading
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks