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

arXiv: https://arxiv.org/abs/2607.04963 · HTML: https://mlsanity.com/paper/2607.04963/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Reinforcement learning for LLM agents, one of 16 directions in LLM agents and tool use.
- Builds on 25 of its 39 references that are in the map.
- First posted on arXiv 6 Jul 2026.

## Placed in

- LLM agents and tool use / [Reinforcement learning for LLM agents](https://mlsanity.com/agents/reinforcement-learning-for-llm-agents/)

## Builds on (most cited first, 25 of 25)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [ReAct: Synergizing Reasoning and Acting in Language Models](https://mlsanity.com/paper/2210.03629/) (2022)
- [Toolformer: Language Models Can Teach Themselves to Use Tools](https://mlsanity.com/paper/2302.04761/) (2023)
- [HotpotQA: A Dataset for Diverse, Explainable Multi-hop Question Answering](https://mlsanity.com/paper/1809.09600/) (2018)
- [RT-2: Vision-Language-Action Models Transfer Web Knowledge to Robotic Control](https://mlsanity.com/paper/2307.15818/) (2023)
- [TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension](https://mlsanity.com/paper/1705.03551/) (2017)
- [Search-R1: Training LLMs to Reason and Leverage Search Engines with Reinforcement Learning](https://mlsanity.com/paper/2503.09516/) (2025)
- [Constructing A Multi-hop QA Dataset for Comprehensive Evaluation of Reasoning Steps](https://mlsanity.com/paper/2011.01060/) (2020)
- [Text Embeddings by Weakly-Supervised Contrastive Pre-training](https://mlsanity.com/paper/2212.03533/) (2022)
- [WebShop: Towards Scalable Real-World Web Interaction with Grounded Language Agents](https://mlsanity.com/paper/2207.01206/) (2022)
- [ALFWorld: Aligning Text and Embodied Environments for Interactive Learning](https://mlsanity.com/paper/2010.03768/) (2020)
- [OSWorld: Benchmarking Multimodal Agents for Open-Ended Tasks in Real Computer Environments](https://mlsanity.com/paper/2404.07972/) (2024)
- [When Not to Trust Language Models: Investigating Effectiveness of Parametric and Non-Parametric Memories](https://mlsanity.com/paper/2212.10511/) (2022)
- [Measuring and Narrowing the Compositionality Gap in Language Models](https://mlsanity.com/paper/2210.03350/) (2022)
- [GPT-4V(ision) is a Generalist Web Agent, if Grounded](https://mlsanity.com/paper/2401.01614/) (2024)
- [Captum: A unified and generic model interpretability library for PyTorch](https://mlsanity.com/paper/2009.07896/) (2020)
- [Group-in-Group Policy Optimization for LLM Agent Training](https://mlsanity.com/paper/2505.10978/) (2025)
- [Navigating the Digital World as Humans Do: Universal Visual Grounding for GUI Agents](https://mlsanity.com/paper/2410.05243/) (2024)
- [ZeroSearch: Incentivize the Search Capability of LLMs without Searching](https://mlsanity.com/paper/2505.04588/) (2025)
- [The Unreasonable Effectiveness of Entropy Minimization in LLM Reasoning](https://mlsanity.com/paper/2505.15134/) (2025)
- [Multimodal Web Navigation with Instruction-Finetuned Foundation Models](https://mlsanity.com/paper/2305.11854/) (2023)
- [Embodied Agent Interface: Benchmarking LLMs for Embodied Decision Making](https://mlsanity.com/paper/2410.07166/) (2024)
- [RLVMR: Reinforcement Learning with Verifiable Meta-Reasoning Rewards for Robust Long-Horizon Agents](https://mlsanity.com/paper/2507.22844/) (2025)
- [Harnessing Uncertainty: Entropy-Modulated Policy Gradients for Long-Horizon LLM Agents](https://mlsanity.com/paper/2509.09265/) (2025)
- [EPO: Entropy-regularized Policy Optimization for LLM Agents Reinforcement Learning](https://mlsanity.com/paper/2509.22576/) (2025)

## Built on by (most cited first, 1 of 1)

- [Matched-Context Group Policy Optimization for LLM Agent Training](https://mlsanity.com/paper/2610.13755/) (2026)

## Abstract (arXiv)

Reinforcement Learning (RL) is the dominant paradigm for training Large Language Model (LLM) agents on long-horizon tasks. However, sparse and delayed rewards often lead to trajectory neglect, in which agents lose focus on the task goal and interaction history at intermediate steps. Prior work has explored step-level supervision using Shannon-entropy-based uncertainty signals, which conflate inherent state complexity with agent confidence and therefore provide unreliable estimates of decision reliability. To address this issue, we propose normalized entropy, which measures confidence deviations relative to an agent's average behavior under a given state, thereby strengthening the association between low-quality actions and trajectory neglect. Building on this insight, we introduce Selective Trajectory-Aware Policy Optimization (STAPO), a hierarchical group-based RL framework. STAPO leverages normalized entropy to locate outlier steps associated with trajectory neglect and optimizes them via a joint mechanism of trajectory-aware reward and trajectory-independent penalty, enhancing trajectory awareness while preserving training stability. Extensive experiments on ALFWorld, WebShop, and Search-Augmented QA demonstrate that STAPO achieves state-of-the-art performance while substantially alleviating trajectory neglect, validating its effectiveness and robustness for agentic tasks.