Token ratios stay controlled while states drift
The paper's starting observation is a mismatch between what existing methods measure and what actually goes wrong. In long-horizon autoregressive rollouts, each token is generated at a state defined by all previously sampled tokens, so the target policy's state distribution depends on the entire history of policy deviations. Small local mismatches can therefore accumulate into a large shift at later positions.
The authors illustrate this with measurements taken from randomly sampled points during GRPO training: the token-level policy ratio stays relatively stable across positions, but the state-level ratio progressively departs from unity toward later positions. That is the gap the paper names. GRPO corrects with token-level importance ratios and GSPO aggregates token ratios into a sequence-level correction, yet neither explicitly corrects the state distribution induced by the preceding trajectory. As response lengths grow, the accumulated shift appears as rising KL divergence between behavior and target policies, and the paper reports eventual training collapse for both baselines.
Why action-level correction cannot cover the trajectory
Importance sampling is the standard tool for off-policy mismatch, and the paper's argument is about where that correction is applied. Treating the whole trajectory as one action, as GSPO does, also has a token-efficiency cost: once the sequence-level ratio leaves the clipping range, every token in the trajectory is clipped, including intermediate tokens whose mismatch may still be well controlled. Conversely, a sequence-level ratio inside the clipping range can still hide substantial state shifts at intermediate positions, because local deviations can average out in the final sequence ratio while the corresponding token updates are still computed at severely shifted states.
The theoretical step is a decomposition. In the appendix derivation, the cumulative importance correction at step t factors into a state-distribution correction, the ratio of target to behavior state densities at that step, multiplied by an action-level ratio. Because LLM state transitions are deterministic, each state uniquely identifies its preceding token sequence, so the state ratio equals the product of token ratios over all earlier positions. Existing methods keep the action term and drop the state term; TAPO keeps both.
A geometric mean keeps the correction from exploding
Keeping the full product correction creates a new problem. Under the paper's modeling assumption that log-ratios are independent Gaussian variables with variance, the variance of the exact product correction grows exponentially with trajectory length t, while the variance of a geometric-mean version, the product raised to the power 1/t, decays as order 1/t. The paper proves both results and also gives a concentration bound showing the geometric-mean correction's log deviation from its mean is sub-Gaussian with variance proxy scaled by 1/t.
This is why TAPO does not simply multiply all token ratios together. The geometric mean preserves the direction of the correction while preventing the estimator's variance from overwhelming the gradient as horizons lengthen. The ablation on Qwen3-8B-Base compares the original trajectory ratio against variants including geometric averaging applied only to the state correction ratio, which is how the paper isolates the estimator's contribution from the correction itself.
Clipping adapts to trajectory length
The second component addresses update size rather than estimator variance. TAPO applies a length-aware dynamic clipping rule intended to prevent overly deviated updates across the trajectory. The ablation study on Qwen3-8B-Base for mathematical reasoning examines this clipping strategy alongside the construction of the trajectory-aware correction ratio, so the paper treats both as design choices to be tested rather than assumed defaults.
The reported training dynamics connect the two components to stability: GRPO and GSPO show declining training scores, fluctuating entropy, and growing KL divergence as policy shift accumulates, while TAPO maintains steadily improving training scores with controlled KL divergence and stable entropy. The paper interprets this as evidence that trajectory-aware correction mitigates the state distribution shift induced by long-horizon updates.
Gains shrink against stronger baselines
Across eight competition-level mathematical reasoning and algorithmic coding benchmarks, TAPO has the best average performance in both domains. Relative to GRPO, the paper reports average improvements of 48.0% on math and 11.8% on coding. The comparison that matters more for judging the method is against stronger correction baselines: versus RIPO, described as the strongest token-level correction method, TAPO improves average performance by 1.2% on math and 2.4% on coding; versus the sequence-level method GSPO, the gains are 16.8% and 9.1%.
These are averages over benchmark suites, not per-task results, and the size of the margin depends heavily on which baseline is chosen. The large GRPO gap and the small RIPO gap describe the same method against different reference points, so readers comparing TAPO to a modern token-level correction should anchor on the 1.2% and 2.4% figures rather than the headline 48%.
Scale results hold, but only two smaller models tested
The generalization experiment trains TAPO on Qwen3-1.7B-Base and Qwen3-4B-Base for mathematical reasoning and evaluates with Avg@8. TAPO achieves the best average performance on both, improving over GRPO by 31.5% and 21.3% respectively, and over GMPO by 13.2% on the 1.7B model and 9.3% on the 4B model.
The scope of this evidence is narrow in two ways. The scale check covers two smaller Qwen3 base models on mathematical reasoning only, while the eight-benchmark comparison spans math and algorithmic coding. The paper's stability claims come from training-dynamics observations on mathematical reasoning. Nothing in the packet tests other model families, other task domains, or longer horizons than those trained, so the correction's benefit should be read as demonstrated within these settings rather than as a general property of long-horizon RL.
Compare against the right baseline first
The paper's practical priority follows from its own comparison structure. If the alternative on the table is GRPO or GSPO, the reported margins are large and the stability argument is the main reason to switch: TAPO keeps KL divergence controlled where the baselines collapse as response lengths grow. If the alternative is a strong token-level correction such as RIPO, the expected average gain is modest, 1.2% on math and 2.4% on coding, and the case for TAPO rests on the trajectory-level correction and clipping design rather than on a large benchmark jump.
The most decision-relevant question the paper leaves readers with is therefore not whether state distribution shift exists, which its own measurements show, but whether the added state correction earns its complexity at the horizon lengths and baseline strength of a given training setup. The evidence supports that trade-off most clearly against GRPO and GSPO, and least clearly against the strongest token-level method.
Sources used in this answer
Trajectory-Aware Policy Optimization: Closing the State Distribution Gap in Long-Horizon LLM RL
TAPO derives a state-distribution correction from the off-policy policy gradient, adds a geometric-mean estimator and length-aware dynamic clipping, and reports average gains over GRPO of 48.0% on math and 11.8% on coding across eight competition-level benchmarks, with smaller gains of 1.2% and 2.4% over RIPO.