Trust Region Masking for Long-Horizon LLM Reinforcement Learning
Yingru Li, Jiacai Liu, Jiawei Xu, Yuxuan Tong, Ziniu Li, Baoxiang Wang
TL;DR
This work addresses the challenge of off-policy mismatch in long-horizon reinforcement learning for large language models, where classical trust-region bounds grow quadratically with sequence length and become vacuous. It derives two tighter bounds—Pinsker-Marginal and Mixed—that scale as $O(T^{3/2})$ and $O(T)$ respectively and depend on the maximum token-level KL divergence $D_{KL}^tok,max$, a sequence-level quantity. To realize practical guarantees, it proposes Trust Region Masking (TRM), which excludes any sequence containing a token that violates the trust region, yielding non-vacuous monotonic improvement guarantees for long-horizon LLM-RL. This approach provides a principled way to achieve reliable improvements in policy optimization for LLMs under realistic off-policy dynamics, including MoE routing and distributed staleness, and offers actionable guidelines for masking criteria and KL estimation.
Abstract
Policy gradient methods for large language models optimize a surrogate objective computed from samples of a rollout policy $π_{\text{roll}}$. When $π_{\text{roll}} \ne π_θ$, there is approximation error between the surrogate and the true objective. Prior work has shown that this off-policy mismatch is unavoidable in modern LLM-RL due to implementation divergence, mixture-of-experts routing discontinuities, and distributed training staleness. Classical trust region bounds on the resulting error scale as $O(T^2)$ with sequence length $T$, rendering them vacuous for long-horizon tasks. We derive two tighter bounds: a Pinsker-Marginal bound scaling as $O(T^{3/2})$ and a Mixed bound scaling as $O(T)$. Crucially, both bounds depend on $D_{kl}^{tok,max}$ -- the maximum token-level KL divergence across all positions in a sequence. This is inherently a sequence-level quantity: it requires examining the entire trajectory to compute, and therefore cannot be controlled by token-independent methods like PPO clipping. We propose Trust Region Masking (TRM), which excludes entire sequences from gradient computation if any token violates the trust region, providing the first non-vacuous monotonic improvement guarantees for long-horizon LLM-RL.