Policy optimization boosts long task performance in text agents
One Act, Many Masks: Effect-Marginal Policy Optimization for Long-Horizon Agents
Machine Learning
Summary
When AI agents respond to complex tasks over many steps, different ways of saying the same action can confuse the learning process. The authors found that linking rewards to specific wording rather than the actual effect causes problems in training. They created a new method called Effect-Marginal Policy Optimization that groups all ways of saying the same action together, making learning more focused on what really matters. This approach improved performance and learning speed in several simulated tasks.
What this means in practice
- •For ai developers: Improve training of AI agents that perform complex multi-step tasks by aligning learning updates with action effects rather than specific wording.
- •For game ai engineers: Enhance long-term decision making for text-driven game agents by optimizing policies across equivalent action descriptions.
Authors
Shaojin Ma, Ruitao Wang, Pengyi Li, Jie Yang, Liangyu Chen, Gangyi Zhao, Hongyao Tang, Yan Zheng, Jianye Hao
Abstract
Recent advances in multi-turn reinforcement learning (RL) have substantially improved the performance of LLM agents on complex, long-horizon interactive tasks. However, despite progress in fine-grained credit assignment, policy optimization still applies these learning signals through response-specific probability ratios. Different textual responses can induce the same executable effect, yet credit differences arising from subsequent outcomes remain tied to their individual textual realizations. We identify this mismatch as credit-update misalignment. Our analysis reveals that credit-text pairing directly shapes policy-update directions, even when the policy context and executable effect are held fixed. Motivated by this finding, we propose Effect-Marginal Policy Optimization (EMPO), which aligns policy optimization with execution-level equivalence through an effect-level surrogate that couples updates across alternative textual realizations. EMPO preserves the original credit signals, integrates with both trajectory-level and step-level credit estimators, and requires no additional critics or rollouts. Experiments on ALFWorld, WebShop, and Sokoban show that EMPO outperforms strong baselines in both task performance and learning efficiency.