Dual-Axis Policy Optimization for LLM Agents: Bayesian Feedback Attribution and Trajectory Mass Normalization
Yingxuan Zhuang, Binhe Yu, Jingxiao Yang, Ruopei Sun, Ziting Li, Cheng Tan, Xuhong Zhang, Jianwei Yin, et al.
cs.AI · stat.ML
Quick verdict
BATON reweights token losses via feedback-likelihood evidence and equalizes trajectory mass, yielding consistent GRPO/GiGPO gains; stats and theory stay thin.
Summary
The paper targets two sources of bias in RL fine-tuning of LLM agents: under-exploitation of environment feedback within a trajectory, and domination of batch aggregation by long trajectories. It proposes BATON, which (i) applies Bayesian Feedback Attribution — for each visited decision it samples one counterfactual action, forms a two-candidate posterior over feedback likelihoods, length-calibrates it via geometric evidence, and uses normalized stop-gradient weights to redistribute token-level loss mass within each trajectory — and (ii) Trajectory Mass Normalization, which replaces token-proportional trajectory weighting with a uniform 1/B share. A proposition shows expected log evidence equals a KL divergence between the executed action's feedback distribution and the pairwise mixture. With GRPO and GiGPO hosts on Qwen2.5-1.5B/7B (ALFWorld, WebShop) and 3B/7B (seven search-augmented QA sets), BATON lifts ALFWorld success from 72.9% to 85.6% (GRPO, 1.5B) and 86.8% to 94.1% (GiGPO), WebShop success from 57.1% to 68.2% and 67.1% to 76.4%, and QA averages by 1.2–3.2 points; the Table 3 ablation shows both axes contribute independently. Learned policies take fewer actions and generate fewer tokens (e.g., GiGPO 1,431 to 746); overhead is 1.10–1.17×. Limitations: no significance tests or stated run counts, a suspicious Bamboogle cell (GiGPO 3B 63.8 vs EnvRL 44.2), no code release, and a thin theory that leaves the single-sample counterfactual posterior's reliability unexamined.
Review complete
Assessment at a glance
4
risk flags identified
● high 0
● medium 3
● low 1
4 strengths · 5 weaknesses · 5 author questions
AI text signal = Binoculars zero-shot detector score over 20 sampled passages (Fast-DetectGPT co-reported) — absence does not prove human authorship, and a flag does not imply misconduct.
Levels reflect verification complexity — not a quality score. Evolve never rates your paper.
Key claims under review
- Policy optimization for LLM agents decomposes into two independently improvable dimensions: intra-trajectory feedback attribution and inter-trajectory objective aggregation.
- Bayesian Feedback Attribution — a detached reweighting derived from a feedback-conditioned two-candidate posterior — improves GRPO and GiGPO on its own (Table 3).
- Trajectory Mass Normalization, replacing token-proportional aggregation with uniform 1/B mass, yields independent gains by removing length bias.
- The combination is consistently the strongest configuration across ALFWorld, WebShop, seven QA datasets and multiple model scales, while shortening trajectories (e.g., GiGPO 1,431→746 generated tokens).
Strengths
Clean dual-axis formalization with mass conservation
Eq. 3-4 factorize the objective into within-trajectory weights (with the constraint that token-weighted weights sum to 1 per trajectory) and outer trajectory mass; the 2×2 ablation in Table 3 maps one-to-one onto this formalism, so the independence claims are actually tested rather than asserted.
Broad, consistent empirical coverage
Two host learners (GRPO, GiGPO), three environments, three model scales. Headline gains are large on embodied/web tasks: ALFWorld GRPO 1.5B 72.9→85.6, GiGPO 7B 90.9→94.9; WebShop GRPO 1.5B success 57.1→68.2. Gains persist over GiGPO, which already provides step-level credit — the strongest evidence the axes add value beyond finer credit assignment.
Behavioral corroboration beyond score tables
Honest motivation and quantified overhead
Weaknesses
Statistical reporting is incomplete
Suspect data cell
Theory is descriptive, not predictive
Missing the closest normalization baseline
No code release; EnvRL underdescribed
Full Review Report
$1.99 · one-timeUnlock the complete preliminary analysis Evolve produced for this paper.
Report preview
Scope & Claims
BATON (arXiv:2609.19830, September 2026) addresses RL fine-tuning of multi-turn LLM agents along two declared axes: Intra-Trajectory Feedback Attribution (how environment feedback shapes credit within a trajectory) and Inter-Trajectory Objective Aggregation (how complete trajectories share batch mass). Four claims are made: (1) the axes are conceptually and empirically separable; (2) Bayesian Feedback Attribution (BFA), a detached reweighting from a feedback-conditioned posterior, helps alone; (3) Trajectory Mass Normalization (TMN), uniform 1/B trajectory weighting, helps alone; (4) the combination is consistently strongest across GRPO/GiGPO hosts, ALFWorld/WebShop/seven QA sets, and three model scales. The motivation is quantified: the longest 25% of trajectories carry 60.3% (ALFWorld) and 42.2% (WebShop) of token-weighted aggregation mass (§1), and the authors appropriately note this imbalance does not by itself establish the behavioral cause of failure.
The remaining sections cover Methodology Assessment, Statistical Evidence, Results & Reproducibility, Limitations & Risks, and Verdict & Recommendations.
- All 4 strengths & 5 weaknesses, analyzed in detail
- 4 risk flags with severity, section & reviewer questions
- Methods & statistical evidence assessment
- 5 questions a peer reviewer would ask
- Full six-dimension evaluation report & recommendations
One-time purchase · this paper only
One-time purchase, this paper only. AI-generated preliminary analysis — not peer review.
Frequently asked questions
What is BATON?
BATON (Bayesian Attribution and Trajectory Objective Normalization) is a reinforcement-learning fine-tuning framework for LLM agents. For each decision inside a trajectory it compares the executed action with one sampled counterfactual action using the model's likelihood of the observed environment feedback, and uses the resulting evidence score to reweight per-token policy losses within the trajectory. Separately, it gives every complete trajectory equal weight (1/B) in the batch objective instead of weighting by token count. It plugs into existing optimizers such as GRPO and GiGPO without changing their rewards, advantages, or rollout procedures.
How much does BATON improve performance?
In the paper's Qwen2.5 experiments, BATON lifts ALFWorld success rate from 72.9% to 85.6% (GRPO, 1.5B) and from 86.8% to 94.1% (GiGPO), WebShop success rate from 57.1% to 68.2% and from 67.1% to 76.4%, and adds roughly 1.2–3.2 points to the average of seven search-augmented QA benchmarks. Results are reported as means with standard deviations, but the paper states no number of runs and performs no significance tests.
Does BATON slow down training?
It adds no extra environment interaction. Bayesian Feedback Attribution requires additional feedback-likelihood evaluations; Trajectory Mass Normalization is only an aggregation change. The paper's Table 5 reports wall-clock overhead of 1.17× for GRPO and 1.10× for GiGPO on ALFWorld with a 1.5B model, with TMN alone costing about 1.01–1.02×.
Is code available for BATON?
The reviewed paper text contains no code or data release statement. The method is specified mathematically and uses public environments (ALFWorld, WebShop, search-augmented QA) and public Qwen2.5 checkpoints, so reproducing it requires reimplementation plus substantial GPU training — the authors' own overhead table implies 18–25 hours per training run at 1.5B scale on ALFWorld.
About this review
This is an AI-generated preliminary review by Evolve.Pub, produced from the openly available text of the paper. It is not peer review and does not represent the authors' views. The canonical record of the paper is at the publisher link above.
- Figures 1-2 are recovered only as garbled layout text; the diagrams themselves were not analyzable, though all numeric tables and Figure 3's values extract cleanly.
- Training hyperparameters (learning rate, batch size, rollout group size, total steps) are not present in the extracted text, limiting reproducibility assessment.
- The EnvRL baseline (Wang et al., 2026) is only cited, not described, so its comparability cannot be fully judged from this text.
Go deeper on this paper
Wrote a paper? Stress-test it before reviewers do.
Upload your manuscript for the same structured scrutiny — reasoning, statistics, evidence, figures, and reproducibility.