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

2026-07-06

AI Summary

This work addresses the trajectory neglect problem in large language model agents during long-horizon tasks, where sparse and delayed rewards often lead to intermediate steps deviating from the task objective and interaction history. To mitigate this issue, the authors propose a hierarchical grouped reinforcement learning framework that leverages normalized entropy to precisely identify anomalous decision steps. The approach jointly optimizes trajectory-aware rewards with trajectory-agnostic penalties and introduces a selective trajectory-aware policy, enhancing the agent’s sensitivity to complete trajectories while preserving training stability. Experimental results on ALFWorld, WebShop, and search-augmented question answering benchmarks demonstrate that the method substantially alleviates trajectory neglect, achieving state-of-the-art performance and confirming its effectiveness and robustness.

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.

Problem

• trajectory neglect
• reinforcement learning
• large language model agents
• sparse rewards
• decision reliability

Innovation

• normalized entropy
• trajectory neglect
• trajectory-aware policy optimization
• hierarchical reinforcement learning
• LLM agent training