Q* Agent: Optimizing Language Agents with Q-Guided Exploration

We sincerely thank the reviewer for spending time reading our paper and giving valuable feedback. We address your concerns below.

Comparison with MCTS

We would like to argue that compared with the widely-used MCTS, our method shows difference in the following perspectives:

1. The reward in the MCTS does not take the depth of the trajectory into account, while Q value represents a discounted value for the future outcome, which is better at modeling long-term value.
2. MCTS is a model-based setup, which requires the environment to output the states so that they can leverage to do backtracking in their framework during the sampling process, e.g., in Agent Q and MCTS. Our method, which adopts a Q-learning framework, is a model-free setup which does not necessarily need to know the function of how the environment outputs the future states. The sampling process is purely guided by feeding historical states and current action to the QNet.

Known Environment Assumption

We would like to clarify that our approach does not rely on constructing a reasoning tree to achieve state transitions. Instead, the reasoning tree structure naturally arises from the characteristics of interactive tasks.

Not enough exploration depth

Thanks for the comments. We respectfully argue that using 3~5 steps as the exploration steps for expansion are already enough for Webshop. Specifically, we do additional ablation studies by sampling 150 instructions and construct reasoning trees with different exploration depths of 3, 5, and 7. The average trajectory rewards are in the table. We can observe that current exploration depth 3 maintains a relatively reasonable averaged trajectory quality. For tasks with longer trajectory length, our method is flexible enough to be adapted to larger exploration depth.

Task-specific Q-value Network

Thanks for the comment. We would like to respectfully argue that our current method treats the QNet as an external module making it more flexible. Also, our current results show, because our QNet is trained on the trajectories collected from the training set of WebShop, and we still observe the performance gain on the test set of WebShop. This actually indicates the generalizability of our proposed method because there are some unseen item types in the test set that have not shown up in the training set, showing our QNet actually captures some general capability of long term value modeling. On the other hand, if we modify the model weights of the language policy, e.g., using RL to optimize the LLM policy with the q function, this will make the language agent overfit more to this task distribution rather than generalizable to other environments.

Question on MDP design using the full observations and actions as current states

We would like to argue that most of the language agent works such as Agent Q[1], language agents are using the full states and actions to serve as the current state. This approach arises from the fundamental differences between language agent tasks and traditional reinforcement learning (RL) environments. In language agent tasks, the observations received from the environment at each step often lack sufficient context to fully represent the state of the environment. This complete interaction history is critical for enabling the agent to reason effectively and make informed decisions.

Experimental Results

Thanks for the comments. We would like to argue that the current improvement on the webshop is fairly good enough. For example, in ETO[2], Agent Lumos[3], and Agent Q[1], they also get around 7~10% performance gain on the WebShop, which already is a good enhancement compared with baselines on the webshop in the language agent domain. Also, we want to claim that tree search algorithm, given the limited resources in the academia, running one iteration from fine-tuning, self-generation, training QNet and using QNet to guide inference already takes three to four days if using 4*A6000. Considering that ALFWorld，ScienceWorld takes long steps to finish, it might take more than one week even if we have one node of A100 to finish one iteration of experiments mentioned above. We are conducting experiments on new agent tasks and making every effort to update the results before the rebuttal deadline. We sincerely thank the reviewer for the understanding.

References

[1] Putta et al. (2024) Agent Q: Advanced Reasoning and Learning for Autonomous AI Agents

[2] Song et al. (2024) Trial and Error: Exploration-Based Trajectory Optimization for LLM Agents

[3] Yin et al. (2023) Agent Lumos: Unified and Modular Training for Open-Source Language Agents

We express our great appreciation for the reviewer's time on reading our paper and giving us insightful suggestions. We address your concerns below.

W1、3、4 Paper writing needs to be refined.

Thanks for the detailed reviews and valuable suggestions to refine our paper! We have revised our paper according to your advice.

W2 The versatility of the method needs to be verified in more agent tasks.

Thanks for the advice. Currently, we are conducting experiments on new agent tasks requiring larger exploration space. Due to the increased trajectory length in these tasks and the need to adjust hyperparameters to adapt to the new tasks, it takes approximately one week to complete a single run on eight A00 GPUs. We are doing our best to obtain these results as an academic lab with limited computational resources, and we will make every effort to update them before the rebuttal deadline. Meanwhile, we address the other points first. We sincerely thank the review for the understanding.

