Stanford CS329A Self-Improving AI Agents | Part 5 | Planning and Multi-Step Reasoning
Multi-step agent tasks require reasoning, action, environmental feedback, and search over alternative trajectories.
Why planning is more than prompting. A multi-step agent must repeatedly connect four operations: reason about the goal and constraints, act through tools, observe what happened, and search or revise the plan. LATS targets this setting by explicitly exploring multiple trajectories instead of accepting the language model’s first plan.
LATS combines chain-of-thought expansion, action outcomes, environmental observations, and reflection in a Monte Carlo tree-search framework.
LATS unifies reasoning, acting, and planning. Candidate actions form branches of a search tree; actions may run in parallel, and their observations return to the context. Outcome-based evaluation and self-reflection then guide later expansion, adding a deliberate exploration–exploitation process beyond a single ReAct trajectory.
A LATS iteration performs selection, expansion, evaluation, simulation, backpropagation, and reflection.
Six-stage search loop. LATS selects a state, expands it with several model-generated actions, evaluates the resulting action–observation states, simulates a promising branch, backpropagates its terminal return, and records a reflection. This makes tool feedback part of the search state rather than an after-the-fact note.
LATS balances exploitation and exploration with a learned state value, self-consistency, UCT, and trajectory returns.
Scoring and updating the tree. The value estimate mixes judge-based promise and sampling frequency. UCT favors both high-value states and underexplored children, preventing a purely greedy search from overlooking useful branches. Terminal success or failure updates every visited state, and reflection summarizes the causal lesson for later searches.
LATS turns additional test-time compute into better multi-hop performance, but its cost and assumption of reversible exploration limit deployment.
Evidence and limitations. LATS shows that structured search can convert inference compute into higher multi-step accuracy across question answering and web shopping. Its modularity is attractive, but every expansion adds calls and latency. More importantly, experimental tree search is unsafe when actions have irreversible real-world consequences, so production systems need simulation, approval, or rollback boundaries.
SPRINT starts from the observation that long reasoning traces contain independent work that need not be generated and executed serially.
Parallelism inside reasoning. Longer deliberation can improve accuracy, but token-by-token serial execution wastes time when branches are independent. SPRINT asks the model to identify those opportunities, produce plans in stages, and expose independent executor work that a runtime can launch concurrently.
SPRINT creates supervision by decomposing existing reasoning traces into plans, executions, dependencies, and parallelizable DAG stages.
Turning traces into a parallel program. SPRINT uses an LLM to convert successful sequential reasoning into a dependency graph. Nodes with no dependency are packed together, then the tagged trace becomes fine-tuning data. At inference, the runtime can recognize the tags, fan work out, synchronize results, and return the combined context for the next planning stage.
SPRINT changes the learned output protocol rather than the transformer architecture, while retaining context for revision and replanning.
No architectural fork is required. Parallel branches are expressed through a structured text protocol, not parallel transformer heads. Execution happens outside the model, and synchronized results return to the ordinary autoregressive context, preserving the ability to replan.
Fine-tuning on parallelized reasoning traces reduced sequential tokens and unexpectedly improved accuracy and out-of-domain performance.
Efficiency and quality can reinforce each other. The structured plan–execute format reduced serial work and also encouraged clearer exploration. Improvements on tasks outside the fine-tuning domain suggest that the model learned a reusable planning behavior rather than memorizing only one benchmark’s solutions.
SPRINT gains depend on task parallelism, synchronization, and load balance, with the largest savings on difficult long-reasoning problems.
Where parallel reasoning helps. Parallelism is task-dependent: early exploration often has wide independent branches, while late reasoning converges. The runtime must synchronize results, tolerate stragglers, and let the model resolve conflicts before answering. Overhead can dominate easy problems, so SPRINT is best viewed as an adaptive strategy for long-horizon reasoning.
SWiRL trains models to perform robust multi-step reasoning and tool use while avoiding live tool execution inside the reinforcement-learning loop.
SWiRL’s design goal. The learner should choose appropriate tools and actions across a trajectory, maintain accuracy as context grows, recover from errors, and decide when enough evidence exists for a final answer. Separating data collection from optimization avoids putting unreliable tools on the critical path of RL training.
SWiRL constructs synthetic trajectories one action at a time and assigns process rewards with an LLM judge.
Offline trajectory annotation. SWiRL repeatedly presents the original task plus accumulated actions and observations, yielding explicit state–action examples. An LLM judge labels the quality of each action, while final-answer correctness supplies an outcome signal. This supports experiments that separate good intermediate decisions from already-successful complete trajectories.
During SWiRL training, the policy proposes the next action from recorded context and receives its stored process reward without executing the tool again.
Training without live environments. Tool calls happen only while building the offline dataset. During RL, the model sees the recorded prefix, proposes an action, and receives a process score. A judge can often determine that a search query or calculator expression is appropriate without knowing the eventual tool output, greatly simplifying the optimization loop.
SWiRL inference alternates tagged reasoning, tool calls, returned observations, and an explicit decision to produce the final answer.
Learned agent loop. Unlike offline RL training, deployment executes the requested calculator, Python, or search action and feeds its observation back to the model. Explicit call and answer tags make termination controllable. The data study then asks which combination of process quality and final correctness produces the most transferable policy.
Process-filtered SWiRL data generalized across datasets and tools, indicating learning of multi-step problem solving rather than one tool’s syntax.
Cross-tool generalization is the central result. Policies trained with one task and tool improved on a different task and tool, and gains grew with more synthetic data. The reported diagnostics show higher process correctness at each step, supporting the interpretation that RL taught a reusable pattern of decomposition, tool invocation, and recovery.
Multi-step reinforcement learning benefited more from process-filtered data than imitation learning, which required fully correct trajectories.
RL and SFT use synthetic traces differently. SFT copies demonstrations, so erroneous outcomes contaminate its target. SWiRL RL instead treats each prefix as another chance to choose a better next action and learn from a local reward. Across the lecture’s three papers, the common theme is to structure search, parallel planning, and step-level feedback so additional computation produces better multi-step behavior.
Enjoy Reading This Article?
Here are some more articles you might like to read next: