Summary
Training AI systems like large language models often uses a method called policy gradient, which assumes the training is done on the same type of data the model makes decisions on. This paper explains how using policy gradient on different, or off-policy, data can still work if certain techniques are applied. The authors show that either a special weighting of decisions across time or adding a type of regularization based on KL divergence with a clever baseline helps manage this mismatch. They also find that adjusting the baseline can make the method more cautious, which is helpful when data coverage is limited. Their work links these ideas to known approaches in offline reinforcement learning, providing a clearer theoretical understanding.
What this means in practice
- •For ai model trainers: Improve post-training of large language models by using off-policy policy gradient methods that handle distribution shifts via KL regularization and baseline adjustments.
- •For offline reinforcement learning engineers: Design offline RL algorithms that incorporate pessimism through baseline shifts to better handle limited data coverage and ensure stable policy updates.
A theory result. No direct application yet.
Abstract
Reinforcement learning has become a central tool for large language model (LLM) post-training, where policy gradient methods are routinely deployed off-policy, even though vanilla policy gradient assumes on-policy sampling. We study when such off-policy deployment is justified, and what roles its core ingredients, including importance sampling, KL regularization, and baselines, play in handling the resulting distribution shift. We show that for a general class of off-policy policy gradient updates, consistency, in the sense that the optimal policy remains a stationary point, requires one of two structures at the optimal policy: trajectory-level importance weighting, which suffers from the well-known curse of horizon, or KL regularization with a baseline equal to the optimal soft value, instantiated by the grouped mean. Focusing on the KL route, we identify it with the group-relative policy gradient (GRPG) update, establish an equivalence to a trajectory-level Bellman residual minimization objective, and obtain a finite-sample regret bound under all-policy coverage. This also gives a theoretical account of a variant of the empirically successful GRPO objective. We further uncover a hidden role of baselines in the off-policy setting: a constant shift in the baseline implements pessimism, the standard remedy in offline RL for partial data coverage. We instantiate this via the asymmetric REINFORCE objective, and connect it to the game-theoretic offline RL framework, showing a regret bound with single-policy coverage. Together, these results offer a unified theoretical view of off-policy policy gradient in LLM post-training and connect it to classical ideas in offline RL.