Trajectory-Aware Policy Optimization: Closing the State Distribution Gap in Long-Horizon LLM RL

arXiv:2610.13020 · cs.LG · Submitted 2026-10-09 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "Trajectory-Aware Policy Optimization".

Jane: The gist:

Tom: First, who's behind it and why it matters.

Title and authors: Tom: So we're looking at this paper today titled Trajectory-Aware Policy Optimization: Closing the State Distribution Gap in Long-Horizon LLM RL. It sounds like they're tackling something pretty specific about how reinforcement learning works when you generate really long sequences.

Jane: Yeah, it’s about these long rollouts, where local mistakes don't just stay local; they build up and mess with the overall state distribution over time. The authors are basically saying that current methods like GRPO and GSPO only fix the mistake at each individual token or sequence level, but they miss this bigger problem where the state distribution shifts.

Lu: Exactly. They show that as you go further down a long trajectory, the state-level policy ratio progressively deviates from one toward later positions, which is a clear sign of that accumulation over time.

Meng: So it's not just about getting one token right; it's about keeping the model in the right context for a whole response. That sounds like it could be a big deal for making AI agents that can actually follow complex instructions over many steps.

Lalam: I see what they’re pointing out. They identify and formalize this accumulation of state distribution shift as a fundamental issue in long-horizon LLM RL, and derive its corresponding correction from the off-policy policy gradient >

The paper's summary: Tom: What’s the actual fix they propose? It seems like they’re trying to derive a specific correction term directly from the policy gradient, not just slapping some general adjustment on top. They call it Trajectory-Aware Policy Optimization or TAPO >

Jane: Right. The core idea is that the exact off-policy correction needed at token t has two parts: a state distribution correction and an action-level correction, and they show how to formally write that as Rt = dπt(stx)dµt(stx)z state correction·π(ot st)µ(ot st)z action correction >

Lu: That derivation is key because it grounds the state distribution correction in the policy gradient itself, rather than guessing what kind of adjustment to make. It connects the required change directly to the optimization process >

Meng: From a deployment standpoint, if you can derive this mathematically from the gradient, it suggests a more principled way to stabilize training that doesn't require just tweaking hyperparameters blindly >

Lalam: And then they address the problem of variance by using a geometric-mean estimator over the preceding trajectory. This makes the correction ratio have a variance decay of O(one/t), which is much better than what you get from the exact product correction where it grows exponentially > <ref:2610.13020#pg1>

The paper's improvements: Tom: Moving into how they make this work practically, they introduce a length-aware dynamic clipping rule. They scale the clipping boundary with one/√t, defined by epsilon l t = epsilon l min + (epsilon l max - epsilon l min)one/√t − one/√T > <ref:2610.13020#pg1>

Jane: That length-aware clipping is important because the geometric averaging makes the correction increasingly concentrated as the trajectory grows. So, they need to adapt how much they clip based on how long the response already is >

Lu: It’s a smart way to manage that concentration. The rule scales epsilon l t with one/√t, meaning early tokens get a wider tolerance while later tokens admit a tighter clipping range > <ref:2610.13020#pg1>

Meng: That makes sense for stability when you're generating something long; you don't want the early steps being overly restricted while the later steps are still highly unstable >

Lalam: And they show that even with these changes, extensive experiments across both mathematical reasoning and algorithmic coding benchmarks demonstrate that TAPO consistently improves performance with enhanced training stability >

Conclusion: Tom: So, to wrap up, the main thing here is that controlling the variance of the complete trajectory correction is essential for effective trajectory-aware policy optimization >

Jane: It moves beyond just fixing local action mismatch to explicitly closing that state distribution gap over long horizons using a geometric-mean estimator and dynamic clipping >

Lu: What this means for us is that we have a theoretically grounded RL algorithm now, one that doesn't just rely on empirical tuning for stability when dealing with long sequences >

Meng: For production, it suggests we can actually push the horizon length further because we know the updates are staying controlled KL divergence and showing steady improvement >

Lalam: I agree. The Trajectory-Aware Policy Optimization framework provides a principled foundation for advancing long-horizon autoregressive policy optimization >

Zhicheng Cai, Weinan Dai, Yuyan Zhou, Wei-Ying Ma, Ya-Qin Zhang, Hao Zhou

