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

Abstract: Reinforcement Learning (RL) is the dominant paradigm for training 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.

• The paper introduces normalized entropy to robustly detect trajectory neglect, refining step-level credit assignment in LLM agents.
• It employs a two-stage STAPO framework that first localizes outlier decisions using IQR criteria and then optimizes them with a combined trajectory-aware reward and penalty.
• Experimental results on ALFWorld and WebShop benchmarks demonstrate significant performance improvements and enhanced scalability with STAPO.

Selective Trajectory-Aware Policy Optimization for LLM Agent Training

Motivation and Problem Formulation

In the context of long-horizon tasks, LLM agents frequently encounter the challenge of sparse and delayed rewards, exacerbating the temporal credit assignment problem. This often culminates in trajectory neglect, where agents lose focus on task goals and the correct management of interaction histories, especially at intermediate decision steps. Existing group-based RL approaches—such as RLOO and GRPO—have improved sample efficiency by estimating advantages via group sampling, yet they lack selective mechanisms to pinpoint and rectify suboptimal, goal-disregarding behaviors at the step level. Prior uncertainty-guided RL fine-tuning, typically leveraging Shannon entropy, fails to disentangle model confidence from state complexity, leading to unreliable outlier detection.

Figure 1: Differential ability of Shannon entropy and normalized entropy to detect trajectory neglect, isolating agent confidence from the intrinsic complexity of the environment.

Normalized Entropy for Outlier Detection

The paper introduces a normalized entropy metric, computed by comparing the Shannon entropy of a sampled step's action distribution against the intra-group statistics (mean and variance) of the same anchor state observed across sampled trajectories. This normalization effectively isolates idiosyncratic uncertainty unrelated to state complexity, providing a more robust signal for detecting decision outliers due to trajectory neglect.

The authors formalize anchor state grouping, ensuring each interaction state is only compared against semantically equivalent contexts. The localization mechanism leverages the interquartile range (IQR) criterion to automatically distinguish outlier steps, circumventing brittle manual thresholds. Empirical analysis across ALFWorld and WebShop environments demonstrates a clear association: detected outliers, as flagged by normalized entropy, correlate strongly with both negative step-level advantage and reduced trajectory attribution scores—precise indicators of trajectory neglect.

Figure 2: Outlier vs. inlier contrast in step-level advantage and attribution score distributions across multiple models on ALFWorld and WebShop; outliers exhibit consistently lower advantage and trajectory focus.

The STAPO Framework

Building on these insights, the proposed Selective Trajectory-Aware Policy Optimization (STAPO) framework integrates outlier detection with a hierarchical policy optimization regime. STAPO divides optimization into two stages:

1. Outlier Localization: Each batch computes normalized entropy for every sampled step; the IQR-based mask flags specific steps/decisions (within an anchor state) exhibiting atypical uncertainty.
2. Selective Optimization: Only outlier steps are subjected to a joint optimization objective that combines a trajectory-aware reward and a trajectory-independent penalty.

The trajectory-aware reward is operationalized as the KL divergence between the model's predicted action probabilities under the full historical prompt and a trajectory-blind version (with goal and history masked). This incentivizes policies that causally depend on the full trajectory context specifically at outlier steps. The penalty regularizes the model under the trajectory-blind prompt towards a reference baseline, deterring shortcut behaviors such as mode collapse that could exploit the auxiliary reward.

Figure 3: STAPO selectively locates outlier decisions via normalized entropy and optimizes them through a combination of trajectory-aware reward and trajectory-independent penalty, preserving overall training stability.

Experimental Evaluation

Extensive experimentation demonstrates STAPO's performance and robustness across multiple domains: ALFWorld and WebShop as primary benchmarks, with additional trials on multi-turn, search-augmented QA. Notable quantitative improvements are reported:

• On ALFWorld, using Qwen2.5-7B-Instruct, STAPO yields an overall success rate of 96.9%, consistently higher than strong baselines like GiGPO.
• On WebShop, STAPO boosts success from 65.0% (GiGPO) to 69.0% using Qwen2.5-1.5B-Instruct.
• In search-augmented QA, STAPO outperforms all competing process RL methods, with a more modest margin, reflecting the relative paucity of trajectory neglect in short-horizon QA tasks.

Ablation studies confirm the superiority of normalized entropy over Shannon entropy for localization, demonstrate that dual masking of goal and history maximizes performance, and that selective application of the trajectory-independent penalty avoids the inefficiencies of coarse global regularization.

Figure 4: Outlier localization sensitivity with respect to IQR coefficient $\lambda$; selective optimization outperforms non-selective training across a spectrum of threshold strictness.

Figure 5: Attribution analysis reveals that lower proportions of outlier steps are tightly coupled with higher task success rates post-STAPO training.

Figure 6: Training dynamics on WebShop show that, even in reward-sparse cold-start scenarios, the trajectory-aware reward $R_{TA}$ furnishes continuous supervisory signal, driving rapid improvement over standard RL baselines.

Implications and Prospective Directions

The findings highlight that selective, context-sensitive policy optimization—anchored by precise detection of trajectory neglect—yields substantial improvements in complex, multi-step LLM agent environments. The normalized entropy mechanism provides a reliable, automatic proxy for focus degradation, overcoming a persistent limitation of entropy-based credit assignment. Importantly, STAPO's selective refinement paradigm demonstrates clear scalability to larger and more capable LLMs, as well as adaptability to different model architectures (e.g., Llama3.1-8B-Instruct and Qwen2.5-14B-Instruct).

From a practical perspective, the framework incurs only marginal additional computational overhead and is orthogonal to ongoing developments in group-based RL, rendering it immediately compatible with existing and emerging RLHF and group sampling methods. This positions STAPO as an efficient and generalizable approach for robustifying agentic LLM behaviors, especially where exploration and long-horizon planning are critical.

Theoretically, the approach suggests a broader principle: leveraging contextually normalized uncertainty as a signal for targeted optimization, opening new avenues in selective intervention within high-dimensional, long-horizon policy learning. Potential future directions include:

• Direct extension to vision-language and general multimodal agent tasks.
• More sophisticated contextual disambiguation of outliers in environments with large or continuous action spaces.
• Deeper exploitation of trajectory-aware interventions for credit assignment in hierarchical RL and meta-RL scenarios.

Conclusion

By introducing normalized entropy as a principled estimator of trajectory neglect and coupling it with a selective, dual-objective policy optimization regime, STAPO achieves robust, state-of-the-art LLM agent performance on challenging long-horizon benchmarks. The work establishes a compelling methodology for tackling temporally extended credit assignment in RL for LLMs, with significant implications for the design of future scalable, context-sensitive agentic learning systems.