Reinforcing LLM Agents via Policy Optimization with Action Decomposition

Authors: muning  wen
ziyu  wan
jun  wang
weinan  zhang
ying  wen
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 16:22:12

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.

Main paper: 
 introduction: 
 Large language models (LLMs) have demonstrated promising capabilities of solving various tasks, from instructions following to complex reasoning and real-world interaction [1] [2] [3] . This growing tasksolving 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 decision-making tasks robustly and efficiently.Recent advances [4] [5] [6] [7] [8] have showcased that the aforementioned challenges can be alleviated in a trial-and-error learning style, namely Reinforcement Learning (RL) [9] . Representatively, GLAM [4] and TWOSOME [5] treat the language actions, i.e. token sequences outputted by a language model, as whole units, and optimize actions' likelihood, calculated as the products of 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 Reward/Advantage > 0, there is a coffee machine, chosen option.Reward/Advantage < 0, irrelevant option, distinguished by common sense."Walk" "to" "Walk" "kitchen" "Walk to" "Walk to kitchen" "Walk" "to" "Walk" "kitchen" "Walk to''Figure 1 : 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 "coffee table is empty". (b) Action-level optimization is uncertain to what extent the key tokens, i.e. P("kitchen"|p, "Walk to"), will be enhanced when optimizing the joint probability P("Walk to kitchen"|p).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 1 , 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"|p). 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-making process [10] [11] [12] . 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 output distribution separately with fine-grained signals from value backups. Empirical success in many single-step tasks, e.g. question answering [13] and alignment [10, 11] , have demonstrated its effectiveness. However, the multi-step nature of sequential decision-making tasks leads to 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 a significant discrepancy has been ignored in previous research and has not been tackled properly for a long time.In this paper, we provide a comprehensive theoretical analysis of the discrepancy and derive the Bellman backup with Action Decomposition (BAD), which guarantees theoretical consistency with optimizing the original MDP. While being possible to seamlessly integrate with a variety of traditional RL methods, in this work, we apply BAD to Proximal Policy Optimization (PPO) [15] , 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 1 , thereby enjoying better interpretability, lower optimization complexity, and higher training efficiency. Meanwhile, it theoretically maintains consistency between the tokenlevel training process for language models and the RL objective of maximizing actions' utilities. We justify our claims by evaluating POAD in both classical sequential decision-making environments with limited action space, i.e Overcooked and VirtualHome [5] , and a 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. 
related works: 
 Language-based 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 [16] [17] [18] . Similarly, LLMs work particularly well in text-based environments [19, 20] .ReAct [21] combines chain-of-thought reasoning [2] with acting-based prompt, efficiently solving hard decision-making 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 [7, 24, 25, 6] .Besides, RL has also been widely applied in fine-tuning LLMs [10, 11, 7, 24, 25, 6] . Particularly, proximal policy optimization (PPO) [15] is the most commonly used RL method for reinforcement learning from human feedback (RLHF), proposing a breakthrough in aligning LLMs with human preference [10, 11] . In classical sequential decision-making scenarios, to 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 as a whole. In this work, instead of treating an action as a whole, 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 a concurrent work, ArCHer [26] , also targets token-level supervision for LLMs in interactive environments, it employs a hierarchical RL framework, using a 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.Sequential Decomposition in RL. Recent years have witnessed an increasing trend of decomposing high-dimension actions to leverage the powerful modeling abilities of sequential models like Transformer [28] in RL problems [29] [30] [31] [32] [33] [34] . While Kuba et al. [35, 36] proposing a sequential decomposition method to provide finer-grained supervision for multi-agent joint actions, Multi-agent Transformer [31] inherits this idea and solves multi-agent RL problems with Transformer. More recently, Q-transformer [34] managed to decompose the Q-functions for high-dimensional actions by representing each action dimension as 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 to intra-action tokens. 
