On-Policy Distillation Frameworks

• On-policy distillation frameworks are techniques that transfer knowledge from a teacher to a student by using data sampled from the student’s own rollouts.
• They mitigate distribution mismatch and improve sample efficiency and stability across reinforcement learning, generative modeling, and language tasks.
• Key design choices include adaptive loss functions, alternating optimization, and self-distillation methods that enable robust policy compression and faster inference.

On-policy distillation frameworks are a class of techniques in machine learning and reinforcement learning for transferring knowledge from a teacher model (or ensemble of models, or privileged source) into a student policy, where the key distinction is that the student is updated with respect to its own (on-policy) distribution of trajectories or outputs rather than a fixed external dataset or offline expert demonstrations. These frameworks address problems of distribution mismatch (exposure bias) between training and inference, provide dense feedback aligned with model usage, and enable efficient model compression, policy acceleration, and even collaborative exploration or self-improvement. On-policy distillation is central in settings ranging from deep reinforcement learning (DRL), sequence modeling, diffusion generative modeling, and LLMs for reasoning and control, especially where sample efficiency, stability, and empirical fidelity to teacher behaviors are paramount.

1. Mathematical Foundations and Loss Constructions

On-policy distillation frameworks unify methodologies where the distillation loss is evaluated with respect to the student’s own trajectory distribution. Consider a family of objectives where the student policy $\pi_S$ (parameterized by $\theta$) and teacher policy $\pi_T$ (possibly fixed, possibly context-shifted or self-conditioned) yield distributions over action sequences or tokens:

$$L_{\mathrm{on}}(\theta) = \mathbb{E}_{x \sim D} \, \mathbb{E}_{y \sim \pi_S(\cdot|x)} \left[ \sum_{t=1}^{|y|} D_{\mathrm{KL}}(\pi_T(\cdot|x,y_{<t}) \,\|\, \pi_S(\cdot|x,y_{<t})) \right]$$

This contrasts with off-policy distillation, where the expectation is taken instead over a static dataset or teacher-generated traces $y^*$, leading to potential distribution mismatch between training and inference. The on-policy formulation can operate in forward-KL, reverse-KL (mode-seeking), or with f-divergences or adaptive bridges (as in Veto (Jang et al., 12 Jan 2026)) to tune gradient behavior and diversity properties. In generalized forms (G-OPD (Yang et al., 12 Feb 2026)), the expectation incorporates a reward-scaling factor and arbitrary reference distributions, reducing standard OPD to a special case of KL-constrained RL.

In diffusion or flow-based modeling, the matching loss can operate at the level of score functions along the student trajectory (e.g., policy-based flow matching in π-Flow (Chen et al., 16 Oct 2025) and One-Step Diffusion Policy (Wang et al., 2024)), and in multimodal generation, additional critic or auxiliary score networks are optimized via on-policy samples under the student’s rollout distribution (Chern et al., 29 Dec 2025).

2. Algorithmic Structures and Training Procedures

On-policy distillation frameworks generally consist of iterative procedures tightly coupling sampling, teacher feedback, and student updates. Key algorithmic patterns include:

• Alternating Optimization: Student parameters are updated by minimizing the on-policy loss, frequently with additional critic or auxiliary networks for stability, as in LiveTalk (Chern et al., 29 Dec 2025) and π-Flow (Chen et al., 16 Oct 2025).
• Two-Stage Approaches: Many frameworks use an initial supervised or off-policy phase to reduce distributional gap (e.g., SFT in VOLD (Bousselham et al., 27 Oct 2025)), followed by an on-policy distillation or RL phase for fine alignment.
• Self-Distillation: In frameworks like OPSD (Zhao et al., 26 Jan 2026) and OPSDC (Sang et al., 5 Mar 2026), models act as both teacher and student by conditioning on privileged context (e.g., answers, concise instructions), performing distillation over their own rollouts.
• Peer/Ensemble Distillation: In DPD (Lai et al., 2020) and online distillation with Decision-Attention (Yu et al., 2024), multiple learners co-evolve and exchange knowledge, using group-derived soft targets.
• Adaptive Objectives: Beta-scheduled bridging (Veto (Jang et al., 12 Jan 2026)) or reward scaling (ExOPD in G-OPD (Yang et al., 12 Feb 2026)) can smoothly interpolate between pure imitation and RL-style exploration.

General training pseudocode includes on-policy sampling, (optional) teacher or auxiliary scoring, per-token or per-state loss computation, and gradient descent steps, sometimes with batch normalization, exponential moving average parameter updates, or KL-masking to ensure stability or support controlled exploration.

3. Principal Applications and Empirical Impact

Reinforcement Learning

