Dual-Axis Policy Optimization for LLM Agents: Bayesian Feedback Attribution and Trajectory Mass Normalization
Summary
The paper introduces BATON, a dual-axis policy optimization framework for LLM agents using Bayesian Feedback Attribution and Trajectory Mass Normalization, demonstrating improved performance in reinforcement learning experiments.
View Cached Full Text
Cached at: 09/18/26, 09:25 AM
# Dual-Axis Policy Optimization for LLM Agents: Bayesian Feedback Attribution and Trajectory Mass Normalization
Source: [https://arxiv.org/html/2609.19830](https://arxiv.org/html/2609.19830)
Binhe YuAffiliation:Zhejiang UniversityJingxiao YangAffiliation:Zhejiang UniversityRuopei SunAffiliation:University of Science and Technology of ChinaZiting liAffiliation:University of New South WalesCheng TanAffiliation:Zhejiang UniversityXuhong ZhangAffiliation:Zhejiang UniversityJianwei YinAffiliation:Zhejiang UniversityJintao ChenAffiliation:Zhejiang UniversityAffiliation:Shanghai Artificial Intelligence Laboratory\*Equal contribution
###### Abstract
Reinforcement learning for LLM agents involves two distinct optimization dimensions: how environment feedback is exploited within a trajectory, and how complete trajectories are aggregated across a batch\. We formulate these dimensions asIntra\-Trajectory Feedback AttributionandInter\-Trajectory Objective Aggregation, and introduceBATON\(BayesianAttribution andTrajectoryObjectiveNormalization\), 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 Trajectory Mass Normalization \(TMN\), which assigns equal optimization mass to complete trajectories\. Experiments with GRPO and GiGPO on ALFWorld, WebShop, and SearchQA show that both axes provide independent gains and that their combination consistently achieves the strongest overall performance across model scales\.
## 1Introduction
Reinforcement learning is increasingly used to train LLM agents that reason, call tools, and act through multi\-turn interaction rather than produce a single response\. Prompted and tool\-augmented agents already operate in embodied, web, software, and conversational environments\([Yao et al\., 2023](https://arxiv.org/html/2609.19830#bib.bib21);[Schick et al\., 2023](https://arxiv.org/html/2609.19830#bib.bib26);[Shridhar et al\., 2021](https://arxiv.org/html/2609.19830#bib.bib38);[Yao et al\., 2022](https://arxiv.org/html/2609.19830#bib.bib39);[Wang et al\., 2022](https://arxiv.org/html/2609.19830#bib.bib40);[Yang et al\., 2024](https://arxiv.org/html/2609.19830#bib.bib37)\), while recent benchmarks expose increasingly long and heterogeneous interaction trajectories\([Liu et al\., 2023](https://arxiv.org/html/2609.19830#bib.bib30);[Ma et al\., 2024](https://arxiv.org/html/2609.19830#bib.bib31);[Zhou et al\., 2023b](https://arxiv.org/html/2609.19830#bib.bib32);[Deng et al\., 2023](https://arxiv.org/html/2609.19830#bib.bib33);[Drouin et al\., 2024](https://arxiv.org/html/2609.19830#bib.bib34);[Lù et al\., 2024](https://arxiv.org/html/2609.19830#bib.bib35);[Yao et al\., 2024](https://arxiv.org/html/2609.19830#bib.bib36)\)\. Beyond prompting and supervised trajectory tuning\([Zeng et al\., 2023](https://arxiv.org/html/2609.19830#bib.bib28)\), recent work increasingly optimizes agent policies directly with reinforcement learning\([Qi et al\., 2024](https://arxiv.org/html/2609.19830#bib.bib17);[Wang et al\., 2025](https://arxiv.org/html/2609.19830#bib.bib16);[Feng et al\., 2025](https://arxiv.org/html/2609.19830#bib.bib1);[Zeng et al\., 2025a](https://arxiv.org/html/2609.19830#bib.bib19);[Luo et al\., 2025](https://arxiv.org/html/2609.19830#bib.bib18)\)\.
Policy optimization for LLMs has rapidly evolved from PPO\-based RLHF\([Schulman et al\., 2017](https://arxiv.org/html/2609.19830#bib.bib7);[Ouyang et al\., 2022](https://arxiv.org/html/2609.19830#bib.bib10)\)toward critic\-free and group\-based objectives such as RLOO and GRPO\([Ahmadian et al\., 2024](https://arxiv.org/html/2609.19830#bib.bib8);[Shao et al\., 2024](https://arxiv.org/html/2609.19830#bib.bib2)\), together with refinements to clipping, reduction, and sequence\-level optimization\([Yu et al\., 2025](https://arxiv.org/html/2609.19830#bib.bib5);[Liu et al\., 2025](https://arxiv.org/html/2609.19830#bib.bib6);[Hu, 2025](https://arxiv.org/html/2609.19830#bib.bib11);[Zheng et al\., 2025a](https://arxiv.org/html/2609.19830#bib.bib12)\)\. In multi\-turn agents, however, optimization acts on structured trajectories in which each action is followed by an environment observation\. This exposes two conceptually distinct questions: how observed feedback should affect learning signals*within*a trajectory, and how complete trajectories should contribute to the objective*across*a batch\. We call these dimensionsIntra\-Trajectory Feedback AttributionandInter\-Trajectory Objective Aggregation\.
This distinction becomes especially visible in variable\-length interaction\. As shown in Figure[1](https://arxiv.org/html/2609.19830#S1.F1), unsuccessful trajectories are often substantially longer than successful ones\. Inspection of prolonged failures reveals repeated ineffective actions, misread state changes, and actions taken before required constraints are satisfied—behaviors consistent with insufficient exploitation of environment feedback\. Meanwhile, global valid\-token averaging creates a separate optimization effect: trajectories containing more valid policy tokens receive proportionally larger scalar mass\. In our training data, the longest 25% of trajectories account for 60\.3% of the aggregation mass in ALFWorld and 42\.2% in WebShop\. Thus, prolonged failures expose two distinct learning challenges: feedback may be insufficiently exploited locally, while long trajectories may simultaneously dominate aggregation globally\. The measured aggregation imbalance does not by itself establish the behavioral cause of failure; it motivates treating the two dimensions separately\.
Figure 1:Motivation of BATON\. Flat token\-level aggregation gives longer trajectories disproportionate optimization mass, motivating separate treatment of feedback attribution and trajectory aggregation\.We introduceBATON\(BayesianAttribution andTrajectoryObjectiveNormalization\), a dual\-axis policy optimization framework\. For Intra\-Trajectory Feedback Attribution, BATON uses Bayesian Feedback Attribution\. Unlike verbal reflection or generic self\-correction\([Shinn et al\., 2023](https://arxiv.org/html/2609.19830#bib.bib22);[Madaan et al\., 2023](https://arxiv.org/html/2609.19830#bib.bib23);[Gou et al\., 2024](https://arxiv.org/html/2609.19830#bib.bib24);[Zhou et al\., 2023a](https://arxiv.org/html/2609.19830#bib.bib25);[Wang et al\., 2023](https://arxiv.org/html/2609.19830#bib.bib27)\), it compares actions sampled under the same pre\-action context by combining their empirical frequencies with the likelihood of the observed next environment response\. The resulting feedback\-conditioned posterior determines a normalized attribution weight that reallocates the host learner’s policy signal across decisions within each trajectory; we further show that its expected attribution gain is governed by the distinguishability of action\-conditioned feedback distributions\. For Inter\-Trajectory Objective Aggregation, Trajectory Mass Normalization \(TMN\) averages valid learning terms within each complete trajectory before averaging trajectories, replacing token\-proportional outer mass with a uniform1/B1/Bshare\. Both components leave the host learner’s reward, advantage estimator, clipping rule, and rollout procedure unchanged\.
We evaluate BATON with GRPO and GiGPO on ALFWorld, WebShop, and search\-augmented question answering, where multi\-turn search policies have recently been optimized with outcome\- and step\-level reinforcement learning\([Jin et al\., 2025](https://arxiv.org/html/2609.19830#bib.bib4);[Sun et al\., 2025](https://arxiv.org/html/2609.19830#bib.bib20);[Zheng et al\., 2025b](https://arxiv.org/html/2609.19830#bib.bib9)\)\. Across model scales, both axes yield independent gains and their combination is consistently strongest\. Improvements also persist with GiGPO, which already provides finer\-grained within\-trajectory credit assignment, and are accompanied by shorter evaluation trajectories and fewer generated policy tokens\.
## 2Preliminaries
#### Problem setup\.
We consider a language\-model agent interacting with an environment to complete multi\-step tasks\. At decision steptt, the agent conditions on the interaction historyht=\(o0,a0,…,ot\)h\_\{t\}=\(o\_\{0\},a\_\{0\},\\ldots,o\_\{t\}\)and generates an actionat∼πθ\(⋅∣ht\)a\_\{t\}\\sim\\pi\_\{\\theta\}\(\\cdot\\mid h\_\{t\}\), after which the environment returns the next observationot\+1o\_\{t\+1\}\. A complete trajectoryτi\\tau\_\{i\}containsTiT\_\{i\}decision steps and receives a task returnRiR\_\{i\}\. Because each textual action may contain multiple policy tokens, we letni,tn\_\{i,t\}denote the number of valid policy tokens at stepttandNi=∑t=1Tini,tN\_\{i\}=\\sum\_\{t=1\}^\{T\_\{i\}\}n\_\{i,t\}the total number of valid policy tokens in trajectoryτi\\tau\_\{i\}\. We optimize over a batch ofBBcomplete trajectories\.
#### Group\-based policy optimization\.
Recent agent RL methods such as GRPO and GiGPO optimize groups of sampled trajectories without requiring a separate value model\. For GRPO, each trajectoryτi\\tau\_\{i\}is assigned a group\-relative advantageA^i\\hat\{A\}\_\{i\}based on the returns of sibling rollouts\. Letri,t,k\(θ\)=πθ\(yi,t,k∣hi,t,k\)/πθold\(yi,t,k∣hi,t,k\)r\_\{i,t,k\}\(\\theta\)=\\pi\_\{\\theta\}\(y\_\{i,t,k\}\\mid h\_\{i,t,k\}\)/\\pi\_\{\\theta\_\{\\rm old\}\}\(y\_\{i,t,k\}\\mid h\_\{i,t,k\}\)denote the likelihood ratio of thekk\-th valid policy token\. Under global valid\-policy\-token averaging, the policy objective is
ℒGRPO\(θ\)=−1S∑i=1B∑t=1Ti∑k=1ni,tmin\[ri,t,k\(θ\)A^i,clip\(ri,t,k\(θ\),1−ϵ,1\+ϵ\)A^i\],\\mathcal\{L\}\_\{\\rm GRPO\}\(\\theta\)=\-\\frac\{1\}\{S\}\\sum\_\{i=1\}^\{B\}\\sum\_\{t=1\}^\{T\_\{i\}\}\\sum\_\{k=1\}^\{n\_\{i,t\}\}\\min\\\!\\left\[r\_\{i,t,k\}\(\\theta\)\\hat\{A\}\_\{i\},\\,\\operatorname\{clip\}\\\!\\left\(r\_\{i,t,k\}\(\\theta\),1\-\\epsilon,1\+\\epsilon\\right\)\\hat\{A\}\_\{i\}\\right\],\\quad\(1\)whereS=∑i=1BNiS=\\sum\_\{i=1\}^\{B\}N\_\{i\}\. We omit separately reduced regularization terms for clarity\.
#### Flat trajectory aggregation\.
More generally, letℓi,t,k𝒜\\ell^\{\\mathcal\{A\}\}\_\{i,t,k\}denote the local policy\-surrogate loss of a host learner𝒜\\mathcal\{A\}, and letℓ¯i𝒜=Ni−1∑t,kℓi,t,k𝒜\\bar\{\\ell\}^\{\\mathcal\{A\}\}\_\{i\}=N\_\{i\}^\{\-1\}\\sum\_\{t,k\}\\ell^\{\\mathcal\{A\}\}\_\{i,t,k\}be its trajectory\-level mean\. A generic trajectory\-weighted objective can be written as
𝒥𝒜\(θ,𝐪\)=∑i=1Bqiℓ¯i𝒜\(θ\),∑i=1Bqi=1,qiflat=Ni∑jNj\.\\mathcal\{J\}\_\{\\mathcal\{A\}\}\(\\theta;\\mathbf\{q\}\)=\\sum\_\{i=1\}^\{B\}q\_\{i\}\\bar\{\\ell\}^\{\\mathcal\{A\}\}\_\{i\}\(\\theta\),\\qquad\\sum\_\{i=1\}^\{B\}q\_\{i\}=1,\\qquad q\_\{i\}^\{\\rm flat\}=\\frac\{N\_\{i\}\}\{\\sum\_\{j\}N\_\{j\}\}\.\(2\)Thus, global valid\-token averaging assigns trajectoryτi\\tau\_\{i\}a scalar contribution proportional to its number of valid policy tokens\.
## 3Method
We proposeBATON, a dual\-axis policy optimization framework for multi\-turn LLM agents\. BATON separates policy optimization into two complementary dimensions:Intra\-Trajectory Feedback Attribution, which allocates learning mass among decisions within a trajectory, andInter\-Trajectory Objective Aggregation, which determines the scalar mass assigned to complete trajectories\. We instantiate the first axis with Bayesian Feedback Attribution \(BFA\) and the second with Trajectory Mass Normalization \(TMN\)\. Both components modify different levels of the objective and can be combined with host learners such as GRPO and GiGPO\.
Figure 2:Overview of BATON\. Bayesian Feedback Attribution evaluates feedback compatibility between visited decisions and observed environment responses, producing normalized stop\-gradient weights within trajectories\. Trajectory Mass Normalization controls the outer aggregation across complete trajectories\.### 3\.1Dual\-Axis Policy Optimization
For a batch of trajectoriesℬ=\{τi\}i=1B\\mathcal\{B\}=\\\{\\tau\_\{i\}\\\}\_\{i=1\}^\{B\}, we decompose policy optimization into within\-trajectory allocation and across\-trajectory aggregation:
ℒ\(θ\)=∑i=1Bqi1Ni∑t=1Ti∑k=1ni,tωi,tℓi,t,k𝒜\(θ\),\\mathcal\{L\}\(\\theta\)=\\sum\_\{i=1\}^\{B\}q\_\{i\}\\frac\{1\}\{N\_\{i\}\}\\sum\_\{t=1\}^\{T\_\{i\}\}\\sum\_\{k=1\}^\{n\_\{i,t\}\}\\omega\_\{i,t\}\\ell^\{\\mathcal\{A\}\}\_\{i,t,k\}\(\\theta\),\(3\)
whereqiq\_\{i\}controls trajectory\-level mass andωi,t\\omega\_\{i,t\}controls the allocation among decisions\. We require
1Ni∑tni,tωi,t=1,\\frac\{1\}\{N\_\{i\}\}\\sum\_\{t\}n\_\{i,t\}\\omega\_\{i,t\}=1,\(4\)
so that attribution redistributes mass inside trajectories without changing their total contribution\.
### 3\.2Bayesian Feedback Attribution
BFA estimates how strongly a visited decision is associated with its observed feedback\. The host learner remains responsible for the signed optimization direction; BFA only adjusts the relative contribution of different decisions\.
For a visited decision\(h,a,o′\)\(h,a,o^\{\\prime\}\), wherehhis the pre\-action context,aais the executed action, ando′o^\{\\prime\}is the observed environment feedback, BFA samples a counterfactual action
a~∼πθold\(⋅\|h\)\.\\tilde\{a\}\\sim\\pi\_\{\\theta\_\{\\rm old\}\}\(\\cdot\|h\)\.\(5\)
The executed and counterfactual actions form a pairwise candidate set\. For each candidate action, BFA evaluates the compatibility with the observed feedback:
logfθold\(o′\|h,a\)=∑m∈𝒱\(o′\)logπθold\(om′\|h,a,o<m′\),\\log f\_\{\\theta\_\{\\rm old\}\}\(o^\{\\prime\}\|h,a\)=\\sum\_\{m\\in\\mathcal\{V\}\(o^\{\\prime\}\)\}\\log\\pi\_\{\\theta\_\{\\rm old\}\}\(o^\{\\prime\}\_\{m\}\|h,a,o^\{\\prime\}\_\{<m\}\),\(6\)
where𝒱\(o′\)\\mathcal\{V\}\(o^\{\\prime\}\)denotes valid feedback\-target token positions\. The likelihood measures model\-based feedback compatibility rather than causal effects of counterfactual execution\.
#### Bayesian evidence\.
Under the pairwise construction,ρ0\(a\|h\)=ρ0\(a~\|h\)=1/2\\rho\_\{0\}\(a\|h\)=\\rho\_\{0\}\(\\tilde\{a\}\|h\)=1/2\. The feedback\-conditioned posterior of the executed action is
ρθold\(a\|h,o′\)=fθold\(o′\|h,a\)fθold\(o′\|h,a\)\+fθold\(o′\|h,a~\)\.\\rho\_\{\\theta\_\{\\rm old\}\}\(a\|h,o^\{\\prime\}\)=\\frac\{f\_\{\\theta\_\{\\rm old\}\}\(o^\{\\prime\}\|h,a\)\}\{f\_\{\\theta\_\{\\rm old\}\}\(o^\{\\prime\}\|h,a\)\+f\_\{\\theta\_\{\\rm old\}\}\(o^\{\\prime\}\|h,\\tilde\{a\}\)\}\.\(7\)
We define the Bayesian evidence ratio
e\(h,a,o′\)=ρθold\(a\|h,o′\)ρ0\(a\|h\)=2fθold\(o′\|h,a\)fθold\(o′\|h,a\)\+fθold\(o′\|h,a~\)\.e\(h,a,o^\{\\prime\}\)=\\frac\{\\rho\_\{\\theta\_\{\\rm old\}\}\(a\|h,o^\{\\prime\}\)\}\{\\rho\_\{0\}\(a\|h\)\}=\\frac\{2f\_\{\\theta\_\{\\rm old\}\}\(o^\{\\prime\}\|h,a\)\}\{f\_\{\\theta\_\{\\rm old\}\}\(o^\{\\prime\}\|h,a\)\+f\_\{\\theta\_\{\\rm old\}\}\(o^\{\\prime\}\|h,\\tilde\{a\}\)\}\.\(8\)
###### Proposition 1\(Bayesian Evidence Interpretation\)\.
Let
m\(o′\|h\)=12fθold\(o′\|h,a\)\+12fθold\(o′\|h,a~\)m\(o^\{\\prime\}\|h\)=\\frac\{1\}\{2\}f\_\{\\theta\_\{\\rm old\}\}\(o^\{\\prime\}\|h,a\)\+\\frac\{1\}\{2\}f\_\{\\theta\_\{\\rm old\}\}\(o^\{\\prime\}\|h,\\tilde\{a\}\)\(9\)
denote the pairwise feedback mixture\. The Bayesian evidence satisfies
e\(h,a,o′\)=fθold\(o′\|h,a\)m\(o′\|h\)\.e\(h,a,o^\{\\prime\}\)=\\frac\{f\_\{\\theta\_\{\\rm old\}\}\(o^\{\\prime\}\|h,a\)\}\{m\(o^\{\\prime\}\|h\)\}\.\(10\)
Moreover,
𝔼o′∼fθold\(⋅\|h,a\)\[loge\(h,a,o′\)\]=DKL\(fθold\(⋅\|h,a\)∥m\(⋅\|h\)\)≥0\.\\mathbb\{E\}\_\{o^\{\\prime\}\\sim f\_\{\\theta\_\{\\rm old\}\}\(\\cdot\|h,a\)\}\[\\log e\(h,a,o^\{\\prime\}\)\]=D\_\{\\mathrm\{KL\}\}\(f\_\{\\theta\_\{\\rm old\}\}\(\\cdot\|h,a\)\\\|m\(\\cdot\|h\)\)\\geq 0\.\(11\)
Proposition[1](https://arxiv.org/html/2609.19830#Thmproposition1)characterizes the statistical quantity underlying BFA evidence: larger evidence corresponds to greater distinguishability between the executed action’s feedback distribution and its policy\-sampled alternative\. The proof is provided in Appendix[A](https://arxiv.org/html/2609.19830#A1)\.
#### Length\-calibrated attribution\.
For decision\(i,t\)\(i,t\), let
ei,t=e\(hi,t,ai,t,oi,t\+1\)\.e\_\{i,t\}=e\(h\_\{i,t\},a\_\{i,t\},o\_\{i,t\+1\}\)\.
Since the feedback likelihood is accumulated over tokens, we use geometric evidence
ri,t=ei,t1/Li,t,r\_\{i,t\}=e\_\{i,t\}^\{1/L\_\{i,t\}\},\(12\)
whereLi,t=\|𝒱\(oi,t\+1\)\|L\_\{i,t\}=\|\\mathcal\{V\}\(o\_\{i,t\+1\}\)\|\. The normalized attribution weight is
ωi,t=ri,t1Ni∑s=1Tini,sri,s\.\\omega\_\{i,t\}=\\frac\{r\_\{i,t\}\}\{\\frac\{1\}\{N\_\{i\}\}\\sum\_\{s=1\}^\{T\_\{i\}\}n\_\{i,s\}r\_\{i,s\}\}\.\(13\)
This normalization gives
1Ni∑tni,tωi,t=1,\\frac\{1\}\{N\_\{i\}\}\\sum\_\{t\}n\_\{i,t\}\\omega\_\{i,t\}=1,\(14\)
preserving each trajectory’s total optimization mass\.
BFA applies detached weights to the host objective:
ℓ¯i𝒜,BFA\(θ\)=1Ni∑t,ksg\[ωi,t\]ℓi,t,k𝒜\(θ\)\.\\bar\{\\ell\}^\{\\mathcal\{A\},\\mathrm\{BFA\}\}\_\{i\}\(\\theta\)=\\frac\{1\}\{N\_\{i\}\}\\sum\_\{t,k\}\\operatorname\{sg\}\[\\omega\_\{i,t\}\]\\ell^\{\\mathcal\{A\}\}\_\{i,t,k\}\(\\theta\)\.\(15\)
Therefore, BFA only rescales host learner updates and does not introduce an independent optimization direction\.
### 3\.3Trajectory Mass Normalization
BFA controls allocation within trajectories, while TMN controls aggregation across trajectories\.
Under global valid\-token averaging, the trajectory coefficient is
qiflat=Ni∑jNj\.q\_\{i\}^\{\\rm flat\}=\\frac\{N\_\{i\}\}\{\\sum\_\{j\}N\_\{j\}\}\.\(16\)
Thus, longer trajectories receive larger objective mass\. TMN instead assigns equal mass to complete trajectories:
qiTMN=1B\.q\_\{i\}^\{\\rm TMN\}=\\frac\{1\}\{B\}\.\(17\)
Because BFA preserves the within\-trajectory mass constraint, the two components are independent: BFA reallocates mass inside trajectories, whereas TMN changes only the outer trajectory aggregation\.
### 3\.4Unified BATON Objective
Combining BFA and TMN gives
ℒBATON\(θ\)=1B∑i=1B1Ni∑t=1Ti∑k=1ni,tsg\[ωi,t\]ℓi,t,k𝒜\(θ\)\.\\mathcal\{L\}\_\{\\rm BATON\}\(\\theta\)=\\frac\{1\}\{B\}\\sum\_\{i=1\}^\{B\}\\frac\{1\}\{N\_\{i\}\}\\sum\_\{t=1\}^\{T\_\{i\}\}\\sum\_\{k=1\}^\{n\_\{i,t\}\}\\operatorname\{sg\}\[\\omega\_\{i,t\}\]\\ell^\{\\mathcal\{A\}\}\_\{i,t,k\}\(\\theta\)\.\(18\)
BFA determines the normalized within\-trajectory allocationωi,t\\omega\_\{i,t\}, while TMN determines the uniform outer trajectory mass1/B1/B\. This formulation yields the controlled2×22\\times 2decomposition in our ablation study: the host learner uses uniform within\-trajectory weights and token\-proportional aggregation; BFA modifies only within\-trajectory allocation; TMN modifies only trajectory\-level aggregation; BATON combines both\.
Table 1:Main results on ALFWorld and WebShop\. We report success rate \(%\) for each ALFWorld task domain and the overall result, and task score and success rate \(%\) on WebShop\. Within each backbone, bold marks the highest mean in each metric column separately within the GRPO\-based and GiGPO\-based method groups\.
## 4Experiments
We evaluate the proposed dual\-axis framework from three perspectives: overall task performance, the independent and joint contributions of its two optimization axes, and the execution behavior of the learned policies\.
### 4\.1Experimental Setup
We evaluate Qwen2\.5\-Instruct models at 1\.5B and 7B scales on ALFWorld and WebShop, and at 3B and 7B scales on search\-augmented question answering\. We report success rate on ALFWorld, task score and success rate on WebShop, and exact match on seven search\-augmented QA datasets together with their unweighted average\. We instantiate the framework with both GRPO and GiGPO as host learners and compare against prompting and reinforcement\-learning baselines, including EnvRL\. Matched variants share the same initial checkpoint, rollout budget, reward and advantage construction, optimization settings, and evaluation protocol, differing only in the component under study\.
### 4\.2Main Results
Table[1](https://arxiv.org/html/2609.19830#S3.T1)shows that the full dual\-axis framework consistently improves both GRPO and GiGPO across ALFWorld and WebShop and at both model scales\. With Qwen2\.5\-1\.5B\-Instruct, the ALFWorld success rate increases from 72\.9% to 85\.6% for GRPO and from 86\.8% to 94\.1% for GiGPO\. On WebShop, the corresponding success rates improve from 57\.1% to 68\.2% and from 67\.1% to 76\.4%, respectively\. The gains remain at the 7B scale: the framework improves ALFWorld from 77\.7% to 86\.9% with GRPO and from 90\.9% to 94\.9% with GiGPO, while also improving their WebShop success rates from 66\.3% to 70\.3% and from 74\.9% to 79\.5%\. These results indicate that jointly improving feedback attribution within trajectories and objective aggregation across trajectories benefits host learners with substantially different local credit constructions\. In particular, the improvements over GiGPO suggest that inter\-trajectory calibration and feedback\-grounded supervision remain useful even when the host learner already employs finer\-grained within\-trajectory credit signals\.
### 4\.3Search\-Augmented Question Answering
We further evaluate BATON on search\-augmented question answering, where agents interact with external retrieval systems during multi\-step reasoning\. Table[2](https://arxiv.org/html/2609.19830#S4.T2)reports the averaged exact match over seven QA domains\. BATON improves both GRPO and GiGPO across model scales\. Compared with the corresponding host learners, BATON improves GRPO from 37\.9% to 41\.1% on Qwen2\.5\-3B and from 42\.8% to 45\.9% on Qwen2\.5\-7B\. Similarly, BATON improves GiGPO from 42\.0% to 43\.2% and from 47\.1% to 48\.7%, respectively\. The results indicate that the proposed dual\-axis optimization remains effective in retrieval\-augmented agent settings\. Complete domain\-level results are reported in Appendix[C](https://arxiv.org/html/2609.19830#A3)\.
Table 2:Main results on search\-augmented question answering\. We report exact match \(%\) on seven evaluation domains and their unweighted average\. NQ, TriviaQA, and PopQA are single\-hop domains; HotpotQA, 2WikiMultihopQA, MuSiQue, and Bamboogle are multi\-hop domains\. Within each backbone, bold marks the highest mean in each metric column separately within the GRPO\-based and GiGPO\-based method groups\.
### 4\.4Component Ablation
Table[3](https://arxiv.org/html/2609.19830#S4.T3)provides a controlled2×22\\times 2decomposition of BATON: the host learner alone, Bayesian Feedback Attribution only, TMN only, and their combination\. Both axes yield substantial gains independently\. On ALFWorld with Qwen2\.5\-1\.5B\-Instruct, GRPO improves from 72\.9% to 84\.1% with Bayesian Feedback Attribution and to 82\.4% with TMN, while BATON reaches 85\.6%\. For GiGPO, the corresponding results are 86\.8%, 93\.4%, 90\.7%, and 94\.1%\. The same complementarity holds across tasks and scales\. On WebShop with the 7B backbone, BATON reaches 70\.3% success with GRPO and 79\.5% with GiGPO, outperforming either component alone\. These results support the dual\-axis formulation: the two components are independently effective and provide complementary gains when combined\.
Table 3:Component ablation of BATON’s two optimization axes\. BFA denotes Bayesian Feedback Attribution for Intra\-Trajectory Feedback Attribution, while TMN denotes Trajectory Mass Normalization for Inter\-Trajectory Objective Aggregation\. Values are reported as means with standard deviations\.
### 4\.5Trajectory Execution Cost
We further examine whether the full dual\-axis framework changes the execution behavior of the learned policies on ALFWorld with Qwen2\.5\-1\.5B\-Instruct\. Figure[3](https://arxiv.org/html/2609.19830#S4.F3)reports the mean number of environment actions and generated policy tokens per evaluation trajectory, including both successful and failed episodes\. For GRPO, the full framework reduces the mean action count from 21\.51 to 16\.39 and the mean generated\-token count from 1,734 to 1,301\. For GiGPO, the corresponding quantities decrease from 13\.51 to 10\.82 actions and from 1,431 to 746 generated tokens\. These results indicate that the performance improvements are accompanied by shorter and less verbose interaction trajectories, suggesting that the learned policies make more efficient use of their interaction budget\.
#### Computational overhead\.
BATON introduces no additional environment interaction\. The additional computation comes from Bayesian Feedback Attribution, which performs feedback likelihood evaluation for attribution weighting, while TMN only modifies the trajectory\-level aggregation rule\. The wall\-clock overhead analysis is provided in Appendix[B](https://arxiv.org/html/2609.19830#A2)\.
Figure 3:Trajectory execution cost on ALFWorld with Qwen2\.5\-1\.5B\-Instruct\. We report mean environment\-action counts and generated policy\-token counts for the host learners, EnvRL variants, and the full dual\-axis framework at update 150, including both successful and failed episodes\. Lower is better\.
## 5Related Work
#### Policy optimization for LLMs and agents\.
PPO\-based RLHF provides the foundation for LLM policy optimization\([Schulman et al\., 2017](https://arxiv.org/html/2609.19830#bib.bib7);[Ouyang et al\., 2022](https://arxiv.org/html/2609.19830#bib.bib10)\), while RLOO and GRPO introduce critic\-free and group\-relative alternatives\([Ahmadian et al\., 2024](https://arxiv.org/html/2609.19830#bib.bib8);[Shao et al\., 2024](https://arxiv.org/html/2609.19830#bib.bib2)\)\. Recent reasoning\-oriented methods improve optimization through modified clipping, reduction, process supervision, and sequence\-level objectives, including DAPO, Dr\. GRPO, REINFORCE\+\+, PRIME, Open\-Reasoner\-Zero, SimpleRL\-Zoo, GSPO, MGAP, TASPO and SARE\([Yu et al\., 2025](https://arxiv.org/html/2609.19830#bib.bib5);[Liu et al\., 2025](https://arxiv.org/html/2609.19830#bib.bib6);[Hu, 2025](https://arxiv.org/html/2609.19830#bib.bib11);[Cui et al\., 2025](https://arxiv.org/html/2609.19830#bib.bib13);[Hu et al\., 2025](https://arxiv.org/html/2609.19830#bib.bib14);[Zeng et al\., 2025b](https://arxiv.org/html/2609.19830#bib.bib15);[Zheng et al\., 2025a](https://arxiv.org/html/2609.19830#bib.bib12);[Yang et al\., 2026a](https://arxiv.org/html/2609.19830#bib.bib43);[Yang et al\., 2026b](https://arxiv.org/html/2609.19830#bib.bib42);[Zhuang et al\., 2026](https://arxiv.org/html/2609.19830#bib.bib41)\)\. For interactive agents, WebRL and RAGEN study reinforcement learning in environmental settings\([Qi et al\., 2024](https://arxiv.org/html/2609.19830#bib.bib17);[Wang et al\., 2025](https://arxiv.org/html/2609.19830#bib.bib16)\), while GiGPO, turn\-level credit assignment, and Agent Lightning investigate finer\-grained trajectory optimization\([Feng et al\., 2025](https://arxiv.org/html/2609.19830#bib.bib1);[Zeng et al\., 2025a](https://arxiv.org/html/2609.19830#bib.bib19);[Luo et al\., 2025](https://arxiv.org/html/2609.19830#bib.bib18)\)\. BATON is complementary to these approaches by separating within\-trajectory feedback attribution from across\-trajectory objective aggregation while preserving the host learner’s optimization structure\.
#### Learning from interaction feedback\.
Interaction feedback has long been used to ground language agents through actions and subsequent observations\([Yao et al\., 2023](https://arxiv.org/html/2609.19830#bib.bib21);[Schick et al\., 2023](https://arxiv.org/html/2609.19830#bib.bib26)\)\. Reflexion, Self\-Refine, CRITIC, LATS, and Voyager exploit verbal, tool\-based, or environmental feedback for iterative improvement\([Shinn et al\., 2023](https://arxiv.org/html/2609.19830#bib.bib22);[Madaan et al\., 2023](https://arxiv.org/html/2609.19830#bib.bib23);[Gou et al\., 2024](https://arxiv.org/html/2609.19830#bib.bib24);[Zhou et al\., 2023a](https://arxiv.org/html/2609.19830#bib.bib25);[Wang et al\., 2023](https://arxiv.org/html/2609.19830#bib.bib27)\)\. AgentTuning and AgentGym further study learning from interaction trajectory collections\([Zeng et al\., 2023](https://arxiv.org/html/2609.19830#bib.bib28);[Xi et al\., 2024](https://arxiv.org/html/2609.19830#bib.bib29)\), while EnvRL introduces state\-prediction and inverse\-dynamics objectives for agent reinforcement learning\([Wang et al\., 2026](https://arxiv.org/html/2609.19830#bib.bib3)\)\. Different from these approaches, Bayesian Feedback Attribution compares the observed feedback compatibility of executed and policy\-sampled alternative actions, and uses the resulting evidence to reweight the host policy objective without replacing its reward or credit estimator\.
#### Interactive agent environments and search\.
Long\-horizon agent learning has been studied in embodied and text environments including ALFWorld, WebShop, and ScienceWorld\([Shridhar et al\., 2021](https://arxiv.org/html/2609.19830#bib.bib38);[Yao et al\., 2022](https://arxiv.org/html/2609.19830#bib.bib39);[Wang et al\., 2022](https://arxiv.org/html/2609.19830#bib.bib40)\), as well as broader evaluation suites such as AgentBench, AgentBoard, WebArena, Mind2Web, WorkArena, and WebLINX\([Liu et al\., 2023](https://arxiv.org/html/2609.19830#bib.bib30);[Ma et al\., 2024](https://arxiv.org/html/2609.19830#bib.bib31);[Zhou et al\., 2023b](https://arxiv.org/html/2609.19830#bib.bib32);[Deng et al\., 2023](https://arxiv.org/html/2609.19830#bib.bib33);[Drouin et al\., 2024](https://arxiv.org/html/2609.19830#bib.bib34);[Lù et al\., 2024](https://arxiv.org/html/2609.19830#bib.bib35)\)\. Recent tool\-use benchmarks and software environments further includeτ\\tau\-bench and SWE\-agent\([Yao et al\., 2024](https://arxiv.org/html/2609.19830#bib.bib36);[Yang et al\., 2024](https://arxiv.org/html/2609.19830#bib.bib37)\)\. Search\-R1, ZeroSearch, and StepSearch optimize multi\-turn retrieval agents with reinforcement learning\([Jin et al\., 2025](https://arxiv.org/html/2609.19830#bib.bib4);[Sun et al\., 2025](https://arxiv.org/html/2609.19830#bib.bib20);[Zheng et al\., 2025b](https://arxiv.org/html/2609.19830#bib.bib9)\)\. These environments motivate explicit treatment of long and heterogeneous interaction trajectories, which BATON addresses through separate attribution and aggregation mechanisms\.
## 6Conclusion
We presentBATON, a dual\-axis policy optimization framework for multi\-turn LLM agents\. BATON combines Bayesian Feedback Attribution, which reallocates learning mass among decisions within trajectories based on feedback compatibility, with Trajectory Mass Normalization, which balances contributions across trajectories\. The framework preserves the host learner’s optimization structure and is compatible with GRPO and GiGPO\. Experiments on ALFWorld, WebShop, and search\-augmented question answering demonstrate consistent improvements across model scales\.
## References
- Ahmadianet al\.\(2024\)A\. Ahmadian, C\. Cremer, M\. Gallé, M\. Fadaee, J\. Kreutzer, O\. Pietquin, A\. Üstün, and S\. HookerBack to basics: revisiting REINFORCE\-style optimization for learning from human feedback in LLMs\.InProceedings of the 62nd Annual Meeting of the Association for Computational Linguistics \(Volume 1: Long Papers\),pp\. 12248–12267\.External Links:[Document](https://dx.doi.org/10.18653/v1/2024.acl-long.662)Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p2.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px1.p1.1)\.
- Cuiet al\.\(2025\)G\. Cui, L\. Yuan, Z\. Wang, H\. Wang, W\. Li, B\. He, Y\. Fan, T\. Yu, Q\. Xu, W\. Chen, J\. Yuan, H\. Chen, K\. Zhang, X\. Lv, S\. Wang, Y\. Yao, X\. Han, H\. Peng, Y\. Cheng, Z\. Liu, M\. Sun, B\. Zhou, and N\. DingProcess reinforcement through implicit rewards\.External Links:2502\.01456,[Link](https://arxiv.org/abs/2502.01456)Cited by:[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px1.p1.1)\.
- Denget al\.\(2023\)X\. Deng, Y\. Gu, B\. Zheng, S\. Chen, S\. Stevens, B\. Wang, H\. Sun, and Y\. SuMind2Web: towards a generalist agent for the web\.External Links:2306\.06070,[Link](https://arxiv.org/abs/2306.06070)Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p1.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px3.p1.1)\.
- Drouinet al\.\(2024\)A\. Drouin, M\. Gasse, M\. Caccia, I\. H\. Laradji, M\. D\. Verme, T\. Marty, L\. Boisvert, M\. Thakkar, Q\. Cappart, D\. Vazquez, N\. Chapados, and A\. LacosteWorkArena: how capable are web agents at solving common knowledge work tasks?\.External Links:2403\.07718,[Link](https://arxiv.org/abs/2403.07718)Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p1.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px3.p1.1)\.
- Fenget al\.\(2025\)L\. Feng, Z\. Xue, T\. Liu, and B\. AnGroup\-in\-group policy optimization for llm agent training\.External Links:2505\.10978,[Link](https://arxiv.org/abs/2505.10978)Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p1.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px1.p1.1)\.
- Gouet al\.\(2024\)Z\. Gou, Z\. Shao, Y\. Gong, Y\. Shen, Y\. Yang, N\. Duan, and W\. ChenCRITIC: large language models can self\-correct with tool\-interactive critiquing\.InInternational Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=Sx038qxjek)Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p4.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px2.p1.1)\.
- Hu \(2025\)J\. HuREINFORCE\+\+: a simple and efficient approach for aligning large language models\.External Links:2501\.03262,[Link](https://arxiv.org/abs/2501.03262)Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p2.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px1.p1.1)\.
- Huet al\.\(2025\)J\. Hu, Y\. Zhang, Q\. Han, D\. Jiang, X\. Zhang, and H\. ShumOpen\-reasoner\-zero: an open source approach to scaling up reinforcement learning on the base model\.External Links:2503\.24290,[Link](https://arxiv.org/abs/2503.24290)Cited by:[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px1.p1.1)\.
- Jinet al\.\(2025\)B\. Jin, H\. Zeng, Z\. Yue, J\. Yoon, S\. Arik, D\. Wang, H\. Zamani, and J\. HanSearch\-R1: Training LLMs to Reason and Leverage Search Engines with Reinforcement Learning\.External Links:2503\.09516,[Link](https://arxiv.org/abs/2503.09516)Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p5.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px3.p1.1)\.
- Liuet al\.\(2023\)X\. Liu, H\. Yu, H\. Zhang, Y\. Xu, X\. Lei, H\. Lai, Y\. Gu, H\. Ding, K\. Men, K\. Yang, S\. Zhang, X\. Deng, A\. Zeng, Z\. Du, C\. Zhang, S\. Shen, T\. Zhang, Y\. Su, H\. Sun, M\. Huang, Y\. Dong, and J\. TangAgentBench: evaluating llms as agents\.External Links:2308\.03688,[Link](https://arxiv.org/abs/2308.03688)Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p1.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px3.p1.1)\.
- Liuet al\.\(2025\)Z\. Liu, C\. Chen, W\. Li, P\. Qi, T\. Pang, C\. Du, W\. S\. Lee, and M\. LinUnderstanding R1\-Zero\-Like Training: A Critical Perspective\.External Links:2503\.20783,[Link](https://arxiv.org/abs/2503.20783)Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p2.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px1.p1.1)\.
- Lùet al\.\(2024\)X\. H\. Lù, Z\. Kasner, and S\. ReddyWebLINX: real\-world website navigation with multi\-turn dialogue\.External Links:2402\.05930,[Link](https://arxiv.org/abs/2402.05930)Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p1.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px3.p1.1)\.
- Luoet al\.\(2025\)X\. Luo, Y\. Zhang, Z\. He, Z\. Wang, S\. Zhao, D\. Li, L\. K\. Qiu, and Y\. YangAgent lightning: train any ai agents with reinforcement learning\.External Links:2508\.03680,[Link](https://arxiv.org/abs/2508.03680)Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p1.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px1.p1.1)\.
- Maet al\.\(2024\)C\. Ma, J\. Zhang, Z\. Zhu, C\. Yang, Y\. Yang, Y\. Jin, Z\. Lan, L\. Kong, and J\. HeAgentBoard: an analytical evaluation board of multi\-turn llm agents\.External Links:2401\.13178,[Link](https://arxiv.org/abs/2401.13178)Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p1.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px3.p1.1)\.
- Madaanet al\.\(2023\)A\. Madaan, N\. Tandon, P\. Gupta, S\. Hallinan, L\. Gao, S\. Wiegreffe, U\. Alon, N\. Dziri, S\. Prabhumoye, Y\. Yang, S\. Gupta, B\. P\. Majumder, K\. Hermann, S\. Welleck, A\. Yazdanbakhsh, and P\. ClarkSelf\-refine: iterative refinement with self\-feedback\.External Links:2303\.17651,[Link](https://arxiv.org/abs/2303.17651)Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p4.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px2.p1.1)\.
- Ouyanget al\.\(2022\)L\. Ouyang, J\. Wu, X\. Jiang, D\. Almeida, C\. L\. Wainwright, P\. Mishkin, C\. Zhang, S\. Agarwal, K\. Slama, A\. Ray, J\. Schulman, J\. Hilton, F\. Kelton, L\. Miller, M\. Simens, A\. Askell, P\. Welinder, P\. Christiano, J\. Leike, and R\. LoweTraining language models to follow instructions with human feedback\.External Links:2203\.02155,[Link](https://arxiv.org/abs/2203.02155)Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p2.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px1.p1.1)\.
- Qiet al\.\(2024\)Z\. Qi, X\. Liu, I\. L\. Iong, H\. Lai, X\. Sun, W\. Zhao, Y\. Yang, X\. Yang, J\. Sun, S\. Yao, T\. Zhang, W\. Xu, J\. Tang, and Y\. DongWebRL: training llm web agents via self\-evolving online curriculum reinforcement learning\.External Links:2411\.02337,[Link](https://arxiv.org/abs/2411.02337)Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p1.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px1.p1.1)\.
- Schicket al\.\(2023\)T\. Schick, J\. Dwivedi\-Yu, R\. Dessì, R\. Raileanu, M\. Lomeli, L\. Zettlemoyer, N\. Cancedda, and T\. ScialomToolformer: language models can teach themselves to use tools\.External Links:2302\.04761,[Link](https://arxiv.org/abs/2302.04761)Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p1.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px2.p1.1)\.
- Schulmanet al\.\(2017\)J\. Schulman, F\. Wolski, P\. Dhariwal, A\. Radford, and O\. KlimovProximal policy optimization algorithms\.arXiv preprint arXiv:1707\.06347\.Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p2.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px1.p1.1)\.
- Shaoet al\.\(2024\)Z\. Shao, P\. Wang, Q\. Zhu, R\. Xu, J\. Song, X\. Bi, H\. Zhang, M\. Zhang, Y\. K\. Li, Y\. Wu, and D\. GuoDeepSeekMath: pushing the limits of mathematical reasoning in open language models\.External Links:2402\.03300,[Link](https://arxiv.org/abs/2402.03300)Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p2.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px1.p1.1)\.
- Shinnet al\.\(2023\)N\. Shinn, F\. Cassano, E\. Berman, A\. Gopinath, K\. Narasimhan, and S\. YaoReflexion: language agents with verbal reinforcement learning\.External Links:2303\.11366,[Link](https://arxiv.org/abs/2303.11366)Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p4.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px2.p1.1)\.
- Shridharet al\.\(2021\)M\. Shridhar, X\. Yuan, M\. Côté, Y\. Bisk, A\. Trischler, and M\. HausknechtALFWorld: aligning text and embodied environments for interactive learning\.InInternational Conference on Learning Representations,External Links:[Link](https://arxiv.org/abs/2010.03768)Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p1.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px3.p1.1)\.
- Sunet al\.\(2025\)H\. Sun, Z\. Qiao, J\. Guo, X\. Fan, Y\. Hou, Y\. Jiang, P\. Xie, F\. Huang, and Y\. ZhangZeroSearch: incentivize the search capability of llms without searching\.External Links:2505\.04588,[Link](https://arxiv.org/abs/2505.04588)Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p5.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px3.p1.1)\.
- Wanget al\.\(2023\)G\. Wang, Y\. Xie, Y\. Jiang, A\. Mandlekar, C\. Xiao, Y\. Zhu, L\. Fan, and A\. AnandkumarVoyager: an open\-ended embodied agent with large language models\.External Links:2305\.16291,[Link](https://arxiv.org/abs/2305.16291)Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p4.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px2.p1.1)\.
- Wanget al\.\(2022\)R\. Wang, P\. Jansen, M\. Côté, and P\. AmmanabroluScienceWorld: is your agent smarter than a 5th grader?\.External Links:2203\.07540,[Link](https://arxiv.org/abs/2203.07540)Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p1.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px3.p1.1)\.
- Wanget al\.\(2026\)Z\. Wang, S\. Li, H\. Peng, S\. Si, Y\. Wang, M\. Sun, and J\. LiEnvRL: learn from environment dynamics in agentic reinforcement learning\.External Links:2606\.17680,[Link](https://arxiv.org/abs/2606.17680)Cited by:[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px2.p1.1)\.
- Wanget al\.\(2025\)Z\. Wang, K\. Wang, Q\. Wang, P\. Zhang, L\. Li, Z\. Yang, K\. Yu, M\. N\. Nguyen, L\. Liu, E\. Gottlieb, M\. Lam, Y\. Lu, K\. Cho, J\. Wu, L\. Fei\-Fei, L\. Wang, Y\. Choi, and M\. LiRAGEN: understanding self\-evolution in llm agents via multi\-turn reinforcement learning\.External Links:2504\.20073,[Link](https://arxiv.org/abs/2504.20073)Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p1.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px1.p1.1)\.
- Xiet al\.\(2024\)Z\. Xi, Y\. Ding, W\. Chen, B\. Hong, H\. Guo, J\. Wang, D\. Yang, C\. Liao, X\. Guo, W\. He, S\. Gao, L\. Chen, R\. Zheng, Y\. Zou, T\. Gui, Q\. Zhang, X\. Qiu, X\. Huang, Z\. Wu, and Y\. JiangAgentGym: evolving large language model\-based agents across diverse environments\.External Links:2406\.04151,[Link](https://arxiv.org/abs/2406.04151)Cited by:[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px2.p1.1)\.
- Yanget al\.\(2026a\)J\. Yang, W\. Gan, Y\. Zhuang, W\. Zhang, J\. Chen, and X\. ZhangReconciling process supervision with outcome\-based credit in agentic policy optimization\.External Links:2608\.31077,[Link](https://arxiv.org/abs/2608.31077)Cited by:[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px1.p1.1)\.
- Yanget al\.\(2026b\)J\. Yang, D\. He, M\. Pan, K\. Yao, G\. Su, W\. Zhang, Y\. Hu, T\. Li, Y\. Li, and X\. ZhangSARE: sample\-wise adaptive reasoning for training\-free fine\-grained visual recognition\.External Links:2603\.17729,[Link](https://arxiv.org/abs/2603.17729)Cited by:[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px1.p1.1)\.
- Yanget al\.\(2024\)J\. Yang, C\. E\. Jimenez, A\. Wettig, K\. Lieret, S\. Yao, K\. Narasimhan, and O\. PressSWE\-agent: agent\-computer interfaces enable automated software engineering\.External Links:2405\.15793,[Link](https://arxiv.org/abs/2405.15793)Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p1.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px3.p1.1)\.
- Yaoet al\.\(2022\)S\. Yao, H\. Chen, J\. Yang, and K\. NarasimhanWebShop: towards scalable real\-world web interaction with grounded language agents\.External Links:2207\.01206,[Link](https://arxiv.org/abs/2207.01206)Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p1.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px3.p1.1)\.
- Yaoet al\.\(2024\)S\. Yao, N\. Shinn, P\. Razavi, and K\. Narasimhanτ\\tau\-bench: a benchmark for tool\-agent\-user interaction in real\-world domains\.External Links:2406\.12045,[Link](https://arxiv.org/abs/2406.12045)Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p1.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px3.p1.1)\.
- Yaoet al\.\(2023\)S\. Yao, J\. Zhao, D\. Yu, N\. Du, I\. Shafran, K\. Narasimhan, and Y\. CaoReAct: synergizing reasoning and acting in language models\.InInternational Conference on Learning Representations,External Links:[Link](https://arxiv.org/abs/2210.03629)Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p1.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px2.p1.1)\.
- Yuet al\.\(2025\)Q\. Yu, Z\. Zhang, R\. Zhu, Y\. Yuan, X\. Zuo, Y\. Yue, W\. Dai, T\. Fan, G\. Liu, L\. Liu, X\. Liu, H\. Lin, Z\. Lin, B\. Ma, G\. Sheng, Y\. Tong, C\. Zhang, M\. Zhang, W\. Zhang, H\. Zhu, J\. Zhu, J\. Chen, J\. Chen, C\. Wang, H\. Yu, Y\. Song, X\. Wei, H\. Zhou, J\. Liu, W\. Ma, Y\. Zhang, L\. Yan, M\. Qiao, Y\. Wu, and M\. WangDAPO: An Open\-Source LLM Reinforcement Learning System at Scale\.External Links:2503\.14476,[Link](https://arxiv.org/abs/2503.14476)Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p2.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px1.p1.1)\.
- Zenget al\.\(2023\)A\. Zeng, M\. Liu, R\. Lu, B\. Wang, X\. Liu, Y\. Dong, and J\. TangAgentTuning: enabling generalized agent abilities for llms\.External Links:2310\.12823,[Link](https://arxiv.org/abs/2310.12823)Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p1.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px2.p1.1)\.
- Zenget al\.\(2025a\)S\. Zeng, Q\. Wei, W\. Brown, O\. Frunza, Y\. Nevmyvaka, and M\. HongReinforcing multi\-turn reasoning in llm agents via turn\-level credit assignment\.External Links:2505\.11821,[Link](https://arxiv.org/abs/2505.11821)Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p1.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px1.p1.1)\.
- Zenget al\.\(2025b\)W\. Zeng, Y\. Huang, Q\. Liu, W\. Liu, K\. He, Z\. Ma, and J\. HeSimpleRL\-zoo: investigating and taming zero reinforcement learning for open base models in the wild\.External Links:2503\.18892,[Link](https://arxiv.org/abs/2503.18892)Cited by:[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px1.p1.1)\.
- Zhenget al\.\(2025a\)C\. Zheng, S\. Liu, M\. Li, X\. Chen, B\. Yu, C\. Gao, K\. Dang, Y\. Liu, R\. Men, A\. Yang, J\. Zhou, and J\. LinGroup sequence policy optimization\.External Links:2507\.18071,[Link](https://arxiv.org/abs/2507.18071)Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p2.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px1.p1.1)\.
- Zhenget al\.\(2025b\)X\. Zheng, K\. An, Z\. Wang, Y\. Wang, and Y\. WuStepSearch: igniting LLMs search ability via step\-wise proximal policy optimization\.InProceedings of the 2025 Conference on Empirical Methods in Natural Language Processing,pp\. 21805–21830\.External Links:[Document](https://dx.doi.org/10.18653/v1/2025.emnlp-main.1106)Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p5.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px3.p1.1)\.
- Zhouet al\.\(2023a\)A\. Zhou, K\. Yan, M\. Shlapentokh\-Rothman, H\. Wang, and Y\. WangLanguage agent tree search unifies reasoning acting and planning in language models\.External Links:2310\.04406,[Link](https://arxiv.org/abs/2310.04406)Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p4.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px2.p1.1)\.
- Zhouet al\.\(2023b\)S\. Zhou, F\. F\. Xu, H\. Zhu, X\. Zhou, R\. Lo, A\. Sridhar, X\. Cheng, T\. Ou, Y\. Bisk, D\. Fried, U\. Alon, and G\. NeubigWebArena: a realistic web environment for building autonomous agents\.External Links:2307\.13854,[Link](https://arxiv.org/abs/2307.13854)Cited by:[§1](https://arxiv.org/html/2609.19830#S1.p1.1),[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px3.p1.1)\.
- Zhuanget al\.\(2026\)Y\. Zhuang, J\. Yang, M\. Pan, C\. Tan, Y\. Cai, S\. Tan, C\. Zhi, X\. Zhang, J\. Yin, and J\. ChenMitigating manifold departure: uncertainty\-aware subspace rectification for trustworthy mllm decoding\.External Links:2606\.09859,[Link](https://arxiv.org/abs/2606.09859)Cited by:[§5](https://arxiv.org/html/2609.19830#S5.SS0.SSS0.Px1.p1.1)\.
## Appendix AProofs of Theoretical Properties
### A\.1Proof of Proposition[1](https://arxiv.org/html/2609.19830#Thmproposition1)
###### Proof\.
Under the pairwise construction, the Bayesian posterior is
ρθold\(a\|h,o′\)=fθold\(o′\|h,a\)fθold\(o′\|h,a\)\+fθold\(o′\|h,a~\)\.\\rho\_\{\\theta\_\{\\rm old\}\}\(a\|h,o^\{\\prime\}\)=\\frac\{f\_\{\\theta\_\{\\rm old\}\}\(o^\{\\prime\}\|h,a\)\}\{f\_\{\\theta\_\{\\rm old\}\}\(o^\{\\prime\}\|h,a\)\+f\_\{\\theta\_\{\\rm old\}\}\(o^\{\\prime\}\|h,\\tilde\{a\}\)\}\.
Since the empirical prior is uniform,
ρ0\(a\|h\)=12\.\\rho\_\{0\}\(a\|h\)=\\frac\{1\}\{2\}\.
Therefore,
e\(h,a,o′\)=ρθold\(a\|h,o′\)ρ0\(a\|h\)e\(h,a,o^\{\\prime\}\)=\\frac\{\\rho\_\{\\theta\_\{\\rm old\}\}\(a\|h,o^\{\\prime\}\)\}\{\\rho\_\{0\}\(a\|h\)\}
can be written as
e\(h,a,o′\)=2fθold\(o′\|h,a\)fθold\(o′\|h,a\)\+fθold\(o′\|h,a~\)\.e\(h,a,o^\{\\prime\}\)=\\frac\{2f\_\{\\theta\_\{\\rm old\}\}\(o^\{\\prime\}\|h,a\)\}\{f\_\{\\theta\_\{\\rm old\}\}\(o^\{\\prime\}\|h,a\)\+f\_\{\\theta\_\{\\rm old\}\}\(o^\{\\prime\}\|h,\\tilde\{a\}\)\}\.
Let
m\(o′\|h\)=12fθold\(o′\|h,a\)\+12fθold\(o′\|h,a~\)\.m\(o^\{\\prime\}\|h\)=\\frac\{1\}\{2\}f\_\{\\theta\_\{\\rm old\}\}\(o^\{\\prime\}\|h,a\)\+\\frac\{1\}\{2\}f\_\{\\theta\_\{\\rm old\}\}\(o^\{\\prime\}\|h,\\tilde\{a\}\)\.
Then,
e\(h,a,o′\)=fθold\(o′\|h,a\)m\(o′\|h\)\.e\(h,a,o^\{\\prime\}\)=\\frac\{f\_\{\\theta\_\{\\rm old\}\}\(o^\{\\prime\}\|h,a\)\}\{m\(o^\{\\prime\}\|h\)\}\.
Taking expectation overo′∼fθold\(⋅\|h,a\)o^\{\\prime\}\\sim f\_\{\\theta\_\{\\rm old\}\}\(\\cdot\|h,a\)gives
𝔼o′∼fθold\(⋅\|h,a\)\[loge\(h,a,o′\)\]\\displaystyle\\mathbb\{E\}\_\{o^\{\\prime\}\\sim f\_\{\\theta\_\{\\rm old\}\}\(\\cdot\|h,a\)\}\[\\log e\(h,a,o^\{\\prime\}\)\]=𝔼o′∼fθold\(⋅\|h,a\)\[logfθold\(o′\|h,a\)m\(o′\|h\)\]\\displaystyle=\\mathbb\{E\}\_\{o^\{\\prime\}\\sim f\_\{\\theta\_\{\\rm old\}\}\(\\cdot\|h,a\)\}\\left\[\\log\\frac\{f\_\{\\theta\_\{\\rm old\}\}\(o^\{\\prime\}\|h,a\)\}\{m\(o^\{\\prime\}\|h\)\}\\right\]=DKL\(fθold\(⋅\|h,a\)∥m\(⋅\|h\)\)≥0\.\\displaystyle=D\_\{\\mathrm\{KL\}\}\(f\_\{\\theta\_\{\\rm old\}\}\(\\cdot\|h,a\)\\\|m\(\\cdot\|h\)\)\\geq 0\.
Thus, the expected log evidence equals the KL divergence between the executed action’s feedback distribution and the pairwise mixture distribution\.
∎
## Appendix BComputational Overhead Analysis
We analyze the additional computational cost introduced by BATON\. All measurements are conducted on ALFWorld with Qwen2\.5\-1\.5B\-Instruct under the same rollout budget, batch size, optimizer settings, and hardware configuration as the main experiments\.
BATON does not introduce additional environment interaction\. The overhead mainly comes from Bayesian Feedback Attribution, which requires additional feedback likelihood evaluation for attribution weighting\. Trajectory Mass Normalization only changes the trajectory\-level aggregation rule and therefore introduces negligible computational cost\.
Table[5](https://arxiv.org/html/2609.19830#A3.T5)reports the total wall\-clock training time and the relative cost normalized by the corresponding host learner\.
## Appendix CAdditional Search\-Augmented QA Results
Table[4](https://arxiv.org/html/2609.19830#A3.T4)reports the domain\-level exact match results on the seven search\-augmented question answering benchmarks\. The averaged score reported in the main paper is computed as the unweighted mean over these domains\.
Table 4:Main results on search\-augmented question answering\. We report exact match \(%\) on seven evaluation domains and their unweighted average\. NQ, TriviaQA, and PopQA are single\-hop domains; HotpotQA, 2WikiMultihopQA, MuSiQue, and Bamboogle are multi\-hop domains\. Within each backbone, bold marks the highest mean in each metric column separately within the GRPO\-based and GiGPO\-based method groups\. BATON rows are shaded in blue\.Table 5:Wall\-clock training overhead of BATON on ALFWorld\. Training time is measured over the complete training process under identical rollout and optimization settings\. Values are normalized by the corresponding host learner\.Similar Articles
BiPACE: Bisimulation-Guided Policy Optimization with Action Counterfactual Estimation for LLM Agents
BiPACE introduces a drop-in advantage estimator that fixes state-action credit mismatch in stepwise group-based RL for LLM agents, using bisimulation-guided state clustering and action counterfactual estimation, achieving significant performance gains on ALFWorld, WebShop, and TextCraft with Qwen2.5 models.
Bayesian-Agent: Posterior-Guided Skill Evolution for LLM Agent Harnesses
Bayesian-Agent presents a framework that treats reusable skills and SOPs as hypotheses, using Bayesian inference to guide agent behavior and improve task performance through posterior-guided harness optimization. It achieves significant improvements on multiple benchmarks with deepseek-v4-flash.
Beyond Trajectory Imitation: Strategy-Guided Policy Optimization for LLM Reasoning
Introduces Strategy-Guided Policy Optimization (SGPO) for LLM reasoning, which replaces trajectory imitation with strategy distillation, improving generalization on math benchmarks.
Bayesian Optimization with Rich Auxiliary Information via LLMs
This paper presents methods to integrate rich auxiliary information into Bayesian Optimization using LLMs, showing superior performance over standard and existing LLM-based approaches in benchmarks and real-world tasks like nuclear fusion.
Agentic Bayesian Optimization through Surrogate-Augmented Autoresearch
This paper introduces agentic Bayesian optimization, where an LLM agent acts as the central decision-maker in the BO loop with a Bayesian backend, enabling online strategy revision and problem reframing. The authors instantiate this in Sara and lenz, demonstrating reliability and performance gains over standard BO and LLM-based baselines.