LATS (Language Agent Tree Search)

An agent framework that binds ToT, ReAct, Reflexion and MCTS into a single tree search to explore candidate solutions.

Concept

LATS (Language Agent Tree Search) is a framework that unifies four techniques, ToT, ReAct, Reflexion and MCTS:

• Tree of Thought (ToT): exploring multiple paths
• ReAct: interacting with external tools
• Reflexion: learning from failure
• Monte Carlo Tree Search (MCTS): optimizing the search

flowchart LR
    S["Selection"] --> E["Expansion"]
    E --> V["Evaluation"]
    V --> M["Simulation"]
    M --> B["Backpropagation"]
    B --> R["Reflection"]
    R --> S

The six core operations

1. Selection: pick the node to explore.
2. Expansion: generate new possibilities.
3. Evaluation: score how promising each one is.
4. Simulation: execute the action and observe the result.
5. Backpropagation: propagate the result back up.
6. Reflection: reflect on failure.

Each step in detail

The techniques it unifies

Performance

The final ICML 2024 paper (arXiv:2310.04406) reports 92.7% HumanEval pass@1 with GPT-4 and a WebShop average score of 75.9 with GPT-3.5. The HumanEval result includes search and test feedback. These tasks use different models and metrics, so the scores cannot be compared directly or interpreted as single-call accuracy.

The coding example and scores below illustrate the procedure. They are not execution records or reproduced results from the paper.

Coding Problem

Problem: "Implement merge sort in Python"

Selection: pick the node with the highest UCB1 score. The recursive implementation approach is selected.

Expansion: generate three implementation approaches: approach 1: plain recursion, approach 2: optimized recursion (in-place), approach 3: iterative implementation.

Evaluation: score how promising each approach is: approach 1: 0.7 (simple but memory-inefficient), approach 2: 0.9 (well balanced), approach 3: 0.6 (complicated).

Simulation (ReAct): implement approach 2.

• Thought: split the array in half and recurse
• Action: write the code def merge_sort(arr): ...
• Observation: code written
• Thought: run the tests
• Action: run merge_sort([3,1,4,1,5,9])
• Observation: [1,1,3,4,5,9], success

Backpropagation: the passing test propagates a score back up, and the UCB1 score along the approach 2 path rises.

Done: return the final code.

Reflection on failure

Simulation failure: running merge_sort([]) raises IndexError: list index out of range

Reflection: "There was no boundary handling for an empty array. An empty-array check has to be added to the recursion base case."

Retry: add the boundary condition if len(arr) <= 1: return arr and the problem is fixed.

The UCB1 algorithm

LATS balances exploration against exploitation using UCB1 (Upper Confidence Bound) from MCTS:

$\text{UCB1} = \bar{x} + C \times \sqrt{\frac{\ln(N)}{n}}$

• (\bar{x}) (mean reward): how good this path has been (exploitation)
• (\sqrt{\frac{\ln(N)}{n}}) (exploration term): a bonus for less-visited paths (exploration)
• (C): the exploration-exploitation balance constant

Compared to earlier techniques

• CoT: single-path search, no tools, no learning, no optimization
• ToT: multi-path search, no tools, no learning, no optimization
• ReAct: single-path search, tools yes, no learning, no optimization
• Reflexion: single-path search, tools yes, learning yes, no optimization
• LATS: multi-path search, tools yes, learning yes, optimization yes

Implementation structure

The code below is unexecuted pseudocode, not the full paper implementation. Real execution requires a search budget, stopping conditions, and environment-state restoration.

class LATS:
    def __init__(self, llm, tools):
        self.llm = llm
        self.tools = tools
        self.memory = []  # Reflexion memory
 
    def solve(self, problem):
        root = Node(state=problem)
 
        while not self.is_solved(root):
            # 1. Selection
            node = self.select(root)
 
            # 2. Expansion
            children = self.expand(node)
 
            # 3. Evaluation
            for child in children:
                child.value = self.evaluate(child)
 
            # 4. Simulation (ReAct)
            best_child = max(children, key=lambda x: x.value)
            result = self.simulate(best_child)
 
            # 5. Backpropagation
            self.backpropagate(best_child, result)
 
            # 6. Reflection (only on failure)
            if not result.success:
                reflection = self.reflect(best_child, result)
                self.memory.append(reflection)
 
        return self.get_solution(root)

The methods in the solve loop map one to one onto the six-step table above. select is Selection, expand is Expansion, evaluate is Evaluation, simulate is Simulation (ReAct), and backpropagate is Backpropagation. Only on failure does reflect run, accumulating the lesson in memory for the next expansion to consult.

Cost and selection

Candidate generation, evaluation, and simulation require additional calls. There is no universal ranking in which LATS is the most accurate or most expensive method on every task. Compare success rate, calls, and latency across search budgets with the same model and problem set.

Related concepts

• Tree of Thought (ToT): multi-path search
• ReAct: tool use
• Reflexion: learning from failure
• Monte Carlo Tree Search: the search algorithm from game AI

Summary

LATS connects actions and environmental feedback to tree search. Its selection-through-reflection procedure explores candidates, but evaluator errors and limited budgets can cause it to miss a solution. Measure the value of additional calls on the target task.

Use a single generation or ReAct as a baseline when evaluating additional LATS search. The benchmark numbers are as reported in the paper, and results will differ depending on the deployment.