On-policy distillation enhances the sample efficiency, robustness, and asymptotic performance of deep RL agents. In "Dual Policy Distillation" (Lai et al., 2020), peer-to-peer on-policy distillation allows mutual improvement, leading to 4–58% gains over standard PPO on continuous control. Decision-Attention (Yu et al., 2024) further extends on-policy distillation to ensemble settings, with group-weighted soft targets providing a +50% boost in Atari PPO returns.

Expected entropy-regularized on-policy distillation provides favorable convergence and stability behavior compared to naïve cross-entropy or entropy-only approaches, especially in large discrete or stochastic action spaces (Czarnecki et al., 2019).

Generative Diffusion and Few-Step Policies

In high-dimensional generative modeling, on-policy trajectory or score matching enables sample-efficient distillation of slow, high-quality diffusion models into fast, single- or few-step generators. In “One-Step Diffusion Policy” (Wang et al., 2024), the distilled student achieves a 41-fold improvement in action rate in real-world visuomotor tasks, matching or surpassing teacher performance. “π-Flow” (Chen et al., 16 Oct 2025) introduces closed-form network-free policy outputs, enabling on-policy ODE integration with no inference overhead, leading to state-of-the-art FID and diversity in fast ImageNet generation.

“LiveTalk” (Chern et al., 29 Dec 2025) demonstrates that improved on-policy distillation, with multimodal condition regularization and critic warm-up, enables real-time avatar video diffusion with 20× speedup and sub-second response latency, while eliminating multimodal artifacts and preserving frame consistency.

LLM Distillation and Reasoning Compression

On-policy distillation frameworks have enabled several advances in LLM training. On-policy self-distillation (OPSD (Zhao et al., 26 Jan 2026), OPSDC (Sang et al., 5 Mar 2026)) compresses privileged information or concise reasoning directly into model weights via reverse-KL losses over model-generated tokens, yielding simultaneous gains in accuracy and efficiency. Context distillation (OPCD (Ye et al., 12 Feb 2026)) allows the transfer of experiential knowledge or system prompt behaviors into a student policy, outperforming off-policy approaches on both in-distribution and OOD tasks.

On-policy prefix distillation (PPD (Zhang et al., 16 Feb 2026)) exploits the observation that dense teacher feedback is concentrated in reasoning prefixes, enabling 2–47× reductions in FLOP with negligible performance loss.

Black-box and Memory-constrained Settings

Generative Adversarial Distillation (GAD (Ye et al., 13 Nov 2025)) enables black-box, on-policy distillation by learning a reward model (discriminator) in parallel, removing the need for token-level teacher supervision and providing adaptive, policy-aligned guidance. On-policy verbal distillation (OVD (Xiong et al., 29 Jan 2026)) further generalizes this concept to replace full-probability teacher feedback with trajectory-level discrete scores, reducing memory cost by orders of magnitude and allowing free policy exploration.

Safety, Modularization, and Adaptive Control

On-policy distillation with adaptive target reformulation (Veto (Jang et al., 12 Jan 2026)) introduces geometric bridges in logit space to stabilize KL objectives. Constrained RL extensions, as in “Training Wheels” (Liu et al., 24 Feb 2025), permit the student model to selectively invoke the teacher at test time, optimizing the trade-off between quality and compute budget.

4. Comparative Analysis: On-Policy vs. Off-Policy and RL

The primary advantage of on-policy distillation is its ability to align student behavior with its own operating distribution, mitigating the exposure bias and error compounding prevalent in off-policy knowledge distillation and supervised fine-tuning. Unlike pure RL, on-policy distillation can deliver dense, low-variance feedback at the token or timestep level, supporting stable mode-seeking (reverse-KL) updates that target high-probability teacher outputs encountered in practice (Zhao et al., 26 Jan 2026, Sang et al., 5 Mar 2026).

Empirical studies consistently show that on-policy distillation outperforms off-policy variants in generalization, sample efficiency, and OOD robustness (Zhang et al., 16 Feb 2026, Ye et al., 12 Feb 2026, Bousselham et al., 27 Oct 2025). Token efficiency is also dramatically improved—OPSD (Zhao et al., 26 Jan 2026) matches or exceeds RL-based methods with 4–8× fewer rollout tokens. Prefix-based on-policy distillation allows for further drastic reductions in computational cost by focusing distillation signal where it matters most (Zhang et al., 16 Feb 2026).

Frameworks such as G-OPD (Yang et al., 12 Feb 2026) show that reward scaling can allow students to outperform even their teacher policies by extrapolating the teacher’s preferences, and modular approaches (peer-to-peer, ensemble, constrained RL) enable further enhancements not available to pure supervised or fixed-policy RL methods.

5. Design Choices, Extensions, and Limitations

Design Choices

