On-Policy Uniform Training in RL

• On-policy uniform training is a unified framework that integrates on- and off-policy updates to enhance sample efficiency, stability, and convergence.
• It employs a perceptron-style loss with clipped surrogate objectives, ensuring reliable and monotonic policy improvement across diverse RL settings.
• The approach enforces a uniform distribution over correct solutions to maintain diversity, using iterative on-policy sampling and coverage improvement techniques.

On-policy uniform training is a theoretical and algorithmic paradigm in reinforcement learning (RL) and preference-based LLM alignment that unifies on-policy and off-policy updates, eliminates the need to distinguish between them, and promotes uniformity over desirable solution sets. This approach leverages core properties of the RL objective, coverage improvement principles in preference learning, and uniform-correct optimization for sample efficiency, convergence, and diversity across application domains such as continuous control and LLM alignment (Hu et al., 2019, Kim et al., 13 Jan 2026, Lochab et al., 1 May 2026).

1. Unified On-Off-Policy Loss and the Perceptron-Style Update

The perceptron-style loss reformulates the standard policy gradient objective to admit uniform updates, regardless of whether data is on-policy or off-policy. Consider a Markov decision process (MDP), with target policy $\pi$, behavior policy $\mu$, and advantage $A^\mu(s,a)$. The objective difference can be represented as

$$\eta(\pi)-\eta(\mu) = \sum_{s}\rho^\pi(s)\sum_{a}\mu(a|s)\left[\frac{\pi(a|s)}{\mu(a|s)}-1\right]A^\mu(s,a),$$

where updates maximize a clipped perceptron surrogate

$$L_{\rm perc}(\pi,\mu) = \mathbb{E}_{s\sim\rho^\pi, a\sim\mu}\left[\min\left\{\left(\frac{\pi(a|s)}{\mu(a|s)}-1\right)A^\mu(s,a), \xi\right\}\right].$$

Updates are only performed when $A^\mu(s,a)(\frac{\pi(a|s)}{\mu(a|s)}-1) < 0$, i.e., when the policy underperforms the behavior policy for given state-action pairs. Discarding other cases ensures each update step improves $\eta(\pi)$ locally, independent of policy proximity, and supports arbitrarily stale (off-policy) data as long as the advantage sign is correct (Hu et al., 2019).

This formulation is mathematically equivalent to the clipped surrogate objective of PPO, with margin $\xi=\epsilon |A^\mu|$, thereby allowing interchangeable use of on-policy and off-policy data with a single, unified loss.

2. Coverage Improvement in On-Policy Preference Learning

In preference learning, coverage improvement formalizes how iterated on-policy sampling enhances the statistical informativeness and sample efficiency of updates. Let $\pi_k$ be the sampling policy at iteration $k$, and $\mu$0 the target. "Coverage" is quantified as a covariance ratio

$\mu$1

with $\mu$2 for covariance matrices.

The principle is that each on-policy update, given sufficient batch size $\mu$3 exceeding a coverage threshold, moves $\mu$4 to a neighborhood of $\mu$5 with strictly better coverage. This yields a geometric contraction for the error and exponential convergence: $\mu$6 whereas static, offline updates are limited by the (much slower) minimax rate $\mu$7 (Kim et al., 13 Jan 2026).

3. Uniform Policy Optimization over Correct Sets

Standard RL with verifiable rewards (RLVR), such as GRPO, is structurally indifferent to the allocation of probability mass over the set of correct solutions $\mu$8: $\mu$9 with the condition that any policy placing all mass on $A^\mu(s,a)$0 is a global maximizer. Stochastic on-policy updates induce a "rich-get-richer" cycle, concentrating mass on a narrow subset of correct solutions and causing diversity collapse as evidenced by reduced Pass@$A^\mu(s,a)$1 and solution-level diversity (Lochab et al., 1 May 2026).

Two optimality criteria—robustness under adversarial erasure and entropy-regularized optimality—both single out the Uniform-Correct Policy: $A^\mu(s,a)$2 Optimizing for this policy ensures maximal resilience to target loss and maximizes diversity among correct responses.

The Uniform-Correct Policy Optimization (UCPO) objective augments RLVR with a conditional uniformity penalty: $A^\mu(s,a)$3 where $A^\mu(s,a)$4 is the conditional policy on $A^\mu(s,a)$5 and $A^\mu(s,a)$6 is uniform over $A^\mu(s,a)$7. This directly counteracts diversity collapse by distributing gradient mass toward underrepresented correct solutions while preserving total advantage (Lochab et al., 1 May 2026).

