Driving · 2026-10-12
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
AI-assisted summary based on the paper. Consult the linked source for methods and limitations. Benchmark numbers appear after editor verification.
TL;DR
The authors identify that long-horizon autoregressive rollouts in LLM RL accumulate state distribution shift from local policy mismatch, which token/sequence-level correction (e.g., GRPO/GSPO) does not address. They propose Trajectory-Aware Policy Optimization (TAPO), which derives the needed state-distribution correction from the off-policy policy gradient, uses a geometric-mean estimator to reduce variance, and applies length-aware dynamic clipping. Experiments show improved performance and more stable training.
Why it matters
Long-horizon LLM RL can collapse because controlling only token/sequence likelihood ratios does not bound the trajectory-induced state distribution shift. TAPO targets this gap directly by incorporating an explicit state-distribution correction into the RL objective, with variance reduction and length-aware clipping aimed at stability as horizon grows.
Method
- Derives an off-policy policy gradient showing that exact correction includes both state-distribution ratio and action likelihood ratio, where state ratio corresponds to preceding-trajectory effects in autoregressive LLMs.
- Proposes TAPO with a trajectory-aware correction ratio and replaces the product with a geometric-mean estimator to reduce variance for long trajectories.
- Introduces length-aware dynamic clipping whose clipping range tightens with preceding trajectory length to prevent overly deviated updates while retaining flexibility early.
Abstract (from arXiv)
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.
Related papers
- FSPO: Policy-Consistent Risk and Pareto-Feasible Control for Budgeted LLM RL Post-Training
- T2SPO: Trajectory-to-Step Policy Optimization for Agentic Reinforcement Learning
- TrajEvolve: Trajectory Evolution for Reinforcement Learning with Hindsight Credit Assignment
- TD-MPC2: Scalable, Robust World Models for Continuous Control
- Reachability-Aware Diffusion Policy Optimization
Publication record
Bibliographic evidence describes where and when a record appeared. It does not validate scientific claims.
First public release: 2026-10-09 · Indexed by Pub-AI: 2026-10-12 UTC
Trajectory-Aware Policy Optimization: Closing the State Distribution Gap in Long-Horizon LLM RL
Preprint posted: 2026-10-09 · Original record ↗
- arxiv
- 2610.13020
Explicit source versions (1)
- v1 · 2026-10-09
- Preprint available
Source date: 2026-10-09 · First detected: 2026-10-25 UTC · Historical baseline
Source evidence (1)
- Official evidence ↗ · atom:published (official metadata)
Abstract · arxiv
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.