Institute for AI Industry Research (AIR) Tsinghua University

cs.LG

Submitted: 2026-10-09

Updated: 2026-10-09

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 90/100

The gist: The gist: Trajectory-Aware Policy Optimization (TAPO) is a principled framework that theoretically derives and incorporates state-distribution correction from the off-policy policy gradient to close

Key concepts

State Distribution Shift: This occurs in long sequences when small errors in the policy at early steps accumulate into large differences in the predicted future states. It means the model starts predicting entirely different sequences than what it was trained on, leading to unstable learning and eventual training collapse.
Off-Policy Policy Gradient Correction: The paper uses the off-policy gradient to mathematically derive the exact correction needed for state distribution mismatch. This correction is split into two parts: one that fixes the state distribution and another that handles action-level mismatches, providing a theoretical basis for stabilization.
Geometric Mean Estimator: To control variance in the trajectory correction term, TAPO replaces simple products of likelihood ratios with their geometric mean. This technique ensures that as the sequence grows, the correction becomes more concentrated and stable, leading to a variance decay of O(1/t).
Length-Aware Dynamic Clipping: This rule prevents overly aggressive updates by scaling the clipping boundary based on trajectory length (1/√t). It allows for looser constraints when processing early tokens and tighter constraints later in the sequence, ensuring updates remain practical across different lengths.

Summary

The gist: Trajectory-Aware Policy Optimization (TAPO) is a principled framework that theoretically derives and incorporates state-distribution correction from the off-policy policy gradient to close the state distribution gap in long-horizon LLM reinforcement learning.

Problem Identification

Long-horizon autoregressive rollouts expose RL to a fundamental vulnerability where local policy mismatch can compound into severe state distribution shift over trajectories. Existing methods like GRPO and GSPO control policy mismatch at the token or sequence level but overlook this trajectory-induced shift, causing training instability and eventual collapse as the horizon grows. This gap is exposed because controlling action-level mismatch does not guarantee control of the state distribution shift induced over long-horizon trajectories.

The Theoretical Correction

The exact off-policy correction consists of two components: a statedistribution correction and an action-level correction. In the context of LLM RL, the state transition is deterministic, meaning the state-distribution ratio reduces to the likelihood ratio of the preceding trajectory. The theoretically required importance correction Rt at token t can be written as Rt = dπt(stx)dµt(stx)z state correction·π(ot st)µ(ot st)z action correction.

Variance Reduction Mechanism

To control the variance of the exact trajectory correction, which grows exponentially with trajectory length, TAPO employs a geometric-mean estimator over the preceding trajectory. This results in a variance decay of O(1/t) for Ret = R 1/t t.

Length-Aware Dynamic Clipping

Geometric averaging makes the correction increasingly concentrated as the trajectory grows, necessitating a length-aware dynamic clipping rule. This rule scales the clipping boundary with 1/√t, defined by ϵl t = ϵl min + (ϵl max − ϵl min)1/√t − 1/√T and similarly for epsilon h.

Final Objective and Results

The final TAPO objective is JTAPO(π) = Ex∼D,γGg=1∼µ(·x)1/o go g Xt=1 min Rbg t Aˆg t, clip Rbg t, 1−ϵl t, 1+ϵh t Aˆg t. Furthermore, as Fig. 2 shows, while GRPO and GSPO eventually suffer from training collapse as the policy shift grows, TAPO maintains a steadily increasing training score with controlled KL divergence.

Generalization and Ablation

TAPO consistently achieves the best performance across both mathematical reasoning and algorithmic coding benchmarks. When generalized to different model scales, TAPO consistently achieves the best average performance on both models. Ablation studies confirm that geometric averaging applied to the full trajectory ratio yields stable training dynamics and sustained score improvement. The results demonstrate that controlling the variance of the complete trajectory correction is essential for effective trajectory-aware policy optimization.

This paper introduces Trajectory-Aware Policy Optimization (TAPO), a principled framework for correcting the policy distribution shift in long-horizon LLM RL, derived from the off-policy policy gradient. The gist: Trajectory-Aware Policy Optimization (TAPO) is a principled framework that theoretically derives and incorporates state-distribution correction from the off-policy policy gradient to close the state distribution gap in long-horizon LLM reinforcement learning.

