STAPO: Selective Trajectory-Aware Policy Optimization for LLM Agent Training
2026-07-06 • Artificial Intelligence
Artificial Intelligence
AI summaryⓘ
The authors focus on a problem in training AI agents, where the agents sometimes lose track of their goals during long tasks because they get unclear feedback. They find that previous methods to measure the agent's uncertainty mixed up how hard a situation is with how confident the agent is, which made the feedback unreliable. To fix this, the authors create a new way to measure confidence that compares current actions to the agent's usual behavior, helping spot when the agent is getting off track. They then build a new training method called STAPO that uses this measure to improve learning on tricky parts of a task, which they show works better on several benchmarks.
Reinforcement LearningLarge Language ModelsTrajectory NeglectNormalized EntropyAgent ConfidencePolicy OptimizationHierarchical RLTrajectory-Aware RewardsSparse RewardsTraining Stability
Authors
Qiuyi Qi, Tian Liang, Mutian Bao, Jinjian Zhang, Dongnan Liu, Wei Zhou, Linjian Mo, Ming Kong, Jie Liu, Feng Zhang, Qiang Zhu
Abstract
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.