Long-horizon RL training stops collapsing mid-run: state-shift correction beats GRPO by 48% on math and 12% on code
Synopsis
In long-horizon autoregressive reinforcement learning, local policy mismatch compounds along the trajectory into state distribution shift, destabilizing training and eventually collapsing it as the horizon grows. Trajectory-Aware Policy Optimization derives a state-distribution correction from the off-policy policy gradient and writes it into the objective, estimates it with a geometric mean to cut variance, and tightens the clipping boundary with the length of the preceding trajectory. Across eight competition-level math and code benchmarks it improves over GRPO by up to 48% on math and 12% on code on average, with training scores rising steadily while KL divergence stays controlled.
Interpretation
Long-horizon autoregressive reinforcement learning gains an explicit correction target: not only the per-action probability ratio, but the state distribution shift induced by the preceding trajectory. GRPO corrects at the token level and GSPO at the sequence level, and both correct policy mismatch through action likelihoods without explicitly accounting for the state distribution induced by the preceding trajectory. As response length grows, the KL divergence between behavior and target policies increases, and GRPO and GSPO eventually suffer training collapse, shown by declining training scores and fluctuating entropy.
The correction comes from a derivation of the off-policy policy gradient, and in practice it replaces the product of per-token likelihood ratios with a geometric mean while adapting the clipping range to trajectory length. The variance of the product of per-token ratios grows exponentially with the length of the preceding trajectory, so a fixed clipping boundary cannot provide a consistent tolerated region across positions. The geometric mean reduces the correction ratio's variance to decay with length, and length-aware dynamic clipping tolerates wider deviation at early tokens and tighter deviation later; training used Qwen3-8B-Base with 8 rollouts, a maximum response length of 20k tokens and 750 steps, on DAPO-Math-17k for math and Eurus-Code-25k for code.
On eight competition-level benchmarks, trajectory-aware correction improves both mathematical reasoning and algorithmic coding while keeping training stable. Against GRPO, average math performance improves by up to 48% and average coding by 12%; against RIPO, the strongest token-level method, math and coding averages improve further by 0.5 and 1.4 points. Math used AIME24, AIME25, BRUMO25 and SMT25, coding used Codeforces, CodeContests, TACO and APPS, with the best checkpoint evaluated 8 times; on Qwen3-1.7B-Base and Qwen3-4B-Base the average gains over GRPO were 31.5% and 21.3%.
Perspective
The report covers policy optimization in long-horizon autoregressive LLM reinforcement learning. Experiments train Qwen3-8B-Base, Qwen3-4B-Base and Qwen3-1.7B-Base, using DAPO-Math-17k for mathematical reasoning and Eurus-Code-25k for algorithmic coding, and evaluate on eight competition-level benchmarks: AIME24, AIME25, BRUMO25 and SMT25 for math, and Codeforces, CodeContests, TACO and APPS for code, taking the best checkpoint and repeating evaluation 8 times. The intended readers are reinforcement learning researchers and engineering implementers training long-horizon reasoning models, focused on how training stability and reasoning scores move together.
The training dynamics show GRPO and GSPO eventually collapsing as policy shift grows, while trajectory-aware correction keeps KL divergence controlled and entropy stable; this comparison comes from the mathematical reasoning training curves. Ablations indicate that geometric averaging over the full trajectory ratio is essential for stable training, that removing clipping causes sudden collapse, that a GRPO-order boundary constrains later tokens insufficiently, and that a GSPO-order boundary constrains early tokens too tightly. The method is validated on math and code tasks with verifiable rewards; behavior on other long-horizon reasoning settings is not given in the material.