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

The Core Problem

Large language model policy optimization gets trapped between two bad options. The standard approach uses a Policy-KL regularizer—a penalty that keeps the model's output distribution close to a baseline—to prevent the model from drifting too far during training. But this regularizer does double duty: it both stabilizes training and limits exploration of better responses. Practitioners either accept constrained behavior or drop the regularizer and lose stability. The authors argue this framing misses the real source of instability.

Where Drift Actually Happens

The insight is that instability doesn't come only from the model changing how it responds to a fixed set of queries. It also comes from the model changing which queries it prefers to see. During reinforcement learning training, the policy learns to request certain types of queries more often based on reward signals. This query distribution—the set of problems the model wants to solve—drifts away from the original training data without any explicit constraint. The authors observe this overlooked drift is a major source of training instability.

Environment-Regularized Policy Optimization (ERPO)

Instead of regularizing response behavior, ERPO regularizes the input side: the distribution of queries the policy implicitly selects during training. The method introduces two components. First, a Query-KL (QKL) term bounds how far the query distribution can shift from the reference distribution of the original training set. This is the core innovation—control drift where it actually happens. Second, ERPO uses per-query weights derived from the reference distribution to bias each update toward queries that were typical in the original data.

Critically, the QKL gradient flows only through query likelihood and never touches the response scoring used by policy-gradient methods. This means the regularizer exerts no direct pressure on response generation, preserving the model's freedom to explore better ways to answer questions. The method integrates cleanly into existing pipelines (GRPO, PPO, REINFORCE) without requiring additional forward passes—it's computationally lightweight.

Results and Impact

On six mathematical reasoning benchmarks, ERPO replaces the standard Policy-KL regularizer while maintaining effective control over query distribution drift. The model achieves stronger accuracy and substantially more stable behavior, especially under high-temperature decoding (more random sampling) and long-horizon training (many iterations). The concrete gains suggest this is not merely a theoretical refinement but a practical improvement that makes training both safer and more effective.

The paper's key contribution is reframing the stability-exploration trade-off. By moving regularization from action space to input space, ERPO decouples two problems that the standard approach conflated. The model can explore response distributions freely while the training procedure itself remains anchored to the reference query distribution. This breaks the false dilemma that has constrained policy optimization in LLMs, opening a path toward faster, more stable learning without sacrificing the ability to find better behaviors.