Tree Search over Agent Trajectories (LATS)

Sources#

• CS329A Self-Improving AI Agents — Part 5: Planning and Multi-Step Reasoning

Summary#

LATS — Language Agent Tree Search unifies reasoning, acting and planning in language models (ICML 2024) — is the first paper in CS329A lecture 5 and this wiki's first treatment of Monte Carlo Tree Search applied to an agent's actions in an environment. The wiki already carried tree search over proofs (Evolutionary Proof Search's P-UCB) and search over inference pipelines (Inference-Time Architecture Search); LATS is the version where the nodes are states of a world the agent is acting on.

Its one-sentence claim: take ReAct's thought/action/observation alternation, stop committing to the first trajectory, and put MCTS around it. Mirhoseini's framing of the problem it targets is a capability complaint that she says still holds: models "can be weak in creating a diversity of solutions or acting upon them," so a single sampled plan under-explores by default. Search is the harness-side fix.

Evidence. CS329A Self-Improving AI Agents — Part 5: Planning and Multi-Step Reasoning (Azalia Mirhoseini solo, delivered 2025-10-06, published 2026-08-03, practitioner-opinion) — a slide walkthrough. The LATS paper is not in raw/; every figure was read off a slide by ASR, and the lecture reports almost no numbers for this paper at all. No COI here: LATS is not her lab's work (it is the only one of the lecture's three papers that isn't).

The six stages#

The lecture enumerates them and walks a maze example — "you're in a dimly lit room, there are two doors" — through each:

1. Selection. Pick the node to expand, by UCT (below), not by best value.
2. Expansion. Sample k candidate actions from the selected node (the example samples three: open the left door, open the right door, inspect the room for clues).
3. Evaluation. Execute each action in the environment, append the observation to the context, and score the resulting state.
4. Simulation. Roll the best child forward — greedily, in the lecture's version — sampling and expanding until a terminal state (success or failure) or the expansion budget runs out.
5. Back-propagation. Push the trajectory's return back up the path, updating each node's value.
6. Reflection. Have the model write, in free text, why the trajectory succeeded or failed, and append that to the context for later expansions. Mirhoseini: this "apparently has been very helpful in overall increasing the quality of this approach."

Stages 1–5 are textbook MCTS. Stages 3 and 6 are where the language model does something a game-playing MCTS cannot.

The value function is two scores added together#

A state's value is the sum of:

• An LLM-as-a-judge score. Prompt a model with the action and its observation and ask, literally, for a number in [0, 1] for how promising this state is. See LLM-as-a-Judge.
• A self-consistency score. Sample many actions from the parent — the lecture's illustration is 50 rather than 3 — bucket them by type, and score an action by the frequency with which it was sampled. In the worked example, state A scores higher because "75% of the time maybe action A was sampled."

The two are summed to give one number per state.

Selection and backup#

UCT (upper confidence bounds applied to trees) is taken over from the MCTS literature unmodified: value plus a weighted exploration term built from the parent's visit count n_p and the node's own visit count n_s, so that a node visited rarely relative to its parent gets a bonus and a heavily-visited one gets discounted. (The lecture describes the term's behaviour rather than writing it out legibly; the exact form is not recoverable from the transcript and is not reconstructed here.) Asked whether the paper compared alternative bandit algorithms, Mirhoseini says no and treats the choice as arbitrary — "their main contribution is that they created a platform that now others can bring in other approaches to optimization."

Backup is a running average: a node's new value is (V_old·(n−1) + return) / n over its visit count. Nothing learned, nothing trained — every component is prompting.

What it diverges from, in the lecture's own comparison#

Mirhoseini positions LATS against two things the course had already taught, and the distinction is sharper than it first looks:

• Against Math-Shepherd (Process vs Outcome Reward Models): there, a trained verifier scores reasoning steps and guides the search. In LATS the scoring is over the outcomes of actions taken in an environment, plus the model's reflection on the trajectory and the observations the environment returned. The unit of credit moves from a token span to a world state.
• Against ReAct (Reasoning–Acting Interleaving (ReAct)): ReAct never revisits. LATS adds the memory of scored alternatives, the ability to back up to a sibling, and the reflection step. In the lecture's phrasing, "we have more and more planning in the process."

Results, such as the lecture reports them#

Two benchmarks, both already in this wiki via ReAct:

• HotpotQA — multi-hop QA requiring retrieval from at least two Wikipedia pages, so multi-step by construction. The reported shape is that accuracy climbs materially with the number of sampled trajectories, and that adding the reflection traces "gives a lot of boosts." No figures survive the ASR. Mirhoseini's summary is the one that matters for this wiki: LATS supplies "a mechanism to translate more compute at test time effectively to better solutions" for multi-step tasks — test-time scaling applied to acting rather than to answering.
• WebShop — buy a product matching a natural-language spec in a simulated storefront. LATS reportedly reaches "really high results, even close to human experts," with no fine-tuning at all. For scale, ReAct alone scores 66.6 against a human expert's 82.1 on the lecture-4 slide; LATS's number is not stated, so the comparison is directional only.

The portability claim is the honest one: everything is prompting over a frozen model, so the method is "very portable and relatively easy to create."

The two limits the lecture concedes, and the one it does not#

Cost was never analysed. Mirhoseini states this outright — every expansion, every rollout, every judge call and every backup adds inference, and "the cost-benefit was not really analysed in the paper." A method whose entire pitch is converting test-time compute into quality, published without a compute axis, is exactly the gap Compute-Controlled Benchmarking exists to name.

Actions must be reversible. This is the deeper one, and it is stated as an assumption the paper did not address: the search executes candidate actions in the environment in order to score them, so exploring a branch means actually taking it. Mirhoseini's example is a model "running a transaction… paying for a service." In any environment without an undo, the expansion stage is not a probe — it is a commitment, and the sibling branches were paid for. LATS is therefore a method for simulated or sandboxed environments, and the lecture's own maze and storefront are both simulators — which puts the sandbox from containment in a role it was not designed for: not damage limitation, but a precondition for the algorithm to be sound at all.

The unstated one: half the value function is a selector this course already showed plateaus. The self-consistency term scores an action by how often it was sampled — per-node majority voting. Lecture 2's central result is that majority voting saturates at 10–50 samples while coverage keeps climbing, because the hardest problems are solved 1–3 times in 10,000, so a frequency-based selector is structurally blind to the rare-but-right branch. LATS sums that blind selector with an LLM judge and steers the whole search with the result. The lecture states both facts three weeks apart and never joins them; the join is this compile's reading, and it predicts the failure mode — search that converges confidently on the modal plan, which is the same complaint about solution diversity that motivated the paper.

Repeated actions, and the tree/graph question#

A student asks what happens when the same action recurs across branches (A, B, A, B) or the same state is reachable by different paths — can the tree be collapsed? Mirhoseini's answer keeps it a tree: repetition under a given parent is captured by the visit count feeding UCT, and "ideally this is a tree that you're forming, not some kind of fully connected graph." So identical states reached via different ancestors are separate nodes, evaluated separately and paid for separately — a known and unpriced inefficiency in this design, and one more reason the missing cost analysis matters.

The same algorithm where the environment is a compiler (2026-08)#

Vamshi and Yang (Reward-Oracle MCTS for Formal Theorem Proving: Sample-Efficient Search and the Need for Kernel-Level Proof Auditing, arXiv 2608.28639, empirical) run structurally the same search in formal mathematics, and the differences are exactly the three places this page says LATS is weak. Their tree's nodes are natural-language proof plans rather than world states; a decomposer expands each selected node into K = 4 candidate next subgoals; a generator writes S complete Lean 4 proof attempts per child; UCB with c = √2 selects; the rollout reward backpropagates as a running mean. Per-model results on Agentic Loops Overtake Bespoke Systems; its place among the benchmark denominators on AI-Driven Formal Proof Search.

The value function is still two scores added together, and the second one is sound. LATS sums an LLM judge's [0,1] score with a sample-frequency term. This system sums an LLM critic's score (normalized from a 0–100 rubric, five samples at τ = 0.3, averaged) with s_k/S — the fraction of the node's proof attempts that the Lean compiler accepted. Same shape, same arity, but the second term is a frequency of success measured by a sound verifier rather than a frequency of sampling measured by the model's own prior. That is the structural fix for the blind selector this page identifies: a branch nobody samples often but that compiles gets credit, where LATS's frequency term would bury it.

Reversibility is free here, which is why the method is deployable at all. This page's sharpest criticism of LATS is that expansion means executing the candidate action, so the search only works in a simulator. Compiling a candidate Lean proof has no side effects: the probe and the commitment are different objects by construction. Formal proof search is therefore the natural habitat of trajectory tree search, and the reason it keeps appearing there (Evolutionary Proof Search's P-UCB is the other instance in this wiki).

The cost analysis LATS never ran. At an identical proof-attempt budget the three-role search consumes 32.8% fewer total inference tokens and 35.8% fewer output tokens than flat sampling — the decomposer and critic calls are short (1,024 and 3 max output tokens) and a generator conditioned on an explicit decomposition writes materially shorter proofs. So in this domain the search is not a test-time-compute trade at all; it is Pareto. The paper also publishes the allocation sweep this page's UCT discussion wants: under a fixed N·K·S = 32, one wide expansion (N=1, K=16) scores 82.1% and four iterations of four (N=4, K=4) scores 84.2% with identical numbers of decomposition calls, critic evaluations and proof attempts — so the gain is the intermediate backpropagation, not the call budget. The narrowest setting (N=16, K=1) buys 0.1 points for 2.8× the wall clock against 1.5×.

Connections#

• CS329A: Self-Improving AI Agents (Stanford) — lecture 5's first paper; the course's move from verification to planning
• Reasoning–Acting Interleaving (ReAct) — the inner loop LATS wraps: ReAct supplies the thought/action/observation unit, LATS supplies the branching, scoring and backtracking ReAct has no way to do
• Intra-Trace Parallel Planning (SPRINT) — lecture 5's second paper and the opposite trade: LATS spends more sequential compute to search alternatives at inference; SPRINT trains the model to spend less by emitting independent plans that run at once. Both are called "planning" in the same lecture and they push in opposite directions on latency
• Offline Multi-Step Tool-Use RL (SWiRL) — lecture 5's third, and the training-time answer to the same problem: LATS searches harder around a frozen model, SWiRL changes the weights so the first trajectory is better
• Process vs Outcome Reward Models — the verifier lineage LATS defines itself against: Math-Shepherd scores reasoning steps with a trained PRM, LATS scores world states with a prompted judge plus a frequency term
• LLM-as-a-Judge — half of LATS's value function is a 0-to-1 judge prompt, with all of that pattern's known reliability limits inherited by the search
• Large-Scale Test-Time Compute — LATS is the acting-agent instance of spending inference budget for quality, and the lecture frames it exactly that way
• The Verifiability Thesis — the search is only as good as the score that steers it, and the self-consistency half is the selector this hub explains the limits of
• Turn-Level Credit Assignment — the same credit problem solved at the other end of the pipeline: LATS back-propagates a terminal return through a tree at inference time with no gradients; TRACE splits one trajectory's return across turns at training time. The backup formulas are cousins; the objects updated are not
• Evolutionary Proof Search — the closest existing relative in the wiki: a population of proof sketches ranked by Elo and selected by P-UCB, i.e. the same UCB-family search with a compiler instead of an LLM judge as the fitness signal
• Inference-Time Architecture Search — search at a different altitude: Archon searches over pipelines offline, LATS searches over trajectories online
• Stopping Under a Noisy Verifier — what the missing cost analysis would have to reckon with: under a noisy scorer, more search is not monotonically better, and the loop's true quality can decline while its reported score rises
• Compute-Controlled Benchmarking — the discipline this paper's headline result is missing
• Blast Radius (Agentic) — why the reversibility assumption is load-bearing: expanding a node means executing it, so the sandbox stops being containment and becomes a correctness precondition
• Kernel-Level Proof Auditing — the formal-math instance of this algorithm and the reason its verifier is trustworthy: same UCB tree, same two-term value function, but the second term is a compiler's accept count rather than a sample frequency — and the page that measures what happens when the harness reporting that count is weaker than the kernel behind it
• Azalia Mirhoseini — the lecturer
• Selection Under a Submission Budget — the need this answers, re-derived from the code side two lectures later. Massive one-shot sampling plateaus on hard competitive-programming problems, and the lecture's own conclusion is that they want decomposition, per-step sampling and backtracking — this page's search, arrived at by exhausting the parallel alternative

Open Questions#

• LATS scores a state with a prompted judge plus a sample-frequency term, and the same course showed frequency-based selection is blind to rare-correct solutions. Does the judge half carry the search on hard problems, or does the frequency half dominate and collapse the exploration the method exists to create? An ablation of the two terms would settle it. Partially answered 2026-09-23 by Reward-Oracle MCTS for Formal Theorem Proving: Sample-Efficient Search and the Need for Kernel-Level Proof Auditing, which runs exactly that ablation in a domain where the second term is a sound verifier rather than a sample frequency: critic-score-only against critic-plus-compile-fraction, same tree, same budget, five seeds. The judge half does not carry the search — dropping the second term costs 3.4 points on DeepSeek-Prover-V2-7B (77.1 → 73.7 at PAB@32) and 2.6 on Goedel-Prover-V2-8B (84.2 → 81.6), and the gap widens rather than closes as budget grows. So the two-term design is load-bearing and the ablation is cheap to run, which is the transferable half. It stays #oq/source because the substituted term is the wrong object: LATS's frequency term measures how often the model proposed an action, this one measures how often a compiler accepted the result, and the rare-correct blindness this question is about lives entirely in the former. Nobody has ablated LATS's own two terms.
• The reversibility assumption confines tree search to simulators. Is there a version that scores a candidate action without executing it — a learned or prompted transition model — or does search over real-world actions reduce to "explore only in a sandbox, then replay the winning trajectory"?

Sources#

• CS329A Self-Improving AI Agents — Part 5: Planning and Multi-Step Reasoning — CS329A Self-Improving AI Agents — Part 5: Planning and Multi-Step Reasoning, Azalia Mirhoseini solo, Stanford Online. Delivered 2025-10-06, published to YouTube 2026-08-03 (practitioner-opinion, YouTube auto-caption transcript, ~11.3k words). The LATS third of the lecture: the trip-planning and maze walkthroughs, the six stages, the two-term value function and its 0–1 judge prompt, UCT selection and the running-average backup, the reflection step, the divergences from Math-Shepherd and ReAct, the HotpotQA and WebShop results, and the conceded limits (no cost analysis, irreversible actions) plus the two student exchanges on bandit alternatives and repeated actions. The LATS paper is not in raw/; the lecture names no authors and almost no numbers, and everything here is ASR-read off slides and hedged accordingly