arXiv：cs.LG· Zhicheng Cai, Weinan Dai, Yuyan Zhou, Wei-Ying Ma, Ya-Qin Zhang, Hao Zhou·· 13 小时前
TAPO：弥合长程 LLM 强化学习中的状态分布差距
Trajectory-Aware Policy Optimization: Closing the State Distribution Gap in Long-Horizon LLM RL
AI 导读
针对长程自回归 rollout 中局部策略失配累积成严重状态分布偏移的问题，研究者提出轨迹感知策略优化（TAPO），从 off-policy 策略梯度理论推导出状态分布修正项并纳入 RL 目标，同时用几何平均估计器降低方差、用长度感知动态裁剪避免过度偏离更新。在八个竞赛级基准上，TAPO 在数学推理和算法编程上分别较 GRPO 提升最高 48% 和 12%，训练稳定性显著增强。
正文
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.
| Subjects: | Machine Learning (cs.LG) |
| Cite as: | arXiv:2610.13020 [cs.LG] |
| (or arXiv:2610.13020v1 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2610.13020 arXiv-issued DOI via DataCite (pending registration) |
Submission history
From: Zhicheng Cai [view email]
[v1]
Fri, 9 Oct 2026 05:22:07 UTC (4,443 KB)
来源：arXiv：cs.LG · arxiv.org