Q1 In Section 5.5, how were the "Averaged reward" and "Reward" obtained during inference? I believe the average step-level reward is unavailable during inference.

We apologize for any confusion caused by our explanation. We would like to clarify that the process rewards are not obtained during inference but are instead computed after Stage 2 in our framework. Specifically, we first construct a reasoning tree during the self-exploration phase in Stage 2. The recorded rewards for each node in the tree are then utilized to construct various process rewards and train process reward models. The results shown in the figure correspond to the Q*Agent-ST (self-training) setup. We have refined our writing in this section and added more explanations.

Q2 What is the advantage of utilizing Bellman equation to estimate Q-values compared to widely-used MCTS?

Compared with the widely-used MCTS, our method shows difference in the following perspectives:

1. The reward in the MCTS does not take account the depth of the trajectory, while the Q value function adopts a discounted value for the future outcome, which is better at modeling long-term value.
2. MCTS is a model-based setup, which requires the environment to output the states so that they can leverage to do backtracking in their framework during the sampling process, e.g., in AgentQ and MCTS. Our method, which adopts a Q-learning framework, is a model-free setup which does not necessarily need to know the function of how the environment outputs the future states. The sampling process is purely guided by feeding historical states and current action to the QNet.

Q3 Statistics of agent tasks

We adopt the open-sourced expert trajectory datasets in ETO[1]. The average turns of interactions of expert trajectories in WebShop, ScienceWorld, and ALFWorld are 4.9, 14.4, and 10.1, respectively. To construct an effective reasoning tree, the number of trajectories that need to be explored grows exponentially with the depth of the tree. This further emphasizes the importance of designing effective node expansion strategy with proper tree pruning.

References

[1] Song et al. (2024) Trial and Error: Exploration-Based Trajectory Optimization for LLM Agents

I appreciate the authors' efforts in providing responses; however, not all of my concerns have been adequately addressed. The paper still lacks clarity in several key areas. For instance, the termination condition for building the inference tree and the temperature used for sampling actions are not provided. Additionally, the ablation studies are insufficient to convincingly demonstrate the effectiveness of each proposed component.

While the paper is well-motivated and tackles a significant challenge in the domain of LLM agents, I am maintaining my current rating due to the unsatisfactory writing quality, the absence of these critical details, the lack of robust evidence supporting the effectiveness of each component, and the unavailability of the source code. These factors collectively impact the paper's quality and overall completeness.

Thank you for your precious feedback. We answer your concern and comments as follows:

(1) the termination condition for building the inference tree and the temperature used for sampling actions are not provided

We have added these details in the revised paper (Line 390~393). “During sampling process, the environment will give termination signal after certain action “Click” or achieve the maximum steps set in advance. Specifically, we set the maximum as 5 for WebShop during self-generation, Q-guided generation and evaluation on the test set.” and we set temperature as 0.7 (Table3 in page16).

(2) comparison with MCTS

We would like to adopt the results from Zhao et al. (2024)[1], which is our concurrent work, to show the effectiveness of our methods. Note that they are using data collected from MCTS to train a Q-value network to guide the action sampling process of language agents. In their experiments, even though they are using a much stronger base-model (GPT-4-turbo), the best average outcome reward for WebShop is only 0.64, which is significantly lower than our best results (0.726) with Llama-2-7B. We believe this observation can prove the superiority of our proposed tree-search process compared with MCTS.

(3) robust evidence supporting the effectiveness of each component

As our key design, we actually have quantitative and qualitative analysis of our proposed QNet compared with other reward function designs. To better showcase the robustness of our proposed method, we conducted experiments on different policies, e.g., models with different levels of behavior cloning, and models with different numbers of parameters. We have done ablations using different process rewards and its inference efficiency (Figure 3.), and we add experiments using a new SFT dataset filtering out trajectories with rewards higher than 0.5 from the original dataset. The results are below. From the table, we can observe that our $Q^{\star}$Agent-I-aug can greatly enhance agent inference by providing effective process reward guidance with perturbation augmented action. Besides, $Q^{\star}$Agent-ST has better performance than ETO, demonstrating its effectiveness in generating high-quality data for self-training. Overall, $Q^{\star}$Agent still performs exceptionally well even when the SFT dataset is with poor quality.

