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

cs.CLAug 24, 2026

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

Authors: Xianlei Zhou, Xiangdi Meng, Yu He, Tianyu Qi, Shuyan Guan, Xianli Zhang, Jian Zhang, Xin Li,

Organizations: AMAP, Alibaba Group · Xi’an Jiaotong University · Beijing Normal University · National University of Singapore

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/AlibabaResearch/ERPO

Figures & tables

Figure 1: KL losses during GRPO training. The Query-KL (dark) rises while the Policy-KL (light) stays low, showing action-only KL does not stabilize the query process.

Figure 2: Comparison of the GRPO and ERPO Computation Flows. Given a query $q$ and its output $o$ , the policy model $\theta$ and the reference model $\theta_{0}$ induce query distributions $\rho_{\theta}(q)$ and $\rho_{\theta_{0}}(q)$ , as well as output distributions $\rho_{\theta}(o)$ and $\rho_{\theta_{0}}(o)$ . The GRPO-like flow (red) applies policy-level KL regularization to the output distributions and optimizes the standard within-query advantage. In contrast, the ERPO-like flow (green) measures query-distribution drift through $D_{\mathrm{KL}}(\text{query})$ for environmental regularization and reweights the advantage at the query level using $w_{B}(q)=f(\rho_{\theta_{0}}(q))$ (Section 4.2 ).

Figure 3: Avg@32 over Sampling Temperatures on Mathematical Reasoning Tasks

Mean Avg@32
Method	AIME24	AIME25	AMC	MATH500	Minerva	Olympiad	Avg.
Base	0.087	0.030	0.246	0.340	0.058	0.099	0.143
GRPO	0.174	0.072	0.398	0.528	0.207	0.266	0.274
ERPO	0.218	0.110	0.478	0.677	0.214	0.316	0.336
Mean Pass@32
Base	0.373	0.206	0.674	0.764	0.349	0.411	0.463

Table 1: Performance comparison across mathematical reasoning benchmarks. Best results per column are highlighted in bold.

Figure 5: Training Dynamics on Long-term RL

Figure 6: Performance Variation Across Training Steps

Base Model	Method	$\boldsymbol{\alpha}$	$w(s)$	Rollout Count	Temperature Metrics
Base Model		$\boldsymbol{\alpha}$	$w(s)$	Rollout Count	0.1	0.6	1	1.5	$\leq 1.0$	1.2–1.5
Baseline	—	—	—	—	52.40	46.80	32.80	0.40	44.44	6.15
Qwen-7B	GRPO	$1\times 10^{-2}$	—	8	66.80	68.40	73.80	0.40	68.80	12.50
	GRPO w/ $w(s)$	$1\times 10^{-2}$	✓	8	78.20	76.80	71.20	7.60	76.14	23.79
	GRPO w/ query-KL	$1\times 10^{-2}$	—	8	81.60	81.60	79.00	2.60	80.90	38.00
	ERPO	$1\times 10^{-2}$	✓	8	79.40	80.60	75.20	8.60	78.74	37.90

Table 2: Performance Comparison Under Different Experimental Settings

Method	Query-KL	Policy-KL	Entropy
GRPO (Avg@3)	0.9679	0.0601	0.5063
GRPO w/ $w_{B(s)}$	0.5933	0.0113	0.2782
GRPO w/Query-KL	0.0041	0.1001	0.5674
ERPO (Avg@3)	0.0828	0.0728	0.4244

Table 3: Influence of Query-KL and Query-Reweighting on Training Stability

Panel of Figure 7 — Figure 7: Pass@1 accuracy and training dynamics under different settings: (a)–(c) Model performance at various temperatures on MATH500; (d) Policy - KL divergence variation with GRPO using only Query - KL.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 8: Variation of Metrics with Temperature

Method	Avg. Train Acc	Avg. Eval@TP1	Avg. Eval@TP2	Gap@TP1	Gap@TP2	Avg. Gap
GRPO	77.3	69.8	70.1	7.5	5.5	6.47
ERPO	77.1	74.1	73.8	3.0	3.3	3.14
Improvement	–	–	–	$\downarrow 4.5$ (60%)	$\downarrow 2.2$ (40%)	$\downarrow 3.33$ (51%)

Appendix

Table 4: Quantification of Reward Hacking via Train–Inference Gap. The table compares the average accuracy during training sampling versus inference decoding. A larger gap indicates severe reward hacking (overfitting to training dynamics). ERPO reduces this gap by

\approx 51\%

