Unified Policy Gradient Estimator (UPGE)

• UPGE is a unified framework that decomposes policy gradient updates into stabilization, reference, advantage, and likelihood components for balanced optimization.
• It generalizes techniques across reinforcement learning, supervised fine-tuning, and hybrid objectives, enabling task-specific adaptations.
• Empirical results show that UPGE achieves improved stability and performance metrics through adaptive loss allocation in large model training.

A Unified Policy Gradient Estimator (UPGE) is a mathematical and algorithmic framework that generalizes policy gradient estimation across reinforcement learning (RL), supervised fine-tuning (SFT), and hybrid post-training paradigms by decomposing the gradient update into four modular components. This formulation enables a rigorous synthesis of optimization techniques for both online reinforcement learning using model-generated rollouts and offline learning from demonstration data in LLMs and other decision-making systems (Lv et al., 4 Sep 2025). The UPGE paradigm abstracts widely adopted optimization methods—such as supervised fine-tuning, trust-region RL, and hybrid algorithms—into a common estimator that can be tailored by selecting appropriate components, providing both a theoretical foundation and a practical route to more effective and stable post-training.

1. Mathematical Formulation and Decomposition

UPGE expresses the gradient update for post-training LLMs or RL agents as

$$\mathbb{E}_{\tau \sim \pi_{\mathrm{ref}}} \left[ \mathbb{1}_{\mathrm{stable}}(\tau, q) \cdot \frac{1}{\pi_{\mathrm{ref}}(\tau|q)} \cdot \widehat{A}_{\mathrm{uni}}(\tau, q) \cdot \nabla_{\theta} \pi_\theta(\tau|q) \right]$$

where:

• $\mathbb{1}_{\mathrm{stable}}(\tau, q)$: stabilization mask, indicating stability (e.g., trust-region or clipping).
• $\pi_{\mathrm{ref}}$: reference policy denominator, sets the data/reweighting regime.
• $\widehat{A}_{\mathrm{uni}}(\tau, q)$: unified advantage estimator, encoding reward or demonstration adherence.
• $\nabla_{\theta} \pi_\theta(\tau|q)$: likelihood gradient.

Each part is interchangeable, allowing algorithm designers to specify the estimator according to task requirements, data sources, and stability constraints. The estimator recovers classical supervised (SFT) or RL approaches via specializations of these choices: for example, SFT reduces to cross-entropy with $\pi_{\mathrm{ref}} = \pi_{\theta}$ and a demonstration-based advantage, while RL with PPO-style trust regions uses on-policy rollouts with clipping ($\mathbb{1}_{\mathrm{stable}}$) and relative advantage estimates (Lv et al., 4 Sep 2025).

Component Table

Component	Typical Choices / Role	Examples
Stabilization mask ($\mathbb{1}_{\mathrm{stable}}$)	Clipping/trust region	PPO/TRPO
Reference denominator ($\pi_{\mathrm{ref}}$)	SFT: $\pi_{\theta}$, RL: $\pi_{\theta,\text{old}}$, Offline: 1	SFT, PPO, offline RL
Advantage ($\widehat{A}_{\mathrm{uni}}$)	RL reward, normalized reward, SFT adherence	Gumbel RL, SFT, mixed
Likelihood gradient ($\nabla_{\theta} \pi_\theta$)	Gradient of policy w.r.t parameters	All

This decomposition makes explicit which aspects of the gradient update are responsible for stability (mask), proper weighting (reference policy), learning signal (advantage), and parameter updates (likelihood gradient).

2. Instantiations: SFT, RL, and Hybrid Objectives

The UPGE encompasses diverse post-training objectives under a unified notation:

• Supervised Fine-Tuning (SFT): Training on demonstration data, with $\pi_{\mathrm{ref}} = \pi_\theta$, advantage set to a demonstration-adherence term, and no stabilization mask; reduces to minimizing cross-entropy.
• Online RL (e.g., PPO/GRPO): Training with model rollouts, $\pi_{\mathrm{ref}} = \pi_{\theta,\text{old}}$ (importance weights), $\widehat{A}$ is reward-based, and $\mathbb{1}_{\mathrm{stable}}$ masks updates violating trust region constraints.
• Offline RL: Training on demonstration/other-model rollouts with $\pi_{\mathrm{ref}} = 1$.
• Mixed/Heteroscedastic Objectives: UPGE enables simultaneous optimization over mixtures of reference policies and advantages, supporting hybrid algorithms.

For example, the advantage estimate is unified as

$$\widehat{A}_{\mathrm{uni}}(\tau, q) = r(\tau|q) + \mu \cdot \mathbb{1}\{\pi_{\mathrm{ref}} = \pi_\beta\} \frac{\pi_\beta(\tau|q)}{\pi_\theta(\tau|q)}$$

combining reward maximization and direct supervision signals.

3. Theoretical Rationale and Stability Properties

The four-component UPGE architecture is theoretically motivated by the dual need for sample efficiency, stability, and bias-variance control:

