PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2610.13020v1cs.LG

TL;DR

Long-horizon rollouts can compound local policy mismatch into state-distribution shift that token- and sequence-level correction do not explicitly address. TAPO derives and applies trajectory-aware correction, with a geometric-mean estimator and length-aware clipping. Across eight mathematical reasoning and coding benchmarks, TAPO outperforms existing RL algorithms, with up to 48% and 12% average improvements over GRPO on math and code, respectively.

  • Problem

    Long-horizon rollouts can compound local policy mismatch into state-distribution shift that token- and sequence-level corrections do not explicitly address.

    Source: §1 · p. 1, §2.3 · p. 4

  • Method

    TAPO derives state-distribution correction from the off-policy policy gradient and uses a geometric-mean trajectory ratio with length-aware dynamic clipping.

    Source: §1 · p. 2

  • Results

    Across eight mathematical reasoning and coding benchmarks, TAPO outperforms existing RL algorithms, with up to 48% and 12% average improvements over GRPO on math and code, respectively.

    Source: §1 · p. 2, §5 · p. 10

  • Takeaways & Limitations

    TAPO provides a principled foundation for long-horizon autoregressive policy optimization.

    Source: §5 · p. 10

    The concentration analysis assumes independent log-ratios with centered deviations that are σ²-sub-Gaussian.

    Source: §3.2.1 · p. 5, §3.2.2 · p. 6

Abstract

from arXiv · show

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.

1 INTRODUCTION

Long-horizon rollouts can accumulate state-distribution shift even when local token-level ratios remain relatively stable, a gap not addressed by token- or sequence-level correction alone. TAPO adds trajectory-aware correction and reports stronger benchmark performance and training stability.

1 INTRODUCTION

  • State-level policy ratios drift from unity toward later positions even as token-level ratios remain relatively stable, revealing accumulated shift over long trajectories.Source: §1 · p. 1, Figure 1 · p. 2
  • GRPO and GSPO correct action-level mismatch at token or sequence granularity but omit the state-distribution shift induced by preceding trajectories.As response lengths grow, this shift is associated with increasing KL divergence and eventual training collapse.Source: §1 · pp. 1–2
  • 48% and 12% average improvements over GRPO on math and code, respectively, accompany stable score growth and controlled KL as GRPO and GSPO collapse with longer responses.The improvements are reported across mathematical reasoning and algorithmic coding; TAPO maintains controlled KL divergence while GRPO and GSPO eventually collapse as response lengths increase.Source: Figure 2 · p. 2, §1 · pp. 2–3
  • TAPO derives the required state-distribution correction from the off-policy policy gradient and incorporates it into the RL objective.Source: §1 · pp. 2–3
  • TAPO uses a geometric-mean trajectory ratio to control variance and length-aware dynamic clipping to restrain overly aggressive updates as horizons grow.The clipping retains flexibility during early exploration.Source: §1 · pp. 2–3

2 PRELIMINARY

The preliminary section models autoregressive generation as a finite-horizon MDP, then describes token-level GRPO and sequence-level GSPO correction. It also explains why sequence-level correction can miss intermediate state shifts and mismatch the token-by-token structure.

2.1 MDP FORMULATION FOR AUTOREGRESSIVE LLMS

  • An autoregressive LLM is modeled as a finite-horizon MDP: the state is the prompt plus preceding tokens, and each generated token is an action.The complete response forms a trajectory whose next state is deterministically induced by the current token.Source: §2.1 · p. 3

2.2 GRPO: TOKEN-LEVEL CORRECTION

  • GRPO applies token-level importance ratios to correct off-policy mismatch and estimates each response's advantage by normalizing its reward within a group.Subsequent variants named in the section retain the token-level importance ratio.Source: §2.2 · p. 3

2.3 GSPO: SEQUENCE-LEVEL CORRECTION

  • GSPO treats a full response as one action and applies sequence-level correction, but this can clip all tokens together or miss shifted intermediate states.The section notes that sequence-level clipping can be inefficient and that intermediate state shifts may remain hidden even when the sequence ratio stays within range.Source: §2.3 · pp. 3–4

3 METHOD

TAPO derives a trajectory-level state-distribution correction from the off-policy policy gradient, then uses variance-controlled correction and length-aware clipping in a unified objective. Its token-specific correction weights each token according to accumulated policy shift along the trajectory.

3.1 FROM ACTION CORRECTION TO STATE CORRECTION

  • The off-policy policy gradient requires both state-distribution correction for preceding trajectories and action-likelihood correction at the current token.In autoregressive generation, each state is induced by the preceding token sequence, so prior policy mismatch can accumulate at later positions.Source: §3.1 · p. 5

3.2 TRAJECTORY-AWARE POLICY OPTIMIZATION

  • The method first derives the theoretically required trajectory-aware correction, then addresses its high variance and length-dependent effective range with complementary mechanisms.Directly multiplying per-token likelihood ratios can produce high variance and unstable optimization, while fixed clipping boundaries can be inadequate as the preceding trajectory length changes.Source: §3.2 · p. 5

3.2.1 VARIANCE-CONTROLLED TRAJECTORY CORRECTION

  • TAPO replaces the product of preceding per-token likelihood ratios with their geometric mean to control variance that can destabilize long-horizon optimization.Under the stated analysis, product-correction variance grows exponentially with trajectory length, while geometric-mean correction variance decays as O(1/t).Source: §3.2.1 · pp. 5–6

