Paper: REX: Rapid Exploration and eXploitation for AI agents (ICLR-2024)

Reviewer comment (Experiments):
The evaluation is somewhat lacking, only evaluated on blocksworld and GSM8K, and ablations were missing (only statistics of # of actions were given), no intuition on how to tune the hyperparameters B,K,C.

Author response:
* The primary rationale behind selecting Blocksworld and GSM8K is rooted in their suitability for evaluating distinct facets: Blocksworld provides an ideal dataset for assessing the planning capability of LLM and its respective technique, while GSM8K offers a suitable dataset for evaluating problem-solving abilities. We consider these two categories of datasets, focusing on Planning and Problem Solving, as strong benchmarks for our evaluation.
* B, K, C values we set based on the grid search technique. We set a reasonable range for all the parameters - $1 \leq C \leq 10000$
, $1 \leq B \leq 1000$, $1 \leq K \leq 1000$ and ran a grid search to approximate the values.

Reviewer comment (Evaluation):
How come RAP scores are missing for every evaluated environment? It is mentioned RAP is not compatible with OAI APIs, maybe supply more results with open-source models?

Author response:
We added RAP to demonstrate its slower performance compared to REX in problem-solving. Reflexion and CoT were used as other advanced techniques for performance comparison.

Reviewer comment (Evaluation):
Do you have some intuition on when is REX-UCL better than REX-UCB and vice-versa? (Paper only included REX-UCL is better at exploration than REX-R).

Author response:
* Accuracy:
    * REX-UCL demonstrates strong performance in mathematical reasoning scenarios, particularly in cases with a significantly expansive action space. For instance, in GSM8K where the number of operands can be extensive, REX-UCL effectively narrows down the action space by effectively suppressing the respective logits.
    * REX-UCB tends to perform well in planning scenarios, where the action space is bounded. 
* Speed:
    * REX-UCB is faster than REX-UCL, below is the number of API calls made by these two techniques for Blocksworld and GSM8K datasets: 
| Method	| Blocksworld |	GSM8K |
|-----------|-------------|-----------|
|REX-UCB|	10	|10|
|REX-UCL|	60	|80|
(For Blocksworld, depth(d)=6, number of passes(n)=10, number of potential actions(m)=4; For GSM8K, depth(d)=8, number of passes(n)=10, number of potential actions(m)=4)

Reviewer comment (Evaluation):
How come when more steps (6) is required, method is worse than Reflexion? Shouldn't it be the opposite, since REX employs search trees. Do you have any intuition?

Author response:
Reflexion utilizes natural language feedback, whereas REX employs a simpler form of feedback, either HIGH/LOW or a scalar value. We posit that in specific scenarios, natural language feedback may convey more valuable information compared to our proposed feedback system. Consequently, Reflexion outperforms REX in the 6-step subset. Our approach is particularly advantageous in real-world situations where the true nature of an action remains ambiguous, such as when selecting an API. In such cases, we believe our strategy holds an advantage. Results in Table 1 are averaged over 10 runs/iterations. While 10 iterations suffice for a 2-step problem to find a solution, arriving at the correct solution quickly may be more challenging for a 6-step problem.

Reviewer comment (Evaluation):
Evaluation does not match techniques on the number of tokens they generate or evaluation calls they make. This makes fair comparisons impossible, as the different methods are afforded different computational budgets for the same objective.

Author response:
The time complexity presented in Table 1 roughly addresses this concern. Below is the number of inference calls made by each of the techniques for every task/problem in Blocksworld and GSM8K datasets:
(For Blocksworld, depth(d)=6, number of passes(n)=10, number of potential actions(m)=4; For GSM8K, depth(d)=8, number of passes(n)=10, number of potential actions(m)=4)
| Method | Blocksworld | GSM8K |
|----------|-----------|----------|
| CoT | 10 | 10 |
| Reflexion | 30 | 30 |
| RAP | 240 | 320 |
| REX-R | 10 | 10 |
| REX-UCB | 10 | 10 |
| REX-UCL | 60 | 80 |

Reviewer comment (Presentation):
In Figure 2, consider adjusting the visuals to make the differences clearer, perhaps by showing where LLM evaluations occur along the search tree to make it clearer that REX-UCL is evaluating the LLM for each step of the depth of the tree.

Author response:
The primary purpose of Figure 2 is to illustrate the transformation of the classical MCTS technique into various REX variations. Figure 2(b) specifically details the rollout strategy for REX-UCL. The phrase 'Repeat d times,' highlighted in red in Figure 2(b), signifies invoking the LLM 'd' times. As each invocation of LLM triggers the simultaneous execution of the Selection and Expansion steps within the REX-UCL strategy, the loop is situated under these blocks. Additionally, the red arrow marks in Figure 2(b) represent the step-by-step invocation of LLM. Further elaboration on the LLM input, specifically the prompt structures for REX-UCB and REX-UCL, is provided in section B of the Appendix.

Reviewer comment (Evaluation):
For Table 1, is it possible to provide uncertainty estimates on the scores provided? (For example, using multiple seeds for different runs of the algorithms.)

Author response:
We have set the temperature to 0.0 to ensure that it does not obscure the effectiveness of the proposed approach and to facilitate a stable comparison. The following stats were derived from 5 runs each with gpt-3.5-turbo-0301 LLM. We applied REX-UCB and REX-UCL to the Blocksworld dataset to demonstrate the observable level of uncertainty. For the results reported in the paper, the same experimental setup was used for other techniques (CoT, Reflexion) as well. The mean(μ) and standard deviation(σ) of number of correct answers across 5 runs is presented below:

| Method | Blocksworld (2 step) | Blocksworld (4 step) | Blocksworld (6 step) |
|----------|-----------|----------|-----------|
|  | No of questions = 30 | No of questions = 56 | No of questions = 114 |
| REX-UCB | μ=24, σ=0.4 | μ=21.6, σ=0.89 | μ=28.8 σ=1.63 |
| REX-UCL | μ=18, σ=0.4 | μ=24.8, σ=0.48 | μ=22.8, σ=0.47 |

**Note:** In the paper, the results reported are average over 10 runs.

Reviewer comment (Evaluation):
What is the intuition for why the REX methods do worse in the largest Blocksworld scenario? They show strong advantages in the 2-step scenario (80% vs 41.7%), but their advantage shrinks in the 4-step scenario (44.6% vs 42.0%) and becomes a deficit in the 6-step scenario (25.4% vs 30.0%).

Author response:
As the complexity of a problem increases (measured in terms of the number of steps required to solve it), the permutations of action sequences, i.e., trajectories, also increase. For instance, in the 6-step Blocksworld problem, which comprises 6 steps involving 4 actions (pick, put, stack, unstack) and several blocks, there are approximately 6 * 4 * 4 (= 96) potential actions to consider simultaneously. This potentially contributes to a decline in the model's performance. Moreover, we conducted all experiments with 10 runs/passes. Increasing the number of runs for more complex problems could potentially enhance performance.

Reviewer comment (Evaluation):
Is there other evidence for the strength of REX on GSM8k-test? The results look very close to CoT for REX-UCB and to Reflexion for REX-UCL. Could the results be subdivided by number of steps required?

Author response:
As already stated, the experimental results were obtained by averaging the results across 10 random runs. This supports the robust and consistent superiority of REX over Reflexion and CoT. As you have pointed out, REX has superior speed when d<=3.

Reviewer comment (Theory):
Reflexion requires `3n` time complexity, while REX-UCL requires `dn`, so on a 3-step problem the two are the same (if understood correctly), while on a 4 or more step problem REX-UCL will be more computationally demanding unless search depth is limited.

Author response:
That's definitely an interesting point. However, it's crucial to consider the task variations across CoT, REX-R, REX-UCB, and REX-UCL. In CoT, REX-R, and REX-UCB, the agent is tasked with predicting the entire solution—comprising both intermediate steps and the final answer—in a single attempt. On the other hand, REX-UCL requires the agent to predict either a single intermediate step or the final answer within one attempt. As a result, although REX-UCL necessitates more inference calls, the total number of generated tokens remains approximately consistent. Consequently, we believe that the computational resources utilized by all these techniques are roughly equivalent.

Reviewer comment (Reproducibility):
What search depth(s) (`d`) were used for the experiments?

Author response:
We have set d=6 for Blocksworld and d=8 for GSM8K.

Reviewer comment (Experiments):
Could the results in Table 1 be replicated with a model that exposes logits and compared? These additional results would strengthen the claims about the core technique and would also enable a comparison to RAP.

Author response:
(Currently working on it. Will share the results soon.)

Reviewer comment (Experiments):
How were `C`, `B`, and `K` selected for the experiments?

Author response:
The above values we set based on the grid search technique. We set a reasonable range for all the parameters - $1 \leq C \leq 10000$, $1 \leq B \leq 1000$, $1 \leq K \leq 1000$ and ran a grid search to approximate the values.

Reviewer comment (Novelty):
The paper presents two algorithms: REX-UCB and REX-R, which differ only in the way the rewards are classified. I wonder whether these could have been unified into one coherent approach. If not, it is important to explicitly mention when one should use which algorithm.

Author response:
As noted in the paper, REX-R equals REX-UCB when setting the constant $C$ in equation (1) to zero. We introduced these two distinct approaches to highlight the importance of the UCB score over absolute reward. Merely relying on feedback (as Reward) isn't enough for optimal performance; achieving superior results requires a nuanced balance between exploring and exploiting the action space.
* REX-UCB has showcased better performance compared to REX-R, indicating our preference for REX-UCB across all scenarios. The introduction of REX-R aimed to emphasize that relying solely on Reward for feedback does not lead to optimal performance, as highlighted in the earlier point.

Reviewer comment (Theory):
Is it assumed that a user can map the UCB scores to HIGH/LOW? This may not always be possible.

Author response:
No, this assumption was never made. In Algorithm 1, line 5, we can see that the action with the highest UCB score gets 'HIGH' and all the remaining actions get 'LOW'. The user is never involved in any of the REX algorithms.

Reviewer comment (Theory):
Can the method deal with non-binary rewards?

Author response:
Certainly! The current REX implementation supports non-binary rewards. Essentially, our approach utilizes UCB scores, meaning if the UCB formula can handle diverse rewards, our approach can as well.

Reviewer comment (Experiments):
Can the tasks described be solved with reinforcement learning alone? If so, this would be a natural baseline to add to the experiments.

Author response:
Not really. We have taken into account benchmark datasets like GSM8K, where questions are presented in natural language. Thus, relying solely on RL techniques is impractical in this scenario.

Reviewer comment (Writing):
The introduction could have included more emphasis about how the contributions (1), (2) and (3) are achieved.

Author response:
In the Introduction, our emphasis has been on establishing the significance and current trends in Autonomous Agents. We consider it vital to convey the motivation behind our project. The achievement of (1), (2), and (3) is outlined in the final paragraph of the Introduction. We delve deeper into these aspects within the "Proposed Methodology" and "Experiments" sections.

Reviewer comment (Writing):
In Section 4.1 - `form’ should be ‘from’. This mistake is also repeated in later parts of the paper.

Author response:
Thank you for helping us identify this typo. We have corrected it.

Reviewer comment (Experiments):
Regarding experiments, I have several concerns. First, I consider that the authors should compare the performance of their LLM-based approaches with SOTA methods that are not based on LLMs.

Author response:
Our paper focuses on demonstrating the efficiency of MCTS, particularly the UCB scoring mechanism, within LLMs. Comparing with non-LLM methods might not be feasible for datasets where questions are in natural language, making it challenging to adapt them to non-LLM approaches. Dataset nature impacts comparison feasibility; transforming natural language to discrete states for non-LLM methods may not be viable for all datasets. Hence we have only considered LLM-based approaches for comparison.

Reviewer comment (Writing):
Maths are sometimes weird to me. For example:  
- Algorithm 1, line 3: $\forall (s,a) \in Q$. $Q$ is not a set of state-action pairs. This part should be $\forall (s,a) \in S \times A$.  
- Algorithm 1, line 14: $\forall (s,a) \in H$. $H$ is a time step.

Author response:
In Algorithm 1, $Q$ includes all possible combinations of $(s, a)$. Equation (1) shows that to compute the UCB score, every $(s, a)$ pair must have a representation in $Q$. To avoid confusion for the reader, we opted to maintain the expression as "$\forall (s, a) \in Q$".

As already defined in the "*Require:*" section of Algorithm 1, $H$ represents sequence-of-intermediate-steps. so in line 14, we are conveying the idea that - for every intermediate step taken in the latest solution, we are updating $Q(s, a)$.

Reviewer comment (Novelty):
[Q1] Could you tell me the similarities and differences between the authors' approach and Tree-of-Thoughts (Yao et al., 2023)? Is it possible to compare the performance of the two methods in the same benchmark? Because the authors use BlocksWorld and GSM8K and Yao et al. (2023) use Game of 24, Creative Writing, and Mini Crosswords, I guess some of them are used for the comparison.

Author response:
* We have addressed the similarities and differences between REX and ToT in the above section.
* (Currently running ToT on Blocksworld and GSM8K, We will share the results soon. )

As a continuation from our previous correspondence, we have completed the experiments using the Tree-of-Thoughts (Yao et al., 2023) technique. The implementation details are as follows: the BFS version of TOT was employed for the experiment, utilizing OAI's GPT-3.5-turbo-0301 LLM. We used value-based action selection in TOT instead of vote-based action selection. Please find the accuracy of various techniques below:

| Technique | Blocksworld (2-step) | Blocksworld (4-step) | Blocksworld (6-step) | GSM8K |
|-------------|-------------------------|-------------------------|-------------------------|-----------|
| ToT | 50% | 3.2% | 0% | 52.6% |
| REX-R | 53.3% | 37.5% | 14.9% | 81.34% |
| REX-UCB | 80% | 39.28% | 25.43% | 82.03% | 
| REX-UCL | 60% | 44.64% | 20.17% | 90.44% |

The Tree-of-Thoughts (ToT) technique demonstrates strong performance when paired with a language model proficient in Self-Evaluation, such as GPT-4. However, its efficacy diminishes when used with less robust models like gpt-3.5. This discrepancy primarily arises from the higher error rate exhibited by gpt-3.5 when selecting a specific action from the sampled set, which undermines the proposed trajectories within the ToT technique. In contrast, REX, relying on systematic reward collection, exhibits comparatively sound performance even with models like gpt-3.5. The code and the experimental setup have been uploaded for your reference.

Reviewer comment (Reproducibility):
[Q2] Although I can guess what are Agent.Solve and CALCULATEUCB to some extent, there is almost no information on what is Agent.ValidateAction. Could you elaborate more?

Author response:
* Details available in Appendix - Section D

Reviewer comment (Experiments):
[Q3] I suspect that the high performance of REX in BlocksWorld is due to the data leak. Did the authors use the ordinary BlocksWorld environments with default settings?

Author response:
* We have confirmed that there's no data leakage in our setup. We created our few-shot examples (3 in total) to form the prompt and used these examples consistently across Reflexion, CoT, and all REX variations. It's impossible for REX to experience data leakage while others do not.