• Objective Selection: Forward-KL, reverse-KL, or adaptive bridges can be chosen based on desired trade-offs between diversity, stability, and mode-seeking behavior. Veto (Jang et al., 12 Jan 2026) shows that linear $\beta$-decay achieves optimal stability.
• Reference Model: G-OPD (Yang et al., 12 Feb 2026) allows arbitrary reference models (including the teacher’s pre-RL base), providing stricter or more relaxed reward signals as desired.
• Policy Initialization: Two-stage pipelines (SFT/cold-start + on-policy) are critical when student and teacher distributions have high initial divergence (Bousselham et al., 27 Oct 2025, Chern et al., 29 Dec 2025).
• Auxiliary Signals: Critic warm-up (20:1 ratio, (Chern et al., 29 Dec 2025)), context scheduling, or KL-masking can stabilize training, particularly when transitioning from pre-training to on-policy fine-tuning or bridging stark supervision gaps.
• Self-Distillation Frequency: For on-policy self-distillation, teacher weight copy interval $M$ (e.g., 50 steps, (Sang et al., 5 Mar 2026)) must be tuned to avoid moving target collapse or insufficient compression.

Limitations

• Compute and Memory: Full-probability or token-level alignment incurs substantial memory costs; approach varies by framework (OVD (Xiong et al., 29 Jan 2026) addresses this for RL).
• Model Capacity: Self-distillation and context assimilation require that the student is sufficiently capable to leverage privileged teacher signals (Zhao et al., 26 Jan 2026).
• Teacher Design: Reward extrapolation requires careful tuning; in strong-to-weak settings, reward correction depends on access to the teacher’s pre-RL variant (Yang et al., 12 Feb 2026).
• Stability: Poor teacher/student alignment can lead to instability, as can unmoderated forward-KL (gradient explosion) or reverse-KL (diversity collapse) (Jang et al., 12 Jan 2026).
• Domain Generalization: Most results are on reasoning, sequence modeling, or specific generative domains; transferability to other modalities or interactive tasks relies on problem-specific adaptations.

Extensions

Potential extensions noted include broader peer graph architectures, curriculum learning, hybrid on/off-policy replay, multi-modal and cross-embodiment transfer, and meta-learning confidence weighting for source selection (Lai et al., 2020, Xiong et al., 29 Jan 2026, Chen et al., 16 Oct 2025). Black-box and verbal feedback approaches are active areas for memory/resource-limited or proprietary teacher settings.

6. Empirical Benchmarks and Performance Table

Framework	Main Application	Key Result/Metric
LiveTalk (Chern et al., 29 Dec 2025)	Real-time multimodal video diffusion	20× speedup, artifact-free multimodal avatar video
DPD (Lai et al., 2020)	RL continuous control (PPO peer distill.)	+4–58% best-of-run return, robust convergence
OneDP (Wang et al., 2024)	Robotic diffusion policy (imitation learning)	41× faster, matches/exceeds teacher success rate
OPSD (Zhao et al., 26 Jan 2026)	LLM math reasoning self-distillation	4–8× token efficiency over RL, best test accuracy
OVD (Xiong et al., 29 Jan 2026)	RL with trajectory score feedback	+25.7pp math, +12.9pp Q&A EM, 48,000× memory savings
π-Flow (Chen et al., 16 Oct 2025)	Image & text-generator ODE policy	SOTA FID/diversity, network-free policy, single evaluation
GAD (Ye et al., 13 Nov 2025)	Black-box LLM distillation	+1–2 GPT-4o points over SeqKD, robust OOD performance
OPCD (Ye et al., 12 Feb 2026)	In-context knowledge distillation (LLM)	+4 points OOD safety; eliminates catastrophic forgetting
OPSDC (Sang et al., 5 Mar 2026)	Reasoning compression (LLM self-distill.)	41–59% token reduction; +9–16 absolute accuracy
PPD (Zhang et al., 16 Feb 2026)	Prefix-only reasoning distillation	2–47× FLOP reduction, matched full-OPD accuracy
Veto (Jang et al., 12 Jan 2026)	Adaptive distillation stabilization	+4.8% GSM8K, +6% HumanEval P@10, superior stability

7. Summary and Outlook

On-policy distillation frameworks have matured into a central tool for stable, efficient, and generalizable knowledge transfer across domains. By correcting the train-test distribution mismatch and leveraging dense, on-policy feedback, these methods deliver supersets of capabilities previously exclusive to RL or supervised imitation. Extensions to resource-constrained, multimodal, and memory-limited settings (e.g., black-box, verbal, or self-distillation) broaden their impact. Continuing research will likely emphasize scalable curriculum design, hybrid control/distillation objectives, robust peer-to-peer transfer, and principled handling of modality gaps, aligning distilled policies and generative models with real-world demands for speed, fidelity, and adaptability.