Reinforcing LLM Agents via Policy Optimization with Action Decomposition
Authors: Muning Wen
Ziyu Wan
Jun Wang
Weinan Zhang
Ying Wen
Venue: NeurIPS 2024 poster (2024)
Keywords: Reinforcement Learning
Language Agent
LLM agent
Large Language Models
Primary area: reinforcement_learning

Abstract: Language models as intelligent agents push the boundaries of sequential decision-making agents but struggle with limited knowledge of environmental dynamics and exponentially huge action space. Recent efforts like GLAM and TWOSOME manually constrain the action space to a restricted subset and employ reinforcement learning to align agents' knowledge with specific environments. However, they overlook fine-grained credit assignments for intra-action tokens, which is essential for efficient language agent optimization, and rely on human's prior knowledge to restrict action space. This paper proposes decomposing language agent optimization from the action level to the token level, offering finer supervision for each intra-action token and manageable optimization complexity in environments with unrestricted action spaces. Beginning with the simplification of flattening all actions, we theoretically explore the discrepancies between action-level optimization and this naive token-level optimization. We then derive the Bellman backup with Action Decomposition (BAD) to integrate credit assignments for both intra-action and inter-action tokens, effectively eliminating the discrepancies. Implementing BAD within the PPO algorithm, we introduce Policy Optimization with Action Decomposition (POAD). POAD benefits from a finer-grained credit assignment process and lower optimization complexity, leading to enhanced learning efficiency and generalization abilities in aligning language agents with interactive environments. We validate POAD across diverse testbeds, with results affirming the advantages of our approach and the correctness of our theoretical analysis. The source code can be accessed directly with this link: https://github.com/morning9393/ADRL.

Reinforcing LLM Agents via Policy Optimization with Action Decomposition
Muning Wen', Ziyu Wan', Jun Wang?, Weinan Zhang lit, Ying Wen
SShnnghai Jiao Tong University, Uniierriity College London
Abstract
Language models as intelligent agents push the boundaries of sequential decision- making agents but struggle with limited knowledge of environmental dynamics and exponentially huge action space. Recent efforts like GLAM and TWOSOME manually constrain the action space to restricted subset and employ reinforcement learning to align agents' knowledge with specific environments. However, they overlook fine-grained credit assignments for antra-action tokens, which is essential for efficient language agent optimization, and rely on human's prior knowledge to restrict action space. This paper proposes decomposing language agent opti- mization from the action level the token level, offering finer supervision for each intra-action token and manageable optimization complexity in environments with unrestricted action spaces. Beginning with the simplification of flattening all actions, theoretically explore the discrepancies between action-level op- timization and this naive token-level optimization. We then derive the Bellman backup with Action Decomposition (BAD) to integrate credit assignments for both intra-action and inter-action tokens, effectively eliminating the discrepancies. Implementing BAD within the PPO algorithm, we introduce Policy Optimization with Action Decomposition (POAD). POAD benefits from a finer- mner-grained credit as- signment process and lower optimization complexity. leading to enhanced learning efficiency and generalization abilities in aligning language agents with interactive environments. We validate POAD across diverse testbeds, with results affirming the advantages of our approach and the correctness of our theoretical analysi
1 Introduction
Large language models (LLMs) have demonstrated promising capabilities of solving various tasks, from instructions following to complex reasoning and real-world interaction [13). This growing task- solving ability underscores their potential as intelligent language agents in interactive environments 4.5) However, despite in-context language generation aiding comprehension of environmental states and action spaces in sequential decision-making tasks, misalignment issues [4] 5 such as generating invalid actions and lacking knowledge of environmental dynamics hinder these agents' ability to complete docision-making tasks robustly and efficiently.
Recent advances 4 have showcased that the aforementioned challenges can be alleviated in trial-and-error learning style, namely Reinforcement Learning (RL) [ Representatively, GLAM [ and TWOSOME (5) treat the language actions, i.e. token sequences outputted by language model, as whole units, and optimize actions' likelihood, calculated as the products conditional probabilities of intra-action tokens. Leveraging the chain rule of probability, they build the bridge between optimizing actions and optimizing tokens, aligning the training process of language models as
Corresponding to Ying Wen <ying. Being.wenCe cenn@stu..du..n......... edu.en.. Weinan Zhang <wnzhang sjtu.edu.en The source code could be accessed directly with this link https://githuu http:/hwww..coooooongn//////////
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

TV
inlevant option. disinguished
Rewancikkvartage +O,
X
by common sense. RowiiriAccvaraage vagg coffee table empty. disting by alignment Rowardikkkvanian O
Task Prompt
order coffee, your next slep is lo
Language Agent
coffee
table
kitehen
Prere coffee machine, chosen option.
Action Prompt
Walk TV;
Walk to coffee table, Walk kitchen.
P("Walk"| p) P("to"| p. "Walk")
P ("kitchen" p, "Walk to")
Enhance Token
Enhance?
Enhanos?
Enhance?
Enhance Action
P("Walk kitchen"
p) ("Walkk"| p) x P(to" p, "Wllk") x P""ktchhen") p, "alW to*)
A Case to demonstrate: (a) the necessity of aligning language agents with environments to exclude the wrong option, since the agent does not initially know that "caffee table is empty". (b) Action-level optimization is uncertain to what extent the key tokens, i.e. IP |P("kichen""|p "Walk to"",, will be enhanced when optimizing the joint probability P(""Walk to kitchen"(p)
Figure I:
next-token predictors with the RL objective of maximizing actions' utilities. However, they still suffer from limitations in optimization and exploration efficiency, due to the uncertainty of credit assignment to intra-action tokens. As shown in Figure when optimizing the distribution over three candidate actions that only differ from the last token, action-level policy optimization strategies in previous works cannot ensure that the probability of the key tokens, i.e. P("kitchen") p. "Walk to") here, will be enhanced precisely when optimizing the joint probability P("Walk to kitchen") pp) Furthermore, optimizing at the action level poses the challenge of overlarge optimization complexity due to exponentially growing action spaces, leading GLAM and TWOSOME to manually constrain action spaces. But they remain incapable in environments whose action spaces cannot be restricted.
A natural attempt to solve these problems is to incorporate the token generation process in each decision step as part of the sequential decision massinn process (10H12). This approach resolves the credit assignment problem and reduces the growth of optimization complexity from multiplicative to additive as the number of intra-action tokens increases, by updating each token's 's output distribution separately with fine- ene-grained signals from value backups. Empirical success in many single-step tasks, e.g. question answering [3] and alignment [10] [1], have demonstrated its effectiveness. However, tasks leads extra difficulties in coordinating the credit assignment process across actions and their constituent tokens. In such scenarios, embedding the intra-action token generation process into the original Markov Decision Process (MDP) [14] results in a new MDP inconsistent with the original one. Intuitively speaking, it introduces an undesired assumption of "in a sentence, tokens appearing later are more important for conveying meaning than those appearing earlier." This assumption, however, is unrealistic due to the nature of linguistic actions. Such significant discrepancy has been ignored in previous research and has not been tackled properly for long time.
the multi-step nature of sequential
In this paper, we provide comprehensive theoretical analysis of the diserepancy and derive the Bellman backup with Action Decomposition (BAD), which guarantees theoretical consistency with optimizing the original MDP. While being possible seamlessly integrate with variety of traditional RL methods, in this work, we apply BAD to Proximal Policy Optimization (PPO) [5)] resulting in the formulation of Policy Optimization with Action Decomposition (POAD). Benefiting from the finer-grained supervision afforded by BAD, POAD mitigates the uncertainty in the credit assignment process described in Figure [ thereby enjoying better interpretability, lower optimization complexity, and higher training efficiency. Meanwhile, it theoretically maintains consistency between the token- level training process for language models and the RL objective of maximizing actions' utilities. We secision-making environments
justify our claims by evaluating POAD in both classical sequential with limited action space, i.e Overcooked and VirtualHome (5). and self constructed data science coding environment featuring an unrestricted action space, i.e. DataSciCoding: results verify POAD's advantages in performance and efficiency over baseline methods, highlighting the significance of BAD, Moreover, we empirically demonstrate that language agents trained with POAD exhibit excellent generalization ability across unseen tasks, without compromising the inherent functionalities of language models. Finally, ablation studies confirm the correctness of our theoretical insights.

---

2
Related Works
Language-baser Decision-Making Agents. Leveraging LLMs' powerful capability and plenty of common knowledge, recent efforts successfully adapt LLMs in decision-making tasks as a policy model in interactive environments. In robotics, LLMs have been employed as high-level planners of control policies I8 Similarly, L.MMs work particularly well in text-based environments [0) ReAct [21 combines chain-of-thought reasoning (2) with satingbbased prompt, efficiently solving hard decision-man makion tasks. Self-Refine [22] and Reflexion [23) further improve language agents' efficiency and robustness via online adaptation through the self-reflection process. In this work, we also apply language agents in sequential decision-making scenarios, i.e. interactive environments.
Fine-tuning Language Models with RL. A body of literature explores the prospects of leveraging strategic planning methodologies to enhance the performance of language agents 24
25 6].
Besides, RL has also been widely applied in fine-tuning LLMs [TO, 7 24
25 6
Particularly,
proximal policy optimization (PPO) 15] is the most commonly used RI. method for reinforcement learning from human feedback (RLHF), proposing a breakthrough in aligning LLMs with human preference 10 [.. In classical sequential decision-making scenarios, align the objective of token-level optimization with action-level optimization, GLAM 4 and TWOSOME [5] estimate the probability of possible actions with the products of the conditional probability of tokens composing the actions and update the action we attempt to decompose actions and explicitly assign precise credit to each intra-action token, while ensuring that its optimality is consistent with updates at the action level. While concurrent work, ArCHer [26]. also targets token-level supervision for LLMs in interactive environments, employs hierarchical RL framework, using Q-network for action-level credit approximation and REINFORCE (27) for token-level backpropagation. However, ArCHer's use of multiple value networks (Q. V, and optional baseline networks) demands extensive manual tuning and can introduce cumulative bias and variance, potentially affecting stability.
whole. In this work, instead of treating an action whole,
Sequential Decomposition in RL. Recent years have witnessed an increasing trend of decom- posing high-dimension actions to leverage the powerful modeling abilities of sequential models like Transformer (28) in RL problems [29]34]. While Kuba et al. [3]] 36) proposing sequential decomposition method to provide finer meer-graieed supervision for multi multi-agent joint actions, Multi-agent Transformer 31 inherits this idea and solves multi-agent RL problems with Transformer. More recently, Q Q-transformer [4] managed to decompose the Q-functions for high-dimensional actions by representing each action dimension separate tokens. For language agents, the language generation process inherently conforms to the pattern of sequential decomposition, which offers a promising avenue for providing finer-grained supervision intra-action tokens.
3 Preliminaries
3.1 Language-augmentted RL
In this work, we assume a textual RL setting that frames sequential decision-making problems with linguistic inputs and outputs as language-augmented Markov Decision Process (MDP) M (V, S, S,A.T. T.. R, 7) [37]] 4. Given V the vocabulary and w E V the tokens, A C V., S C V are the action space and state space respectively, i.e. actions and states are sequences of tokens. T T:S. S x A + S is the state transition function. R S x A + R is the reward function that only responds to complete actions, and 7 is the discounted factor that typically less than 1. At time step t E N, a language agent receives a textual state E S from an interactive environment as input ( where |ae] ==] are tokens in it. Then, textual action a; will be grounded to a specific API call or command in the environment (38). After execution, the language agent receives a reward r. = R(s,, ax) along with the next state 51+1, based on the transition function T. Following this process with trajectories of maximum timestep T, the agents earn a discounted cumulative return of R I 7 rr, which aimed be maximized by RI. algorithms.
and generates an action d, E A in an auto-regressive manner, i.c. d, = (w) ...., w.t.....!!!! [oe] lee]
denotes the number of tokens in the action string and [w]]]e]
r',,
3

---

