Multi-domain On-Policy Distillation

• MOPD is a technique that trains a student model using its own on-policy data while distilling expertise from multiple domain-specific teacher models.
• The method employs domain-based sampling, importance weighting, and tailored loss functions to merge outputs from different expert policies.
• Empirical results show that MOPD improves sample efficiency and mitigates negative interference, leading to faster convergence and robust multi-task performance.

Multi-domain On-Policy Distillation (MOPD) refers to a family of techniques that leverage on-policy data to distill the behaviors of multiple specialized teacher models ("domain experts") into a single student model capable of robust multi-domain performance. These methods are situated at the intersection of reinforcement learning, policy optimization, and large-scale supervised learning, with applications ranging from control and reasoning to generative modeling and multimodal systems. The defining principle is that student updates are computed on trajectories generated by the student itself (on-policy), while supervision is supplied by frozen teacher models sampled across distinct domains, in contrast to earlier off-policy or joint training approaches.

1. Foundational Principles and Mathematical Formulation

MOPD generalizes standard policy distillation to settings with multiple task- or domain-specific experts, and formalizes distillation as an on-policy RL objective. In the general case, consider $M$ domains with data distributions $\{D_i\}_{i=1}^M$ and teachers $\{\pi^*_i\}$. The core objective averaged over domains is:

$$\mathcal{J}_{\rm MOPD}(\theta) = \frac{1}{M} \sum_{i=1}^M \mathbb{E}_{x \sim D_i,\, y \sim \pi_\theta(\cdot|x)} \left[ \beta \log \frac{\pi^*_i(y|x)}{\pi_{\rm ref}(y|x)} - \mathrm{KL}\left(\pi_\theta(\cdot|x)\|\,\pi_{\rm ref}(\cdot|x)\right) \right]$$

Here, $\beta$ is a reward-scaling factor; $\pi_{\rm ref}$ serves as a reference policy (most often the student initialization). The student $\pi_\theta$ is trained by sampling from its own policy (i.e., on-policy rollouts), while teacher logit distributions provide dense token- or action-level supervision for each trajectory and domain (Yang et al., 12 Feb 2026).

In the continuous-action regime, as in motion control, the corresponding objective often uses mean-squared error between mean actions or full Gaussian KL, sampled from on-policy student rollouts and labeled by the appropriate domain expert (Berseth et al., 2018).

Algorithmic implementations introduce importance sampling, truncated importance weights for stability, and domain-based sampling to cover all domain distributions in a balanced or weighted manner (Yang et al., 19 Mar 2026).

2. Domain Specialization, Teacher Selection, and Expert Merging

A central theme in MOPD is the consolidation of multiple expert policies, each independently trained on a distinct task or environment, into one robust student model. Teacher policies are frozen at their strongest checkpoint—either via early stopping, best-validation-score selection, or specialized RL pipelines (e.g., math reasoning, RL-from-human-feedback) (Yang et al., 19 Mar 2026). In the Nemotron-Cascade 2 pipeline, teachers include SFT-only, RL-from-verifiable-rewards (RLVR), and RLHF domain experts, each supplying reference outputs for distillation on their respective prompt sets.

Teachers may also be derived from separate RL optimization trajectories towards their respective reward structures, followed by merging using the above expectation over domains. The multi-domain context often requires distinct output heads or conditioning on domain identifiers to accommodate different action/state/reward spaces (Rusu et al., 2015, Chern et al., 29 Dec 2025).

Expert selection and orchestration strategies involve:

• Uniform or data-proportional sampling of domains during training.
• Context-sensitive replay buffer management and labeling to ensure coverage of all domains.
• Cyclic or curriculum-based domain traversal to favor smooth knowledge transfer and monotonic improvement (e.g., cyclic distillation for sim-to-real control) (2207.14561).

3. On-Policy Data Generation, Supervision, and Optimization

MOPD strictly leverages student-generated trajectories to address train-test distribution mismatch and mitigate exposure-bias—student states can diverge rapidly from the teacher's off-policy trajectories in high-dimensional or multi-domain regimes. Supervision is thus always provided on data encountered by the student policy itself.

Key procedure elements:

• For each batch, a domain is sampled, a prompt/state is drawn from the corresponding distribution, and a student trajectory is rolled out.
• The teacher policy for this domain supplies per-token or per-action logit targets for the visited states, defining the loss.
• Token-level importance corrections (student train/inference ratio) are applied to stabilize gradients (Yang et al., 19 Mar 2026).
• Reward extrapolation (scaling $\beta>1$) may be used to push students beyond teacher boundaries when teacher improvements over baseline are reliably positive (Yang et al., 12 Feb 2026).
• Sampling and weighting schemes vary by domain size, trajectory length, or application requirements.

Tables summarizing variants in MOPD instantiations:

Domain	Teacher Type	Data Source
Math	SFT / RL expert	RL math blend, AceReason-Math
RLVR	RL expert	Nano-v3 RL blend
RLHF	RLHF expert	HelpSteer3, safety-blend

(Yang et al., 19 Mar 2026)

4. Algorithmic Implementations and Practical Considerations

The implementation follows a modular structure, with separate modules for:

