PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Xianlei Zhou, Xiangdi Meng, Yu He, Tianyu Qi, Shuyan Guan, Xianli Zhang, Jian Zhang, Xin Li, Qika Lin, Jun Liu

arXiv:2608.23311v1cs.CL

TL;DR

Policy optimization for LLMs must balance stability against exploration while policy-induced query distributions drift from their pre-RL reference. ERPO moves regularization to the input side with Query-KL and reference-derived query weighting, preserving response exploration; across six mathematical reasoning benchmarks, it improves accuracy and stability during high-temperature, long-horizon training.

  • Problem

    LLM policy optimization faces a stability–exploration dilemma, while policy-induced query distributions can drift unchecked from their pre-RL reference.

  • Method

    ERPO regularizes the policy-induced query distribution with Query-KL and uses a dataset-static reference-derived per-query weight in GRPO/PPO/REINFORCE-style pipelines.

  • Results

    Across six mathematical reasoning benchmarks, ERPO outperforms GRPO on Avg@32 accuracy, with gains of up to 14.9% and an overall average improvement of 6.2%.

  • Takeaways & Limitations

    ERPO provides KL divergence control without explicit policy regularization while delivering superior performance and evidence of stability during long-term, multi-temperature training.

  • Takeaways & Limitations

    Experiments focus primarily on mathematical reasoning benchmarks and Qwen-family models, so transfer to broader settings remains to be validated.

Abstract

from arXiv · show

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

1 Introduction

LLM policy optimization faces accumulating instability and distribution shift while balancing stability against exploration. ERPO addresses unchecked policy-induced query drift with input-side regularization and reference-derived query weighting, preserving response exploration.

  • Long-horizon policy optimization can accumulate noise and distribution shift, causing oscillations and occasional collapses.
  • The policy-induced query distribution drifts from its pre-RL reference even when the training corpus remains fixed.
  • Under a fixed Policy-KL budget, Query-KL rises while Policy-KL stays flat, leaving query-process drift uncontrolled.
  • ERPO adds Query-KL between current and reference query distributions and a reference-derived per-query weight.
  • Across diverse reasoning benchmarks, ERPO consistently improves accuracy and supports multi-temperature stability evaluation.

2 Related Works

Prior work addresses RLVR rewards and language-model training stability, but explicit management of input query distributions remains underexplored. ERPO is presented against this gap in output-focused regularization.

  • RLVR uses automatically verifiable outcomes as scalable reward signals for tasks such as mathematical reasoning, code generation, and logical problem solving.
  • Process reward models extend RLVR by evaluating and rewarding intermediate reasoning steps.
  • RLVR can introduce high variance from sparse rewards and potential reward hacking, motivating more stable training methods.
  • Language-model RL stability is challenged by discrete actions, large parameter spaces, and complex reward landscapes.
  • Existing approaches largely focus on output regularization, leaving explicit input query-distribution management underexplored.

3 Preliminaries

The preliminaries define an RLVR setting with fixed training queries and policy-generated responses, then formalize policy-induced query distributions and their drift from a pre-RL reference.

  • In RLVR, an LLM policy produces responses to queries and a verifier assigns scalar rewards.
  • The training objective averages per-query expected reward over the training query distribution, and GRPO is used as the concrete estimator.
  • Training queries are shaped by corpus composition, curricula, difficulty filters, and active sampling schedulers.
  • The policy-induced query distribution ρθ is defined by the model’s autoregressive likelihood of each query and remains distinct from fixed ρtrain.
  • As parameters update, ρθ can drift from ρθ0; Query-KL is used to quantify and control this environment drift.

4 Method

ERPO treats the pre-RL policy-induced query distribution as a reference environment, combining Query-KL regularization with dataset-static query reweighting. Its input-side penalty avoids direct response-distribution pressure and is compatible with common policy-gradient estimators.

  • ERPO bounds drift between the current query distribution ρθ and pre-RL reference ρθ0 while leaving response-side exploration unconstrained.
  • The method assumes, as an empirical premise, that preserving alignment with ρθ0 better retains prior-training generalization than unconstrained drift.
  • ERPO computes query likelihoods autoregressively, caching reference log-likelihoods once and reading current log-likelihoods from the policy-gradient forward pass.
  • Query-KL penalizes divergence between ρθ and ρθ0, while a dataset-static reference-derived weight directs updates toward queries typical under ρθ0.
  • QKL acts through query likelihood ℓθ(q), excludes the response score function, and therefore exerts no direct gradient pressure on response behavior.
  • ERPO is estimator-agnostic and can instantiate GRPO, PPO, or REINFORCE through appropriate policy-gradient surrogate choices.

