Trust Region Masking for Long-Horizon LLM Reinforcement Learning

Abstract

Policy gradient methods for Large Language Models optimize a policy $\pi_\theta$ via a surrogate objective computed from samples of a rollout policy $\pi_{roll}$. However, modern LLM-RL pipelines suffer from unavoidable implementation divergences, such as backend discrepancies, Mixture-of-Experts routing discontinuities, and distributed training staleness. These factors cause an off-policy mismatch ($\pi_{roll} \neq \pi_\theta$), leading to approximation errors between the surrogate and the true objective. We demonstrate that classical trust region bounds on this error scale as $O(T^2)$ with sequence length $T$, rendering them vacuous for long-horizon tasks. To address this, we derive two new bounds: a Pinsker-Marginal bound scaling as $O(T^{1.5})$ and a Mixed bound scaling as $O(T)$. Crucially, both bounds depend on $D_{KL}^{max}$, the maximum token-level KL divergence across the sequence. As a sequence-level term, the divergence cannot be controlled by previous token-independent methods like PPO clipping. We propose Trust Region Masking (TRM), which masks entire sequences that violate the trust region. TRM reduces the vacuity of traditional bounds, offering a principled framework to mitigate training collapse in long-horizon LLM-RL.