Dual-Axis Policy Optimization for LLM Agents: Bayesian Feedback Attribution and Trajectory Mass Normalization
Yingxuan Zhuang and colleagues separate two optimization axes in agent RL, how feedback is exploited within a trajectory and how trajectories are aggregated across a batch, and address each with BATON.
Ask this paper
Bayesian Feedback Attribution handles the first axis. A feedback-conditioned posterior over sampled actions distributes credit within a trajectory instead of spreading the terminal reward uniformly.
Trajectory Mass Normalization handles the second. Each complete trajectory receives equal optimization mass, so long trajectories stop dominating the batch update by token count.
Both axes help independently. Experiments with GRPO and GiGPO on ALFWorld, WebShop and SearchQA show separate gains from each, and the combination is strongest across model scales.
The framing is reusable. Naming the two dimensions makes existing agent RL methods comparable by which axis they modify, which most papers leave implicit.
Abstract
Reinforcement learning for LLM agents involves two distinct optimization di- mensions: how environment feedback is exploited within a trajectory, and how complete trajectories are aggregated across a batch. We formulate these dimen- sions as Intra-Trajectory Feedback Attribution and Inter-Trajectory Objec- tive Aggregation, and introduce BATON (Bayesian Attribution and Trajectory Objective Normalization), a dual-axis policy optimization framework. BATON instantiates the first axis with Bayesian Feedback Attribution, which constructs a feedback-conditioned posterior over sampled actions, and the second with Trajec- tory Mass Normalization (TMN), which assigns equal optimization mass to com- plete trajectories. Experiments with GRPO and GiGPO on ALFWorld, WebShop, and SearchQA show that both axes provide independent gains and that their combi- nation consistently achieves the strongest overall performance across model scales.