How it works

Identification of State Distribution Shift: The paper identifies and formalizes the accumulation of state distribution shift as a fundamental issue in long-horizon LLM RL. This shift occurs because local policy mismatch at preceding positions accumulates along the trajectory and can induce substantial state-distribution mismatch at later positions.
Derivation of Correction: The off-policy policy gradient is derived to reveal the required correction, which consists of two parts: a statedistribution correction and an action-level correction.
Variance Control via Geometric Mean: To stabilize the variance of the exact trajectory correction, the product of per-token likelihood ratios is replaced by its geometric mean. This yields a variance decay of O(1/t) for Ret = R 1/t t.
Length-Aware Clipping: To make the correction practical, a length-aware dynamic clipping rule is introduced. This rule scales the clipping boundary with 1/√t, ensuring that early tokens require a wider tolerance while later tokens admit a tighter clipping range.

This framework provides a principled middle ground between local action correction and global sequence correction, explicitly bridging the state-distribution gap in long-horizon autoregressive generation and enabling consistently stable training. The results demonstrate that explicitly correcting the trajectory-induced state distribution shift not only improves reasoning performance but also substantially enhances training stability. Ultimately, TAPO provides a principled foundation for advancing long-horizon autoregressive policy optimization.

The paper introduces Trajectory-Aware Policy Optimization (TAPO), a principled framework for correcting the policy distribution shift in long-horizon LLM RL, derived from the off-policy policy gradient. The gist: Trajectory-Aware Policy Optimization (TAPO) is a principled framework that theoretically derives and incorporates state-distribution correction from the off-policy policy gradient to close the state distribution gap in long-horizon LLM reinforcement learning.

Improvements for AI systems

textbfTrajectory-Aware Policy Optimization (TAPO) Implementation for Long-Horizon RL Stability: The improved system can now perform long-horizon autoregressive rollouts with significantly enhanced training stability compared to GRPO and GSPO, as TAPO maintains controlled KL divergence and stable improvement whereas others suffer from growing KL divergence and eventual training collapse. This allows the system to reliably generate coherent, multi-step responses without catastrophic forgetting or divergence during extended reasoning tasks.
textbf Mechanistic State Distribution Correction: The system will explicitly correct for the accumulation of state distribution shift by deriving the required correction from the off-policy policy gradient, as shown in Theorem 1 and Equation (7). This allows the model to better handle long sequences where local policy mismatch can compound into severe state distribution shift over trajectories, leading to more accurate reasoning across many tokens.
textbf Variance Reduction via Geometric Averaging: The system will utilize a geometric mean estimator for variance reduction as defined in Equation (9) and proven in Theorem 2, which yields a correction ratio where the variance of the geometric-mean correction decays as O(1/t), compared to the exact product correction whose variance grows exponentially. This ensures that the policy updates remain stable even for very long trajectories, mitigating high variance and unstable optimization.
textbf Length-Aware Dynamic Clipping: The system will employ a length-aware dynamic clipping mechanism defined by Equation (13) to prevent overly aggressive updates across the trajectory. By adapting the clipping boundaries to the length of preceding trajectory, it ensures that early tokens require a wider tolerance, whereas later tokens admit a tighter clipping range, preventing instability as response lengths grow.

Abstract

Reinforcement learning (RL) has become a dominant paradigm for enhancing LLMs' reasoning capabilities. Yet long-horizon autoregressive rollouts introduce a fundamental limitation: local policy mismatch can compound into severe state distribution shift over trajectories. Existing LLM RL methods, such as GRPO and GSPO, control policy mismatch at the token or sequence level, but overlook this trajectory-induced shift, causing training instability and eventual collapse as the horizon grows. To close the state distribution gap, we propose Trajectory-Aware Policy Optimization (TAPO), a principled framework that theoretically derives the state-distribution correction from the off-policy policy gradient and incorporates it into the RL objective. TAPO further employs a geometric-mean estimator for variance reduction, together with a length-aware dynamic clipping that prevents overly deviated updates across the trajectory. Extensive experiments demonstrate that TAPO significantly surpasses existing LLM RL algorithms across eight competition-level benchmarks, achieving up to 48% and 12% improvement over GRPO on mathematical reasoning and algorithmic coding, with substantially enhanced training stability.