3.2 Action-level Policy Optimization
We begin by briefly reviewing the process of action-level policy optimization which is widely adopted in several State-of-the-wee methods that align language agents with environments via RL algorithms 39]. It facilitates seamless integration between any textual environment and conventional RL algorithms and thus is an ideal starting point for our analysis.
The possibly achieved episodic return following policy л given action and state is usually evaluated by state-action value function 2x a) or state value function Vx (s). Then, language agent updates its policy л according credits calculated on the value functions, defined
Qn(s, a) A BurTTTrr |80 = 8, 8,00 a].
(1)
(s) & BBI.rrTT..rr RR| = a].
(2)
where = -w" ...., w(ae ) -i While Ahn et al. [16] builds the connection between the likelihoods of actions and tokens through the chain rule of probability as
Jol
(aas)) = II (w |s,
w -
(3)
5=1
recent approaches like GLAM 4 and TWOSOME (5) leverage similar ideas and optimize action-level likelihoods with RL methods directly. When considering optimizing for π (als). Equa-
tions [ and are aligned with the definition of the value function in traditional RL settings, allowing
them to updated with traditional Bellman backup 9
Qr(st, 0c) R(8,, at) + max (0(54)1,
(4)
(s( + R(st, a.) + (s-1).
(5)
Moreover, it is noteworthy that Equation calculates the likelihood of action a in exponentially
growing language space as |a] increases, i.e. LA| = | al, Exploration and optimization in such huge action space are typically intractable for RL methods. Therefore, in the settings of GLAM and TWOSOME, the feasible action space is significantly restricted and smaller than the entire language
space, i.e. |A| < ||||al Taking TWOSOME as example, optimizes the likelihood of action concerning the feasible action space with Equation to mask out invalid outputs.
exp(log x [(lss)/((a))
TrwosoME
(6)
E' exp(log (((a)s)
[aa indicates the number of tokens or words in the action prompt, utilized as normalization term to mitigate the effects of varying action lengths.
Underpinned by Equation 3 GLAM and TWOSOME ensure the consistency between token-level optimization for language models and action-level optimization in an RL manner, without the need to explicitly assign credits for intra-action tokens. However, the jointness of the objective causes difficulties associated with the uncertainty in the credit assignment process 44 41 as shown in Figure after assigning credit to an action, it's unsure whether key tokens in this action have been identified, and how much they are influenced. Thus, conducting RL training at the action level introduces uncertainty, which may lead to an inefficient learning process for the language agent.
4 From Actions to Tokens: Naive Token-level Policy Optimization
4.1 Naive Token-level Policy Optimization
To address the unclear credit assignment issue described in Figure and Section 3.2 our target is to provide finer-grained supervision for each token during update while maintaining consistency in the optimality with action-level optimization, i.e. maximizing agents' cumulative returns. For arbitrary subsets of actions w with j S |a], we define token value functions for supervising policy update
(x w 1:3-1 w²) é [80 = 8, wn
1:1-1
w 1:3-1 w/ = w²)
Va w 111 R 18 = 8, wn 11111 = w 1::--
(8)
A natural approach to approximate Qu (s, 1:j-1 w
and VA (s, 1 is conceptualizing the token generation process as part of the MDP, where each token is treated micro action. This enables

---

back- Bacck-roppggtingg credit among all tokens to furnish detailed supervision. Such an idea is borrowed from the modeling process of RLHF in general language tasks 10 42]. In this way, the token-level Bellman backup corresponding Equations and can be expressed
0+ Tw max with Qr (8t, w " wrtt)) R(st, + max Qa(4441, wh+)),
j < |ae]
Qs(ss, 10-1
w/ w!?
(9)
if = |ax|
0 + Var Va (st, wi ( R(s + To (8
if [o||
(10)
if |
To facilitate subsequent theoretical analysis, we separate the discount factor y into intra-action 7v and inter-action "Yo despite their numerical equivalence here. The above backups can be interpreted as applying RL algorithms on modified reward function, which maintains action-level feedback while introducing extra O feedback for tokens within an action, except for the last token. Intuitively, this approach seems feasible and decomposes the action-level reward signal R(s, st. ar) ) intra action tokens, thus alleviating the uncertainty in credit assignment [31] and reducing optimization complexity. However, it must be noted that the optimization of Equations 2 and TT are inconsistent with those of Equations 4 and for sequential pacision-making tasks since it introduces a new MDP M that diverges from the origin one. We will analyze this discrepancy in the following section.
4.2 The Discrepancy
To maintain the consistency between action-level updates and token-level updates, we should ensure that optimizing over tokens gives the same optimality as optimizing the whole action, which is analogous to the multi-dimensional action optimization setting in traditional [L That is, given the deterministic nature of linguistic action generation, and an optimal polity म² after sufficient training following the token-level Bellman backups, ideal optimal token value functions should
satisfy (st,
Qr (st. 0t) and V w, 1t)) Vz- (s:) for Vj < |ae]. To quantify the discrepancy between action-level optimization and token-level optimization with the shape of Equations
1:j-1 w/ !:ii!
9 and 10 we expand the value backup process over each token starting from arbitrary
j < |az] as The following equations (For derivation see Appendix A).
Qr (st, (/,, 15-11 w() = R(se, a+) + Yo max Q (11+1,0++1)
(1I)
at
%o
(all-i))(sss + 7a(1 (1l111lo Jedttle441-3-)
max Q (8141, act1
Discrepancy between Equation 4
and
Va (Bt, w(:)) R(ss,a)) + Ja V+ (#++1)
VaVa
Iael Ieells R[(s, 1 a.)
(1 7u - VVV (sc+1) (12)
Discrepancy between Equation and TO
With Equation I and 12
we observe following significant insights:
The discrepancy of the Bellman optimality between action-level optimization and token-level optimization diminishes The [0, increases, achieving consistency when "iw- 1.
Given You < 1, the discrepancy increases the number of intra-action tokens, |arl, increases. (ii)
Based on the first insight, we will derive our main approach, namely the Bellman backup with Action-Decomppsittonn in the next section. The second insight can be intuitively understood as follows: The diverged MDP M attaches the assumption that "words later a sentence are more expressive than earlier ones", which is unrealistic for representing the semantics of linguistic actions.
5

---

Turn
TV
Completedl
Tum
TV
Completed)
Q:11
Smootop
mmeetop
4)
rp
of
Faled!
Tum
otf
TV
Figure 2: Visual comparison of the differences between action-level Bellman backup (left) and our BAD (right), given the goal turn on the TV, where q is the action or token value estimations, 6, qt-1 or ~ q¹ - -q-- represent the credit assigned to corresponding actions and tokens respectively for policy update, e.g. the advantage value[43]. To facilitate understanding, a step-by- step breakdown of the right figure is provided in Appendix
Action-Decomposition Reinforcement Learning
5.1 Bellman Backup with Action-Decomposition
According the first insight, can modify Equations 9 and 10 Decioo-DDeoomosittion (BAD) as
proposing the Bellman backup with
max Qr w R(sc;as) + max (x (11+11,4411))
if < |as]
Qr(st, w. 1: w/)
(13)
if |a.]
(se, w 1 R(ss, az) + +VVa($241;0) 0),
if < |ad if Jael
(s,, w
(14)
(For proof of optimization consistency see Appendix B Training language agents with BAD provides finer-grained supervision for credit backpropagation, eliminating uncertainty in credit assignment and thus enjoying better interpretability and efficiency of the RL training process. In addition, it theoretically ensures consistency between the token-level training process for language models and the RL objective of maximizing actions' utilities. Figure 2 visually compares the differences between action-level Bellman backup and our BAD and demonstrates how BAD precisely assigns credit to each token.
Another advantage of BAD is the ability to decompose the optimization in an intractable action space of size ((IV into |a] times optimizations in token spaces of size O(VV)) reducing the complexity of RL problem with language agent to O[[a x [V], thus rendering the problem more manageable. Moreover, BAD can be seamlessly integrated into various existing RL methods, including off- off-policy algorithms c.g. DQN[44], on-policy ones like Actor-Critic445 and PPO[I5]. Furthermore, we have also provided a version the Soft Q-function in Appendix B.3 support various entropy- RI. algorithms like 4AC|46 47].
5.2 Policy Optimization with Action Decomposition
In this section, we integrate the BAD into the widely used PPO and propose specific method called Policy Optimization with Action Decomposition (PÓAD), while the integration with other algorithms will be reserved for future works. POAD sequentially decomposes the policy update granularity from the action level to the token level. To approximate the token value function, we introduce a critic network with parameters e whose objective is to minimize the empirical Bellman error of tokens by
L(0) = Σ
(R().) + VV 0) Vole. wi 2
[ae]
[Va(er, u (1 Va w. ",0)) (15)
losl
Inter-action credit assignment
Intra action credit assignment

---

where 0 represents the non non-differentiable parameter of the target network, updated at intervals. The policy network's 's parameters are denoted as o, and optimization follows the clipping PPO objective.
lad
L(o)
Σ
min
(16)
1=
(wr st, w
1 1
ratio! (6)
Toora (wg] |st, w 1) and A/ A(w,] wi 1)
where A is an estimate of the advantage value for each token with the generalized advantage estimation (GAE) 48 To capture more details about POAD, we draw a pseudo-code in Appendix D
6 Experiments
In this section, we show the superiority of POAD in performance, efficiency, and generalization abilitics with different testbeds. Moreover, we conduct meaningful ablations on Ya and Yu to verify the theoretical analysis in Section 4.2 Finally, examine models trained with POAD and baseline methods to investigate their impact on the model's original language abilities. For in-depth analysis, conduct case study in Appendix to validate the effectiveness of BAD in terms of token-level credit assignment. We deploy LLaMA2-7B [49] for Overcooked and VirtualHome, and CodeLLaMA- 7b (50) for DataSciCoding, fine-tuned with Low Rank Adaptation (LoRA) [5] with I Nvidia A100 GPU.
6.1 Environmental Setup
We first evaluate method two classical sequential recision-making environments with limited action space: Overcooked [3] and VirtualHome [5], where the action space consists of approximately 10 possible actions per state, each includes 5-10 tokens. Then evaluate method in data science coding and debugging environment with unrestricted action space: DataSciCoding, where agents generate actions (up 128 tokens) freely. More detailed descriptions can be found in Appendix
Overcooked and VirtualHome. Overcooked challenges agents to prepare dishes such as tomato salad and tomato-lettuce salad in a 7x7 grid kitchen using linguistic actions like Chop, Get-Tomato, and Go- Cro-Cutting-Bouee Board, with rewards for correct deliveries and penalties for incorrect ones or time wastage. Meanwhile, VirtualHome simulates household tasks like reheating pancakes in Food Preparation and organizing an evening of Entertainment, with actions such as walk to the living room and turn the TV, rewarding agents only upon task completion in a partially observable setting.
DataSciCoding. We develop DataSciCoding to automate data science coding tasks with unrestricted action space, currently adopting 3 Kaggle datasets and 3 OpenML datasets (52) with details in Appendix E.I In each task, agents aim to implement the most effective classifier with the scikit-learn module, striving to achieve the highest possible ROC AUC score (53] on test sets. As the prompts provided to the agents contain no detailed information about task datasets, agents are required to interactively modify and debug their code based feedback from the runtime environment until works, thus aligning the task datasets. Agents receive ROC AUC scores E [0. 1] as rewards for workable codes and -1 as penalties for run failed. Adopting the same evaluation metrics as CAAFE [54], for each dataset and code, we evaluate 5 repetitions, each with a random 50% 50% train-test split (55), and record the average ROC AUC score across these splits.
6.2 Baseline Methods
For Overcooked and VirtualHome, we compare POAD's performance with Naive Token-Level Policy Optimization (NTPO) mentioned in Section 4.1 i.e. integrating Equation T0 with Yor % into PPO. and TWOSOME [5]-ther the current tt method on Overcooked and VirtualHome. Besides, we also incorporate ArCHer [26] as a baseline in VirtualHome with comparative analysis, since positioned as an intermediate between NTPO and POAD for token-level credit assignment, theoretically enjoying reduced discrepancy compared to NTPO. We demonstrate the difference in the backup processes between POAD and these baselines as well as the optimality after convergence in Appendix C
7

---

Solad
Trmor: Salat
Preparstion
Fntetstment
Eo coti Batum
basate ROUF
1010
Figure 3: Performance comparisons on Overcooked (first two) and VirtualHome (last two).
For the DataSciCoding environment, in case TWOSOME is inapplicable due to the unrestricted action space, we examine POAD's performance along with NTPO. Mcanwhile, we also compare POAD with CAAFE [54] novel AutoML framework in which humans collaborate with large language models for data science. In CAAFE, humans implement machine learning models, c.g. classifiers, while large language models generate the code for feature engineering. CAAFE represents the current state-of-the the- performance with collaboration between humans and powerful closed-source LLMs such as GPT-3.5 and GPT-4 on the given task datasets we used.
6.3 Main Results
6.3.1
Classical Sequential Decision-Making Tasks with Restricted Action Space
As shown in Figure where curves are averaged over seeds with the shadow showing their standard deviation, the performance drop of NTPO when comparing with POAD verifies the existence and negative impact of the discrepancy analyzed in Section 4.2 While POAD can achieve the same (or even better) convergence performance compared to TWOSOME which further verifies the consistency between token-level optimization with BAD and action-level optimization. In addition. the training curves of POAD are more stable (enjoying smaller standard deviations) than TWOSOME and converge much faster than all other baselines, indicating POAD's stability and efficiency by integrating our BAD. In complex Entertainment tasks, ArCHer is second only to POAD, aligning with our theoretical expectations. However, in tasks such as Food Preparation where the methods' performance gap was less pronounced, ArCHer performed poorly, potentially due to instability in its system that involved multiple value networks.
6.3.2 Data Science Coding Tasks with Unrestricted Action Space
According to the training curves in Figure 4 our method POAD significantly outperforms NTPO both in the convergence speed and the final score. Compared to the results sequential decision- making tasks in Section [6.31]] the performance gap between POAD and NTPO on DataSciCoding is consistently larger, due the much longer action length |az]. most 128 tokens for each action in DataSciCoding. Such results are consistent with our second insights in Section 4.2 and empirically highlight the importance of proper intra-action credit assignment, which, our Bellman Backup with Action Decomposition (BAD) is designed for.
In Figure 4 POAD-Bes Best means the performance of the best code discovered during POAD's training process with Codel.LaMA-7B. We compare it with the best performance achieved by the state-of-the- art AutoMI. framework CAAFE (54 with GPT-4 model. In this experiment, aim to prove that even small-scale language models can also provide hetter outcomes than large-scale models, what is needed is just a stable and efficient training algorithm POAD and only 2-3 hours on Nvidia A100 (Details of wall-time on each task are shown in Appendix H). A more detailed Comparison between POAD and CAAFE with both GPT-3.5 and GPT-4 can be found in Appendix G
6.4 Docablaoyyy Task Generalization
LLMs' ppcabbocayylary feature enables language agents to transfer their learned skills into unseen similar tasks, expanding the capabilities of decision mmaiing agents. We compare the generalization performance of language agents trained by POAD, TWOSOME and NTPO in the Food Preparation task with the original base model LLaMAA-7B. Table shows that token-level policy optimization methods achieve better generalization performance unseen tasks. And our POAD outperforms the other basclines in seven of eight tasks.
8

---

health Insurance
pharyngtts
spaceship
RDC MUC Score
RBC ALC
AUI
VTPC
AmPP HIRD
KTPO
FORD
PAND
Bast
Environment Steps
Environment Steps
Environment Steps
breast,
ROC-AUC Score
NCC- BOC--UC Sccre
0.-
RDC
1800
3050
2013
Environment Steps
Environmentt Steps
Environment Steps
Figure 4: While TWOSOME does not support open action space tasks, we compare the average performance between POAD and NTPO on the DataSciCoding benchmarks, as well as POAD-Best the performance best code explored by POAD during the training phase and CAAFE with GPT-4.
Table 1: Comparison of generalization performance on eight unseen tasks, with episodic returns averaged over 100 episodes. In these tasks, replace the pancake in the original Food Preparation task with other foods like cheese, hamburger, apple pie and pizza, or replace the (pancake, microwave) at the same time with (dishes, dishwasher) (clothes, washing machine) for greater differences.
Methods
Cheese
Hamburger
Apple Pie
Pizza
Washing Plate
Laundry
ILaMA2-7B
0.1351
0.1342
0.1656
0.1409
0.0527
0.0344
TWOSOME
0.7119
0.7058
0.7304
0.7047
0.7031
0.6038
NTPO
0.7428
0.7476
0.7141
0.7355
0.7491
0.5687
POAD
0.7553
0.7602
0.7650
0.7625
0.7075
0.7014
Preparation
Preparation
Preparation
Egicodic Retur
Eploodic Raturn
Raturn
POwo
MTFD-0 65
TWOSOONE-O TIVOSONE-1.
NTFC-3 60
NTFC-3.10
FOAD-CL05
NTFC-0.50
FOAD-10
ARDHEL-2AD
1006
2000
3000
AMDO
2000 X000
500
1000
1300
2000
2000
400
Enviromment Steps
Environment Steps
Environment Steps
% [1.0,0.95] for both TWOSOME and POAD (left), and Yoe (0.95, 0.9, 0.8, 0.5) for NTPO while keeping the >hh = 0.95 unchanged (middle). In the left figure, Setting Ta 1.0 led decreased performance and convergence for TWOSOME and POAD, validat- ing necessity of 7/a < 1.0. While in the right figure, the increasingly larger performance gap between POAD and NTPO, as Ywu decreases, verifies the theoretical analysis Section 4.2 The right one shows the performance change after applying theoretical insights to enhance ArCHer's performance, i.c. ARCHER-BAR BBA with Tw 1.0.
Figure 5: Ablation
6.5 Ablations: the Impact of %a and To to the Discrepancy
To verify our analysis in Section
we conduct ablations on Tw to further investigate how large the discrepancy would be by using different values of Yor in NTPO. Therefore, we deploy NTPO on the Food Preparation task with Yw E {0.95, 0.9, 0.8, 0.5), while we keep %4 0.95 which is consistent with our main experiments. Besides, we also show the performance with 7a E (1.0, 0.95) show the necessity of setting Ya strictly less than 1-0, and thus the necessity to separate inter-action tokens and intra-action tokens. Further, our theoretical analysis is also compatible with ArCHer, motivating us to apply insights in Section 4.2 to enhance ArCHer's performance.

---

The left side of Figure 5 demonstrates an increasingly larger performance gap between POAD (Tu 1) and NTPO as The decreases. These results empirically showcase the discrepancy between naive token-level credit assignment and our BAD, which is consistent with our first theoretical insight Section 4.2 Moreover, in the middle of Figure 5 for both TWOSOME and our proposed method POAD, setting
1.0 performs much worse than % = 0.95, indicating that the discrepancy between NTPO and POAD can not be solved by simply setting "a = Th = 1.0. Furthermore, the right one in Figure 2 shows improved results that validate the effectiveness of applying our insights to ArCHer. However, due to the inherent challenges excessive network complexity and difficult hyper-parameter tuning, ArCHer-BAl BBD still falls short of matching POAD's performance.
6.6 Impact on Original Language Ability
To investigate the impacts of online RL fine-tuning on LLMs' original capabilities, we evaluate the models trained by POAD, TWOSOME and NTPO widely used NLP benchmarks (56) which are also re- ported in Tan et al. [5] and Touvron et al. [49] These models are trained Food Preparation. Table 2 demonstrates these models' zero-shot performance, compared with the original LLaMA2-7B model. The results show significant losses of general ability like natural language understanding after aligning with the embodied environment, even sometimes bringing minor improvements.
Table Zero-shot performance on Language Model Evaluation Harness [56]. with details in Appendix |
Methods
ARC C
HellaSwag
PIQA
MMLU
LLaMA2-7B
0.44
0.57
0.78
0.41
TWOSOME
0.44
0.58
0.78
0.41
NTPO
0.44
0.58
0.78
0.41
POAD
0.45
0.59
0.78
0.41
7 Conclusion
In this work, we propose the Bellman backup with Action Decomposition (BAD), theoretically eliminating the discrepancy between naive token-level policy optimization and action-level policy optimization for language agents. Integrating BAD with PPO, we propose our method of Policy Optimization with Action Decomposition (POAD), providing finer -reiieg supervision for each intra-action token and ensuring theoretical consistency between the token-level training nature of language models and the RI. objective of maximizing actions' utilities. Empirical experiments and thorough ablations showcase the effectiveness of BAD as well as the superiority of POAD in learning efficiency and gencralization abilities, over strong action-level baseline TWOSOME.
Limitation and Future Work.
The existing limitation of POAD is on the requirement for easily attainable some environments. To mitigate this,
quantitative reward function, which we envisage integrating POAD with self-rewarding [57 58 or hindsight relabeling [59].
Social Impact. The advancements in RL for language agents can significantly enhance decision- making processes in various domains such as healthcare, finance, and autonomous systems. Improved decision-making lead better outcomes, increased efficiency, and reduced errors. However, we acknowledge that when optimizing agents using our method, language agents may potentially resort to unscrupulous means maximize rewards, which could lead to potentially harmful results. Thus, we advocate for more comprchensive consideration when designing the reward function, or combining it with consty-con RL methods mitigate these risks.
Acknowledgment
The SJTU team is partially supported by National Key R&D Program of China (20222D011488) 14804), Shanghai Municipal Science and Technology Major Project 20021HZZZZ0012) and National Natural Science Foundation of China (62322603, 62076161, 62 62106141). Muning Wen is supported by Wu Wen Jun Honorary Scholarship, Al Institute, Shanghai Jiao Tong University.
10

---

References
[1] Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Yi Tay, William Fedus, Yunxuan Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, et al. Scaling instruction-finetune language models. Journal of Machine Learning Research, 25(70): 2024.
[2] Jason Wei, Xuezhi Wang. Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoe V Le, Denny Zhou, et al. Ofain-of-thouge prompting elicits reasoning in large language models. Advances neural information processing systems, 35:24824-2488 2022.
[3] Xidong Feng, Yicheng Luo, Ziyan Wang, Hongrui Tang, Mengyue Yang, Kun Shao, David Mguni, Yali Du, and Jun Wang. Chessgpt: Bridging policy learning and language modeling. Advances in Neural Information Processing Systems, 36, 2024.
[4] Thomas Carta, Clément Romac, Thomas Wolf, Sylvain Lamprier, Olivier Sigaud, and Pierre- Yves Oudeyer. Grounding large language models in interactive environments with online reinforcement learning. In International Conference on Machine Learning, pages 3676-3713. PMLR, 2023.
[5] Weihao Tan, Wentao Zhang, Shanqi Liu, Longtao Zheng, Xinrun Wang, and Bo An. True knowl- edge comes from practice: Aligning large language models with embodied environments via reinforcement learning. In The Twelfth International Conference on Learning Representations, 2024. URL https://openi:
[6] Xidong Feng, Ziyu Wan, Muning Wen, Ying Wen, Weinan Zhang, and Jun Wang. Alphazeno- like tree-search can guide large language model decoding and training. arXiv preprint 2309v:2309.17179 2023.
[7] Noah Shinn, Federico Cassano, Ashwin Gopinath, Karthik Narasimhan, and Shunyu Yao. Reflexion: Language agents with verbal reinforcement learning. Advances in Neural Information Processing Systems, 36, 2024.
[8] Fengshuo Bai, Hongming Zhang, Tianyang Tao, Zhiheng Wu, Yanna Wang, and Bo Xu. Picor: Multi-task deep reinforcement learning with policy correction. Proceedings of the AAAI Conference on Artificial Intelligence, 37(6))6728-66 6736, Jun. 2023. doi: 10. 1090//99. v37i6. 25825. BRRpstts:///fp aaai. oow/oorgggee//eeeeooeoo index. php/ AAA./aaticll/eiieeeee.ee.aaaeaee 25825
[9] Richard Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
[10] Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35:27730-27744, 2022.
[11] Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model is secretly a reward model. Advances in Neural Information Processing Systems, 36, 2024.
[12) Rafael Rafailov, Joey Hejna, Ryan Park, and Chelsea Finn. From r to q": : Your language model secretly a q-function. arXiv preprint arXiv:2404. 12358, 2024.
[13] Zhilin Yang, Peng Qi. Saizheng Zhang, Yoshua Bengio, William W Cohen, Ruslan Salakhut- dinov, and Christopher D Manning. Hotpotqa: A dataset for diverse, explainable multi-hop question answering. arXiv preprint (tt::1800009000 2018.
[14] Frédérick Garcia and Emmanuel Rachelson. Markov decision processes. Markov Decision Processes in Artificial Intelligence, pages 1-38, 2013.
[15] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal 00077006477.55 2017.
policy optimization algorithms. arXiv preprint
[16] Michael Ahn, Anthony Brohan, Noah Brown, Yevgen Chebotar, Omar Cortes, Byron David, Chelsea Finn, Chuyuan Fu, Keerthana Gopalakrishnan, Karol Hausman, et Do i can, not as say: Grounding language in robotic affordances. arXiv preprint arXiv:2204.01cm 2022.
11

---

[17] Wenlong Huang, Fei Xia, Ted Xiao, Harris Chan, Jacky Liang, Pete Florence, Andy Zeng, Jonathan Tompson, Igor Mordatch, Yevgen Chebotar, et al. Inner monologue: Embodied reasoning through planning with language models. In 6th Annual Conference on Robot Learning. 2022.
[18] Danny Driess, Fei Xia, Mehdi SM Sajjadi, Corey Lynch, Aakanksha Chowdhery, Brian Ichter, Ayzaan Wahid, Jonathan Tompson, Quan Vuong, Tianhe Yu, et al. Palm-e: An embodied multimodal language model. arXiv preprint 2tt.:::23330033377 2023.
[19] Mohit Shridhar, Xingdi Yuan, Alexande Côté, Yonatan Bisk, Adam Trischler, and Matthew Hausknecht. Alfworld: Aligning text and embodied environments for interactive learning. arXiv preprint arXiv:2010.03766 2020.
[20] Xiao Liu, Hao Yu, Hanchen Zhang, Yifan Xu, Xuanyu Lei, Hanyu Lai, Yu Gu, Hangliang Ding, Kaiwen Men, Kejuan Yang, al. Agentbench: Evaluating Ilms as agents. In The Twelfth International Conference on Learning Representations, 2023.
[21] Shunyu Yao, Jeffrey Zhao, Dian Yu, Nan Du, Izhak Shafran, Karthik R Narasimhan, and Yuan Cao. React: Synergizing reasoning and acting in language models. In The Eleventh International Conference on Learning Representations, 2022.
[22] Aman Madaan, Niket Tandon, Prakhar Gupta, Skyler Hallinan, Luyu Gao, Sarah Wiegreffe, Uri Alon, Nouha Dziri, Shrimai Prabhumoye, Yiming Yang, al. Self-refine: Iterative refinement with self- feedback. In sevethh Conference on Neural Information Processing Systems, 2023.
[23] Noah Shinn, Federico Cassano, Ashwin Gopinath, Karthik R Narasimhan, and Shunyu Yao. Reflexion: language agents with verbal reinforcement learning. In shirth-seventh Conference on Neural Information Processing Systems, 2023.
[24] Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Tom Griffiths, Yuan Cao, and Karthik Narasimhan. Tree of thoughts: Deliberate problem solving with large language models. Ad- vances in Neural Information Processing Systems, 36, 2024.
[25] Shibo Hao, Yi Gu, Haodi Ma, Joshua Jiahua Hong. Zhen Wang, Daisy Zhe Wang, and Zhit- ing Hu. Reasoning with language model is planning with world model. arXiv preprint arXiv:2305. 14992, 2023.
[26] Yifei Zhou, Andrea Zanette. Jiayi Pan, Sergey Levine, and Aviral Kumar. Archer: Training language model agents via hierarchical multi-turn rl. In Forty-first International Conference on Machine Learning, 2024.
[27] Ronald J Williams.
Simple statistical gradient-following algorithms for connectionist reinforce- ment learning. Machine learning, 8:229-256. 1992.
[28] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez. Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[29] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling, Advances in neural information processing systems, 34: 1508441777777
2021.
[30] Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems, 34:1273- 1286, 2021.
[31] Muning Wen, Jakub Kuba, Runji Lin, Weinan Zhang, Ying Wen, Jun Wang, and Yaodong Yang. Multi-agent reinforcement learning is sequence modeling problem. Advances Neural Information Processing Systems, 35: 1659-- 2022.
12

---

[32] Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel
Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, et al. A generalist agent. arXiv preprini M0Xiv:22005001775 2022.
[33] Ying Wen, Ziyu Wan, Ming Zhou, Shufang Hou, Zhe Cao, Chenyang Le, Jingxiao Chen, Zheng Tian, Weinan Zhang, and Jun Wang. On realization of intelligent decision-making the real world: A foundation decision model perspective. arXiv preprini arXiv:22122 22.69999... 2022.
[34] Yevgen Chebotar, Quan Vuong, Karol Hausman, Fei Xia, Yao Lu, Alex Irpan, Aviral Kumar, Tianhe Yu, Alexander Herzog, Karl Pertsch, et al. Q-transformer: Scalable offline reinforcement learning via autoregressive g-functions. In Conference Robot Learning, pages 3909-3928. PMLR, 2023.
[35] Jakub Grudzien Kuba, Muning Wen, Linghui Meng. Haifeng Zhang, David Mguni, Jun Wang, Yaodong Yang, et al. Settling the variance of multi-agent policy gradients. Advances Neural Information Processing Systems, 34:13458-13470, 2021.
[36] Jakub Grudzien Kuba, Ruiqing Chen, Muning Wen, Ying Wen, Fanglei Sun, Jun Wang, and Yaodong Yang. Trust region policy optimisation in multi-agent reinforcement learning. In International Conference on Learning Representations, 2021.
[37] Martijn Van Otterlo and Marco Wiering. Reinforcement learning and markov decision processes. In Reinforcement learning: State-of-the-arr, pages 3-42. Springer, 2012.
[38] Timo Schick, Jane Dwivedi- Yu, Roberto Dessi, Roberta Raileanu, Maria Lomeli, Eric Hambro, Luke Zettlemoyer, Nicola Cancedda, and Thomas Scialom. Toolformer: Language models can teach themselves to use tools. Advances in Neural Information Processing Systems, 36, 2024.
[39] Filippos Christianos, Georgios Papoudakis, Matthieu Zimmer, Thomas Coste, Zhihao Wu, Jingx- uan Chen, Khyati Khandelwal, James Doran, Xidong Feng, Jiacheng Liu, al. Pangu-agent A ine-tunable generalist agent with structured reasoning. arXiv preprint arXiv:2312.1 14878, 2023.
[40] Zhipeng Chen, Kun Zhou, Wayne Xin Zhao, Junchen Wan, Fuzheng Zhang, Di Zhang, and Ji-Rong Wen. Improving large language models via fine-grained reinforcement learning with minimum editing constraint. arXiv preprint arXiv:2401. 100081, 2024.
[41] Gregor Bachmann and Vaishnavh Nagarajan. The pitfalls of next-token prediction. arXiv preprint arXiv:2403.0668 06966, 2024.
[42] Nisan Stiennon, Long Ouyang, Jeffrey Wu, Daniel Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul F Christiano. Learning summarize with human feedback. Advances Neural Information Processing Systems, 33:3008-3021 2020.
[43] Mohammad Babaeizadeh, Iuri Frosio, Stephen Tyree, Jason Clemons, and Jan Kautz. Re- inforcement learning through asynchronous advantage actor-critic on gpu. arXiv preprint arXiv:1611.06256, 2016.
[44] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 518(7540):529-5 -333, 2015.
[45] Vijay Konda and John Tsitsiklis. Actor-critic algorithms. Advances in neural information processing systems, 12, 1999.
[46] John Schulman, Xi Chen, and Pieter Abbeel. Equivalence between policy gradients and soft q-learning arXiv preprint arXiv:1704.0664 0004040 2017.
[47] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off- policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning, pages 1861-1870. PMLR, 2018.
[48] John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High- dimensional continuous control using generalized advantage estimation. arXiv preprint MxXi::556.02438.. 2015.
13

---

[49] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale. Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiao- qing Ellen Tan, Binh Tang. Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliaan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Rohert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023.
[50] Baptiste Rozière, Jonas Gehring, Fabian Gloeckle, Sten Sootla, Itai Gat, Xiaoqing Ellen Tan, Yossi Adi, Jingyu Liu, Romain Sauvestre. Tal Remez, Jérémy Rapin, Artyom Kozhevnikov, Ivan Evtimov, Joanna Bitton, Manish Bhatt, Cristian Canton Ferrer, Aaron Grattafiori, Wenhan Xiong. Alexandre Défossez, Jade Copet, Faisal Azhar, Hugo Touvron, Louis Martin, Nicolas Usunier, Thomas Scialom, and Gabriel Synnaeve. Code llama: Open foundation models for code, 2023.
[51] Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.09665 09685, 2021.
[52] Joaquin Vanschoren, Jan N Van Rijn, Bernd Bischl, and Luis Torgo Openml: networked science in machine learning. ACM SIGKDD Explorations Newsletter, 15(2):49-60, 2014.
[53] Sarang Narkhede, Understanding auc-roc curve. Towards data science, 26(1):220-227, 2018.
[54] Noah Hollmann, Samuel Müller, and Frank Hutter. Large language models for automated data science: Introducing caafe for context-aware automated feature engineering. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[55] Xavier Bouthillier, Pierre Delaunay, Mirko Bronzi, Assya Trofimov, Brennan Nichyporuk, Justin Szeto, Nazanin Mohammadi Sepahvand, Edward Raff, Kanika Madan, Vikram Voleti, et al. Accounting for variance in machine learning benchmarks. Proceedings of Machine Learning and Systems, 3:747-769, 2021.
[56] Leo Gao, Jonathan Tow, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Kyle McDonell, Niklas Muennighoff, et al. A framework for few-shot language model evaluation. Version vO, O.I. Sept, page 8, 2021.
[57] Weizhe Yuan, Richard Yuanzhe Pang, Kyunghyun Cho, Sainbayar Sukhbaatar, Jing Xu, and Jason Weston. Self-rewarding language models. arXiv preprint arXiv:2401.10024, 2024.
[58] Runze Liu, Fengshuo Bai, Yali Du, and Yaodong Yang. Meta-reward-net: Implicitly dif- ferentiable reward learning for preference-based reinforcement learning. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neu- ral Information Processing Systems, volume 35, pages 12270-22284 Curran Associates, Inc., 2022. URL URL neurips. (be8c334001888003337774448888883
ptt..eeee//pppp/22///1iiiee Paper Conffr--ooeeeeeeeeee pdff
[59] Alexander Li, Lerrel Pinto, and Pieter Abbeel. Generalized hindsight for reinforcement learning. Advances in neural information processing systems, 33:7754-7767, 2020.
14

---

A Derivation for the Discrepancy
A.1 Q-Function
The optimization each token following Equation starting from arbitrary |az], could be expanded as
Qx (st, w ux{) Yw max R (st, wit) ut"" 11
(17)
wa" u
Tho max Qx- (82, wh wott wp! wgtt)
The max
wi+
mn?
= 1alll-
R(s,, az) + Va 1all-
max Qx (8L+1, w(+1)
"ii1
R(s ae) + Ta max wht1
The max Qa (5t+1, w w(+1))
1/1
R(st: + Yar
max Q (8t+1, w w w2+11) w
Tw
R(st, az) + Na
11
max Q (S w/ 1:last1l) 1: 1:04111
|e-l-/ Tw
R(st, 01) + 7a 7a71a441
max Q (82+1, a++1
= R(st, + max Q (1++1,
a
(s..a)
larl- (all-))((sr,e ar) 7a -alll-llllllaal------1 max Qu (5441,0011) as+1
Yen
Discrepancy between Equation and
A.2 V-Function
Similar Appendix A.I the optimization over each token following Equation 10 arbitrary j < lael. could be expanded as
starting from
Vn- (st, w/ 1 You Va- (st, (
(18)
Yin Nou Va- (Se, w (:8t2)
= - R(ss, ar) + 7a7 ma -i (8(+1)
RR(ss at) + Ya Vx. (8e+1)
(1 w ladl (all--))((ss az) + 7a - 7 -arl- Vi-. (5c_1)
(*))
Discrepancy between Equation and
B Proof of Optimization Consistency
To show that optimizing value functions with BAD still ensures the same optimality with action-level optimization, we show that optimizing the value functions for each token equivalent optimizing for the full action.
15

---

B.1 Q Function
The optimization each token with the optimal policy n² following Equation 13 starting from arbitrary j < arl, could be expanded as
Qx- (st, wj" 1:3-1, ww/) max Qx (st, 1i wi!" )1
w
max 2x (st, wat 11 wg s+1 un w/
R(s;, az) + y max ou (30+1, w/+22)
w/+1
(1) max wh41
max Q (81+1, w u(+1))
w
R(s), +7 max Qx (st+1, w/t1 w2t1))
Qr (8:+1, W (
-R(s,, 0z) +7 max 1il441l
R(ss, az) + +y max on (81+1.
Qs,
(19)
B.2 V Function
The optimization over each token following Equation 14 starting from arbitrary j < lasl, could be expanded as
Va (St, we 1) )- Vx- (st, :
Vx (86, we 1
R(se, + VV..
Vx- (8t),
(20)
where Vx+ (st, wt 11 and Va (se, wn (aall are equivalent to BuI 1Q -1 w/)] and
Emj ~r (Q(4++1, (18441,0/4111
B.3 Soft Q function in Regruay-Re RL
As an extension of BAD, we also provide the soft Bellman backup with Action- Aetion-Decompositionn (sBAD), support Entrooy-RReglaarin methods like SAC (by setting the reference policy म to uniform distribution).
[ 1 E (i (Q, w/ w/ xuf+t)]
if |
Qr(s w 2:5-1 (,
R(st, as) + (B [Q+ (8 w(+ ) 3D (Dw..////a(aa..... (s+1).
if |a.]
(21)
We show that optimizing the soft Q-function with sBAD at the token level is consistent with optimizing for the full action. We adopt (L((a))) = D [ [|]]] (s), where the KL(aas) indicate it is a action level KL divergence, KL(w)s) indicate token level KL divergence. Given optimal stochastic policy म², the vanilla soft Bellman backup for full actions is:
Eg ~ [Q(st, = Earrr R(si, an) (Ba Q(st, an) - (KL ((44211)))+11)
(22)
Ba [R(se ax) + B Bu (0(8111:00411 7: (l4441188441)
16

---

Following our sBAD, the update over each token within an action (j |a|) is:
Bu (Q(s,, wt))
Eu Bu (Q wt, uz]] - SKL(w?lss wx)]
= But up (Q, w/ wf}1 wt ~x- [KL (st, w!)
= E w! NT (Q w
1:acl-1 wa [m))
lee1-1
fE lo u0 / -M
KL(W?" +2 Se, )
= Euj (Q(8).
lael-1
(L(
Se, w
j=1
(23)
Minus 3KL(w))as (x) from both sides, we have:
Ew! oar (Qs))((s,, wi)] SKL(w] (s) Euixx [Q(ss: az)l - (KL(a))st)
(24)
where KLL(ae(s))= E)ad (W st,, mi (
Then, the update across the current action a and the next action a++1 with our sBAD:
Euewx [Q(se: Enewx
[R(s, )1) (Ea w (0(5++1, w(+1)]
EE cur [KL (8(+1)))
E Barw.
R((,, a;)] y (Q(5)+1,
-3B [L (s411)]
(25)
Equation 25 enjoys the same shape as Equation 22 thus optimizing the soft Q function following sBAD is equivalent the action-level optimization. We prove the consistency between sBAD and the traditional soft Bellman updates over full actions.
C Comparison of the Optimality between Three Backup Approaches
We provided visual comparison of the differences between four different backup approaches with the optimal policy after convergence Figure
Compinted
Complnted
Comp Intect
Complnted
Walk
TV
Wolk
Walk
Gosl: Tum TV
Gowi: Tam on TV
Goal: Turn the TV
Govl: Tum on TV
TWOSOME
NTPO
ARCHER
POAD
Figure 6: Visual comparison of the differences between four different backup approaches with the optimal policy after convergence. We show that optimizing the Q-function for each token with BAD (POAD) is equivalent to backup for the full action (TWOSOME), while others are not.
17

---

D Pseudo Code of POAD
Algorithm Policy Optimization with Action Decomposition
1: Input: LLM p. Learning rate a, MDP M = AV,S,A,T, R,,)
2: Initialize: {p. 0); Va (p,0);:0 < 0
3: Initialize: Data buffer D +- 0
4: for each epoch do
5: for 0 to
do
6: Collect St-
7: Og ~ TT(-(1)).
8: St+1 a())s)))))
9: D+ DU ((1, at, R(s,, az), 8141))-
end for
11: Sample a mini-batch B from D.
12:
for each time step / in B do for each token in dz do
13:
if j < | then
Vtorg - Va( Wi 1:5+1
15:
16:
else if j url then
17:
Vtarg <- R(st, at) + 7 (V((++1,
18:
end if
19:
Vo(se, u.
1:3)
20:
Estimate Ae with GAE or Utarg end for
21:
22:
end for
VVBBB (v [barg)?] Φ + Φ - aVoEs(L(d))
24:
25:
0+0
end for
E Detailed Description of Experimental Environments
Figure visually shows the Overcooked and VirtualHome tasks.
Overcooked is proposed as typical deicision beicision-making environment in Tan et al. 5 There are two tasks in which a language agent is aiming to make and serve tomato salad and tomato-letiuce salad respectively, with the provided ingredients and tools the tasks, the agent needs to explore and learn the correct order to cook the dish with the provided macro-actions, such as Chop, Ger-Tomato, and Go- Pn--Cutting-oatte observable. The agent only observes the objects within 5x5 square centered the agent. The reward involves +0.2 for chopping a correct ingredient, +1 terminal reward for delivering the correct dish, -0.1 for delivering any wrong item, and -0.001 for every time step.
7x7 grid world as kitchen. To solve
The environment is partially
VirtualHome is a simulated physical household enviromment proposed in Tan et al. (.. In this environment, an agent is placed in fully-furnisher household with various objects. There are two tasks in environment, microwave in the kitchen. In the generalization tasks, we replace the pancake with other foods like hamburger and pizza, or replace the (pancake, microwave) at the same time with (dishes, dishwasher) or (clothes, washing machine), to evaluate agents' generalization abilitics for similar but unseen tasks. The second one is to plan to have some entertainment while watching TV, for example picking up chips and milk in the kitchen, bringing them to the living room, turning the TV, sitting on the sofa and enjoying. In order to solve the tasks, the agent need to use macro-actions to interact with the environment such walk the living room, furn the TV and sit the sofa. The environment is partially observable. Both tasks adopt a sparse reward setting, only when the task is finished, will the agent receive +1 reward.
first one is to find the cold pancake on the table and heat it with the
DataSciCoding an environment we developed for data science coding tasks with open action space. We adopt 3 Kaggle datasets and 3 OpenML datasets [52) with details in Appendix E.I In each task,
18

---

Tomato Salad
Tomata-leltuce Salad
(c) Food Preparation
(d) Entertainment
Figure 7: Visual demonstrations of Overcooked and VirtualHome tasks [5]
agents aim to build the most effective classifier with the scikit-learn module, striving to achieve the highest possible ROC AUC score [53] on test sets. As the prompts provided the agents contain no detailed information about task datasets, agents are required to align their knowledge to different datasets, that is, they must continuously modify and debug their code based on feedback from the runtime environment until it works. Therefore, is a sequential decision-making process. Through this interaction, they align their approach to the specific task dataset. When the codes are executed successfully, a done signal will occur immediately, along with an ROC AUC score E [0, 1J. If the codes run wrong, the agent receives constant -1.0 penalty. Adopting the same evaluation metrics as Hollmann et al. (54), for each dataset, we evaluate repetitions, each with random 50% 50% train-test split 55 and record the average ROC AUC score across these splits.
E.1 Details of Datasets Used in DataSciCoding Tasks
These datasets are collected from Kaggle and OpenML. respectively, with pruning approaches consistent with Hollmann al. 54
Table 3: Details of task datasets in DataSciCoding, where [K] denotes Kaggle.
Data set
#Features
#Samples
#Classes
Pharyngitis [K]
19
512
Health Insurance [K]
13
2000
2
Spaceship Titanic [K]
13
2000
2
Airlines
7
2000
2
Balance Scale
125
Breast-w
9
69
2
F Case Study in Token-Level Credit Assignment
In this section, we conduct a case study for in-depth analysis to validate the effectiveness BAD in terms of teven-level credit assignment. We selected the last three states from the Food Prepa-
19

---

Adv of each Laken at first state
Adv of each Loken at last state
Heal the Pancake with Microweve
POXD
-44
NTPO
Figure 8: Case Study; Demonstration of the token-level credit assignment learned by the BAD at two states (left two). And a comparison of the volume of credit assignment for key tokens between different methods (right one), where TWOSOME indicates the credit assigned to entire actions instead of specific tokens.
ration task, which form complete subtask: Heat the Pancake with Microwave. This serves as a simple yet practical case for our analysis, with transitions: T(5g. "open the microwave") -
81, T(so, "close the microwave") + 50. T(si, "put the pancake into the microwave") + 52. (82,, "open the microwave") + 82, T(22, "close the microwave") - success with reward 1.0. Ac- cording to these transitions, the optimal trajectory to complete this subtask is (open the microwave, put the pancake into the microwave, close the microwave). Additionally, the maximum step length for the task is 5; if the task is not completed within 5 steps, it is considered failure, resulting in reward signal of -1.0.
Based on this, we first sample 1,000 examples using random policy, then train a critic to convergence with three different backups: action-level Bellman backup (TWOSOME), naive token-level Bellman backup (NTPO), and BAD (POAD). In Figure 3, we recorded the credit assignment results with each critic for each token in positive and negative actions under the first and last states.
The first two in Figure 8 show the credit assigned to each token in a symmetric (positive, negative) action pair by the BAD-Iearned critic in two different states, results confirm that the BAD critic can effectively assign most of the credit to the key tokens while rarely influencing other irrelevant tokens. The right one illustrates the volume of credits assigned key tokens by POAD and NTPO, compared with the credit assigned to the entire action by TWOSOME, showing that POAD enjoys far smaller discrepancy than NTPO.
G Detailed Comparison between POAD and CAAFE
Table 4 shows the performance of the best code discovered during POAD's training process with CodeLLaMA-7B, comparing with CAAFE with GPT-3.5 and GPT-4.
Table 4: The performance of the best code discovered during POAD's training process with CodeI.LaMA-7B, comparing with CAAFE with GPT-3.5 and GPT-4. [K] denotes that the dataset is collected from Kaggle, while others collected from OpenML.
Task
POAD-Best CAAFE(GPT-3.5)
CAAFE(GPT-4)
[
0.5939
0.5745+0.02
0.5748
pharyngitis(K]
0.7282
0.6976-00.3
0.7078110.4
spaceship-titanice
0.862810.01
0.838310.02
0.8405
airlines
0.664
0.61910.04
0.620
balance-scale
0.9651110.03
0.844
0.88210.26
breast-w
0.9981+-0.11
0.9809+0.02
0.9809+00.02
H Wall-time of POAD on DataSciCoding Environment
Table Sisows the average wall-time of training LLaMA2-7b with LoRA and POAD on all DataSci- Doing tasks.
20

---

Table 5: Average wall- for POAD training with each dataset with 1 Nvidia A100.
Task
Wall-Time
Environmental Steps
health-insurance([]]
2h 34m
10k
pharyngitis{K? [
2h 11m
10k
spaceship-titante titanic[K]
3h 5m
10k
airlines
2h 56m
10k
balance-scale
Ih 43m
10k
breast-w
2h 6m
10k
I Details in Generalization Experiments
Table 6: Comparison of generalization performance on eight unseen tasks, with episodic returns averaged over 100 episodes (The fewer timesteps consumed, the higher the returns). In these tasks, we replace the pancake in the original Food Preparation task with other foods like cheese, hamburger, apple pie and pizza, or replace the (pancake, microwave) at the same time with (dishes, dishwasher) or (clothes, washing machine) for greater differences. In parentheses the success rate of completing the task within 50 timesteps.
Methods
Cheese
Hamburger
Apple Pie
Pizza
Washing Plate
Laundry
LLaMA2-7B 0.1351(0.55)
0.1342(0.55)
0.1656(0.61)
0.1409(0.55)
0.0527(0.27)
0.0344(0.15
TWOSOME
0.7119(10)
0.7058(1.0)
0.7304(1.0)
0.7047(1.0)
0.7031(1.0)
0.6038(1.0
NTPO
0.7428(1.0)
0.7476(1.0)
0.7141(1.0)
0.7355(1.0)
0.7491(1.0)
0.5687(1.0)
POAD
0.7553(1.0)
0.7602(1.0)
0.7650(1.0)
0.7625(1.0)
0.7075(1.0)
0.7014(1.0)
Evaluation on NLP Benchmarks
To investigate the impacts of fine-tuning different RL methods on the LLM's origi- nal abilities, evaluate the models trained
Table 7: Zero-shot performance on Common Sense
Reasoning tasks
by POAD, NTPO and TWOSOME on widely used NLP benchmarks(56) All models are trained the Food Prepa- ration task. The four benchmarks ARC_Challenge HellaSwag. PIQA and MMLU, with results ARC @RC_Challenge, HellaSwag and PIQA are displayed in Ta- ble and the results of MMLU are dis- played in Table g All results are calculated following the default configurations in Gao et al. [56]
Methods
ARC_ C
HellaSwag
PIQA
LLaMA2-7B
0.4352
0.5713
0.7807
TWOSOME
0.4445
0.5819
0.7786
NTPO
0.4428
0.5825
0.7824
POAD
0.4471
0.5856
0.7797
21

---

Table 8: Details of Zero-shor shot performance on Massive Multitask Language Understanding tasks
Tasks
LLaMA2-7B
TWOSOME
NTPO
POAD
Dostract_algebie
0,3100
0,2900
0,3300
0.2900
anatomy
0,4296
0,4296
0,4444
0,4593
astronomy
0,4474
0.4145
0.4342
0.4079
business_etho etthietthcc
0,4300
0.4300
0.4400
0.4500
0.4868
0.4566
0.4792
0.4604
college. biology
0.4236
0.4097
0,4236
0.4306
college ceeemistry
0.2900
0.3000
0.2700
0.2700
college computer science
0,2900
0.3000
0,3100
0,3000
college_mathematte _mathematiisss
0,3700
0,3900
0,4000
0,3900
college_medicine
0.4451
0.4104
0,4220
0.4393
jollege_physics
0,2451
0.1961
0.2157
0.2157
0.5400
0.5100
0.5300
0.5
seceettual_phsiieee
0.3532
0.3660
0.3489
0,3489
econometrics
0,2193
0.2105
0.2018
0.1930
electrical_ ennirieriigginerrinngggog
0.3724
0.3586
0,3931
0.3655
elementary mathematics cemeentary_matthematiis
0,2460
0.2619
0.2672
0,2593
formal Jogic
0,3413
0,3571
0.3175
0,3175
global_fach
0.3100
0.2900
0.3400
0,2800
0.4516
0.4548
0.4452
0.4548
0.3300
0.3054
0.3350
0.3153
high school computer science
0.3800
0.4000
0,4100
0.4100
high school european history
0.6000
0.5879
0.5636
0.5879
high school geography
0.4343
0.4545
0.4444
0.4495
high_schhe nigh_school___ Lthool_govermmet and politics
_
0,5389
0,5181
0,5233
0,5181
high_seo nigh_school__ee wal__sschoooeeeee macrocconomics
0.3769
0.3718
0.3821
0.3821
high_school mathemmatics
0,2704
0,2593
0,2481
0,2630
high_schoo sshool microeconomics
0.3739
0,3739
0,3655
0,3487
high_school_pee schsol_physics
0.2781
0.3179
0.2848
0.2914
fect_school___ee bosccooosyyh
0.5248
0.5376
0.5138
0.5174
gead_sshool_staee
0.2731
0.2824
0.2778
0.2546
high_school
school us. history
0,5196
0,5441
0,5343
0,5049
high school world history
0,5612
0.5823
0,5696
0.5696
human aging
0.4350
0.4484
0.4395
0.4260
human_see seanantyexulltty
0.5649
0.5267
0.5344
0.5191
international_N
0.5620
0.5785
0.5620
0.5702
jurisprudence
0,4722
0.4722
0,4630
0.4537
logical fallacies
0.4785
0.4663
0.4663
0.4847
machine learning
0.3929
0.3839
0.3571
0.3571
management
0.4466
0.4272
0.4272
0.4175
marketing
0.6026
0.6282
0,6239
0.5940
medical genetics
0.4900
0,5000
0,5000
0.4800
miscellaneous
0.5428
0,5428
0,5581
0,5504
moral _iiruu_ssputes
0.4480
0.4249
0.4451
0.4364
moral scenarios
0.2402
0.2413
0.2559
0.2425
nutrition
0.4771
0.4869
0.4739
0.4739
philosophy
0.4887
0.4695
0,4952
0.4823
prehistory
0.4660
0,4568
0.4784
0.4691
professional accounting
0.3546
0.3511
0.3617
0.3511
professional_ie law
0.3096
0.3181
0.3136
0.3123
professional_meddine medicine
0,4118
0.4228
0.3713
0.4044
iiffessional_ssychollee
0.4199
0.4248
0.4167
0.4085
public_relations relations
0.4182
0.4636
0.4273
0.4364
security_sture studies
0.4735
0.4694
0.4571
0.4857
sociology
0.6020
0.6119
0.6269
0.6169
us. foreign policyy
0.6700
0.6400
0,6500
0.6400
virology
0.4217
0,4337
0.4157
0.4036
world religions
0.6140
0.6023
0.5965
0.5965
22

---

K Hyper-Parameters Settings of Experiments
During experiments, the implementations TWOSOME are consistent with their official repositories, reproducing the original best-performing status. We first list the grid search in Table[9] Then, we show the each method and environment in Table 10
candidates used for that achieve the best performance for
Table 9: Hyper-Parameters Parameters candidates for grid search in Overcooked, VirtualHome, and DataSciCod- ing environments.
hyper-paramettrs
candidates
critic Ir actor ppo epochs num mini batch
le-3.5e-4.1e445555.ee1..
5e-4, 1e-4.11-4,55-,,1-55556,, 1e-6,-e--
1,5
2,,.8,16,22
gamma entropy coef max grad norm
0.99,0.95
0.01,0.001.0.0011
10,0.5
Table 10: Hyper-parameters used for POAD in Overcooked tasks.
hyper-parameters
value
hyper-parameters
value
hyper-parameters
value
critic
1e-5
actor Ir
5e-7
ppo epochs
5
batch size
128
num mini batch
2
gamma
0.99
rollout threads
4
entropy coef
0.01
max grad norm
0.5
PPO clip
0.2
value coef
0.5
KL threshold
0.02
Table 11: Hyper-parameters used for NTPO in Overcooked tasks.
hyper-parameters
value
hyper-parameters
value
hyper-parameters
value
critic Ir
le-5
actor Ir
5e-7
ppo epochs
5
batch size
128
num mini batch
4
gamma
0.99
rollout threads
4
entropy coef
0.01
max grad norm
0.5
PPO clip
0.2
value coef
0.5
KL threshold
0.02
Table 12:
used for TWOSOME in Overcooked tasks.
paraeerarameters
value
hyper-parameters
value
hyper-parameters
value
critic
le-5
actor
5e-7
ppo epochs
batch size
128
num mini batch
4
gamma
0.99
rollout threads
4
entropy coef
0.01
max grad norm
0.5
PPO clip
0.2
value coef
0.5
KL threshold
0.02
23

---

Table 13: Hyper-parameters used for POAD in VirtualHome tasks.
prpee-parameters
value
nyper-parameters
value
nyper-parametent
value
critic Ir
5e-5
actor Ir
lc-6
ppo epochs
5
batch size
128
num mini batch
2
gamma
0.95
rollout threads
4
entropy coef
0.0001
max grad norm
0.5
PPO clip
0.2
value coef 0.5
KL threshold
0.02
Table 14: Syype--arametery used for NTPO in VirtualHome tasks.
paper-parameters
value
hyper-parameters
value
hyper-parameteer
value
critic
5e-5
actor
le-6
ppo epochs
5
batch size
128
num mini batch
2
gamma
0.95
rollout threads
4
entropy coef
0.01
max grad norm
0.5
PPO clip
0.2
value coef
0.5
KI. threshold
0.02
Table 15: Hyper-parameters used for TWOSOME in VirtualHome tasks.
hyper-parameters
value
hyper-parameters
value
hyper-parameters
value
critic Ir
5e-5
actor Ir
le-6
ppo epochs
batch size
128
num mini batch
4
gamma
0.95
rollout threads
4
entropy coef
0.01
max grad norm
0.5
PPO clip
0.2
value coef
0.5
KL threshold
0.02
Table 16: Hyper-parameters used for POAD in DataSciCoding tasks.
Hlper-parameterr
value
hyper-parameteer
value
hyper-parameters value
critic Ir
5e-5
actor Ir
1e-6
ppo epochs
batch size
128
num mini batch
4
gamma
0.95
rollout threads
4
entropy coef
0.01
max grad norm
0.5
PPO clip
0.2
value coef
0.5
KL threshold
0.02
Table 17: Hyper-parameters used for NTPO in DataSciCoding tasks.
hyper-parameters
value
hyper-parameters
value
hyper-parameters
value
critic
5e-5
actor Ir
le-6
ppo epochs
batch size
128
num mini batch
4
gamma
0.95
rollout threads
4
entropy coef
0.01
grad norm
0.5
PPO clip
0.2
value coef
0.5
KL threshold
0.02

---

L Step-by-Step Breakdown of BAD
In this section, facilitate the understanding of how BAD precisely assigns credit each token, step-by-step breakdown of the credit assignment process with BAD is shown in Figures 13
TV
Tum
TV
Falled!
Palee!
T
Goal
Goal
(OONEEXII
Complnted
Compltedd
TV
Tum
(a) Step 1
(b) Step
Figure 9: Step 1: Receiving feedback rt and ry for positive and negative trajectories. Step 2: Propagating rf and ry to Q("TV"lo, "Turn on") and Q("TV"|oo, "Tumn off").
Turn
Faled!
Faled!
ry
Goal
[
Goal
[DONTEXI1]
of
Completed
Complsed
Turn
on
TV
Tum
C0
(a) Step
(b) Step
Figure 10: Step and 4: Back-propagatingr? rf' and ry Q("on" 0z, "Turn") and ("foff"" or, "Turn"". and then to Q( Q("Tum""lor )ot) in both trajectories continuously.
x
7
Goal
Goal
JCONTENTI
Mean r. of
G- -
di-i 0
Completed!
Completed!
xxx
7
(a) Step 5
(b) Step
Figure 11: Step 5: Back- propagating to the previous action with - = Mean x ...r x q.: since both trajectories share the same pre-action. Step 6: Similarly, Q( Q("Tur"'(o.) is also shared among trajectories, thus Q(*Tum"(oo) Mean (g g0 0 as well.
25

---

7
07
0
02
Goal
of
Mosn
Goal
[CUNTEXI]
Meun
rf of
G--1
Mr
Completed
Completed!
7xx
7
(a) Step 7
Step
Figure 12: Step 7: Modifying of and qf in both trajectories with the same Q("Tum""lo) ol:)). Step 8: Starting to pack-propagate again for credits assigned to each token for optimization with & --'²-1
for q'-
O; this process is similar to the calculation for advantage values.
xx
xxx
Tum
off
TV
Tum
TV
Fi lell
Feeldl
4e
Goal
(PONTTXII]
Mean
Goal
JOONEXXII
g
Mean 1
Comploted!
Complotedi
Tum
TV
TV
(a) Step 9
(b) Step 10
Figure 13: Step 9: Calculating the credits for key tokens, where 5, = gf 9f and 57. = 4
90-
Step 10: Calculating the credits for tokens with 0 qf z | and 59. z0 / Till now, we precisely emphasized key tokens while keeping irrelevant tokens with 0 credits.
26

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
Section
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how be expected to generalize other settings.
much the results
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Section
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
27

---

Justification: Section 3.1 and Appendix A QB
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: Section 6. and Appendix K
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
28

---

Answer: [Yes]
Justification:
Code is included in the zip file with paperr-prameterss Appendix K
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips...
cessionPolccy for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/ /pipp/.cc/pulli///uuii wwwtc//ggui////eeeeceoocc CoodeSubeiiie ssi asionnolli onnolicy)))) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Section 6.1 and Appendix Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: Secton 6
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
29

---

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of enrors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Section 6
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g.. preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: /ttp:::////uis..eeeeeeeee httpp//csccppp..t.ttttttacteeeee
Answer: [Yes]
Justification: This work has no ethical violation.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This work has societal impact since it is theoretical research.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to
30

---

generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative socictal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: This paper poses such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: Section 6.1
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, wwppppswpptseeee whworee.wtcche....c.. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's is creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and the documentation provided alongside the assets?
31

---

Answer: [NA]
Justification:
This paper does not release new assets.
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details the dataset/coddemoodll as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: This paper does not involve crowdsourcing research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does not involve crowdsourcing research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
32


Reviews:
[{"review_id": "C8GfomiGo3", "reviewer": "Reviewer_H4uQ", "summary": "This paper introduces Policy Optimization with Action Decomposition (POAD), a novel method for reinforcing language agents by optimizing at the token level rather than the action level. The authors derive a theoretical framework called Bellman backup with Action Decomposition (BAD) to address discrepancies between action-level and token-level optimization. Implementing BAD within PPO, POAD offers finer-grained credit assignment and lower optimization complexity. Experiments across various environments demonstrate POAD's improved learning efficiency and generalization compared to baseline methods, supporting the theoretical analysis and advancing language agent optimization in interactive settings.", "questions": "1. How did you tune the baseline?\n2. If you use full scale finetuning instead of Lora, what do you expect? Could you perform on experimentation on this?\n3. How would this method work on standard NLP task?", "limitations": "Yes, author addressed the limitations.", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The paper seems to provide a theoretical foundation for its approach, analyzing the discrepancies between action-level and token-level optimization and deriving the Bellman backup with Action Decomposition (BAD) to address these issues. This theoretical grounding lends good analysis to the proposed method.\n\n2.  The authors implement their theoretical insights into a practical algorithm (POAD) and demonstrate its effectiveness across diverse testbeds, including environments with both restricted and unrestricted action spaces. The empirical results show improved learning efficiency and generalization abilities compared to baseline methods, validating the practical utility of the theoretical contributions.", "weaknesses": "Limited testing environment and lack of baseline. The author only show two baselines for comparison. More baselines should be include for comprehensive analysis. \n\nAnother main concern is that the finetuning is only done on using Lora, when finetune the model using the full capacity, the effect is unknown.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### We thank Reviewer H4uQ for his/her constructive comments that will surely turn our paper into a better shape.\n\n> **Q1** More baselines should be included for comprehensive analysis.\n\n**A1** We appreciate the feedback from H4uQ and also extend our gratitude to Reviewers 671a and 9Wzo for recommending the recent relevant method, ArCHer [1]. ArCHer aims to provide fine-grained token-level supervision for optimizing language agents on interactive (multi-turn) tasks. Upon further investigation, we found that its key idea for token-level credit assignment plays an intermediate role between NTPO and POAD. Thus, **in our revised manuscript, we have incorporated ArCHer as a new baseline in our experiments with comparative analysis.** \n\nPlease refer to the **Figures 1 and 2: New Baseline, ArCHer and Corresponding Analysis** section in meta-responses for detailed comparison between ArCHer and our methods, as well as preliminary experimental results. We believe the extra comparisons and analyses involving ArCHer will better illustrate the advantage of our method and the research trajectory in the community toward providing finer-grained supervision for language agents. \n\n> **Q2** If you use full-scale finetuning instead of Lora, what do you expect? Could you experiment with this?\n\n**A2** We conducted an ablation comparing our method with the action-level baseline, TWOSOME, on both full-scale and LoRA fine-tuning settings, with results referring to **Figure 4 (a) in our meta-responses**.\n1. Comparing POAD's performance under full-scale fine-tuning and LoRA fine-tuning, we observed consistent performance, indicating that POAD remains effective under full-scale fine-tuning. However, due to increased memory consumption with full-scale fine-tuning, we significantly reduced the batch size during training, leading to slightly reduced stability in the training process. This accounts for the observed larger fluctuations in the curves during full-scale fine-tuning compared to LoRA fine-tuning.\n2. Contrasting the two methods under full-scale fine-tuning, POAD demonstrated greater advantages over the action-level baseline method compared to Lora fine-tuning. Potential reasons for this include: a) full-scale fine-tuning's heightened sensitivity, requiring higher precision and granularity of supervised signals compared to Lora; b) the further increase in demand for precision of supervised signals due to the drastically reduced batch size.\n\nWe greatly appreciate the reviewer for this valuable suggestion. These experimental results further underscore the importance of our approach, highlighting that in scenarios involving full-scale fine-tuning, fine-grained and precise supervision signals are more critical compared to situations involving LoRA fine-tuning.\n\n> **Q3** How did you tune the baseline?\n\n**A3** Regarding the baseline, such as TWOSOME, we integrated its implementation into our code framework based on the official repository. Throughout this process, for fair experimental comparisons, we ensured the following:\n1. Key components of the baseline method aligned with those in the original repository.\n2. Be able to replicate the performance reported in the paper using the original hyper-parameters list in their papers.\n3. Regarding the code framework, apart from necessary adaptations specific to the method, we maintained consistency in other aspects.\n4. When tuning hyperparameters, we utilized the same search ranges to maintain consistency (we reused the search ranges as defined in TWOSOME).\n\n> **Q4** How would this method work on standard NLP tasks?\n\n**A4** Theoretically, our method can also be applied to standard NLP tasks such as QA, writing, coding, etc (DataSciCoding in our experiments is a coding environment where agents can generate code freely). Notably, the most suitable application scenarios would be those language tasks involving multi-turn interactions (which language agents usually target). In multi-turn situations, our method's advantage in decomposing inter-action and intra-action credit assignments can be fully leveraged. For other single-turn tasks, in the worst-case scenario, our method would be similar to standard RL approaches with terminal rewards derived from environmental feedback or a reward model.\n\n---\n[1] Zhou, Yifei, et al. \"Archer: Training language model agents via hierarchical multi-turn rl.\" arXiv preprint arXiv:2402.19446 (2024).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Puyd7mjece", "reviewer": "Reviewer_JxkT", "summary": "The paper investigates LLM agents: RL agents where sampling actions from a policy means sampling a sequence of tokens mapping to an action from a (suitably conditioned) large language model. Authors notice a problem with previous implementations of this idea: since actions are typically described as a sequence of more than one token, this introduces issues around credit assignment between tokens contributing to an action. However, fixing it naively, treating each token as an individual action, leads to a problem: when time discounting is present in the MDP, the new (intra-token) MDP is no longer equivalent to the old one. The proposed solution is to only discount \"fully formed\" actions. This idea is then integrated with standard PPO implementation, and empirically tested against SOTA on three RL environments, showing a meaningful improvement. In addition, the paper presents evidence that another fix of simply removing the discount factor (ie setting it to 1) degrades the performance, and through ablation study proves that it is indeed the intra-token discounting removal that improved the final performance.", "questions": "1. There is a well-known equivalence between time-discounted MDP and adding an auxiliary state $\\s_0$ with zero reward and transition probabilities $T(s, s_0) = \\gamma$ for every other state $s$. As far as I understand, it fits into the paradigm in the paper by postulating that additional transitions are added from completed actions only - do authors have any additional insight into this?\n2. Did authors thought about incorporating the discount factor trick into ReAct? Since the impact of $\\gamma$ is rather large (from experimental evidence), it seems that distinguishing between thoughts (potentially incurring low discount penalty) and true env actions might be a potentially useful direction.", "limitations": "Yes.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper is written clearly, and investigates an important topic. The experiments convinced me that indeed, the method works as described and improves the performance of LLM agents. I also highly appreciated the investigation into removing time discounting completely and the impact of varying intra-token discount factor, and into the possible degradation of NL capability of the model after fine-tuning.", "weaknesses": "I am sceptical with regards to the formal treatment of POMDPs in the paper. First, the definition is missing the initial state (/initial distribution). Then, there are numerous issues: the reward function is defined as operating on states and actions, but is then used as taking an observation and an action. Expectations in equations (1) and (2) are ill-defined if conditioned only on observations: they either have to be conditioned on states, or the observation has to somehow induce a probability distribution over states. The definition of Bellman updates for value function update in eq (5) and (10) refer to the next action/next token, which haven't been defined, and also uses $|a_t|$ which is ambiguous for a given sequence $w_t^{1:j}$. The equations (11) and (12) refer to the next observation $o_{t+1}$ which is a random variable (first, because POMDP, second, because the transition function is not assumed to be deterministic).\n\nI understand that the method was implemented and works fine, and the intuition behind it is clear, so I believe that these mistakes are fundamentally not that serious. However, they confuse and frustrate careful reader, and have to be fixed. (In comparison, TWOSOME paper uses simple MDP formalism, while GLAM paper sweeps the formalism of POMDPs under the rug by informally talking about \"approximations\").\n\n---\n\nSmall notes:\n - typo in \"policy\" in line 173\n - Figure 2 was very hard to read, and referring the reader to the appendix to understand it - which is, I think, necessary - seems to warrant additional work to make it more readable\n - \"NTPO\" typo in the description of Figure 4\n - LHS and RHS subfigures are swapped in Figure 5", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### We thank Reviewer JxkT for his/her constructive comments that will surely turn our paper into a better shape.\n\n> **Q1** The unclear formalism and notation errors.\n\n**A1** We apologize for the confusion caused to the reviewers and readers. Indeed, we recognize that the use of POMDP in our formulation was unnecessary, as our paper does not delve into discussions related to partial observability. The inaccurate use of POMDP and the resulting omissions have inadvertently introduced additional confusion and distraction for the readers. To address this issue, in the updated version of our manuscript, we have:\n\n1. referring to TWOSOME, replaced the POMDP formulation with MDP and adjust all related notations accordingly. We believe this change will make our formulation and derivation more clear and concise.\n2. after these modifications, thoroughly reviewed all notations and equations to ensure that all relevant variables are properly defined and clearly explained.\n\n> **Q2** There is a well-known equivalence between time-discounted MDP and adding an auxiliary state  $s_0$ with zero reward and transition probabilities $T(s,s_0)=\\gamma$ for every other state $s$. As far as I understand, it fits into the paradigm in the paper by postulating that additional transitions are added from completed actions only. Do authors have any additional insight into this?\n\n**A2** Yes, we agree with the reviewer's insight about the equivalence between adding the additional transitions for only completed actions, i.e., action-level, and our BAD paradigm. Meanwhile, we appreciate the reviewer providing us with another perspective to understand why removing the intra-action discount is necessary. Specifically, in environments with linguistic actions, the dynamics between intra-action tokens within an action are strictly deterministic. That is, given a token $w^j$ for $(s,w^{1:j−1})$ , we can only transition to $(s,w^{1:j})$. In this case, if we do not remove the intra-action discount, it would be equivalent to adding an extra transition for intra-action tokens, which contradicts the deterministic nature of linguistic action generation.\n\nAdditionally, we can consider the implications of setting different discount factors for intra-action and inter-action tokens from another perspective. Originally, the discount factor was used to balance short-term and long-term rewards, with a large discount factor typically applied in scenarios where actions within the horizon are strongly correlated. In our work, given the particularity of intra-action tokens—where meaningful environmental feedback can only be obtained if intra-action tokens can be combined into a valid action—it is crucial to encourage agents to generate highly correlated intra-action tokens. Therefore, it is reasonable to consider setting the intra-action discount factor to 1, then agents will evaluate each of its intra-action tokens based on the action-level feedback. This approach aligns with the core idea of solving specific MDPs using time-varying discount factors, as seen in works like [1].\n\n\n> **Q3** Did the authors think about incorporating the discount factor trick into ReAct? \n\n**A3** Yes, we recognize the potential benefits of integrating different discount factors into distinct components of the language agent's outputs, such as actions and thoughts. Based on the following intuitions:\n\n1. The thoughts do not directly affect the final state-action transition—only the final action does—instead, thoughts guide the generation to converge to specific low-level actions (they play different roles in interactive tasks). \n2. We can remove the gamma in intra-action tokens since words later in an action may not be more expressive than earlier ones. However, thoughts usually present a progressive structure, i.e. content closer to the final action may be more closely related to the final action, which makes it reasonable to distinguish them from final actions, e.g. with a low discount factor.\n\nFor now, we have observed that discount factors significantly influence the final policy of language agents learned from RL algorithms. Thus, considering separate outputs with varying intentions in a ReAct-like agent, we believe applying different discount factors to these components could be beneficial, whether through heuristic methods or learning-based approaches. For instance, [2] proposed a gradient-based meta-learning approach to optimize the discounting ratio of future rewards, while [1] solves specific MDPs using time-varying discount factors. Both could be promising directions to explore in the future.\n\n> **Q4** Figure 2 should be more readable.\n\n**A4** Thanks very much for the reviewer's suggestion. We are working hard to improve the readability of this figure. **Figure 4 (b) in the meta-responses** is a current version of this diagram, hopefully more readable than the previous one. Besides, to facilitate understanding the effects of BAD shown in this figure, we also conducted a case study (See the **Figure 3: Case Study** section in meta-response) to visually demonstrate the token-level credit assignment results learned by the BAD.\n\n> **Q5** Typos and subfigures swapping issue.\n\n**A5** Thanks to the reviewer’s careful examination, we have addressed these typos and adjusted the figures in our updated version.\n\n---\n[1] Gan, Jiarui, et al. \"Markov decision processes with time-varying geometric discounting.\" Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 37. No. 10. 2023.\n\n[2] Xu, Zhongwen, Hado P. van Hasselt, and David Silver. \"Meta-gradient reinforcement learning.\" Advances in neural information processing systems 31 (2018).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xhjgxl8MIo", "reviewer": "Reviewer_671a", "summary": "The paper proposes a novel approach to optimizing language agents in reinforcement learning (RL) environments, addressing the challenges of limited environmental knowledge and vast action spaces. Traditional methods like GLAM and TWOSOME optimize language actions as whole units, leading to inefficiencies in credit assignment and optimization complexity. This paper introduces Policy Optimization with Action Decomposition (POAD), which decomposes actions to the token level, allowing for finer-grained supervision and manageable optimization complexity. The authors derive the Bellman backup with Action Decomposition (BAD) to ensure theoretical consistency and integrate it with the PPO algorithm. They validate POAD across various testbeds, demonstrating its advantages in learning efficiency, generalization, and theoretical correctness.", "questions": "1. Lines 201-203, the authors mention that the action space is reduced from exponential to polynomial compared to works like GLAM and TWOSOME. However, from Equation 16, it can be seen that the algorithm learned by BAD still samples each token from the entire vocabulary, which does not seem to change the action space?\n2. Can the authors demonstrate the token-level credit assignment learned by the BAD method through a case study?\n3. BAD has a significant correlation with some existing token-level optimization methods [2][3] and hierarchical optimization methods [3]. Can the authors make a more in-depth comparison between BAD and these methods? Is there any equivalence between BAD and these methods? Is it just extending existing methods, such as [2][3], from single-turn to multi-turn?\n\n---\n\n[1] Zhou, Yifei, et al. \"Archer: Training language model agents via hierarchical multi-turn rl.\" arXiv preprint arXiv:2402.19446 (2024).\n\n[2] Rafailov, Rafael, et al. \"From $ r $ to $ Q^* $: Your Language Model is Secretly a Q-Function.\" arXiv preprint arXiv:2404.12358 (2024).\n\n[3] Zeng, Yongcheng, et al. \"Token-level Direct Preference Optimization.\" ICML 2024", "limitations": "The author has already discussed the limitations of the method and its potential impacts in the paper, and provided possible solutions.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- This paper introduces a novel method of decomposing actions to the token level, which provides finer-grained supervision and reduces optimization complexity.\n- The derivation of the Bellman backup with Action Decomposition (BAD) ensures theoretical consistency, addressing a missing gap in previous research.\n- The paper provides many empirical validation across different testbeds, demonstrating the effectiveness and efficiency of POAD compared to baseline methods.", "weaknesses": "1. This paper lacks a thorough comparison and analysis of some important pioneering works, such as LLM decoding from a control or reinforcement learning perspective and hierarchical reinforcement learning.\n2. This paper emphasizes that BAD can perform efficient and reasonable credit assignment at the token level, but the experimental section lacks in-depth analysis, such as case studies.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### We thank Reviewer 671a for his/her constructive comments that will surely turn our paper into a better shape.\n\n> **Q1** BAD has a significant correlation with some existing token-level optimization methods [2][3] and hierarchical optimization methods [1]. Can the authors make a more in-depth comparison between BAD and these methods? Is there any equivalence between BAD and these methods? Is it just extending existing methods, such as [2][3], from single-turn to multi-turn? \n\n**A1** We appreciate the reviewer for recommending these relevant works. Firstly, [2] and [3] primarily focus on single-turn RL problems, performing decomposition and token-level optimization on single-step outputs. While these works, including [3] here, emphasize the importance of fine-grained token-level supervision, they didn’t take policy consistency into account (which has a minor impact on single-turn tasks). Our work, however, concentrates on multi-turn scenarios, where the discrepancy grows with the number of turns, making it an unavoidable challenge in multi-turn RL. Moreover, one of the baselines in our paper - the naive token-level Bellman backup behind NTPO - can be viewed as a direct extension of [2] and [3] from single-turn to multi-turn. BAD, then, is an improvement that further addresses the inconsistency issue. This highlights both the differences and connections between our method and [2][3]. Our experimental results demonstrate that resolving the inconsistency problem, beyond simply extending, has a significant improvement on the ultimate performance of language agents in multi-turn RL problems. **We will incorporate these comparisons in our Related Work section.**\n\nAs for ArCHer [1], we are very grateful to the reviewer for recommending this timely work. ArCHer aims to provide fine-grained token-level supervision for optimizing language agents in interactive (multi-turn) tasks. Upon further investigation, we found that its key idea for token-level credit assignment plays an intermediate role between NTPO and POAD. Thus, **in our revised manuscript, we have incorporated ArCHer as a new baseline in our experiments with comparative analysis.** \n\nPlease refer to the **Figures 1 and 2: New Baseline, ArCHer and Corresponding Analysis** section in meta-responses for detailed comparison between ArCHer and our methods, as well as preliminary experimental results. Thanks again for the reviewers’ valuable recommendation, we believe these comparisons and analyses involving ArCHer will better illustrate the advantage of our method as well as the research trajectory in the community toward providing finer-grained supervision for language agents. \n\n> **Q2** This paper emphasizes that BAD can perform efficient and reasonable credit assignment at the token level, but the experimental section lacks in-depth analysis, such as case studies. Can the authors demonstrate the token-level credit assignment learned by the BAD method through a case study?\n\n**A2** We are very grateful for this insightful suggestion and **we have conducted a case study in updated manuscript for in-depth analysis to validate the effectiveness of BAD**. Please refer to the **Figure 3: Case Study** section in meta-responses for detailed settings and experiemental results. \n\nBy comparing the credit assigned to each token in a symmetric (positive, negative) action pair by the BAD-learned critic in different states, our experimental results confirm that the BAD critic can effectively assign most of the credit to the key tokens while rarely influencing other irrelevant tokens. By investigating the volume of credits assigned to key tokens by POAD and NTPO, compared with the credit assigned to the entire action by TWOSOME, results show that the POAD enjoys a far smaller discrepancy than NTPO.\n\n> **Q3** Lines 201-203, the authors mention that the action space is reduced from exponential to polynomial compared to works like GLAM and TWOSOME. However, from Equation 16, it can be seen that the algorithm learned by BAD still samples each token from the entire vocabulary, which does not seem to change the action space?\n\n**A3** We apologize for the confusion. BAD does not change the action space. In lines 200-203, we intended to convey that, by using BAD, the critic can provide more fine-grained token-level supervision signals, i.e. the token-level advantage values as shown in Equation 16, for the language agent's policy updates, which reduces the complexity of policy optimization from exponential to polynomial. Specifically, before using BAD, we have an action-level signal, which is used to optimize a policy with an action space of $O(|V|^{|a|})$. After using BAD, we have $|a|$ signals, each guiding an optimization with a token space of $O(|V|)$. Essentially, BAD decomposes a large problem into several smaller problems, thereby reducing the complexity of solving it. Thank you for your comment, we acknowledge that our initial statement might have caused some confusion, and **we will further clarify our statement in lines 200-203 based on this discussion to avoid any unnecessary confusion.**\n\n---\n[1] Zhou, Yifei, et al. \"Archer: Training language model agents via hierarchical multi-turn rl.\" arXiv preprint arXiv:2402.19446 (2024).\n\n[2] Rafailov, Rafael, et al. \"From 𝑟 to 𝑄∗: Your Language Model is Secretly a Q-Function.\" arXiv preprint arXiv:2404.12358 (2024).\n\n[3] Zeng, Yongcheng, et al. \"Token-level Direct Preference Optimization.\" ICML 2024", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "I2fkINkFVm", "reviewer": "Reviewer_9Wzo", "summary": "The paper proposes Bellman backup with Action Decomposition (BAD) and its realization on PPO (POAD), which aims to train with token-level policy for more fine-grained credit assignment on the crucial part of the language response made by the language agent. To address the issue of distorted Q-value by the intra-sentence discount factor, the authors propose a novel Bellman backup rule that applies discount factor and reward only at the end of each sentence (high-level action), and provide the realization, POAD, on PPO. Experiments have shown that the proposed method works better than several baselines on decision-making and coding tasks.", "questions": "I have a question for the authors: in line 200-202, the author claims that \"another advantage of BAD is ... reducing the complexity of RL problem with language agent to $O(|a|\\times|V|)$. While it is true that the action space at each step reduces, the episode length is greatly increased from a few steps to potentially hundreds or thousands of steps with much sparser reward, which could be potentially harder. Could the author explain intuitively why the benefit of reduced action space outweighs the problem of extended episode length with sparse reward?\n\nAlso, it will be great if the author can do the following two things (which corresponds to point 2) and 3) in weakness):\n\n1. conduct an ablation showing the intuition in Fig. 2 is actually true, i.e., for the sentence \"turn on / off TV\", the Q-value actually changes more on the word \"on / off\";\n\n2. compare with ArCHer [1], which is a sentence-level RL method, to further show that token-level Bellman update is better.  \n\n**Reference**\n\n[1] Y. Zhou et al. ArCHer: Training Language Model Agents via Hierarchical Multi-Turn RL. In ArXiv, 2024.", "limitations": "The authors mentioned the problem of acquiring quantitative reward function as the limitation, which is a valid concern in some LLM tasks such as reasoning with chain-of-thought, and gave reasonable directions for solution. The paper did not mention any societal impact, for which I would encourage the authors to include a paragraph (the paper is not \"theoretical\"; it proposes empirical solutions with experiments).", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-written and easy to follow. The paper does a good job in stating its motivation; Figure 1 clearly shows why token-level training is necessary, while Eq. 11 and 12 in Sec 4.2 helps much in understanding the inherent problem of token-level agent to be addressed in the rest of the paper. The appendix shows a good amount of details that helps the understanding of the paper, including step-by-step breakdown, hyperparameter tables, wall-clock training time and pseudocode.\n\n2. The proposed method is tested on a variety of test cases, including decision-making and coding, which shows the wide range of potential application for the proposed method. There is also ablation showing the generalizability (i.e., not overfitting) of the trained model, which is a desired property of the LLM community.", "weaknesses": "While working better than sentence-level method TWOSOME, the paper is still not convincing enough in showing that token-level critic is indeed better than sentence-level critic for three reasons: 1) there are potential disadvantages such as increased training complexity; 2) the paper did not show results that \"key token\" identified by token-level critic actually matters; 3) there are other sentence-level critic RL methods, and TWOSOME might not be the best sentence-level critic RL method. See questions section for details.\n\n**Minor issues:**\n\nThe colors of the lines are inconsistent throughout the paper. For example, in Fig. 5, POAD is orange in the left subfigure, but it is purple and blue in the right subfigure.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### We thank Reviewer 9Wzo for his/her constructive comments that will surely turn our paper into a better shape.\n\n> **Q1** Could the author explain intuitively why the benefit of reduced action space outweighs the problem of extended episode length with sparse reward?\n\n**A1** Thank you for your thoughtful question. We believe this discussion is valuable and will significantly help readers better understand our work. **We plan to add a dedicated discussion section on this topic in the revised version.**\n\nAt first glance, action decomposition indeed extends the episode length, seemingly exacerbating the sparse reward problem. However, we can analyze this issue more deeply from different perspectives: \n1. Efficiency of credit backpropagation. The increased episode length due to action decomposition could lead to inefficient training if using a 1-step TD, as credit would only backpropagate one step per training epoch. Fortunately, BAD can be combined with n-step TD or GAE to achieve efficient propagation comparable to action-level Bellman backup within a single training epoch. For this reason, our POAD implementation is based on GAE. We appreciate the reviewer's reminder and will explicitly state in future versions that BAD is best used in conjunction with GAE or N-step TD methods to avoid confusion and misuse.\n2. Credit decay. While the extended episode length could cause exponential decay of credit strength based on the discount factor, which makes it difficult for sparse reward signals to act on actions/tokens in the early stages of an episode, BAD removes the intra-action discount factor. This theoretically prevents the exponential decay of signal strength caused by action decomposition. Our experimental results, where POAD significantly outperforms NTPO, confirm this point.\n3. Credit estimation variance. Even with addressing the efficiency and decay issues, increased horizons often lead to higher variance in credit estimation. However, in linguistic action environments, the dynamics between intra-action tokens are strictly deterministic. In this case, we are reasonable to expect that the additional variance introduced by action decomposition would be much less severe than that in traditional RL problems with increased episode length.\n\nWe hope this discussion intuitively explains why the benefit of reduced action space outweighs the problem of extended episode length with sparse reward when using BAD. Besides, our experiments on the VirtualHome also empirically confirm this. When the agent only receives a binary reward signal upon task completion or reaching the maximum step limit (50 steps), our POAD still outperforms the action-level method, TWOSOME, despite such a sparse reward setting.\n\n> **Q2** Conduct an ablation showing the intuition in Fig. 2 is actually true.\n\n**A2** We are very grateful for this insightful suggestion and **we have conducted a case study/ablation for in-depth analysis to validate the effectiveness of BAD**. Please refer to the **Figure 3: Case Study** section in meta-responses for detailed settings and experiemental results. \n\nBy comparing the credit assigned to each token in a symmetric (positive, negative) action pair by the BAD-learned critic in different states, our experimental results confirm that the BAD critic can effectively assign most of the credit to the key tokens while rarely influencing other irrelevant tokens. By investigating the volume of credits assigned to key tokens by POAD and NTPO, compared with the credit assigned to the entire action by TWOSOME, results show that the POAD enjoys a far smaller discrepancy than NTPO.\n\n> **Q3** compare with ArCHer to further show that token-level Bellman update is better.\n\n**A3** We are very grateful to the reviewer for recommending this timely work. Although ArCHer [1] adopts sentence-level critics, it also aims to provide fine-grained token-level supervision for optimizing language agents in interactive (multi-turn) tasks. Upon further investigation, we found that its key idea for token-level credit assignment plays an intermediate role between NTPO and POAD. Thus, in our revised manuscript, **we have incorporated ArCHer as a new baseline in our experiments with comparative analysis.**\n\nPlease refer to the **Figures 1 and 2: New Baseline, ArCHer and Corresponding Analysis** section in meta-responses for detailed comparison between ArCHer and our methods, as well as preliminary experimental results. Thanks again for the reviewers’ valuable recommendation, we believe these comparisons and analyses involving ArCHer will better illustrate the advantage of our method as well as the research trajectory in the community toward providing finer-grained supervision for language agents. \n\n> **Q4** Minor Issue: The colors of the lines are inconsistent throughout the paper.\n\n**A4** Thanks to the reviewer for his/her careful examination, we have adjusted the related figures and unified the color used in our updated version.\n\n> **Q5** Lack of societal impact.\n\n**A5** We appreciate the reviewer's suggestion. In our updated version, we have added the following paragraph addressing the societal impact:\n\n**Social Impact.** The advancements in RL for language agents can significantly enhance decision-making processes in various domains such as healthcare, finance, and autonomous systems. Improved decision-making can lead to better outcomes, increased efficiency, and reduced errors. However, we acknowledge that when optimizing agents using our method, language agents may potentially resort to unscrupulous means to maximize rewards, which could lead to potentially harmful results. Thus, we advocate for a more comprehensive consideration when designing the reward function, or combining it with safety-constrained RL methods to mitigate these risks.\n\n---\n[1] Zhou, Yifei, et al. \"Archer: Training language model agents via hierarchical multi-turn rl.\" arXiv preprint arXiv:2402.19446 (2024).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
# Meta Responses
We are delighted to receive positive feedback from all the reviewers, and thank the reviewers for their valuable suggestions.
## Answers to some common questions and New results (see the PDF attached here)
For convenience, we provide detailed answers to some of the common questions here.
### Figures 1 and 2: New Baseline, ArCHer and Corresponding Analysis
We appreciate the feedback from reviewer H4uQ about involving extra baselines and also extend our gratitude to reviewers 671a and 9Wzo for recommending the recent relevant method, ArCHer [1]. Upon further investigation, we found that its key idea for token-level credit assignment plays an intermediate role between NTPO and POAD. Thus, **we have incorporated ArCHer as a new baseline in our experiments with comparative analysis.**

Specifically, from a theoretical perspective: Figure 1 demonstrates the key processes of TWOSOME, NTPO, ArCHer, and POAD, in terms of credit assignment. ArCHer, positioned as an intermediate between NTPO and POAD, theoretically enjoys reduced discrepancy compared to NTPO (although the discrepancy issue is not explicitly discussed in its literature). Importantly, our theoretical analysis regarding BAD is compatible with ArCHer, allowing us to apply insights gained from our theoretical analysis to optimize ArCHer.

From a practical perspective: ArCHer is implemented based on a hierarchical RL framework, utilizing an off-policy Q-network for credit backpropagation at the action-level (referred to as utterance-level in the original text, unified as action-level here for clarity). This off-policy Q-network enhances sample efficiency, especially when pre-collected data is available. At the token-level, ArCHer directly backpropagates action-level credits to each intra-action token using REINFORCE. Notably, in practice, ArCHer introduces a V-network at the action-level to facilitate computing action-level advantage values, using these advantage values (rather than Q-values) as terminal rewards at token-level, which may introduce additional inconsistency. Moreover, the use of multiple value estimation networks (including an optional token-level baseline value model besides the Q and V network) increases manual tuning efforts and may lead to cumulative bias and variance, potentially impacting stability.

We have made every effort to test ArCHer's performance in the VirtualHome environment, yielding preliminary experimental results as shown in Figure 2. In complex Entertainment tasks, ArCHer is second only to POAD, aligning with our theoretical expectations. However, in tasks such as Food Preparation where the methods’ performance gap was less pronounced, ArCHer performed poorly, likely due to instability in its system.

Furthermore, we applied insights from our theoretical analysis to enhance ArCHer's performance, i.e. removing the intra-action gamma, as depicted in ArCHer-BAD in Figure 2 (right one), showing improved results that validate the effectiveness of our insights. Nevertheless, due to the inherent challenges of excessive network complexity and difficult hyper-parameter tuning, ArCHer still falls short of matching POAD's performance.

### Figure 3: Case Study
Thanks for the insightful suggestions from reviewers 671a and 9Wzo, **we have conducted a case study for in-depth analysis to validate the effectiveness of BAD in terms of token-level credit assignment.**

**Case study settings.** We selected the last three states from the Food Preparation task, which form a complete subtask: "Heat the Pancake with Microwave". This serves as a simple yet practical case for our analysis. The corresponding transitions are as follows:

$$\mathcal{T}(s_0,\text{``open the microwave"})\rightarrow s_1$$

$$\mathcal{T}(s_0,\text{``close the microwave"})\rightarrow s_0$$

$$\mathcal{T}(s_1,\text{``put the pancake into the microwave"})\rightarrow s_2$$

$$\mathcal{T}(s_2,\text{``open the microwave"})\rightarrow s_2$$

$$\mathcal{T}(s_2,\text{``close the microwave"})\rightarrow \text{success with reward 1.0}$$

According to these transitions, the optimal trajectory to complete this subtask is (open the microwave, put the pancake into the microwave, close the microwave). Additionally, the maximum step length for the task is 5; if the task is not completed within 5 steps, it is considered a failure, resulting in a reward signal of -1.0.

Based on this, we first sampled 1,000 examples using a random policy. We then trained the critic to convergence using three different backup methods: TWOSOME (action-level Bellman backup), NTPO (naive token-level Bellman backup), and POAD (BAD). In Figure 3, we recorded the credit assignment results for each token in positive and negative actions under the First State and Last State conditions for the critics trained with these three methods.

The first two subfigures in Figure 3 show the credit assigned to each token in a symmetric (positive, negative) action pair by the BAD-learned critic in two different states, results confirm that the BAD critic can effectively assign most of the credit to the key tokens while rarely influencing other irrelevant tokens. The right one illustrates the volume of credits assigned to key tokens by POAD and NTPO, compared with the credit assigned to the entire action by TWOSOME, showing that POAD enjoys a far smaller discrepancy than NTPO.

### Figure 4 (a): Full-scale Fine-Tuning
We greatly appreciate the reviewer H4uQ for the valuable suggestion, we conducted an ablation comparing our method with the action-level baseline, TWOSOME, on both full-scale and LoRA fine-tuning settings, with results shown in Figure 4 (a). 

### Figure 4 (b)
Thanks for the suggestion from JxkT. we are working hard to improve the readability of the Figure 2 in our original paper. Figure 4 (b) here is a current version of this diagram.

---
[1] Zhou, Yifei, et al. "Archer: Training language model agents via hierarchical multi-turn rl." arXiv preprint arXiv:2402.19446 (2024).

Meta-review: This paper investigates reinforcement learning settings where sampling actions corresponds to sampling *sequences* of tokens from a suitably conditioned large language model. To address this challenge, the authors designed a method known as Policy Optimization with Action Decomposition (POAD) that allows agents' policies to be optimized at the token level rather than the action level. The paper introduces a theoretical framework called Bellman backup with Action Decomposition (BAD). BAD was explicitly designed to tackle discrepancies between action-level and token-level optimization. This is a key and necessary feature in this setting since actions are noew described as *sequences* of tokens, which introduces issues relating to credit assignment between all tokens linked to a particular action. 

The reviewers raised a few concerns. Reviewer JxkT brought up issues with the notation and formalization adopted by the authors. To the best of my knowledge, the authors acknowledged these issues and addressed them appropriately. In response to these updates and improvements, JxkT thanked the authors for their thorough response and for addressing the points they brought up. JxkT also appreciated that the authors extended their empirical analyses to account for concerns mentioned by other reviewers—specifically, they included comparisons with additional baseline methods and provided further and detailed experiments. This latter limitation was also the main concern mentioned by H4uQ. Based on the rebuttal and the post-rebuttal discussion, this concern seems to have been adequately addressed from H4uQ's perspective.

Furthermore, reviewers 671a and 9Wzo recommended that POAD be compared with a recently proposed relevant method (ArCHer), which the authors did. All results were included in a pdf file attached to the rebuttal. The reviewers were satisfied with the new experiments (671a: "*(...) many thanks to the authors for thoroughly supplementing the key experiments and discussions*"; 9Wzo: "*Thanks for the detailed response. I think the authors have addressed most of my concerns, including empirical comparison with ArCHer, societal impact, and whether Fig. 2 is true*").

Although all reviewers brought up important concerns, and the consensus seemed to be that this paper is not necessarily strongly novel, it does provide the community with interesting insights on how to (at least partially) tackle the issue of credit assignment in settings where RL agents sample actions from an LLM. Based on the feedback provided by the reviewers, I do not see any critical flaws, missing comparisons or analyses, or methodological issues with this paper. More importantly, the extensive and fruitful discussion between the authors and the reviewers, as well as the reviewers' acknowledgment of the authors' contributions and findings, makes me confident that the NeurIPS community will benefit from the insights presented in this work.
Decision: Accept (poster)