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

🕒 Published (v1): 2026-07-06 11:50 UTC · Source: Arxiv · Venue: ACL 2026 · link

Why this paper was selected

ACL 2026; STAPO tackles trajectory neglect in sparse-reward RL—core agent training failure mode

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

TL;DR¶

STAPO addresses trajectory neglect — the tendency of RL-trained LLM agents to lose track of task goals and interaction history at intermediate steps of long-horizon tasks. It introduces normalized entropy to reliably locate such steps and applies selective optimization via a trajectory-aware reward plus a trajectory-independent penalty only at those steps. On ALFWorld and WebShop, STAPO achieves state-of-the-art success rates, outperforming prior process-RL methods like GiGPO.

Problem¶

Group-based RL algorithms (GRPO, RLOO, GiGPO) for LLM agent training suffer from trajectory neglect: with sparse, delayed outcome rewards, agents produce low-quality actions at intermediate steps because they stop attending to the task goal $g$ and interaction history $h_t$. Existing remedies that use Shannon entropy as a step-quality proxy are unreliable — a high-entropy step may simply reflect a genuinely complex state (many plausible actions), not a lack of agent confidence, so raw entropy conflates state complexity with decision reliability.

Method¶

Normalized Entropy. For action $a_t^{(i)}$ at anchor state $\tilde{s}$, STAPO z-scores its Shannon entropy $H(a_t^{(i)})$ within the intra-group distribution over all trajectories visiting $\tilde{s}$:

\[H_n(a_t^{(i)}) = \frac{H(a_t^{(i)}) - \mu\bigl(\{H(a_t^{(j)})\mid (a_t^{(j)},\cdot)\in G_S(\tilde{s})\}\bigr)}{\sigma\bigl(\{H(a_t^{(j)})\mid (a_t^{(j)},\cdot)\in G_S(\tilde{s})\}\bigr) + \epsilon}\]

This decouples state complexity from confidence by measuring deviation relative to the state-specific baseline.

Outlier Localization. An IQR-based mask $M_t \in \{0,1\}$ flags steps whose $H_n$ falls outside $[Q_1 - \lambda\cdot\text{IQR},\; Q_3 + \lambda\cdot\text{IQR}]$ (default $\lambda=1.5$, Tukey's fence), without manual thresholding.

Selective Optimization. Flagged steps receive two additional terms beyond the base clipped-policy objective:

Trajectory-aware reward (maximize): KL divergence between the full-context policy and the trajectory-blind policy (goal and history stripped): $$R_{TA}(c_t^{(i)}, a_t^{(i)}) = D_{\mathrm{KL}}\!\left[\pi_\theta(a_t^{(i)}\mid c_t^{(i)}) \;\|\; \pi_\theta(a_t^{(i)}\mid (c_t^{(i)})_{\setminus T})\right]$$
Trajectory-independent penalty (minimize): KL from the reference policy evaluated on the masked prompt, preventing reward hacking via mode collapse: $$P_{TI}(c_t^{(i)}, a_t^{(i)}) = D_{\mathrm{KL}}\!\left[\pi_\theta(a_t^{(i)}\mid (c_t^{(i)})_{\setminus T}) \;\|\; \pi_{\mathrm{ref}}(a_t^{(i)}\mid (c_t^{(i)})_{\setminus T})\right]$$

Advantage. A two-level advantage combines episode-level ($A_E$, intra-group normalization of outcome rewards) and step-level ($A_S$, intra-anchor-state normalization of cumulative rewards): $A(a_t^{(i)}) = A_E(\tau_i) + \omega A_S(a_t^{(i)})$.

The full objective $J_{\mathrm{STAPO}}$ applies the selective terms $\alpha R_{TA} - \gamma P_{TI}$ only where $M_t=1$ and the base group-RL objective everywhere.

Key Contributions¶

Normalized entropy: a state-relative uncertainty metric that decouples inherent state complexity from agent confidence, yielding reliable trajectory-neglect signals.
STAPO framework: hierarchical group-based RL with (i) dynamic outlier localization via IQR on normalized entropy and (ii) joint selective optimization using trajectory-aware reward + trajectory-independent penalty.
Empirical validation: attribution analysis confirms outlier steps co-locate with reduced trajectory attention and lower action quality; ablations isolate each component's contribution.

Results¶

ALFWorld (Qwen2.5-7B): STAPO 96.9% overall success vs. GiGPO 90.8% and GRPO 77.6%; surpasses GPT-4o (48.0%) and Gemini-2.5-Pro (60.3%).
ALFWorld (Qwen2.5-1.5B): STAPO 92.2% vs. GiGPO 86.7%.
WebShop success rate (Qwen2.5-1.5B): STAPO 69.0% vs. GiGPO 65.0% vs. EMPG 60.8%.
WebShop success rate (Qwen2.5-7B): STAPO 77.6% vs. GiGPO 72.8%.
Search-augmented QA (Qwen2.5-7B avg): STAPO 48.4% vs. GiGPO 47.2%; modest ~1% improvement attributed to shorter interaction horizons.
Ablation (ALFWorld/WebShop): Normalized entropy localization (92.2/69.0) > Shannon entropy (89.6/67.2); full goal+history masking > history-only > goal-only > no masking.

Limitations¶

Gains over GiGPO are modest on shorter-horizon tasks (search-augmented QA ~1%), as trajectory neglect is naturally less severe there.
Anchor state grouping requires exact-match keys for structured observations and similarity-based keys (with a tunable threshold) for free-text — the similarity threshold introduces a hyperparameter sensitive to task type.
Selective optimization requires two forward passes per outlier step (full context and masked), adding computational overhead.
Evaluated only on text-based environments; generalization to multimodal or tool-rich agentic settings is not demonstrated.

Relevance to Agentic AI / LLM Agents¶

STAPO directly addresses one of the core failure modes of RL-trained LLM agents — loss of goal and context coherence over long interaction horizons — making it highly relevant to the broader challenge of credit assignment in multi-turn agent training. The normalized entropy signal is a lightweight, model-internal diagnostic that doesn't require process-reward labeling or auxiliary critic models, positioning it as a practical add-on to existing group-based RL pipelines (GRPO, GiGPO, DAPO). The trajectory-aware reward is conceptually related to attention-forcing and context-grounding objectives explored in chain-of-thought and retrieval-augmented settings. This work advances the process-RL line of research (alongside GiGPO, EMPG, RLVMR) that seeks fine-grained, step-level credit signals for LLM agents operating in long-horizon, interactive environments.