To give a stronger support to the empirical results, we run the experiments on the sciworld in a limited time. We have achieved better performance than SFT, PPO, BoN, RFT, and comparable to ETO.

As shown in the table, our q-guided exploration has already achieved good enough performance on both Webshop and Sciworld. We believe compared with PPO and DPO-based methods like ETO, our method is simpler and more flexible without further modifying the original model weights. In practice, we found tuning DPO and PPO is hard (the parameter is sensitive) and cost a lot of compute.

Since ETO does not release their DPO checkpoints, we apply our checkpoints on other well-trained checkpoints like llama-13B-instruct to see whether it can improve the checkpoints.

(4) the unavailability of the source code

We will definitely release the code upon acceptance. We want to claim that it is not required to provide the source code at the submission stage.

Thank you so much for your feedback! We truly appreciate the time you took to provide detailed comments and discuss with us!

References: [1] Zhai et al (2024). Enhancing Decision-Making for LLM Agents via Step-Level Q-Value Models.

We would like to express our thanks to the reviewer for the detailed reviews! We answer the questions below.

W1 Experiments on other agent tasks

Thanks for the advice. Currently, we are conducting experiments on new agent tasks requiring larger exploration space. Due to the increased trajectory length in these tasks and the need to adjust hyperparameters to adapt to the new tasks, it takes approximately one week to complete a single run on eight A100 GPUs. We are doing our best to obtain these results as an academic lab with limited computational resources, and we will make every effort to update them before the rebuttal deadline. Meanwhile, we address the other points first. We sincerely thank the review for the understanding.

W2 Adding suggested baselines

We appreciate your insightful feedback regarding the inclusion of related works on inference-time self-improvement in language agents. We also acknowledge the importance of positioning our work relative to existing works. However, we respectfully argue that the baselines you have mentioned (e.g. LATS, Reflexion) rely on an external LLM or VLM to provide verbal feedback rather than scalar rewards. This places them in a fundamentally different category of work. Besides, the focus of our work is fundamentally different as we aim to achieve self-improvement with minimal reliance on external resources.

W3 Further explanation on the prompting setting

We would like to clarify that the benefits of ReAct in the SFT dataset and 1-shot evaluation setup are orthogonal to our method. For fair comparison, all the experimental results in the table are evaluated under the same 1-shot prompt, and all the training-based methods including ETO, RFT, PPO, QAgent-ST are based on the same SFT dataset in ReAct style.

Q1 Replace the illustrative example in Figure 2 with an example more close to the experimental study in this paper.

Thanks for the suggestion. We have changed the example to one more close to the WebShop environment.

Q2 Use $Q^\star$ instead of Q* to denote the optimal state-action value function.

Thanks for the advice. We have revised the notion in our paper to denote the estimated optimal state-action value function.

Q5 Is there any reason results of Q*[1] and the enhanced Q-value model[2] were omitted?

As we have stated in answer to Q1, due to the absence of an open-source codebase for the two mentioned concurrent works, we are unable to reproduce their results within our setup and data framework. We provide a concise comparison with Zhai et al. (2024), as we both involve experiments on WebShop using the same SFT dataset from the ETO[3]. This comparison is detailed in Lines 424-427 of our paper: Zhai et al.[2] used Llama-3.1-8b-instruct as their base agent and achieved a final reward of 60. Our $Q^{\star}$Agent used llama2-7b-chat as base agent and $Q^{\star}$Agent-I-aug achieved a higher final reward of 72.6.

Q6 The performance of $Q^{\star}$Agent when the SFT dataset shows poor performance

Thanks for the insightful suggestion! We add experiments using a new SFT dataset filtering out trajectories with rewards higher than 0.5 from the original dataset. The results are below. From the table, we can observe that our $Q^{\star}$Agent-I-aug can greatly enhance agent inference by providing effective process reward guidance with perturbation augmented action. Besides, $Q^{\star}$Agent-ST has better performance than ETO, demonstrating its effectiveness in generating high-quality data for self-training. Overall, $Q^{\star}$Agent still performs exceptionally well even when the SFT dataset is with poor quality.