5 Experiments

Experiments on six mathematical reasoning benchmarks show that ERPO improves accuracy over GRPO and better controls query-distribution drift, especially during high-temperature and long-horizon training. Ablations indicate that Query-KL and reference-derived query reweighting contribute to performance and stability.

  • 5.2 Main Results: ERPO improves Avg@32 accuracy by 6.2% overall and by up to 14.9% across six benchmarks and sampling temperatures from 0.1 to 1.5.The comparison averages performance over sampling temperatures and covers AIME24, AIME25, AMC, MATH500, Minerva, and OlympiadBench.
  • 5.2 Main Results: ERPO surpasses GRPO by 6.2% in Avg@32, 3.64% in Pass@32, and 5.69% in Pass@1.The reported gains hold across the evaluation metrics used in the main comparison.
  • 5.3 Training Dynamics: In GRPO, query-level KL divergence becomes an order of magnitude greater than policy-level KL divergence as the query distribution drifts from the reference.ERPO instead constrains the query distribution and adjusts loss using problem probability, limiting reference divergence while allowing policy-level exploration.
  • 5.3 Training Dynamics: After approximately 400 steps, GRPO degradation begins at high temperatures and spreads across temperatures, whereas ERPO’s deterioration is substantially smaller.ERPO can still experience extended-training collapse, but it consistently outperforms GRPO without an explicit policy-based KL term.
  • 5.4 Analysis: Replacing policy-based KL with Query-KL yields a 15.9% average improvement over GRPO in the MATH500 ablation.The query-based KL setting has the best overall performance without modifying other hyperparameters.
  • 5.4 Analysis: Reference-derived query reweighting reduces policy KL divergence while supporting training, and 16 samples per query raise average Pass@1 to 74.6%.Higher sampling counts also improve high-temperature stability without noticeably increasing divergence from the reference model.
  • 5.4 Analysis: Increasing regularization strength, such as α = 5 × 10^-2, further improves overall performance, although the authors did not exhaustively search for the optimal α.The reported experiments retain the default value for a relatively fair comparison.

6 Conclusion

ERPO regularizes the policy-induced query distribution with query-level KL divergence and reference-derived per-query reweighting, while avoiding explicit policy regularization. Across mathematical reasoning benchmarks, it provides KL control and improved stability during long-term training.

  • ERPO bounds drift of the policy-induced query distribution ρθ from the pre-RL reference ρθ0 using query-level KL divergence.It also reweights advantages with a dataset-static reference-derived per-query weight.
  • Across multiple experiments, GRPO consistently failed to capture the desired output format, so reported results use answer accuracy only.
  • ERPO achieves KL divergence control without explicit policy regularization while delivering superior performance.
  • ERPO examines sampling across temperatures to provide additional evidence of stability during long-term reinforcement-learning training.

Limitations

The evaluation is concentrated on mathematical reasoning benchmarks and Qwen-family models, with transfer to broader application settings left open. The method also depends on query-likelihood estimates and lacks an exhaustive regularization-coefficient sweep.

  • ERPO’s transfer to instruction-following, dialogue, code-generation, and multilingual settings remains to be validated.Experiments focus primarily on mathematical reasoning benchmarks and Qwen-family models.
  • ERPO relies on estimating query-level likelihoods or prevalence statistics, whose quality and computational cost may vary with data-selection mechanisms and model scale.
  • The study does not conduct an exhaustive sweep over the regularization coefficient, leaving systematic hyperparameter analysis for future work.
  • The supplied implementation text specifies reasoning-output formatting requirements for training examples.
  • ERPO reduces the performance gap between sampling temperatures while increasing the likelihood of sampling correct outputs.

C Analysis of Reward Hacking