4. Framework and Algorithmic Structure

On-policy uniform training is characterized by the following iterative structure (Kim et al., 13 Jan 2026):

1. Iterated on-policy sampling: At each round, sample from the current (possibly mixed) policy.
2. Local loss minimization: Solve a convex surrogate (DPO, UCPO, or a reward-distillation variant) for the next policy.
3. Coverage improvement: Prove that updates move to regions with uniformly better coverage and thus greater statistical informativeness.
4. Geometric contraction: Establish exponential error decay when batch size exceeds the generalized coverage threshold.
5. Optimal design samplers (optional): For domains where coverage is initially poor, construct a preferential G-optimal mixture policy to guarantee coverage, allowing convergence in as few as two rounds.

Below is a canonical on-policy UCPO/GRPO update structure for RLVR: $$\eta(\pi)-\eta(\mu) = \sum_{s}\rho^\pi(s)\sum_{a}\mu(a|s)\left[\frac{\pi(a|s)}{\mu(a|s)}-1\right]A^\mu(s,a),$$2 (Lochab et al., 1 May 2026)

5. Practical Implications and Experimental Evidence

On-policy uniform training unifies and extends the strengths of on- and off-policy methods. Key practical consequences include:

• Sample efficiency: Off-policy reuse via replay buffers and V-trace, together with the uniform perceptron loss, enables order-of-magnitude improvements in data efficiency versus strictly on-policy algorithms (Hu et al., 2019).
• Stability and convergence: Clipped objectives and uniform-update conditions yield stability guarantees grounded in monotonic improvement and mirror descent theory.
• Diversity in RLVR/LM alignment: UCPO consistently improves Pass@$A^\mu(s,a)$8 and equation-level diversity (up to +10% absolute gain on AIME24@64, +45% higher equation-level diversity on mathematical reasoning benchmarks), with only minor variation in Pass@1, and negligible computational overhead (Lochab et al., 1 May 2026).

Experiments with on-policy DPO-based preference learning show robust, monotonic improvement, exponential contraction toward the target, and sharp separation from offline methods once past the coverage threshold (Kim et al., 13 Jan 2026). In RL control, the unified PPO+IMPALA implementation demonstrated strong performance and fast convergence for both classic control (pendulum) and real-world quadrotor tasks with minimal tuning (Hu et al., 2019).

6. Limitations and Extensions

Limitations of on-policy uniform training center on:

• Variance of importance weights: For UCPO, high values of the uniformity weight $A^\mu(s,a)$9 can lead to unstable importance sampling weights, although $$\eta(\pi)-\eta(\mu) = \sum_{s}\rho^\pi(s)\sum_{a}\mu(a|s)\left[\frac{\pi(a|s)}{\mu(a|s)}-1\right]A^\mu(s,a),$$0 remains practical (Lochab et al., 1 May 2026).
• Reward model dependence: For preference/reward distillation, the approach presumes access to a reliable reward or preference oracle. Extensions to fuzzy or non-binary rewards require further development.
• Non-verifiable settings and exploration: The core methodologies presume verifiable targets or reliable preference signals. Integration with advanced exploration to expand $$\eta(\pi)-\eta(\mu) = \sum_{s}\rho^\pi(s)\sum_{a}\mu(a|s)\left[\frac{\pi(a|s)}{\mu(a|s)}-1\right]A^\mu(s,a),$$1, or adaptive annealing of regularization weights, remains an area for further research.

A plausible implication is that mixing G-optimal design samplers with on-policy sampling can overcome poor initial coverage in settings with large or structured output spaces, ensuring rapid convergence from the outset (Kim et al., 13 Jan 2026).

7. Synthesis and Theoretical Consolidation

On-policy uniform training represents an overview of the perceptron-style unification from RL control (Hu et al., 2019), the coverage improvement principle and optimal design from preference-based learning (Kim et al., 13 Jan 2026), and conditional uniformity in RLVR for diversity (Lochab et al., 1 May 2026). The overarching framework is:

• Always update on freshly-sampled or uniformly covered policy batches.
• Use clipped, margin-augmented, or conditional-KL losses to enforce monotonic policy improvement and within-target diversity.
• Leverage geometric contraction guarantees and optimal-design constructions where coverage bottlenecks would otherwise slow convergence.

These principles yield a unified, theoretically-grounded roadmap for efficient, stable, and diversity-preserving policy optimization across distinct RL and alignment settings.