What a search budget cannot fix after early pruning

A tidy first step can be the wrong first step

Imagine organizing a fictional archive. One strategy immediately creates attractive folder names, another inspects file contents, and a third checks the naming rules before moving anything. A superficial evaluator prefers the first strategy because the directory looks organized quickly. The final objective, however, is correct placement. Search is useful precisely when the intermediate signal and the terminal result disagree.

Language Agent Tree Search, especially Sections 3.2 and 4.2, combines tree search with language-model action proposals, value estimates, environment feedback, and reflections. Its selection rule includes an exploration bonus: a promising but rarely visited branch can compete with a frequently visited branch. The example below isolates that allocation decision. It does not implement the paper's language-model components.

Make the budget count something observable

The toy tree has three first-step strategies, each with three terminal plans. Their heuristic scores are 9, 6, and 2, but their terminal rewards are respectively [2, 3, 4], [5, 6, 7], and [0, 8, 10]. These numbers are deliberately constructed, not measured file-system outcomes. Each revealed leaf costs one simulated evaluation. Inspecting the supplied root heuristics is free for every policy, so no method receives hidden extra evaluations.

Greedy pruning keeps only the highest-ranked branch. A width-two beam keeps the best two. The UCT-style allocator keeps all three and chooses between their empirical means plus sqrt(2 log(total visits) / branch visits). Unvisited branches are tried first. Once a branch's three leaves have been examined, it leaves the frontier. This is a small branch-allocation skeleton, not a general MCTS implementation with arbitrary depth.

Consider the third branch after its first evaluation: it has reward zero, which looks discouraging. Eliminating it permanently hides both eight and ten. Continuing to explore costs time, but provides information the heuristic did not contain. With a nine-evaluation allowance, exhaustive coverage can expose the ten-point leaf; a pruned search cannot recover a discarded branch merely by receiving a larger allowance.

Read rewards and work together

The program reports both best terminal reward and evaluations actually used. Budget is a ceiling, not a requirement to waste calls after the retained frontier is exhausted. Comparing only rewards would miss that a width-one method stops after three leaves. Comparing only elapsed time would mix Python overhead with the fictional tool cost. For a latency calculation, assign an explicit cost such as 40 milliseconds per evaluation: six evaluations then represent 240 milliseconds of serial simulated work.

What would make the conclusion stronger?

The current fixture favors exploration by construction. Reverse the heuristic ranking, vary branch width, and include tasks where every extra evaluation is redundant. Keep those test trees separate from any trees used to choose exploration strength. Report success thresholds as well as best reward, because a tiny reward improvement may not complete the task. Finally, a reversible in-memory tree is a modeling convenience: real tool actions may not be undoable, and replaying a conversation does not restore external state.

Reverse the heuristic to test the story

The extension enumerates all six possible orders of the three root branches. At a nine-evaluation budget the UCT-style allocator finds reward ten in every order because it eventually exhausts all leaves. Greedy and beam-two depend on which branches survive their initial pruning. This is a coverage statement, not evidence that the exploration bonus is optimal.

At smaller budgets, order and the sequence of rewards within a branch matter. A branch with an early zero can consume exploration effort before its later high rewards become visible. The complete budget sweep therefore reports both the best value found and actual evaluations used. A useful planner comparison needs those paired measurements rather than a maximum reward detached from its search cost.

"""What a search budget cannot fix after early pruning
Deterministic, standard-library experiment. RESULTS contains the full data.
"""
import math

leaves = ((2, 3, 4), (5, 6, 7), (0, 8, 10))
heuristic_order = (0, 1, 2)

def search(mode, budget):
    allowed = heuristic_order[:{'greedy': 1, 'beam2': 2, 'uct': 3}[mode]]
    visits, totals = [0] * 3, [0.0] * 3
    best, used = None, 0
    for _ in range(budget):
        frontier = [a for a in allowed if visits[a] < len(leaves[a])]
        if not frontier:
            break
        if mode == 'uct':
            unseen = [a for a in frontier if visits[a] == 0]
            def score(a):
                bonus = math.sqrt(2 * math.log(max(1, used)) / visits[a])
                return totals[a] / visits[a] + bonus
            arm = unseen[0] if unseen else max(frontier, key=score)
        else:
            arm = min(frontier, key=lambda a: (visits[a], a))
        reward = leaves[arm][visits[arm]]
        visits[arm] += 1
        totals[arm] += reward
        used += 1
        best = reward if best is None else max(best, reward)
    return best, used

for budget in (2, 4, 6, 9):
    row = {mode: search(mode, budget) for mode in ('greedy', 'beam2', 'uct')}
    print('budget', budget, 'best_reward,used', row)
assert search('uct', 9) == (10, 9)
assert search('greedy', 9) == (4, 3)


sweep=[dict(budget=b,mode=m,best_reward=search(m,b)[0],used=search(m,b)[1]) for b in range(1,10) for m in ('greedy','beam2','uct')]
# Exhaustive three-branch allocation recovers all nine leaves regardless of ordering.
original_order=heuristic_order
permuted=[]
from itertools import permutations
for order in permutations(range(3)):
    heuristic_order=order
    permuted.append(dict(order=list(order),greedy=search('greedy',9)[0],beam2=search('beam2',9)[0],uct=search('uct',9)[0]))
heuristic_order=original_order
assert all(r['uct']==10 for r in permuted)
RESULTS=dict(leaves=leaves,heuristics=[9,6,2],sweep=sweep,order_sensitivity=permuted)
print('Heuristic order permutations:',len(permuted),'; full-budget UCT maximum always 10')

budget 2 best_reward,used {'greedy': (3, 2), 'beam2': (5, 2), 'uct': (5, 2)}
budget 4 best_reward,used {'greedy': (4, 3), 'beam2': (6, 4), 'uct': (6, 4)}
budget 6 best_reward,used {'greedy': (4, 3), 'beam2': (7, 6), 'uct': (7, 6)}
budget 9 best_reward,used {'greedy': (4, 3), 'beam2': (7, 6), 'uct': (10, 9)}
Heuristic order permutations: 6 ; full-budget UCT maximum always 10