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

The fifth CS329A lecture brings together LATS, SPRINT and SWiRL. Each addresses multi-step reasoning, but they intervene at different points: explore alternative action trajectories, expose independent work that can run together, or train a model to choose better actions from collected contexts.

That distinction prevents a vague conclusion that “more planning” solves every task. Searching more branches costs computation. Parallel execution needs genuine independence. Training on a judged action depends on whether the judge’s signal corresponds to useful behavior.

The common objective is to make the next step better informed without losing sight of the whole task. The lecture’s examples show how to organize that process and where its assumptions become limiting.

Language Agent Tree Search combines reasoning, acting and planning with ideas from Monte Carlo tree search. Instead of generating one plan and simply executing it, the system samples possible actions, obtains observations and expands selected trajectories.

The maze example starts with alternatives such as opening either door or inspecting the room. Executing an action adds an observation to context. A state value combines a language-model judgment of promise with a self-consistency signal based on how often an action was sampled.

Both components are estimates. A frequently proposed action is not automatically correct, and a confident judge can be wrong. Their role is to guide which trajectory to examine next, not to certify a successful outcome before the environment has supplied one.

UCT balances exploitation of a high-value state with exploration of less-visited children. Relative visit counts encourage the search to inspect a branch that a purely greedy policy might neglect.

After a simulated trajectory succeeds, fails or exhausts its budget, its return updates values along the visited path. This backpropagation updates search statistics; it is not a neural-network gradient update to the language model’s weights.

LATS also appends a reflection on the trajectory to inform later attempts. The lecture reports benefits in the studied multi-hop question-answering and simulated-shopping tasks. It does not claim that the chosen exploration formula has been proved globally optimal for arbitrary language-agent problems.

The lecture explicitly identifies two limits. Tree expansion and repeated evaluation add cost, while the cited work does not fully analyze the cost-benefit trade-off. Strong benchmark results alone do not establish an economical deployment.

The second limit is irreversibility. A simulated action may be easy to revisit; a real transaction or payment may not be. Exploring several possible branches cannot mean casually executing every consequential real-world action.

That is an assumption of the method’s environment, not a detail a planner can ignore. A tree representation also does not automatically solve every repeated-state or arbitrary-graph problem. The task’s state, action and rollback properties determine what kind of search is appropriate.

SPRINT begins with the observation that long reasoning traces can contain independent work. A model may consider separate subproblems or calculations that do not need to wait for one another.

The described data pipeline segments existing reasoning traces, annotates planning and execution pieces, builds a dependency graph and packs independent steps together. Supervised fine-tuning then teaches a reasoning model to emit plans and execution structure in that form.

The underlying language model still predicts tokens. The special structure lets an external execution process branch into independent work and synchronize results back into context. It is not a claim that a single next-token stream has spontaneously become several hardware-parallel streams without orchestration.

At inference time, independent plans can be issued before waiting for each execution to finish. Their results return to context, after which the model can plan again or revise its direction.

33:10View full image

The lecture reports improvements in sequential-token measures and accuracy for the studied models and tasks, including transfer beyond the training dataset. Sequential tokens describe a critical-path opportunity; they are not identical to total generated tokens, actual elapsed time or total infrastructure cost.

Stragglers remain possible, and the available width depends on the task. For short or easy reasoning, planning overhead can even make the method worse than its baseline. Parallelism is valuable where independence exists and the runtime can exploit it, not merely because a prompt contains several numbered plans.

A student asks whether individually plausible branches can contradict one another when combined. The response leaves responsibility with the final synthesis: the model sees returned context and is expected to resolve conflicts before answering.

That expectation is a requirement, not a guarantee. Sequential reasoning can also contain contradictions, while parallel execution adds dependency and synchronization questions that must be checked.

The lecturers observe broader exploration earlier and more focused work later in the studied traces. They present further reinforcement learning and runtime optimization as opportunities, not completed results of every extension. The mechanism should be evaluated as a whole, including how it recognizes when branches were not actually independent.

SWiRL addresses multi-step reasoning and tool use while avoiding live tool execution during the reinforcement-learning update. Tools can be slow, fail or complicate training, so the method first collects trajectories and environment responses offline.

Iterative prompting asks a model to reason, call a tool or produce a final answer one step at a time. The next prompt includes prior actions and their observed responses. A language-model judge scores the proposed steps, creating process signals alongside final-outcome information.

Thus “no tools during training” has a specific scope. Tool calls still occur in data collection, and real tool execution is used when the trained system runs at inference. It does not mean the method learned tool behavior without any tool observations.

During the described RL procedure, a collected prefix supplies prior context and observations. The model proposes a next action, and a judge can evaluate whether that action is sensible without executing the newly proposed tool call inside the update.

60:50View full image

For an age-comparison question, asking for the first person’s age can be a reasonable search query even before the judge knows the returned date. This supplies a process reward about the action’s quality, not evidence that a live tool produced a correct answer.

That separation makes the training procedure more tractable, but preserves a boundary: a plausible query can still retrieve poor information or fail when executed. The action judge and the environment answer different questions, so their signals should remain distinguishable in evaluation.

The lecture reports that process-filtered data can help SWiRL’s RL training more than retaining only trajectories with correct final outcomes. Its explanation is that a model can learn to take a better new action from a useful prefix even when the original trajectory ended poorly.

Supervised fine-tuning has a different relationship to the same data. Since it imitates provided trajectories, incorrect endings can be harmful; in the reported comparison, filtering for both process and outcome was more useful for that objective.

These are findings and interpretations from the described experiments, not permission to treat all incorrect training data as beneficial. Data selection and the optimization objective must be considered together, including how reliable the process judge is.

SWiRL’s reported experiments move between mathematical tasks with calculation tools and multi-hop questions with search. The lecture interprets gains on held-out tasks as evidence that the model learned something about multi-step action selection beyond one particular tool.

It also reports improved process scores after training. Those scores are informative measurements under an evaluator, not independent proof that every reasoning step is logically valid or that any new tool will work without adaptation.

Across LATS, SPRINT and SWiRL, planning becomes concrete when one can identify what changes: the explored branch, the execution schedule or the learned action policy. The resulting system still needs checks on final answers, resource use and the assumptions that made its intermediate feedback meaningful.