Language Agent Tree Search

Real environment

Implement the Language Agent Tree Search pattern in this codebase, using the LLM client and language we already use.

## Goal

Find a sequence of actions that solves a task in an environment you can reset or branch, by running Monte Carlo tree search where a model proposes actions, the environment executes them, and a model scores the outcomes.

## How it works

1. Represent the search as a tree of nodes. Each node holds an environment state, the trajectory that reached it, a visit count and a running value; the root is the task.
2. Select: starting at the root, repeatedly pick the child with the highest UCT score (value plus an exploration bonus that shrinks with visits) until you reach a node that has not been expanded. Do this in code, not with a model.
3. Expand: call the model with the Expand prompt k times (default 2) to sample different next actions from that node, including every saved reflection in the prompt.
4. Act: run each sampled action in the environment from that node's state, such as executing the code against the visible tests, and record the observation as a new child node.
5. Simulate (multi-step tasks only): roll a new child forward to a terminal state by repeating expand and act with one sample per step. Skip this when one action is a complete answer, such as a whole program.
6. Evaluate: if the environment says a child solved the task, return its trajectory. Otherwise score it with the Value prompt; blend that score with a self-consistency score or, where tests exist, the test pass rate, and use the blend as the child's value.
7. Backpropagate: walk from each new child up to the root, adding one to each node's visit count and folding the child's value into each node's running average.
8. Reflect: for each child that ends in a terminal failure, call the model with the Reflection prompt, passing its trajectory and the environment's feedback, and save the reflection for every later expansion. Then start the next rollout at the select step.

## Prompts

### Expand prompt

```markdown
You are solving the task below step by step. Propose the next action.

<task>{{task}}</task>
<trajectory>{{trajectory_so_far}}</trajectory>

{{#if reflections}}
Earlier attempts failed. Lessons from them:
<reflections>{{reflections}}</reflections>
Do not repeat those mistakes.
{{/if}}

Think briefly about what is still missing, then give exactly one action
in this format: {{action_format}}
```

### Value prompt

```markdown
Judge how close this attempt is to solving the task.

<task>{{task}}</task>
<trajectory>{{trajectory}}</trajectory>
<latest_observation>{{observation}}</latest_observation>

Reason in two or three sentences about what is right and what is still wrong.
Then end with one line: SCORE: a number from 0 to 1, where 1 means solved.
```

### Reflection prompt

```markdown
This attempt failed. Work out why, so the next attempt avoids it.

<task>{{task}}</task>
<trajectory>{{trajectory}}</trajectory>
<feedback>{{environment_feedback}}</feedback>

In one or two sentences, name the specific mistake and what to do instead.
Do not write a solution.
```

## Stop rule

Stop as soon as the environment confirms a solution, or after a rollout budget (default 8). On the budget, return the trajectory of the highest-value node with a status saying no solution was confirmed.

## Requirements

- The environment must be resettable or snapshotable, because siblings start from the same state. Never run this over actions that cannot be undone.
- Cap rollouts (default 8) and children per expansion (default 2), and log the model calls and environment runs for each rollout.
- Parse the value prompt's score into a number between 0 and 1; if it does not parse, retry once, then score the child 0.
- Keep the tree, with visit counts and values, in a form you can dump to a log, so a run can be inspected after it ends.
- Add a test with a stub LLM and a fake test runner that fails both children of the first rollout and passes the child from the second; assert reflect was called twice, both lessons appear in every expand prompt of rollout 2, and the passing code is returned.

World model (RAP)

Implement the Language Agent Tree Search pattern (World model (RAP) variant) in this codebase, using the LLM client and language we already use.

## Goal

Plan a sequence of actions toward a goal without touching a real environment, by running Monte Carlo tree search where a model proposes actions, a model predicts the resulting state, and a reward scores progress.

## How it works

1. Represent the search as a tree whose nodes are predicted states, each with a visit count and a running value; the root is the starting state from the task.
2. Select: from the root, repeatedly pick the child with the highest UCT score until you reach a node that has not been expanded. Do this in code.
3. Expand: call the model with the Action proposal prompt to propose up to k (default 3) actions that are legal in the node's state.
4. Simulate: for each action, call the model with the World model prompt to predict the next state, and add it as a child. No real action is taken.
5. Reward: score each child with the Reward prompt, combined with any check in code (for example, how many goal conditions the predicted state already meets).
6. Backpropagate each child's reward up to the root, then start the next rollout; when a predicted state meets the goal, return the actions on its path.

## Prompts

### Action proposal prompt

```markdown
<goal>{{goal}}</goal>
<state>{{state}}</state>

List up to {{k}} different actions that are legal in this state and could
move toward the goal, one per line. Use only these action forms: {{action_forms}}
```

### World model prompt

```markdown
You simulate the environment. Given a state and an action, predict the new state.

<state>{{state}}</state>
<action>{{action}}</action>

If the action is not legal in this state, reply ILLEGAL and why.
Otherwise list every fact in the new state, one per line, and nothing else.
```

### Reward prompt

```markdown
<goal>{{goal}}</goal>
<state>{{state}}</state>
<plan_so_far>{{actions}}</plan_so_far>

Is this a good step toward the goal? Answer Yes or No on the first line,
then one sentence of reasoning.
```

## Stop rule

Stop when a predicted state meets every goal condition, or after a rollout budget (default 10). On the budget, return the highest-value path with a status saying the goal was not reached.

## Requirements

- Keep states as structured text (a list of facts) so a code check can test goal conditions instead of trusting the model.
- Cap rollouts (default 10), depth (default 6) and actions per expansion (default 3).
- Validate each predicted state against simple rules in code where you have them (a block cannot be in two places) and drop children that break them.
- Treat the plan as a proposal: if it will run in the real world, execute it step by step and replan when an observed state differs from the predicted one.
- Add a test with a stub LLM whose world model reports the goal met after the fourth action on one path; assert the returned plan is exactly the four actions on that path, in order.