Beyond the Stability-Exploration Dilemma: Environmental Regularization for LLM Policy Optimization

AI summary

Overview

Research area: Natural Language Processing — reinforcement learning post-training of large language models (specifically Reinforcement Learning with Verifiable Rewards, RLVR), with a focus on training stability and exploration.

Technical level: Advanced. The paper assumes familiarity with policy gradient methods, KL regularization, GRPO/PPO/REINFORCE estimators, and reward-verifiable training pipelines.

Scope: The paper proposes Environment-Regularized Policy Optimization (ERPO), which replaces the standard action-side Policy-KL regularizer with a Query-KL (QKL) term that bounds drift in the policy-induced query distribution, plus a reference-derived per-query reweighting term, and evaluates it on six mathematical reasoning benchmarks with Qwen2.5-Math-7B and Qwen2.5-32B models.

What This Paper Is About

Policy optimization for LLMs currently relies on an action-side Policy-KL regularizer to keep training stable, but this creates a double bind: keeping Policy-KL constrains response behavior and consumes the action-side exploration budget, while dropping it leaves optimization without explicit drift control. The authors argue that a key under-controlled source of instability is the drift of the model's own likelihood over training queries away from its pre-RL reference distribution — an "environment non-stationarity" that action-only KL does not address. ERPO's goal is to bound this input-side drift while leaving the response distribution free to explore.

Key Contributions

1. Query-environment control: The paper treats the model's query likelihood as the regularizable input-side statistic, combining Query-KL (QKL) — which bounds drift of the policy-induced query distribution ρ_θ from a pre-RL reference ρ_θ0 — with a dataset-static, reference-derived per-query weight that reduces estimator variance and tames high-temperature behavior.
2. Estimator-agnostic instantiation: The method adds a QKL term and a reference-derived per-query weight on top of GRPO/PPO/REINFORCE-style pipelines with minimal changes and no additional forward passes.
3. Stability evaluation protocol: The authors assess RL stability via multi-temperature sampling paired with a multi-metric suite (Pass@k, Pass@1, Avg@k), spanning sampling temperatures from 0.1 to 1.5.
4. Empirical gains: ERPO replaces the standard Policy-KL regularizer while consistently improving accuracy across diverse mathematical reasoning benchmarks.

Main Findings

• Query drift is real and unchecked by action-side KL. Under a fixed Policy-KL budget during GRPO training, the batch-estimated Query-KL rises steadily while Policy-KL stays nearly flat, showing that constraining only the action distribution does not stabilize the query process (Figure 1).
• Accuracy improvements on six benchmarks. Figure 3 reports Avg@32 accuracy averaged over sampling temperatures from 0.1 to 1.5, with ERPO gains of up to 14.9% and an overall average improvement of 6.2% over GRPO.
• Benchmark-level results (Table 1, mean over reported metrics): On Mean Avg@32, Base scores 0.143, GRPO 0.274, and ERPO 0.336. On Mean Pass@32, Base is 0.463, GRPO 0.575, ERPO 0.611. On Mean Pass@1, Base is 0.149, GRPO 0.275, ERPO 0.332. Reported deltas over GRPO are 6.2% in Avg@32, 3.64% in Pass@32, and 5.69% in Pass@1.
• Query-KL is the primary driver of gains; query weighting is complementary. Replacing policy-based KL with query-based KL yields the best overall performance, with an average improvement of 15.9% over GRPO, whereas query weighting (QW) mainly stabilizes training.
• Long-horizon stability. When training steps are scaled to 1K, GRPO remains stable for sampling temperatures below 1.0 until approximately 240 steps (epoch=15), but shows pronounced degradation in the high-temperature regime after 400 steps, which then propagates to all temperatures. ERPO shows a modest decline with substantially smaller deterioration, and its performance even improves within the high-temperature range. ERPO is not entirely immune to collapse (sudden entropy increase and loss of sampling capability) but consistently outperforms vanilla GRPO.
• Reduced reward hacking / train–inference gap. Averaged over six checkpoints (Steps 40, 80, 120, 160, 200, 240), GRPO shows an average Train–Eval gap of 6.47%, while ERPO reduces it to 3.14% — a reduction of approximately 51%. GRPO's evaluation accuracy drops sharply at Step-240 (from approximately 75% to 58.4%) while training accuracy stays high.
• Stability mechanics (Table 3, Avg@3): GRPO shows Query-KL 0.9679, Policy-KL 0.0601, entropy 0.5063; GRPO with w_B(s) shows 0.5933, 0.0113, 0.2782; GRPO with Query-KL shows 0.0041, 0.1001, 0.5674; ERPO shows 0.0828, 0.0728, 0.4244.
• Regularization strength matters. Increasing α to 5×10⁻² improves overall performance (Qwen-7B, n=8: 79.00 mean accuracy at temperatures ≤1.0 and 43.35 in the 1.2–1.5 range) versus α=1×10⁻² (78.74 and 37.90) and α=5×10⁻³ (59.94 and 39.30). The authors note they did not exhaustively search for the optimal α.
• More rollouts help. Increasing samples per query to n=16 gives the best performance, with average Pass@1 rising to 74.6%, and significantly improves sampling stability at high temperatures (Qwen-7B n=16: GRPO 39.75 versus ERPO 66.25 in the 1.2–1.5 range).
• Scale: On Qwen-32B, GRPO scores 81.62 (≤1.0) and 57.20 (1.2–1.5), while ERPO scores 84.60 and 82.80.
• Transfers to other RLVR algorithms. Applied to DAPO and RLOO, the method yields absolute improvements of 10.24% and 2.28% at temperatures below 1.0, respectively.
• Removing all KL constraints fails. The authors report that completely removing all KL divergence constraints resulted in training failing to converge.