, demonstrating robust generalization.

GRPO (Baseline)
Model / Metric	Step-0	Step-40	Step-80	Step-120	Step-160	Step-200	Step-240	Avg.
Train Acc	44.48	76.41	76.09	78.29	78.95	79.44	76.70	72.91
Eval@TP1	31.20	73.20	76.00	73.40	72.20	73.60	58.40	65.43
Eval@TP2	30.60	74.00	74.80	74.20	76.60	75.80	66.20	67.46
ERPO (Ours)
Train Acc	44.55	75.63	76.53	78.66	80.63	78.41	81.46	73.70

Appendix

Table 5: Trajectory of Train vs. Inference Accuracy. Detailed performance recorded at 40-step intervals. Note the divergence in GRPO at Step-240, where Eval accuracy drops significantly while Train accuracy remains high—a clear sign of reward hacking. ERPO maintains consistency throughout.

Figure 9: Likelihood Relationship between Query and Response

Base Model	Method	D_KL	Rollout	Temperature
Base Model	Method	D_KL	Rollout	0.1	0.6	1	1.5	$\leq$ 1.0	1.2-1.5
Baseline	–	–	–	52.4	46.8	32.8	0.4	44.44	6.15
Qwen-7B	GRPO	Policy	8	66.8	68.4	73.8	0.4	68.8	12.5
	GRPO	Policy	16	73.0	79.2	75.0	10.6	75.22	39.75
	ERPO	Query	8	79.4	80.6	75.2	8.6	78.74 (+9.94)	37.9
	ERPO	Query	16	80.4	78.8	74.4	56.2	77.82 (+2.6)	66.25

Appendix

Table 6: Performance Comparison Under Different Experimental Settings

RL for Language Model Reasoning

Reinforcement learning for language model reasoning uses reward-driven training to improve multi-step problem solving, tool use, and reliable answers.

See all 1,087 papers

On KL-Regularized Policy Optimization
Oct 6, 2026Yifan ZhangRL for LLM Agents KL-Regularized RL
Princeton University
expo: Exploration-prioritized policy optimization via adaptive kl regulation and gaussian curriculum sampling
May 11, 2026Mingxiong Lin, Zhangquan Gong, Maowen Tang +6Reinforcement Learning KL-Regularized RL
OPPO AI Center
SCOPE-RL: Stable and Quantitative Control of Policy Entropy in RL Post-Training
Oct 9, 2025Chen Wang, Zhaochun Li, Jionghao Bai +4Reinforcement Learning RL Post-Training
College of Software, Nankai University · Zhongguancun Academy · Beijing Institute of Technology +2

Oct 6, 2026J/K move · Enter open · S save

On KL-Regularized Policy Optimization

Yifan Zhang

Princeton University

RL for LLM Agents KL-Regularized RL Off-Policy RL Asynchronous RL Critic-Free RL

Open paper AskPDF arXiv

Asynchronous reinforcement learning (RL) for large language model (LLM) agents trains one policy on trajectories generated by another: rollouts come from stale checkpoints, and the inference engine's probabilities differ from the trainer's even at identical parameters. Standard remedies either clip importance ratios, which biases the update, or, as in GRPO, sample a group of responses per prompt, which is costly when episodes are long. We propose KL-Regularized Policy Optimization (KLPO), a framework that anchors the KL regularizer at the sampler. The regularized improvement step then has a closed-form Gibbs solution, and KLPO fits its log-ratio optimality condition by least squares on the sampler's own trajectories, so the sampler probability enters through a log-ratio and no importance weights are needed. Profiling out the regression intercept replaces the intractable log-partition function with the signal's sampler mean plus a sampler-to-trainer KL divergence. For token-level policy mirror descent targets, we show that the resulting gradient can be computed from terminal returns without a critic, via sampler-centered scores or a single trajectory residual, even under stochastic tool outputs. We further prove that independent Monte Carlo estimates of the KL term keep these gradients unbiased, derive the exact KL gap of cheaper top-

K

and binary approximations, and show that SPPO, GPO, REBEL, and BPO arise as special cases of KLPO. The result is a critic-free update that uses one rollout per prompt and requires neither a learned normalizer nor a group of responses.

Abstract

RL for Language Model Reasoning

Explore similar work

On KL-Regularized Policy Optimization

expo: Exploration-prioritized policy optimization via adaptive kl regulation and gaussian curriculum sampling

SCOPE-RL: Stable and Quantitative Control of Policy Entropy in RL Post-Training