3.2.2 LENGTH-AWARE DYNAMIC CLIPPING

  • Geometric averaging makes the exact trajectory correction increasingly concentrated as trajectories grow, so a fixed clipping boundary cannot maintain a consistent tolerated proximal region across token positions; a length-aware clipping rule addresses this.The clipping range spans [1, T], tightening from a relatively permissive ϵmax boundary toward a conservative ϵmin boundary as the trajectory grows.Source: §3.2.2 · p. 6

3.2.3 TRAJECTORY-AWARE POLICY OPTIMIZATION

  • TAPO combines variance-controlled trajectory correction with length-aware dynamic clipping in a unified objective, while retaining the local policy-gradient direction for each token.The trajectory-level ratio corrects accumulated distribution shift; the current action likelihood provides the policy-gradient direction.Source: §3.2.3 · pp. 6–7
  • TAPO gives each token a trajectory-aware, token-specific correction, between GRPO’s local likelihood ratio and GSPO’s uniform sequence-level correction.Each token is weighted according to accumulated policy shift along the trajectory leading to its current state.Source: §3.2.3 · p. 7

4 EXPERIMENTS

Experiments compare TAPO with representative RL algorithms across mathematical reasoning and coding, then test its performance across model scales and ablate its correction-ratio and clipping choices.

4. EXPERIMENTS

  • Experiments cover eight competition-level benchmarks across math and coding, comparing ten representative RL algorithms on models of different scales.Evaluation uses the best checkpoint and reports avg@8 over eight repeats; training uses Qwen3-8B-Base with DAPO-Math-17k and Eurus-Code-25k.Source: §4 · p. 7, §4.1 · pp. 7–8

4.2 MAIN RESULTS

  • TAPO achieves the best performance across eight math and coding benchmarks, improving average performance over GRPO by 48.0% on math and 11.8% on coding.It also exceeds RIPO by 1.2% and 2.4%, respectively, and GSPO by 16.8% and 9.1%.Source: §4.2 · p. 8
  • In mathematical reasoning training, TAPO maintains rising scores with controlled KL divergence and stable entropy, while GRPO and GSPO eventually collapse as policy shift grows.The comparison is reported in Figure 2.Source: §4.2 · p. 8

4.3 GENERALIZATION ACROSS MODEL SCALES

  • On Qwen3-1.7B and Qwen3-4B math tasks, TAPO has the best average performance, improving over GRPO by 31.5% and 21.3%, respectively.Table 2 compares algorithms across math benchmarks using Avg@8.Source: §4.3 · p. 8, Table 2 · p. 9

4.4 ABLATION STUDIES

  • Full-trajectory geometric averaging yields stable training and sustained score improvement; truncating large ratios or averaging only the state correction does not prevent instability or collapse.The ablation compares the original trajectory ratio, a GMPO-style truncation, state-only geometric averaging, and full-ratio geometric averaging.Source: §4.4 · pp. 8–9, Figure 3 · p. 9
  • No clipping leads to sudden collapse; GRPO-order clipping underconstrains later tokens, while GSPO-order clipping overconstrains earlier tokens and slows score improvement.The lower panel of Figure 3 compares clipping strategies; the results support adapting the boundary to trajectory length.Source: §4.4 · pp. 9–10, Figure 3 · p. 9

5 CONCLUSION

The paper concludes that TAPO offers a principled approach to policy-distribution shift in long-horizon LLM RL, with favorable results across math and coding benchmarks.

5 CONCLUSION

  • TAPO is derived from the off-policy policy gradient to correct policy-distribution shift in long-horizon LLM RL.The paper presents it as a principled framework for trajectory-aware policy optimization.Source: §5 · p. 10
  • Across eight math and coding benchmarks, TAPO substantially outperforms existing RL algorithms.Source: §5 · p. 10
  • The paper presents TAPO as a principled foundation for advancing long-horizon autoregressive policy optimization.Source: §5 · p. 10

A PROOF OF THEOREM1

The proof derives an off-policy policy gradient for autoregressive trajectories, separating state-distribution correction from action-level correction.

A PROOF OF THEOREM1

  • The proof treats an LLM trajectory as a finite-horizon process with deterministic next states formed by appending each sampled token.Under this property, a state identifies its preceding token sequence.Source: A PROOF OF THEOREM1 · pp. 13, 15
  • Importance sampling rewrites the target-policy objective using trajectories sampled from the behavior policy, then differentiates that objective with respect to policy parameters.The behavior policy is fixed during optimization, so only target-policy probabilities depend on the parameters.Source: A PROOF OF THEOREM1 · p. 13
  • The step-wise derivation conditions on the current state and action, integrates over future trajectories, and expresses the gradient using an advantage function.The resulting cumulative importance correction decomposes into state-distribution and action-level corrections.Source: A PROOF OF THEOREM1 · pp. 13–15

B PROOF OF THEOREM 2

The proof analyzes the variance of exact and geometric-mean trajectory corrections under an independent Gaussian log-ratio assumption.

B PROOF OF THEOREM 2

  • The variance analysis assumes independent, identically distributed Gaussian log-ratios with variance σ2.Source: B PROOF OF THEOREM 2 · p. 15
  • For the exact product correction, the proof applies log-normal moment calculations to show that variance grows exponentially with trajectory length.It obtains the distribution by taking logarithms of the product and summing the independent Gaussian log-ratios.Source: B PROOF OF THEOREM 2 · pp. 15–16
  • The proof also formulates a geometric-mean correction by scaling the summed log-ratios by 1/t and analyzes its variance behavior.Source: B PROOF OF THEOREM 2 · p. 16

C PROOF OF THEOREM 3

The proof centers log-ratio variables, applies sub-Gaussian reasoning, and uses a geometric-mean correction to complete the result.

C PROOF OF THEOREM 3

Related papers

Loading 2610.13020v1…