• Sampling/rollout: Two views of student policy for inference and gradient computations ensure correct on-policy ratios.
• Teacher lookup: Domain-specific teacher policies are indexed by the sampled domain.
• On-policy loss computation: Reverse-KL or mean-squared error computed per-token/action across the trajectory.
• Stabilization mechanisms: Importance weight truncation $w_t = 1$ if $r_t \in [E_{\min}, E_{\max}]$, else 0.
• Optimization: Typically AdamW, carefully tuned learning rates, and batch sizes sufficient to maintain trajectory diversity.

Pseudocode examples for these procedures explicitly appear in (Yang et al., 19 Mar 2026) and (Yang et al., 12 Feb 2026). For instance, the MOPD step in Nemotron-Cascade 2 (Python-like) rolls out student trajectories per domain, computes teacher-student logit differences, applies importance weights, and accumulates token-level loss (Yang et al., 19 Mar 2026).

Empirically, convergence in large-scale LLMs or multimodal models is reported within 30–50 optimization steps, with dense supervision delivering superior wall-clock and sample efficiency compared to sparse RL rewards or outcome-based optimization (Yang et al., 19 Mar 2026, Yang et al., 12 Feb 2026).

5. Theoretical and Empirical Benefits

MOPD avoids negative interference and catastrophic forgetting typical in multi-task RL and off-policy joint training. Explicit results include:

• Rapid recovery of capabilities lost due to reward-overfitting or specialized RL in prior pipeline stages (Yang et al., 19 Mar 2026).
• Surpassing single-domain experts in both math and code reasoning by reward extrapolation (e.g., $+1.4\%$ code, $+1.3\%$ math for $\beta=1.25$ over OPD/teacher baselines (Yang et al., 12 Feb 2026)).
• Sample efficiency: MOPD reaches peak reward or accuracy 3–6 times faster than leading RL or off-policy methods in control, reasoning, and sim-to-real transfer (2207.14561, Yang et al., 19 Mar 2026).
• Robustness across multiple domains, model scales, and application settings, confirmed by ablations on curriculum, entropy filtering, mixture reward clipping, modality scaling, etc. (Chern et al., 29 Dec 2025, Ko et al., 11 Mar 2026).

Specific applications extend MOPD to real-time video diffusion (Chern et al., 29 Dec 2025) and fine-grained agentic tool use (Ko et al., 11 Mar 2026).

6. Extensions and Stabilization: Relaxed and Cyclic Variants

Instabilities in standard on-policy distillation (e.g., heavy-tailed likelihood ratios, negative transfer) have motivated the development of extensions:

• REOPOLD introduces mixture-based reward clipping, entropy-guided token-level filtering, and a curriculum from exploration to refinement stages, stabilizing multi-domain training and further boosting sample efficiency (achieving $6.7\times$–$12\times$ improvement in sample usage over recent RL methods) (Ko et al., 11 Mar 2026).
• Cyclic policy distillation (CPD) organizes local policies in sub-domains (e.g., physical parameter ranges), mixes neighbor policies with monotonic improvement constraints, and finally distills into a global policy (2207.14561). This approach accelerates and stabilizes domain randomization and sim-to-real transfer without negative transfer at domain boundaries.

Empirical ablations confirm the necessity of each regularization or scheduling mechanism, and suggest that domain-specific tuning (curation, learning-rate, guidance strengths per modality) can have disproportionately large impact on both quality and stability (Chern et al., 29 Dec 2025).

7. Architectural and Application Scope

MOPD supports diverse architectural paradigms:

• Discrete-action RL (multi-head DQN, controller architectures) (Rusu et al., 2015)
• Continuous control with mean/action-value network distillation (Berseth et al., 2018)
• LLMs and MoE transformers (with per-domain heads or task-context conditioning) (Yang et al., 19 Mar 2026)
• Autoregressive, block-wise causal diffusion for multimodal video generation (Chern et al., 29 Dec 2025)
• Hybrid architectures where teacher and student networks differ in both size and modality coverage

Domains addressed include Atari and Mujoco games, math/code/creative tasks, bipedal and dexterous motion, sim-to-real robotics, multimodal video synthesis, and agentic tool-use (Rusu et al., 2015, 2207.14561, Yang et al., 19 Mar 2026, Chern et al., 29 Dec 2025, Ko et al., 11 Mar 2026).

• "Policy Distillation" (Rusu et al., 2015)
• "Cyclic Policy Distillation: Sample-Efficient Sim-to-Real Reinforcement Learning with Domain Randomization" (2207.14561)
• "LiveTalk: Real-Time Multimodal Interactive Video Diffusion via Improved On-Policy Distillation" (Chern et al., 29 Dec 2025)
• "Learning beyond Teacher: Generalized On-Policy Distillation with Reward Extrapolation" (Yang et al., 12 Feb 2026)
• "Nemotron-Cascade 2: Post-Training LLMs with Cascade RL and Multi-Domain On-Policy Distillation" (Yang et al., 19 Mar 2026)
• "Progressive Reinforcement Learning with Distillation for Multi-Skilled Motion Control" (Berseth et al., 2018)
• "Scaling Reasoning Efficiently via Relaxed On-Policy Distillation" (Ko et al., 11 Mar 2026)