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

Idea

STAPO 先在相同 anchor state 的 actions 中对 entropy 做 normalization，定位异常步骤，再选择性地奖励 policy 对 task goal 与 history 的依赖，并限制 trajectory-blind policy 的漂移。其目标是缓解 long-horizon agent 在中途忽视既有轨迹的 trajectory neglect。

Research Problem

Sparse terminal reward 很难指出哪一步忽略了目标或历史，raw entropy 又同时反映 state complexity 与 model confidence，容易把正常困难状态当作错误。论文研究如何构建更可比较的 outlier signal，并只对相关步骤增加 trajectory-aware optimization，避免全局约束破坏正常决策。

Insight

同一 state 内的相对 confidence 比跨 state 的绝对 entropy 更有解释力；但单纯最大化 full-context 与 masked-context distributions 的差异，可能通过恶化 masked policy 获利。因此作者同时约束 masked policy 接近 reference。Entropy outlier 与 neglect 的联系属于经验 proxy，而非逐步错误的完美标签。

Method

1. 将跨 rollouts 的 observations 按 exact-match 或 semantic similarity 分配到唯一 anchor-state group；计算每个 action 的 token-average entropy，再以该 state group 的均值和标准差转成 normalized entropy。
2. 对当前 batch 的 normalized entropy 使用 IQR interval 标记异常值，包含高低两侧 outliers；只有这些步骤开启 selective module，其他步骤仍执行基础 policy optimization。
3. 删去 goal 与 history 得到 trajectory-blind prompt；Trajectory-Aware Reward 最大化 full/masked policy 的 KL difference，Trajectory-Independent Penalty 则最小化 masked current/reference policy 的 KL，抑制通过破坏 masked behavior 获得奖励。
4. 基础 objective 结合 episode-level reward advantage 与同 anchor state 的 step-level return advantage，并使用 clipped policy update 与全局 KL regularization；selective terms 按 outlier mask 加入，在 ALFWorld、WebShop 和 search QA 中检验。