Dual-Axis Policy Optimization for LLM Agents: Bayesian Feedback Attribution and Trajectory Mass Normalization
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.
Related
- Box-iLQR: Constraint-Aware Feedback via Finite-Barrier Trajectory Optimization
- Multi-Objective Reinforcement Learning Using Routed Ensembles and Trajectory Attribution
- Surprise Reduction and Nullification in Bayesian and Inverse Bayesian Inference under Ambiguous Prediction-Error Attribution
- Online Material-Labeled Environment Reconstruction via Bayesian Multipath Attribution for Low-Altitude ISAC
- Search-Trajectory Patterns as Feedback for Adaptive Metaheuristic Configuration Control
BibTeX
@misc{zhuang2026dualaxis,
title = {{Dual-Axis Policy Optimization for LLM Agents: Bayesian Feedback Attribution and Trajectory Mass Normalization}},
author = {Yingxuan Zhuang and Binhe Yu and Jingxiao Yang and Ruopei Sun and Ziting Li and Cheng Tan and Xuhong Zhang and Jianwei Yin and Jintao Chen},
howpublished = {arXiv},
year = {2026},
month = sep,
doi = {10.48550/arxiv.2609.19830},
eprint = {2609.19830},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2609.19830},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Dual-Axis Policy Optimization for LLM Agents: Bayesian Feedback Attribution and Trajectory Mass Normalization
AU - Yingxuan Zhuang
AU - Binhe Yu
AU - Jingxiao Yang
AU - Ruopei Sun
AU - Ziting Li
AU - Cheng Tan
AU - Xuhong Zhang
AU - Jianwei Yin
AU - Jintao Chen
PY - 2026
DA - 2026/09/17
PB - arXiv
DO - 10.48550/arxiv.2609.19830
UR - https://arxiv.org/abs/2609.19830
AB - 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.
ER -