GRPO exhibits reward hacking: training accuracy remains high while evaluation behavior diverges, including a sharp late-stage evaluation drop. ERPO substantially narrows the train–evaluation gap and regularizes training against spurious reward patterns.

  • GRPO shows reward hacking when high training rewards coexist with plateauing or degrading evaluation performance during later optimization.
  • Training and evaluation accuracy were monitored every ten optimization steps, with evaluation run under TP1 and TP2 settings.
  • 6.47% is GRPO’s average train–evaluation gap, while evaluation accuracy falls from approximately 75% to 58.4% at Step-240 despite high training accuracy.
  • Approximately 51% is ERPO’s reduction of the average train–evaluation gap, from 6.47% to 3.14%.
  • ERPO’s reduced gap indicates that training improvements translate more reliably to inference performance by limiting exploitation of spurious reward patterns.

D Correlation analysis between query and response probabilities

The analysis finds a positive relationship between prompt and response NLL and uses query-likelihood quantities already available in policy-gradient training. ERPO’s query-level modifications therefore reuse existing computation while remaining compatible with several policy-gradient estimators.

  • Over 8K training questions were sampled at temperature 1.0 for separate prompt and response negative log-likelihood computation.
  • For 95% of training samples with prompt NLL below 300, prompt and response NLL have a correlation coefficient close to 1.
  • Low-probability responses in positive samples contribute to increased entropy during training.
  • The query-likelihood appendix defines autoregressive sequence likelihood for token sequences and requires a single forward pass for ℓθ(q).
  • The cached reference likelihood table is fixed throughout training and treated as a constant during backpropagation.
  • ERPO reuses the policy-gradient estimator’s per-step forward pass for current query likelihoods, adding no extra forward cost.
  • ERPO’s query-level KL term and outer per-query reweighting are orthogonal to the selected action-level estimator and leave its inner loop unchanged.
  • Applying ERPO requires precomputing reference likelihoods, replacing the outer query weight, and adding an estimated query-KL term without architectural changes.

H Proof of Proposition 1: Structural Decoupling of Regularization and Exploration

Proposition 1 establishes that Query-KL regularization acts only through query likelihoods, creating structural decoupling from response-side exploration.

  • Structural conclusion: The gradient of Rquery(θ) flows through the query log-likelihood and is orthogonal to the response policy score function.This is the formal statement of Proposition 1.
  • Gradient derivation: The reference query distribution is θ-independent, so its log-gradient vanishes during differentiation.The derivation uses ∇θ log ρθ0(q) = 0 because ρθ0 does not depend on θ.
  • Gradient derivation: The probability-normalization identity makes the corresponding distributional term vanish because ρθ sums to one for every θ.This follows from differentiating the normalization condition.
  • Structural conclusion: The Query-KL gradient depends on θ through ℓθ(q) and ∇θℓθ(q), with no response-policy score factor.Thus QKL acts on the input environment rather than directly on response probabilities.

I Derivation of Query Reweighting

ERPO derives a reference-aligned query objective and implements it with a cached, bounded query weight based on reference-model likelihoods. The resulting weight is static during training, favors reference-typical queries, and limits reweighting variance.

  • Reference-aligned objective: ERPO replaces empirical query sampling with a reference-aligned objective based on the cached reference-induced query prior ρθ0.The objective is expressed under the training distribution through importance weighting.
  • Importance weighting: Under approximately uniform sampling from N queries, the ideal importance weight differs only by a shared constant from the relative query weighting.The shared factor N does not affect relative weighting among queries.
  • Practical query weight: ERPO uses a dataset-static weight monotone in cached reference score ℓθ0(q), assigning larger weights to queries favored by the reference model.The reference score is computed once over the training set and reused throughout optimization.
  • Training objective: The ERPO loss includes query-level regularization controlled by strength parameter α.Rquery(θ) denotes the query-level regularization term.
  • Practical query weight: The cached reference-model log-probability serves as a monotone proxy for relative query preference rather than an unbiased density-ratio estimate.The approximation preserves ordering instead of estimating absolute query density.
  • Practical query weight: The normalized weight is clipped to a fixed range to reduce importance-reweighting variance and prevent any single query from dominating optimization.Normalization centers the scale around the dataset average while preserving monotonicity.
Loading 2608.23311v1…