Elamur Research⌕
☰
home / papers / 2026 / STAPO: Selective Trajectory-Aware Policy Optimization for LLM Agent Training

STAPO: Selective Trajectory-Aware Policy Optimization for LLM Agent Training

Paper · 2026 · Qiuyi Qi, Tian Liang, Mutian Bao et al.

Summary

Agents built on LLMs and trained with reinforcement learning for long-horizon tasks get sparse, delayed rewards. Partway through an episode they often lose track of the task goal and the earlier interaction, a failure the authors call trajectory neglect. Earlier step-level methods use Shannon entropy as an uncertainty signal, but raw entropy mixes up how complex a state is with how confident the agent is. The authors propose normalized entropy, which standardizes a step's entropy against other actions sampled from the same anchor state. Steps falling outside an IQR-based interval are flagged as outliers. STAPO, a hierarchical group-based RL framework, then optimizes only these flagged steps. It pairs a trajectory-aware reward, derived by contrasting full inputs with versions where the goal and history are masked, with a trajectory-independent penalty against reward hacking. On ALFWorld, WebShop and Search-Augmented QA, STAPO outperforms baselines such as GiGPO and EMPG and reduces the share of outlier steps. It also provides learning signal when outcome rewards are absent, giving a targeted and inexpensive way to improve credit assignment for agents.

Key points

  • Outlier detection: Normalized entropy is computed within anchor-state groups, building on GiGPO's anchor grouping. Outliers are detected per batch with the IQR rule, where λ is typically 1.5. Raising λ from 0.0 to 3.0 improved results, suggesting that optimizing fewer, more confidently identified outliers helps.
  • Ablations:
    • Masking both the goal and the history works best. History-only masking beats goal-only masking.
    • Simply increasing a global KL coefficient hurt ALFWorld accuracy (88.3% → 84.6%). Applying the penalty only at outlier steps performed best.
  • Attribution analysis: Analysis with Qwen2.5-1.5B-Instruct shows that performance falls as the proportion of outlier steps rises. STAPO yields far fewer outliers than GiGPO.
  • Cold start: On WebShop with LLaMA-3.1-8B-Instruct, GiGPO stays at zero success. STAPO's trajectory-aware reward gives non-zero guidance early in training (roughly steps 0–50), when outcome rewards are absent.
  • Limitations:
    • The extra forward passes add only about 5.7% overhead.
    • The main experiments use text-only agents, with pilot vision-language results in the appendix.
    • Normalization degrades to raw entropy when a state appears only once in a batch.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🎓Zhejiang UniversityHangzhou · China172 works · 12,988 citations🎓City University of Hong KongHong Kong · China45 works · 3,014 citations🏢Ant GroupHangzhou · China42 works · 1,131 citations

 Fields, topics and concepts

📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 worksArts and Humanities · 16,039 worksLanguage and Linguistics · 15,891 works

Cite

Cite as Qiuyi Qi, Tian Liang, Mutian Bao, Jinjian Zhang, Dongnan Liu, Wei Zhou, Linjian Mo, Ming Kong, Jie Liu, Feng Zhang and Qiang Zhu. 2026. STAPO: Selective Trajectory-Aware Policy Optimization for LLM Agent Training. Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers).

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.18653/v1/2026.acl-long.1308

@inproceedings{qi2026stapo,
  title = {{STAPO: Selective Trajectory-Aware Policy Optimization for LLM Agent Training}},
  author = {Qiuyi Qi, Tian Liang, Mutian Bao, Jinjian Zhang, Dongnan Liu, Wei Zhou, Linjian Mo, Ming Kong, Jie Liu, Feng Zhang and Qiang Zhu},
  booktitle = {Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)},
  year = {2026},
  doi = {10.18653/v1/2026.acl-long.1308},
  url = {https://doi.org/10.18653/v1/2026.acl-long.1308}
}

References