language-augmented rl: 
 In this work, we assume a textual RL setting that frames sequential decision-making problems with linguistic inputs and outputs as a language-augmented Markov Decision Process (MDP) M = (V, S, A, T , R, γ) [37, 4] . Given V the vocabulary and w ∈ V the tokens, A ⊂ V N , S ⊂ V N are the action space and state space respectively, i.e. actions and states are sequences of tokens. T : S × A → S is the state transition function. R : S × A → R is the reward function that only responds to complete actions, and γ is the discounted factor that typically less than 1. At time step t ∈ N, a language agent receives a textual state s t ∈ S from an interactive environment as input and generates an action a t ∈ A in an auto-regressive manner, i.e. a t = (w are tokens in it. Then, textual action a t will be grounded to a specific API call or command in the environment [38] . After execution, the language agent receives a reward r t = R(s t , a t ) along with the next state s t+1 , based on the transition function T . Following this process with trajectories of a maximum timestep T , the agents earn a discounted cumulative return of R γ = T t=0 γ t r t , which is aimed to be maximized by RL algorithms.We begin by briefly reviewing the process of action-level policy optimization which is widely adopted in several state-of-the-art methods that align language agents with environments via RL algorithms [4, 5, 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 Q π (s, a) or state value function V π (s). Then, a language agent updates its policy π according to credits calculated on the value functions, defined asEQUATIONEQUATIONwhere a = (w 1 , . . . , w |at| ) = w 1:|a| . While Ahn et al. [16] builds the connection between the likelihoods of actions and tokens through the chain rule of probability asEQUATIONrecent approaches like GLAM [4] and TWOSOME [5] leverage similar ideas and optimize action-level likelihoods with RL methods directly. When considering optimizing for π(a|s), Equations 1 and 2 are aligned with the definition of the value function in traditional RL settings, allowing them to be updated with traditional Bellman backup [9] Qπ(st, at) ← R(st, at) + γ maxEQUATIONEQUATIONMoreover, it is noteworthy that Equation 3 calculates the likelihood of action a in an exponentially growing language space as |a| increases, i.e. |A| = |V| |a| . Exploration and optimization in such a 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| ≪ |V| |a| . Taking TWOSOME as an example, it optimizes the likelihood of action a concerning the feasible action space with Equation 6 to mask out invalid outputs. L(a) indicates the number of tokens or words in the action prompt, utilized as a 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 [40, 41] -as shown in Figure 1 , 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 
naive token-level policy optimization: 
 To address the unclear credit assignment issue described in Figure 1 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 1:j with j ≤ |a|, we define token value functions for supervising policy update asEQUATIONVπ(s, wEQUATIONA natural approach to approximate Q π (s, w 1:j-1 , w j ) and V π (s, w 1:j-1 ) is conceptualizing the token generation process as part of the MDP, where each token is treated as a micro action. This enables back-propagating 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 to Equations 4 and 5 can be expressed as Qπ(st, w Vπ(st, w 1:j t ) ←EQUATION)To facilitate subsequent theoretical analysis, we separate the discount factor γ into intra-action γ w and inter-action γ a , despite their numerical equivalence here. The above backups can be interpreted as applying RL algorithms on a modified reward function, which maintains action-level feedback while introducing extra 0 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 t , a t ) to 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 9 and 10 are inconsistent with those of Equations 4 and 5 for sequential decision-making tasks since it introduces a new MDP M that diverges from the origin one. We will analyze this discrepancy in the following section.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 RL [34] . 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 Q π * (s t , w 1:j-1 t , w j t ) = Q π * (s t , a t ) and V π * (s t , w 1:j t ) = V π * (s t ) for ∀j < |a t |. To quantify the discrepancy between action-level optimization and token-level optimization with the shape of Equations 9 and 10, we expand the value backup process over each token starting from arbitrary j < |a t | as the following equations (For derivation see Appendix A).Qπ * (st, w 1:j-1 t , w j t ) = R(st, at) + γa maxa tQπ * (st+1, at+1)EQUATION-(1 -γ |a t |-j w )R(st, at) + γa(1 -γ |a t |+|a t+1 |-j-1 w ) max a t+1Qπ * (st+1, at+1)Discrepancy between Equation 4 and 9, Vπ * (st, w 1:j t ) = R(st, at) + γaVπ * (st+1)V π * (s t ) -(1 -γ |a t |-j w )R(st, at) + γa(1 -γ |a t |-j w )Vπ * (st+1)Discrepancy between Equation 5 and 10,With Equation 11 and 12, we observe following significant insights:(i) The discrepancy of the Bellman optimality between action-level optimization and token-level optimization diminishes as γ w ∈ [0, 1] increases, achieving consistency when γ w = 1.(ii) Given γ w < 1, the discrepancy increases as the number of intra-action tokens, |a t |, increases.Based on the first insight, we will derive our main approach, namely the Bellman backup with Action-Decomposition, in the next section. The second insight can be intuitively understood as follows: The diverged MDP M attaches the assumption that "words later in a sentence are more expressive than earlier ones", which is unrealistic for representing the semantics of linguistic actions.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, δ t = q t -q t-1 or δ j = q j -q j-1 represent the credit assigned to corresponding actions and tokens respectively for policy update, e.g. the advantage value [43] . To facilitate understanding, a step-bystep breakdown of the right figure is provided in Appendix L.5 Action-Decomposition Reinforcement Learning 
bellman backup with action-decomposition: 
 According to the first insight, we can modify Equations 9 and 10, proposing the Bellman backup with Action-Decomposition (BAD) asEQUATIONEQUATION(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. Moreover, BAD can be seamlessly integrated into various existing RL methods, including off-policy algorithms e.g. DQN [44] , on-policy ones like Actor-Critic [45] and PPO [15] . Furthermore, we have also provided a version of the Soft Q-function in Appendix B.3 to support various entropy-regularized RL algorithms like SAC [46, 47] .In this section, we integrate the BAD into the widely used PPO and propose a specific method called Policy Optimization with Action Decomposition (POAD), 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 ϕ whose objective is to minimize the empirical Bellman error of tokens byL(θ) = 1 T T -1 t=0 1 |at| R(st, at) + γVθ(st+1, ∅) -Vθ(st, w 1:|at| t ) 2Inter-action credit assignmentEQUATIONIntra-action credit assignment , (15) where θ represents the non-differentiable parameter of the target network, updated at intervals. The policy network's parameters are denoted as ϕ, and optimization follows the clipping PPO objective.EQUATIONratio j t (ϕ) = π ϕ (w j t |st, w 1:j-1 t ) π ϕ old (w j t |st, w 1:j-1 t ), and Âj t = Â(w j t |st,w 1:j-1 t ),where Âj t 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. 
experiments: 
 In this section, we show the superiority of POAD in performance, efficiency, and generalization abilities with different testbeds. Moreover, we conduct meaningful ablations on γ a and γ w to verify the theoretical analysis in Section 4.2. Finally, we examine models trained with POAD and baseline methods to investigate their impact on the model's original language abilities. For in-depth analysis, we conduct a case study in Appendix F 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) [51] with 1 Nvidia A100 GPU.We first evaluate our method on two classical sequential decision-making environments with limited action space: Overcooked [5] and VirtualHome [5] , where the action space consists of approximately 10 possible actions per state, each includes 5-10 tokens. Then we evaluate our method in a data science coding and debugging environment with unrestricted action space: DataSciCoding, where agents generate actions (up to 128 tokens) freely. More detailed descriptions can be found in Appendix E.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-Cutting-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 on 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.1. 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 on feedback from the runtime environment until it works, thus aligning the task datasets. Agents receive ROC AUC scores ∈ [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.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 10 with γ w = γ a into PPO, and TWOSOME [5] -the current state-of-the-art method on Overcooked and VirtualHome. Besides, we also incorporate ArCHer [26] as a baseline in VirtualHome with comparative analysis, since it is 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 For the DataSciCoding environment, in case TWOSOME is inapplicable due to the unrestricted action space, we examine POAD's performance along with NTPO. Meanwhile, we also compare POAD with CAAFE [54] -a novel AutoML framework in which humans collaborate with large language models for data science. In CAAFE, humans implement machine learning models, e.g. classifiers, while large language models generate the code for feature engineering. CAAFE represents the current state-of-the-art 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.As shown in Figure 3 , where curves are averaged over 3 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.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 on sequential decisionmaking tasks in Section 6.3.1, the performance gap between POAD and NTPO on DataSciCoding is consistently larger, due to the much longer action length |a t |, at 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 a proper intra-action credit assignment, which, our Bellman Backup with Action Decomposition (BAD) is designed for.In Figure 4 , POAD-Best means the performance of the best code discovered during POAD's training process with CodeLLaMA-7B. We compare it with the best performance achieved by the state-of-theart AutoML framework CAAFE [54] with GPT-4 model. In this experiment, we aim to prove that even small-scale language models can also provide better 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.LLMs' open-vocabulary feature enables language agents to transfer their learned skills into unseen similar tasks, expanding the capabilities of decision-making 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 LLaMA2-7B. Table 1 shows that token-level policy optimization methods achieve better generalization performance in unseen tasks. And our POAD outperforms the other baselines in seven of eight tasks. To verify our analysis in Section 4.2, we conduct ablations on γ w to further investigate how large the discrepancy would be by using different values of γ w in NTPO. Therefore, we deploy NTPO on the Food Preparation task with γ w ∈ {0.95, 0.9, 0.8, 0.5}, while we keep γ a = 0.95 which is consistent with our main experiments. Besides, we also show the performance with γ a ∈ {1.0, 0.95} to show the necessity of setting γ a 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 (γ w = 1) and NTPO as γ w decreases. These results empirically showcase the discrepancy between naive token-level credit assignment and our BAD, which is consistent with our first theoretical insight in Section 4.2. Moreover, in the middle of Figure 5 , for both TWOSOME and our proposed method POAD, setting γ a = 1.0 performs much worse than γ a = 0.95, indicating that the discrepancy between NTPO and POAD can not be solved by simply setting γ a = γ w = 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 of excessive network complexity and difficult hyper-parameter tuning, ArCHer-BAD still falls short of matching POAD's performance. To investigate the impacts of online RL fine-tuning on LLMs' original capabilities, we evaluate the models trained by POAD, TWOSOME and NTPO on widely used NLP benchmarks [56] which are also reported in Tan et al. [5] and Touvron et al. [49] . These models are trained in Food Preparation. Table 2 demonstrates these models' zero-shot performance, compared with the original LLaMA2-7B model. The results show no significant losses of general ability like natural language understanding after aligning with the embodied environment, even sometimes bringing minor improvements. 
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-grained supervision for each intra-action token and ensuring theoretical consistency between the token-level training nature of language models and the RL 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 generalization abilities, over strong action-level baseline TWOSOME.Limitation and Future Work. The existing limitation of POAD is on the requirement for a quantitative reward function, which is not easily attainable in some environments. To mitigate this, 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 decisionmaking 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. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: This paper poses no such risks. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a 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 is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or 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? 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 or dataset.• The authors should state which version of the asset is used and, if possible, include a 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 (e.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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [NA] Justification: This paper does not release new assets. Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS 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. 
 
 Appendix: A.1 Q-FunctionThe optimization over each token following Equation 9, starting from arbitrary j < |a t |, could be expanded asEQUATION= γ w maxw j+1 t γ w max w j+2 t Q π * (s t , w 1:j t , w j+1 t , w j+2 t ) = γ |at|-j w R(s t , a t ) + γ a γ |at|-j w max w 1 t+1 Q π * (s t+1 , w 1 t+1 ) = γ |at|-j w R(s t , a t ) + γ a max w 1 t+1 γ w max w 2 t+1 Q π * (s t+1 , w 1 t+1 , w 2 t+1 ) = γ |at|-j w R(s t , a t ) + γ a γ w max w 1:2 t+1 Q π * (s t+1 , w 1 t+1 , w 2 t+1 ) = γ |at|-j w R(s t , a t ) + γ a γ |at+1|-1 w max w 1:|a t+1 | t+1 Q π * (s t+1 , w 1:|at+1| t+1 ) = γ |at|-j w R(s t , a t ) + γ a γ |at+1|-1 w max at+1 Q π * (s t+1 , a t+1 ) = R(s t , a t ) + γ a max at+1 Q π * (s t+1 , a t+1 ) Q π * (st,at) -(1 -γ |at|-j w )R(s t , a t ) + γ a (1 -γ |at|+|at+1|-j-1 w ) max at+1 Q π * (s t+1 , a t+1 )Discrepancy between Equation 4 and 9.Similar to Appendix A.1, the optimization over each token following Equation 10, starting from arbitrary j < |a t |, could be expanded asEQUATION= γ w γ w V π * (s t , w 1:j+1 t ) = γ |at|-j w R(s t , a t ) + γ a γ |at|-j w V π * (s t+1 ) = R(s t , a t ) + γ a V π * (s t+1 ) V π * (st) -(1 -γ |at|-j w )R(s t , a t ) + γ a (1 -γ |at|-j w )V π * (s t+1 )Discrepancy between Equation 5 and 10.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 is equivalent to optimizing for the full action.The optimization over each token with the optimal policy π * following Equation 13, starting from arbitrary j < |a t |, could be expanded asEQUATIONThe optimization over each token following Equation 14, starting from arbitrary j < |a t |, could be expanded asEQUATIONwhere V π * (s t , w 1:j-1 t) and V π * (s t , w1:|at| t ) are equivalent to E w j t ∼π * [Q(s t , w 1:j-1 t , w j t )] and E w 1 t+1 ∼π * [Q(s t+1 , w 1 t+1 )].As an extension of BAD, we also provide the soft Bellman backup with Action-Decomposition (sBAD), to support Entropy-Regularized methods like SAC (by setting the reference policy π to a uniform distribution), asEQUATIONWe show that optimizing the soft Q-function with sBAD at the token level is consistent with optimizing for the full action. We adopt KL(a|s) = D KL [π * ||π](s), where the KL(a|s) indicate it is a action level KL divergence, KL(w|s) indicate a token level KL divergence. Given an optimal stochastic policy π * , the vanilla soft Bellman backup for full actions is:EQUATION= E at∼π * [R(s t , a t )] + γE at,at+1∼π * [Q(s t+1 , a t+1 )] -γβE at∼π * [KL(a t+1 |s t+1 )].Following our sBAD, the update over each token within an action (j < |a|) is: E w 1 t ∼π * [Q(s t , w 1 t )] = E w 1 t ∼π * [E w 2 t ∼π * [Q(s t , w 1 t , w 2 t )] -βKL(w 2 t |s t , w 1 t )] = E w 1 t ,w 2 t ∼π * [Q(s t ,Minus βKL(w 1 t |s t ) from both sides, we have:EQUATIONwhere KL(a t |s t ) = |at| j=1 KL(w j t |s t , w 1:j-1 t).Then, the update across the current action a t and the next action a t+1 with our sBAD:EQUATION)Equation 25 enjoys the same shape as Equation 22, thus optimizing the soft Q function following sBAD is equivalent to the action-level optimization. We prove the consistency between sBAD and the traditional soft Bellman updates over full actions.We provided a visual comparison of the differences between four different backup approaches with the optimal policy after convergence in Figure 6 .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. Sample a mini-batch B from D.12:for each time step t in B do 13:for each token j in a t do θ ← θ -α∇ θ E B [(v -v targ ) 2 ] 24: ϕ ← ϕ -α∇ ϕ E B [L(ϕ)] 25:θ ← θ 26: end for E Detailed Description of Experimental Environments Overcooked is proposed as a typical deicision-making environment in Tan et al. [5] . There are two tasks in which a language agent is aiming to make and serve a tomato salad and tomato-lettuce salad respectively, with the provided ingredients and tools in a 7×7 grid world as a kitchen. To solve the tasks, the agent needs to explore and learn the correct order to cook the dish with the provided macro-actions, such as Chop, Get-Tomato, and Go-Cutting-Board. The environment is partially observable. The agent only observes the objects within 5×5 square centered on 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.VirtualHome is a simulated physical household environment proposed in Tan et al. [5] . In this environment, an agent is placed in a fully-furnished household with various objects. There are two tasks in the environment, the first one is to find the cold pancake on the table and heat it with the 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 abilities 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 on 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 as walk to the living room, turn on the TV and sit on 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.DataSciCoding is 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.1. In each task, 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 to 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, it 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 ∈ [0, 1]. If the codes run wrong, the agent receives a constant -1.0 as a penalty. Adopting the same evaluation metrics as Hollmann et al. [54] , for each dataset, we evaluate 5 repetitions, each with a random 50% -50% train-test split [55] , and record the average ROC AUC score across these splits.These datasets are collected from Kaggle and OpenML respectively, with pruning approaches consistent with Hollmann et al. [54] . In this section, we conduct a case study for in-depth analysis to validate the effectiveness of BAD in terms of token-level credit assignment. We selected the last three states from the Food Prepa- ration task, which form a complete subtask: Heat the Pancake with Microwave. This serves as a simple yet practical case for our analysis, with transitions: T (s 0 , "open the microwave") → s 1 , T (s 0 , "close the microwave") → s 0 , T (s 1 , "put the pancake into the microwave") → s 2 , T (s 2 , "open the microwave") → s 2 , T (s 2 , "close the microwave") → 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 sample 1,000 examples using a 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-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.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 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 is the success rate of completing the task within 50 timesteps. To investigate the impacts of fine-tuning different RL methods on the LLM's original abilities, we evaluate the models trained by POAD, NTPO and TWOSOME on widely used NLP benchmarks [56] . All models are trained in the Food Preparation task. The four benchmarks are ARC_Challenge, HellaSwag, PIQA and MMLU, with results on ARC_Challenge, HellaSwag and PIQA are displayed in Table 7 and the results of MMLU are displayed in Table 8 . All results are calculated following the default configurations in Gao et al. [56] . During experiments, the implementations of TWOSOME are consistent with their official repositories, reproducing the original best-performing status. We first list the hyper-parameter candidates used for grid search in Table 9 . Then, we show the hyper-parameters that achieve the best performance for each method and environment in Table 10 -17. Step 3 and 4: Back-propagating r + t and r - t ′ to Q("on"|o t , "Turn") and Q("off"|o t , "Turn"), and then to Q("Turn"|o t ) in both trajectories continuously. = Mean{γ × q 0 t , γ × q 0 t ′ } since both trajectories share the same pre-action. Step 6: Similarly, Q("Turn"|o t ) is also shared among two trajectories, thus Q("Turn"|o t ) = Mean{q 0 t , q 0 t ′ } = 0 as well. Step 8: Starting to back-propagate again for credits assigned to each token for optimization with δ j = q j -q j-1 for j > 0; this process is to the calculation for advantage values. Step 9: Calculating the credits for key tokens, where δ 1 t = q 1 t -q 0 t and δ 1 t ′ = q 1 t ′ -q 0 t ′ .Step 10: Calculating the credits for tokens j = 0 with δ 0 t = q0 t -q |at-1| t-1 and δ 0 t ′ = q0 t ′ -q |at-1| t-1 . Till now, we precisely emphasized key tokens while keeping irrelevant tokens with 0 credits.Justification: Code is included in the zip file with hyper-parameters in Appendix K Guidelines:• The answer NA means that paper does not include experiments requiring code. 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 a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to 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 to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise 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. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).† Corresponding to Ying Wen <ying.wen@sjtu.edu.cn>, Weinan Zhang <wnzhang@sjtu.edu.cn>. The source code could be accessed directly with this link https://github.com/morning9393/ADRL. 


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: Limited testing environment and lack of baseline. The author only show two baselines for comparison. More baselines should be include for comprehensive analysis. 

Another main concern is that the finetuning is only done on using Lora, when finetune the model using the full capacity, the effect is unknown.
questions: questions: 1. How did you tune the baseline?
2. If you use full scale finetuning instead of Lora, what do you expect? Could you perform on experimentation on this?
3. How would this method work on standard NLP task?
limitations: limitations: Yes, author addressed the limitations.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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).

I 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").

---

Small notes:
 - typo in "policy" in line 173
 - 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
 - "NTPO" typo in the description of Figure 4
 - LHS and RHS subfigures are swapped in Figure 5
questions: 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?
2. 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: limitations: Yes.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
2. 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.
questions: 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?
2. Can the authors demonstrate the token-level credit assignment learned by the BAD method through a case study?
3. 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?

---

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

[2] Rafailov, Rafael, et al. "From $ r $ to $ Q^* $: Your Language Model is Secretly a Q-Function." arXiv preprint arXiv:2404.12358 (2024).

[3] Zeng, Yongcheng, et al. "Token-level Direct Preference Optimization." ICML 2024
limitations: limitations: The author has already discussed the limitations of the method and its potential impacts in the paper, and provided possible solutions.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.

**Minor issues:**

The 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.
questions: 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?

Also, it will be great if the author can do the following two things (which corresponds to point 2) and 3) in weakness):

1. 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";

2. compare with ArCHer [1], which is a sentence-level RL method, to further show that token-level Bellman update is better.  

**Reference**

[1] Y. Zhou et al. ArCHer: Training Language Model Agents via Hierarchical Multi-Turn RL. In ArXiv, 2024.
limitations: 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).

Comments: title: Thanks for upgrading the score!
Thank you very much for your recognition of our efforts!
We kindly inquire if reveiwer has any further comments or concerns regarding our responses
We kindly inquire if reveiwer has any further comments or concerns regarding our responses
Thank you for the time and effort you dedicated to re-evaluating our work!
Response
Thank you for your deep engagement
Response to Reviewer 9Wzo
comment: We would like to say thank you sincerely for your recognition of our efforts. Your constructive feedback has been invaluable in guiding our revisions and enhancing the quality of our work. Thanks a lot!
thanks the author, I have increased my score
Thank you very much for your recognition of our efforts. Your constructive feedback has been invaluable in guiding our revisions and enhancing the quality of our work. Thanks a lot!
Many thanks to the authors for thoroughly supplementing the key experiments and discussions. I have raised the score to 6.
Dear reviewer H4uQ,

We sincerely appreciate the time and effort you dedicated to our work and we have made every effort to address your concerns in our rebuttal submission. Given that the discussion deadline is less than 24 hours away, we wanted to kindly inquire if you have any further comments or concerns regarding our responses. Your insights are incredibly important to us, and we are eager to ensure that all of your points have been adequately addressed. If you feel that our responses have addressed your concerns, we would greatly appreciate your consideration in reflecting this in the final evaluation of the manuscript.

Thank you once again for your time and support throughout this process. We look forward to hearing from you soon.

Best regards, 

Authors
Dear reviewer 671a,

We sincerely appreciate the time and effort you dedicated to our work and we have made every effort to address your concerns in our rebuttal submission. Given that the discussion deadline is less than 24 hours away, we wanted to kindly inquire if you have any further comments or concerns regarding our responses. Your insights are incredibly important to us, and we are eager to ensure that all of your points have been adequately addressed. If you feel that our responses have addressed your concerns, we would greatly appreciate your consideration in reflecting this in the final evaluation of the manuscript.

Thank you once again for your time and support throughout this process. We look forward to hearing from you soon.

Best regards, 

Authors
We sincerely appreciate the time and effort you dedicated to re-evaluating our work. Your constructive feedback has been invaluable in guiding our revisions and enhancing the quality of our work. Thanks a lot!
I thank the authors for providing a thorough response, and for committing to fixing the issues pointed out in the review. Based on this, but also others reviewers feedback regarding the baselines comparison, and authors providing additional experiments, I decided to increase my score by one point.
We would like to say thank you for your deep engagement with our work, especially the direction you have suggested helps us make the paper more complete and understandable. We will check the ArCHer repo again and refine our statement continuously. Thanks a lot!
Thanks for the detailed response. I think this is a minor issue anyway, since the proposed method has already been compared with ArCHer-BAD. The ArCHer paper has github repo available; I suggest the authors to double-check the repo to make sure the figure is correct. In conclusion, I decide to keep my score.
Thank you very much for your recognition of our efforts. 
 
Yes, ArCHer adopts sentence-level critics exclusively and does not utilize token-level critics. However, the absence of token-level critics does not always imply that there is no discount factor applied within sentences. In Section 3.4 of the original ArCHer paper, the part of "Low-level token actor," authors stated: "We update the token-level policy with the policy gradient computed via REINFORCE (also known as the Monte Carlo Policy Gradient)". This indicates that rather than approximating token-level credits with critic networks, they employ Monte Carlo estimation to backpropagate token-level credits from sentence-level advantage values, $A(s_c, a_t)$, which serve as terminal rewards (and thus bypass the need for token-level critic networks). 

We also noted that in one instance provided by ArCHer, specifically in Equation (3), they applied a vanilla (relatively old) version [2] of REINFORCE that does not involve a discount factor (they didn't explain why). However, it is important to highlight that in most contemporary practices involving Monte Carlo estimation or REINFORCE, the use of a discount factor is standard for managing horizons and balancing bias-variance (e.g. $G_t=\sum_{k=0}^\infty \gamma^k r_{t+k+1}$, see Equation (13.8) in Section 13.3 and Equation (3.8) in Section 3.3 of [3]). Furthermore, considering that the authors claim ArCHer is a framework that can integrate with other reinforcement learning approaches, we decided to retain the discount factor in the REINFORCE process as depicted in Figure 1 of the global PDF for a more general representation.

Nevertheless, we have considered ArCHer without a discount factor, which aligns with our insights presented in the paper. We conducted an ablation study on this aspect, as illustrated in Figure 2 of the global PDF, i.e. the ArCHer-BAD variant in the right one. In this case, the performance of ArCHer should theoretically match that of POAD, as they share the same theoretical results of credit assignment. Unfortunately, due to the inherent challenges associated with extra network complexity and the difficulties of hyperparameter tuning, ArCHer still falls short of achieving the performance levels of POAD.

Thank you so much again for this valuable feedback, we will add this clarification along with Figure 1 of the global PDF in our later version.

---
[2] Williams, Ronald J. "Simple statistical gradient-following algorithms for connectionist reinforcement learning." Machine learning 8 (1992): 229-256.

[3] Sutton, Richard S., and Andrew G. Barto. Reinforcement learning: An introduction. MIT press, 2018.
Thanks for the detailed response. I think the authors have addressed most of my concern, including empirical comparison with ArCHer, societal impact, and whether Fig. 2 is true. However, there is one problem: is the updated figure for ArCHer in the global pdf correct? Judging from their Fig. 2 in the arXiv version, they never use a token-level critic and thus should not have discount rate inside a sentence. In fact, in page 9 they wrote: "Informally, Theorem 1 shows that the error in estimating advantages using the token-level critic is $\gamma\sqrt{L}$ larger than the utterance-level critic (in the worst case)". Thus, I think ArCHer's figure should be the same as POAD.