• Stabilization mask (e.g., as in PPO clipping) prevents destructive updates by omitting gradients where the current policy diverges excessively from the reference, directly inspired by established trust region policy optimization mechanics.
• Reference denominator controls the estimator's variance and importance-sampling properties, making it possible to train safely on both on-policy and off-policy data.
• Advantage estimator generalizes reward assignment, incorporating SFT adherence for effective demonstration exploitation or group-normalized rewards for variance reduction and equitable credit assignment.
• Likelihood gradient ensures that optimization is performed with respect to actual model parameters, preserving correctness in the presence of off-policy or mixed data.

This architecture makes it possible to interpolate between exploration (favoring RL reward-driven updates) and exploitation (favoring demonstration-based SFT), facilitating robust and balanced learning in diverse environments.

4. Hybrid Post-Training: Algorithmic Realization

The Hybrid Post-Training (HPT) algorithm operationalizes UPGE by dynamically composing RL and SFT losses:

$$\mathcal{L} = \alpha \cdot \mathcal{L}_{\mathrm{RL}} + \beta \cdot \mathcal{L}_{\mathrm{SFT}}$$

where coefficients $\alpha, \beta$ are selected on a per-instance basis in response to model competence, as measured by verifier scores on sampled rollouts (Lv et al., 4 Sep 2025).

• If model performance $P > \gamma$ on a given input, the update uses only RL signals (exploration and refinement).
• If $P \leq \gamma$, SFT dominates to "pull" the model toward correct demonstration-based reasoning.
• This logic enables HPT to allocate training signal adaptively, preserving existing reasoning patterns while incrementally exploiting new reward signals or demonstration knowledge.

Empirical results across six math reasoning and out-of-distribution benchmarks demonstrate that HPT, through UPGE, achieves higher accuracy, better Pass@$k$ performance, and more stable training trajectories in LLMs of various scales.

5. Empirical Evidence and Benchmarking

In extensive evaluation, HPT consistently surpasses baselines:

• Outperforms pure SFT, RL (GRPO), sequential SFT$\to$RL, and previously proposed hybrid algorithms such as LUFFY and SRFT.
• On Qwen2.5-Math-7B and other models, attains $7+$ point gains over the best baseline for Pass@1 and large-$k$ metrics across both in-distribution and out-of-distribution settings.
• Dynamic allocation of SFT and RL results in smooth training curves, higher sequence entropy, and preservation of answer lengths—demonstrating balanced exploration and exploitation.

These findings substantiate the claim that the modular UPGE structure supports both effective demonstration exploitation and stable reward-driven exploration in LLM post-training.

6. Implications, Extensions, and Limitations

The UPGE framework supplies a rigorous umbrella for modern post-training algorithms. Notably:

• Any training procedure that can be cast by selecting components of the UPGE template (stabilization mask, denominator, advantage, likelihood gradient) is immediately compatible.
• The architecture permits future extension to additional variance-reduced, off-policy-corrected, risk-sensitive, or meta-learning estimators by appropriate component selection.
• A plausible implication is that the UPGE approach can be generalized beyond LLMs to other domains where multiple training data sources and stability-bias-variance tradeoffs must be balanced within a single optimization routine.

Potential caveats include the need for principled selection of stabilization thresholds, proper normalization of reference policies and advantage estimators, and the integration of UPGE with highly heterogeneous reward functions or extremely long-horizon dependencies.

7. Relevant Formulas and Schematic

Key expressions central to the UPGE are:

• Unified policy gradient estimator:

$$\nabla_\theta \mathcal{J}_\mu(\theta) = \mathbb{E}_{\tau \sim \pi_{\mathrm{ref}}(\cdot|q)} \left[ \mathbb{1}_{\mathrm{stable}}(\tau, q) \cdot \frac{1}{\pi_{\mathrm{ref}}(\tau|q)} \cdot \widehat{A}_{\mathrm{uni}}(\tau, q) \cdot \nabla_\theta \pi_{\theta}(\tau|q) \right]$$

• Unified advantage estimator:

$$\widehat{A}_{\mathrm{uni}}(\tau, q) = r(\tau|q) + \mu \cdot \mathbb{1}\{\pi_{\mathrm{ref}} = \pi_\beta\} \cdot \frac{\pi_\beta(\tau|q)}{\pi_\theta(\tau|q)}$$

• Hybrid loss:

$$\mathcal{L} = \alpha \cdot \mathcal{L}_{\mathrm{RL}} + \beta \cdot \mathcal{L}_{\mathrm{SFT}}$$

Diagrams such as Figure 1 (“Illustration of the Unified Policy Gradient Estimator”) (Lv et al., 4 Sep 2025) visualize the data flow: selection of data source, determination of reference policy, computation of unified advantage, application of stabilization, and backpropagation via the likelihood gradient.

---

In conclusion, the Unified Policy Gradient Estimator supplies a powerful, theoretically grounded abstraction that reconciles online RL, SFT, and mixed training regimes. By making explicit the role of stabilization, reference policy selection, advantage computation, and likelihood-respecting parameter updates, it provides a modular framework for stable, efficient, and generalizable post-training optimization in LLMs and related RL systems.