Q7 Further explanation on the setting of $Q^{\star}$Agent-ST and $Q^{\star}$Agent-I

Thanks for the question. Both $Q^{\star}$Agent-ST and $Q^{\star}$Agent-I use QNet to choose actions with the highest Q-value at each generation step. $Q^{\star}$Agent-I is using QNet during the searching process at each step, using Q(current_state, current_action) to predict the q-value for each action and to select the action with the highest q-value after sample three actions for each step. $Q^{\star}$Agent-ST also uses QNet to do such step-wise search and collect sampled trajectories to train the language agent.

Minor revisions on notations in the Algorithm 2 of Appendix A.3.

Thanks for pointing this out. We have rewritten notations and restructured the Algorithm 2 in Appendix A.3 to provide a more clear explanation. We kindly invite the reviewer to refer to our revised paper.

References

[1] Wang et al. (2024) Q*: Improving Multi-step Reasoning for LLMs with Deliberative Planning

[2] Zhai et al. (2024) Enhancing Decision-Making for LLM Agents via Step-Level Q-Value Models

[3] Song et al. (2024) Trial and Error: Exploration-Based Trajectory Optimization for LLM Agents

We appreciate the reviewer’s insightful feedback. We address your concerns below one by one.

Q1 More detailed comparison with Q*[1] and the enhanced Q-value model[2].

We would like to clarify first that SFT for behavioral cloning is not our unique contribution. We use SFT to provide our language agent with capabilities to perform reasoning and actions in this environment. Also, we respectfully argue that both Q*[1] and the enhanced Q-value model[2] are concurrent works. The comparisons are listed below:

1. Q*[1] has not been tested on interactive agent tasks which are more complicated and have broader real-world applications. Also, it does not include a self-training stage to further enhance the language agents.

2. We introduce several strategies to reduce the time complexity of tree search to address the heavy searching cost in interactive agent tasks. Also, we propose using LLM to perturb the context to further boost the diversity of the exploration content.

3. For training the process reward model, unlike the Q-value model[2] that leverages DPO, our framework adopts MSE loss to learn the Q-value. This approach is not only much easier to optimize but also more robust to the hyper-parameters.

4. Both Q*[1] and the enhanced Q-value model[2] have not released their code, so we cannot reproduce their methods in our setup. However, since Zhai et al.[2] also experimented on WebShop and used the same SFT dataset in ETO[3] as $Q^{\star}$Agent, we gave a concise comparison in L424-427 of our paper: Zhai et al.[2] used Llama-3.1-8b-instruct as their base agent and achieved a final reward of 60. Our $Q^{\star}$Agent used llama2-7b-chat as base agent and $Q^{\star}$Agent-I-aug achieved a higher final reward of 72.6.

Q2 Whether pruning zero-branches will disregard partially correct trajectories

First we would like to clarify that our approach does not prune or discard branches with a final reward of zero that have already been explored, which could be partially correct trajectories. Designing the pruning strategy is intended to optimize exploration focusing on expanding branches that are more likely to yield non-zero rewards within a limited time. Therefore, while we cease further expansion at nodes deep within a trajectory that culminates in zero reward, we retain the trajectory up to that point for analysis and potential use in subsequent training iterations.

To provide further clarity on how $Q^{\star}$Agent constructs the reasoning tree, we have added Algorithm 3 in Appendix A.4. of our revised version for your reference.

Q3 In Q-guided self-training, why does the $Q^{\star}$Agent algorithm not use the dataset collected during step 2 in Figure 1?

We do not use the trajectories in step 2 because there are many step-level overlaps between the self-explored trajectories used to train the QNet and Q-guided explored trajectories. Also, we believe self-training requires data with higher quality, and q-guided explored trajectories have higher quality than the self-explored one, so we do not use the data collected during step 2 for the current version.

Q4 The 'exploration' terminology may be confusing

Thanks for pointing this out. We revised the terminology used in step2 and step4 to “self-guided generation” and “Q-guided generation” respectively.

Thank you for your detailed response.

However, I am still confused about the definitions of -ST and -I.

Additionally, it is not clear that Q*Agent performs better than other baselines, both empirically and theoretically, including ETO.

As a result, I am on the borderline and would not be surprised if this work gets accepted, but I will maintain my current score.