Beyond the Stability-Exploration Dilemma: Environmental Regularization for LLM Policy Optimization
2026-08-24 • Computation and Language
Computation and Language
AI summaryⓘ
The authors address a problem in training large language models where controlling how much the model changes versus exploring new behaviors is tricky. Instead of regularizing the model's responses directly, they propose ERPO, which controls the shift in the input queries during training. This keeps the training process stable without limiting exploration of new answers. They tested ERPO on reasoning tasks and found it improves accuracy and stability compared to traditional methods.
Policy OptimizationLarge Language ModelsPolicy-KL RegularizationExploration-Exploitation Trade-offQuery DistributionReinforcement LearningPolicy GradientGRPOPPOMathematical Reasoning Benchmarks
Authors
Xianlei Zhou, Xiangdi Meng, Yu He, Tianyu Qi, Shuyan Guan, Xianli Zhang, Jian Zhang, Xin Li, Qika Lin, Jun Liu
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