44 references, 21 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, and 1 others. 2023. GPT-4 technical report. arXiv preprint arXiv:2303.08774.
  2. Shivam Agarwal, Zimin Zhang, Lifan Yuan, Jiawei Han, and Hao Peng. 2025. The unreasonable effectiveness of entropy minimization in llm reasoning. arXiv preprint arXiv:2505.15134. find
  3. Arash Ahmadian, Chris Cremer, Matthias Gallé, Marzieh Fadaee, Julia Kreutzer, Olivier Pietquin, Ahmet Üstün, and Sara Hooker. 2024. Back to basics: Revisiting reinforce style optimization for learning from human feedback in LLMs. In Proceedings of the 62nd Annual Meeting of the Association for Compu… find
  4. Ganqu Cui, Yuchen Zhang, Jiacheng Chen, Lifan Yuan, Zhi Wang, Yuxin Zuo, Haozhan Li, Yuchen Fan, Huayu Chen, Weize Chen, and 1 others. 2025. The entropy mechanism of reinforcement learning for reasoning language models. arXiv preprint arXiv:2505.22617. find
  5. Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, and 1 others. 2019. Natural questions: a benchmark for question answering research. Transactions of the Association for Computational Lin…
  6. Manling Li, Shiyu Zhao, Qineng Wang, Kangrui Wang, Yu Zhou, Sanjana Srivastava, Cem Gokmen, Tony Lee, Erran Li Li, Ruohan Zhang, and 1 others. 2024. Embodied agent interface: Benchmarking LLMs for embodied decision making. Advances in Neural Information Processing Systems, 37:100428–100534. find
  7. Lang Feng, Zhenghai Xue, Tingcong Liu, and Bo An. 2025. Group-in-group policy optimization for llm agent training. arXiv preprint arXiv:2505.10978. find
  8. Alex Mallen, Akari Asai, Victor Zhong, Rajarshi Das, Daniel Khashabi, and Hannaneh Hajishirzi. 2022. When not to trust language models: Investigating effectiveness of parametric and non-parametric memories. arXiv preprint arXiv:2212.10511. find
  9. Hiroki Furuta, Kuang-Huei Lee, Ofir Nachum, Yutaka Matsuo, Aleksandra Faust, Shixiang Shane Gu, and Izzeddin Gur. 2024. Multimodal web navigation with instruction-finetuned foundation models. In The Twelfth International Conference on Learning Representations. find
  10. Ofir Press, Muru Zhang, Sewon Min, Ludwig Schmidt, Noah A Smith, and Mike Lewis. 2022. Measuring and narrowing the compositionality gap in language models. arXiv preprint arXiv:2210.03350. find
  11. Boyu Gou, Ruohan Wang, Boyuan Zheng, Yanan Xie, Cheng Chang, Yiheng Shu, Huan Sun, and Yu Su. 2025. Navigating the digital world as humans do: Universal visual grounding for GUI agents. In The Thirteenth International Conference on Learning Representations. find
  12. Xanh Ho, Anh-Khoa Duong Nguyen, Saku Sugawara, and Akiko Aizawa. 2020. Constructing a multi-hop qa dataset for comprehensive evaluation of reasoning steps. arXiv preprint arXiv:2011.01060.
  13. Timo Schick, Jane Dwivedi-Yu, Roberto Dessì, Roberta Raileanu, Maria Lomeli, Eric Hambro, Luke Zettlemoyer, Nicola Cancedda, and Thomas Scialom. 2023. Toolformer: Language models can teach themselves to use tools. Advances in Neural Information Processing Systems, 36:68539–68551.
  14. John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. 2017. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347. find
  15. Yuxiang Ji, Ziyu Ma, Yong Wang, Guanhua Chen, Xiangxiang Chu, and Liaoni Wu. 2025. Tree search for llm agent reinforcement learning. arXiv preprint arXiv:2509.21240. — [arXiv:2509.21240](https://arxiv.org/abs/2509.21240) find
  16. Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, YK Li, Y Wu, and 1 others. 2024. DeepSeek-Math: Pushing the limits of mathematical reasoning in open language models. arXiv preprint arXiv:2402.03300.
  17. Bowen Jin, Hansi Zeng, Zhenrui Yue, Dong Wang, Hamed Zamani, and Jiawei Han. 2025. Search-R1: Training LLMs to reason and leverage search engines with reinforcement learning. arXiv preprint arXiv:2503.09516.
  18. Noah Shinn, Federico Cassano, Ashwin Gopinath, Karthik Narasimhan, and Shunyu Yao. 2024. Reflexion: Language agents with verbal reinforcement learning. Advances in Neural Information Processing Systems, 36.
  19. Mandar Joshi, Eunsol Choi, Daniel S Weld, and Luke Zettlemoyer. 2017. TriviaQA: A large scale distantly supervised challenge dataset for reading comprehension. arXiv preprint arXiv:1705.03551.
  20. Mohit Shridhar, Xingdi Yuan, Marc-Alexandre Cote, Yonatan Bisk, Adam Trischler, and Matthew Hausknecht. 2021. ALFWorld: Aligning text and embodied environments for interactive learning. In International Conference on Learning Representations.
  21. Narine Kokhlikyan, Vivek Miglani, Miguel Martin, Edward Wang, Bilal Alsallakh, Jonathan Reynolds, Alexander Melnikov, Natalia Kliushkina, Carlos Araya, Siqi Yan, and 1 others. 2020. Captum: A unified and generic model interpretability library for pytorch. arxiv. arXiv preprint arXiv:2009.07896, 2:5. — [arXiv:2009.07896](https://arxiv.org/abs/2009.07896) find
  22. Wouter Kool, Herke van Hoof, and Max Welling. 2019. Buy 4 reinforce samples, get a baseline for free! In ICLR 2019 Workshop. find
  23. Hao Sun, Zile Qiao, Jiayan Guo, Xuanbo Fan, Yingyan Hou, Yong Jiang, Pengjun Xie, Fei Huang, and Yan Zhang. 2025. ZeroSearch: Incentivize the search capability of llms without searching. arXiv preprint arXiv:2505.04588.
  24. Richard S Sutton and Andrew G Barto. 2018. Reinforcement Learning: An Introduction. MIT press. find
  25. Weihao Tan, Wentao Zhang, Xinrun Xu, Haochong Xia, Gang Ding, Boyu Li, Bohan Zhou, Junpeng Yue, Jiechuan Jiang, Yewen Li, and 1 others. 2024. Cradle: Empowering foundation agents towards general computer control. In NeurIPS 2024 Workshop on Open-World Agents. find
  26. Gemini Team, Rohan Anil, Sebastian Borgeaud, Jean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, Katie Millican, and 1 others. 2023. Gemini: A family of highly capable multimodal models. arXiv preprint arXiv:2312.11805.
  27. Harsh Trivedi, Niranjan Balasubramanian, Tushar Khot, and Ashish Sabharwal. 2022. MuSiQue: Multi-hop questions via single-hop question composition. Transactions of the Association for Computational Linguistics, 10:539–554. find
  28. Guanzhi Wang, Yuqi Xie, Yunfan Jiang, Ajay Mandlekar, Chaowei Xiao, Yuke Zhu, Linxi Fan, and Anima Anandkumar. 2024a. Voyager: An open-ended embodied agent with large language models. Transactions on Machine Learning Research. find
  29. Jiawei Wang, Jiacai Liu, Yuqian Fu, Yingru Li, Xintao Wang, Yuan Lin, Yu Yue, Lin Zhang, Yang Wang, and Ke Wang. 2025a. Harnessing uncertainty: Entropy-modulated policy gradients for long-horizon llm agents. arXiv preprint arXiv:2509.09265. — [arXiv:2509.09265](https://arxiv.org/abs/2509.09265) find
  30. Junyang Wang, Haiyang Xu, Haitao Jia, Xi Zhang, Ming Yan, Weizhou Shen, Ji Zhang, Fei Huang, and Jitao Sang. 2024b. Mobile-Agent-v2: Mobile device operation assistant with effective navigation via multi-agent collaboration. Advances in Neural Information Processing Systems, 37:2686–2710. find
  31. Liang Wang, Nan Yang, Xiaolong Huang, Binxing Jiao, Linjun Yang, Daxin Jiang, Rangan Majumder, and Furu Wei. 2022. Text embeddings by weakly-supervised contrastive pre-training. arXiv preprint arXiv:2212.03533.
  32. Ziliang Wang, Xuhui Zheng, Kang An, Cijun Ouyang, Jialu Cai, Yuhang Wang, and Yichao Wu. 2025b. StepSearch: Igniting LLMs search ability via stepwise proximal policy optimization. arXiv preprint arXiv:2505.15107. — [arXiv:2505.15107](https://arxiv.org/abs/2505.15107) find
  33. Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, and 1 others. 2022. Chain-of-thought prompting elicits reasoning in large language models. Advances in Neural Information Processing Systems, 35:24824– 24837.
  34. Tianbao Xie, Danyang Zhang, Jixuan Chen, Xiaochuan Li, Siheng Zhao, Ruisheng Cao, Toh Jing Hua, Zhoujun Cheng, Dongchan Shin, Fangyu Lei, Yitao Liu, Yiheng Xu, Shuyan Zhou, Silvio Savarese, Caiming Xiong, Victor Zhong, and Tao Yu. 2024. OSWorld: Benchmarking multimodal agents for open-ended tasks in…
  35. Wujiang Xu, Wentian Zhao, Zhenting Wang, Yu-Jhe Li, Can Jin, Mingyu Jin, Kai Mei, Kun Wan, and Dimitris N Metaxas. 2025. Epo: Entropy-regularized policy optimization for llm agents reinforcement learning. arXiv preprint arXiv:2509.22576. — [arXiv:2509.22576](https://arxiv.org/abs/2509.22576) find
  36. Zhilin Yang, Peng Qi, Saizheng Zhang, Yoshua Bengio, William W Cohen, Ruslan Salakhutdinov, and Christopher D Manning. 2018. HotpotQA: A dataset for diverse, explainable multi-hop question answering. arXiv preprint arXiv:1809.09600.
  37. Shunyu Yao, Howard Chen, John Yang, and Karthik Narasimhan. 2022. WebShop: Towards scalable real-world web interaction with grounded language agents. Advances in Neural Information Processing Systems, 35:20744–20757.
  38. Shunyu Yao, Jeffrey Zhao, Dian Yu, Nan Du, Izhak Shafran, Karthik R Narasimhan, and Yuan Cao. 2023. ReAct: Synergizing reasoning and acting in language models. In The Eleventh International Conference on Learning Representations.
  39. Qiying Yu, Zheng Zhang, Ruofei Zhu, Yufeng Yuan, Xiaochen Zuo, Yu Yue, Tiantian Fan, Gaohong Liu, Lingjun Liu, Xin Liu, and 1 others. 2025. DAPO: An open-source LLM reinforcement learning system at scale. arXiv preprint arXiv:2503.14476.
  40. Chaoyun Zhang, Liqun Li, Shilin He, Xu Zhang, Bo Qiao, Si Qin, Minghua Ma, Yu Kang, Qingwei Lin, Saravan Rajmohan, and 1 others. 2024. UFO: A UI-focused agent for windows OS interaction. arXiv preprint arXiv:2402.07939.
  41. Zhuosheng Zhang and Aston Zhang. 2024. You only look at screens: Multimodal chain-of-action agents. In Findings of the Association for Computational Linguistics ACL 2024, pages 3132–3149.
  42. Zijing Zhang, Ziyang Chen, Mingxiao Li, Zhaopeng Tu, and Xiaolong Li. 2025. Rlvmr: Reinforcement learning with verifiable meta-reasoning rewards for robust long-horizon agents. arXiv preprint arXiv:2507.22844. — [arXiv:2507.22844](https://arxiv.org/abs/2507.22844) find
  43. Boyuan Zheng, Boyu Gou, Jihyung Kil, Huan Sun, and Yu Su. 2024. GPT-4V (ision) is a generalist web agent, if grounded. arXiv preprint arXiv:2401.01614.
  44. Brianna Zitkovich, Tianhe Yu, Sichun Xu, Peng Xu, Ted Xiao, Fei Xia, Jialin Wu, Paul Wohlhart, Stefan Welker, Ayzaan Wahid, and 1 others. 2023. RT-2: Vision-language-action models transfer web knowledge to robotic control. In Conference on Robot Learning, pages 2165–2183. PMLR. find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/qi-2026-stapo-selective-trajectory-aware-policy/