Methodology in Plain English

The starting observation is that during RL training, even though the training corpus is fixed, the model's own likelihood assigned to each training question changes as its parameters update. This shifts what the authors call the policy-induced environment — the distribution over queries — and amplifies gradient variance. Standard Policy-KL only constrains the model's output distribution, not this input-side drift.

ERPO makes three moves:

1. Measure query likelihood. For each query, the model computes an autoregressive sequence log-likelihood, using the same forward pass the policy-gradient estimator already performs. The pre-RL reference model's log-likelihoods are computed once over the entire training set and cached.
2. Add a Query-KL penalty. This penalizes divergence between the current query distribution and the pre-RL reference distribution. Because the gradient of this term flows strictly through the query likelihood and not through the response score function used by policy-gradient estimators, it exerts no direct gradient pressure on the response distribution — so exploration is preserved.
3. Reweight queries. Each per-query update is weighted by a bounded, dataset-static weight derived from the cached reference log-likelihood (w(q) ∝ ℓ_θ0(q)), biasing updates toward queries that are typical under the reference and downweighting gradients from low-probability queries that tend to produce low-probability responses and higher gradient variance.

The two components plug into any policy-gradient surrogate of the standard form, recovering GRPO, PPO, or REINFORCE through different choices of the estimator terms. Training used the EasyR1 framework on Qwen2.5-Math-7B and Qwen2.5-32B, with MATH Level 3–5 problems (approximately 8.5K examples), maximum sequence length 3K tokens, eight responses per problem at inference temperature 1.0, rollout batch size 512, update batch size 128, 240 training steps (scaled to 1K for the long-horizon study), token-level loss, and a default KL divergence coefficient of 0.01. Evaluation used AIME24, AIME25, AMC, MATH500, Minerva, and OlympiadBench, aggregated over sampling temperatures from 0.1 to 1.5. The authors state a working assumption (A1) — treated as a premise, not a theorem — that keeping the model-induced query distribution aligned with its pre-RL reference preserves generalization better than unconstrained drift.

Why This Matters

Impact on research. The paper reframes LLM RL stability as an input-side, environment-level control problem rather than purely an action-side trust-region problem, connecting LLM post-training to classic non-stationary and robust RL literature and to covariate-shift lessons from imitation learning. It offers a drop-in replacement for Policy-KL that the authors report works across GRPO, DAPO, and RLOO.

Real-world applications (each a plausible direction given the paper's setting):

• Post-training mathematical reasoning models where long-horizon RL runs and inference-time high-temperature sampling both matter.
• Stable RLVR pipelines for any task with automatically verifiable rewards, such as code generation or logical problem solving — though the authors note transfer to code generation remains to be validated.
• Serving scenarios that need high-temperature or diverse sampling without the performance collapse observed in the baseline after extended training.
• Diagnosing and mitigating reward hacking, given the measured reduction in the train–inference accuracy gap.

Industry relevance. The method is described as adding no additional forward passes beyond the underlying policy-gradient estimator, which matters for large-scale training budgets. The work comes from AMAP (Alibaba Group) with collaborators at JD.com, Xi'an Jiaotong University, Beijing Normal University, and the National University of Singapore, and source code is released at https://github.com/AlibabaResearch/ERPO. The authors' explicit temperature-controlled, multi-metric evaluation protocol also addresses a comparability problem they identify in prior reporting practice.

Future Directions

• Broader domains. The authors state that transfer to instruction-following, dialogue, code-generation, and multilingual settings remains to be validated, since experiments focus on mathematical reasoning benchmarks and Qwen-family models.
• Better likelihood estimation. ERPO relies on estimating query-level likelihoods or prevalence statistics during training; the quality and computational cost of these estimates may vary with the data-selection mechanism and model scale.
• Systematic hyperparameter analysis. The authors did not conduct an exhaustive sweep over the regularization coefficient α, leaving this for future work.
• Understanding the residual constraint and collapse. The paper notes a residual constraint from coupling among LLM parameters, and reports that ERPO is not entirely immune to collapse during extended training, leaving the mechanisms of that collapse as an open question.

Target Audience

Researchers and engineers working on RL post-training of large language models — particularly those implementing RLVR pipelines with GRPO, PPO, REINFORCE, DAPO, or RLOO — as well as practitioners concerned with long-horizon training stability, reward hacking, and reliable behavior under high-temperature decoding. Readers without a background in policy gradient methods and KL regularization will find the method sections demanding.

Authors’ abstract

Policy optimization (PO) for Large Language Models faces a stability--exploration trade-off, currently mediated by an action-side Policy-KL regularizer. This puts practitioners in a double bind: keeping Policy-KL constrains response behavior and consumes the action-side exploration budget, while dropping it leaves the optimization without an explicit drift control. We argue for an alternative that breaks the dilemma by moving regularization to the input side. As training progresses, the distribution over training queries induced by the current policy drifts unchecked from its pre-RL reference distribution. Concretely, Environment-Regularized Policy Optimization (ERPO) introduces a Query-KL (QKL) term that bounds this query distribution shift, together with a dataset-static reference-derived per-query weight that biases each per-query update toward queries typical under the reference. The QKL gradient flows strictly through the query likelihood; the response score function used by policy-gradient estimators does not appear in the QKL term, so QKL exerts no direct gradient pressure on the response distribution---exploration is preserved. ERPO plugs into GRPO/PPO/REINFORCE-style pipelines without additional forward passes. On six mathematical reasoning benchmarks, ERPO replaces the standard Policy-KL regularizer while achieving effective control over query distribution drift, delivering stronger accuracy and substantially more stable behavior under high-temperature decoding and long-horizon training.Our source code are available at https://github.com/alibaba/ERPO