Stanford CS329A Self‑Improving AI Agents ｜ Part 5 ｜ Planning and Multi‑Step Reasoning — Annotated Lecture Notes & Key Takeaways

Course Information

Institution: Stanford

Original Course: Stanford CS329A Self‑Improving AI Agents ｜ Part 5 ｜ Planning and Multi‑Step Reasoning

Instructor Bio: This session is co-instructed by **Aakanksha Chowdhery** and **Azalia Mirhoseini**, faculty and researchers at Stanford University and Google DeepMind. Aakanksha Chowdhery earned her PhD from Stanford University and is a Senior Researcher at Google DeepMind. She led end-to-end training of the 540B PaLM model and drove pre-training and scaling efforts for multiple generations of Gemini MoE models. She has also contributed core components to PaLM-E, Med-PaLM, and the Pathways infrastructure that underpins Google’s large language model ecosystem, with deep expertise in dense and Mixture-of-Experts architectures and large-scale model training. Azalia Mirhoseini is an Assistant Professor of Computer Science at Stanford University and founder of the Scaling Intelligence Lab. She also serves as a Senior Staff Scientist at Google DeepMind and co-founded Ricursive Intelligence. Previously she held research roles at Google Brain and Anthropic, contributing to the development of the Claude and Gemini model families. Her pioneering work spans MoE architectures, test-time compute scaling, and deep reinforcement learning for chip design (AlphaChip), with a core research focus on recursive self-improving AI systems.

Course Description: This lecture dives into planning and multi-step reasoning capabilities that enable agents to tackle complex, long-horizon tasks. It covers chain-of-thought reasoning, tree-of-thought search, hierarchical planning frameworks, and task decomposition strategies. It also discusses multimodal reasoning workflows, agent orchestration for multi-stage pipelines, and the challenges of maintaining coherence and correctness across extended reasoning chains.

This lecture focuses squarely on planning and multi‑step reasoning, a core skill large‑language‑model‑based agents still struggle to get right. The presenter walks through three recent research papers: LATS, Sprint, and Swirl. Each piece of work tackles distinct real‑world pain points that surface when LLMs need to complete tasks requiring chained thought, tool calls, and exploration of multiple possible paths. Coming out of the talk, it’s clear the field is moving past basic chain‑of‑thought prompts. Researchers are borrowing established ideas from reinforcement learning and tree search to make agents more capable. Even so, plenty of practical trade‑offs have yet to be worked out.
 

LATS: Language Agent Tree Search

 

LATS comes from an ICML paper published the year before this lecture. Its main goal is to weave reasoning, action‑taking, and planning together within one unified language‑agent system (0:31). The speaker opened with a holiday‑planning example to illustrate what multi‑step tasks actually demand.
 

Take putting together a trip to Hawaii. Three different kinds of work feed into the final plan. First is reasoning: the agent works through practical constraints such as available budget or overall trip goals. Next comes acting, where the model goes out to gather new information, for instance by browsing travel‑related community posts. Then there is iterative search and refinement. After pulling in real‑world information, the model revises its initial ideas. It might swap out potential destinations or add new activities to the itinerary. Real‑world agent tasks cycle repeatedly through these three phases.
 

The speaker pointed out a well‑known shortcoming of standard LLMs. Given a broad prompt, a model can generate one single plan and follow it through. But it rarely naturally explores multiple alternative routes. It tends to lock into one trajectory without testing other ways to reach a correct answer (2:33). LATS tries to fix this gap by bringing Monte‑Carlo Tree Search techniques, long used in reinforcement‑learning research, into LLM workflows.
 

MCTS allows an agent to branch off into many potential actions, assign scores to each new branch, and use feedback from its environment to refine future planning. The Hawaii‑trip example helps make this concrete. When prompted to plan the vacation, the model suggests several possible actions. Each action receives a calculated score. Branches with higher scores get expanded further within the growing search tree. Multiple actions can even run side‑by‑side. The system balances leaning into high‑promise paths and checking less‑explored options, which is exactly what the UCT formula handles (5:33).
 

The lecturer drew clear comparisons against two older frameworks. Math‑Shepard used a separate verification model to score pure reasoning choices during inference. LATS uses a different scoring logic. It judges outcomes after an action runs, and adds a reflection step where the model assesses how well a given path actually worked out (6:07). LATS also builds on top of ReAct. ReAct already alternates reasoning steps with tool‑driven environment interaction. LATS adds more sophisticated tree‑based planning plus formal reflection once a path finishes.
 

LATS operates in six sequential phases: selection, expansion, evaluation, simulation, back‑propagation, and reflection (8:11). The presenter used a text‑based maze scenario to walk through these steps. Imagine starting inside a dim room with two doors. The model thinks up three possible moves: open the left‑hand door, open the right‑hand door, or search the room for useful clues. Each action runs inside a simulated environment, and the resulting observation gets added to the working context.
 

Evaluation creates a composite state score built from two separate inputs. One input treats the LLM itself as a judge. Reading an action paired with its resulting observation, the model assigns a score between 0 and 1 showing how promising that new state feels. The second input relies on self‑consistency, measured by sampling frequency. Actions the model tends to suggest more often receive higher weighting (10:40). These two values combine to produce one single score for each new state.
 

During simulation, the system expands the highest‑scoring path until hitting an end condition. That could mean reaching success, hitting total failure, or burning through a pre‑set compute budget. Once a trajectory completes, back‑propagation updates every node along that path using the final reward from that run. The UCT metric picks which node gets expanded next. UCT balances strong average reward against visit frequency. This keeps the system from writing off paths that look weak early on but might turn out well later (13:42).
 

Reflection wraps up the LATS workflow. Whether a trajectory succeeds or fails, the model writes short notes explaining what drove that outcome. The lecture mentioned this reflection component noticeably boosts overall performance (16:57).
 

Researchers tested LATS on HotpotQA, a multi‑hop question‑answering dataset requiring facts pulled from several different Wikipedia pages. They also ran tests on WebShop, which simulates real‑world online‑shopping workflows. With no fine‑tuning at all, just extra computation during inference, LATS hit performance numbers close to human experts on WebShop (18:22).
 

From there, the presenter shifted to the framework’s real‑world limitations. Every tree expansion and back‑propagation step eats up additional compute resources. The original paper never systematically analyzed how performance scales against cost. Another major catch: LATS assumes agent actions stay safe and reversible. If the agent triggers high‑stakes, irreversible real‑world actions like processing payments, this framework cannot safely operate (19:44).
 

In‑class discussion brought up unresolved questions. One student noted UCT comes from multi‑armed‑bandit research and asked whether authors tested alternative bandit algorithms. The speaker replied they had not. LATS mostly provides a working platform; other exploration formulas could easily be plugged in for follow‑up research (21:28). Another student worried about repeating actions inside the tree, for example looping through actions A‑B‑A‑B. The lecturer explained LATS constructs a tree structure rather than a full graph. Repeated actions under the same parent node raise visit counts and adjust UCT scores accordingly (22:52).
 

Sprint: Building Parallel Reasoning Trajectories

 

Sprint is a 2025 paper that had not yet had its official conference talk at the time of this lecture (23:32). Its starting observation comes from studying state‑of‑the‑art reasoning‑focused models. When given harder problems, these models produce longer reasoning chains, and longer thinking generally correlates with better final answers. Still, large chunks of that extended thinking contain independent subtasks that do not need to run one right after another (24:07).
 

Right now LLMs generate reasoning in strict sequential token order. Even when two subtasks have no dependencies on each other, the model still finishes one before starting the next. Sprint tries to unlock parallel execution for these independent reasoning pieces. Its goal is cutting down sequential token volume without hurting, and ideally improving, overall task accuracy.
 

Sprint works through fine‑tuning instead of changing underlying model architecture. Its training‑data creation process was quite interesting. Researchers took high‑quality reasoning outputs from models such as DeepSeek‑R1. They fed these reasoning traces to another powerful LLM like GPT‑4o. That annotator model labelled sections of text as either planning work or execution work. It also mapped out task dependencies, creating a directed acyclic graph (DAG) showing which steps rely on prior outputs and which can run side‑by‑side (28:08).
 

Once the DAG existed, steps were grouped into batches ready for parallel execution. These annotated traces became supervised fine‑tuning examples. The base model learned special tags to separate planning blocks from their corresponding parallel‑execution blocks. When running inference, the model outputs these tagged plans. An external runtime reads those tags, kicks off independent execution tasks simultaneously, collects all outputs, merges everything back into context, and hands the full updated context back to the LLM for its next round of planning (31:08).
 

I found one piece a little tricky to follow. The model still writes out plan descriptions sequentially. The parallel speed‑up happens on the execution side, after the model has written multiple independent plans in one pass. It does not rewrite the transformer to generate tokens in parallel internally.
 

Experimental results showed two pleasant surprises. Sprint‑fine‑tuned models reduced the number of sequential tokens needed. Unexpectedly, accuracy also rose roughly three‑and‑a‑half percentage points on math benchmarks. Thinking in this structured, parallel‑aware way improved reasoning quality all on its own (37:46). Encouragingly, this ability transferred outside training data. Trained purely on math examples, the model still posted better results on GPQA‑Diamond and Countdown, datasets it had never seen during fine‑tuning (38:30).
 

Several students asked thoughtful questions during Q&A. One wondered how the system keeps track of which execution result belongs to which original plan, especially when the agent needs to revise old plans. The lecturer pointed out the full context stays visible for the model at every step. Fine‑tuning teaches the model to generate revised or corrected plans whenever needed (34:40).
 

Another student brought up conflicting outputs generated from separate parallel branches. The lecturer acknowledged parallel branches can produce contradictory intermediate outputs, but noted this risk also exists for fully sequential reasoning. At the end of a working cycle, all intermediate outputs sit within context, and the model has to sort out contradictions before producing its final answer (39:30).
 

Someone else asked about load balancing. If different parallel subtasks require vastly different effort, some will finish fast while others lag behind, creating bottlenecks. The lecture said no rigorous analysis existed for this problem. One practical heuristic they mentioned was folding very trivial steps into the planning block to build larger chunks suitable for parallel work (42:06).
 

Potential parallel speed‑ups also depend heavily on the task itself. Problems requiring only short reasoning offer little room for improvement. Sprint delivers the biggest gains on hard problems with many independent subtasks (45:21). The lecturer added future research could layer reinforcement learning or GRPO over Sprint’s supervised fine‑tuning setup to get even stronger generalization (49:52).
 

Swirl: Multi‑Step RL Without Live Tool‑Calls During Training

 

Swirl is the third paper covered, presented here as a preview ahead of its upcoming conference talk in Montreal. Swirl targets multi‑step reasoning and tool‑use learning. It tries solving a well‑known pain point for reinforcement‑learning‑based agent training (50:29).
 

Lots of real‑world multi‑step work requires repeated external‑tool calls: search engines, calculators, Python interpreters. Running live tool invocations directly inside RL training loops creates major headaches. Tools can crash, time‑out, or return noisy outputs. This slows training down and introduces instability (52:20). Many existing RL pipelines only assign reward based on whether the final answer is correct. They offer almost no feedback judging whether each intermediate reasoning or tool‑call step itself was sensible. Swirl aims to deliver step‑wise rewards for process quality while avoiding live‑tool runs inside the RL phase.
 

Swirl’s workflow splits into two distinct phases: offline synthetic‑data generation followed by reinforcement‑learning fine‑tuning.
 

During offline data generation, a high‑quality seed model generates complete multi‑step trajectories. Real tools run at this stage, and the system saves every prompt, reasoning step, tool call, and value returned by that tool. A judge LLM then assigns a score to each individual step. Importantly, the judge evaluates how sensible the tool query itself is, not whether the tool’s eventual output was right or wrong. For instance, given a question asking which public figure is older, the judge can tell searching for person A’s birth year is a logical move, without needing to know their actual birth year (1:02:21).
 

Once trajectories carry per‑step judge scores, researchers filter data using different rules. Process filtering keeps trajectories where nearly every single step earned high marks from the judge. Outcome filtering keeps any path that ends with a correct final answer, no matter how many mis‑steps happened along the way (57:40).
 

The reinforcement‑learning fine‑tuning phase never runs live tools. Training feeds the model partial trajectories stopping at some prior state. The model is asked to produce the next action. Rewards come from pre‑computed judge scores saved within the dataset. Tool outputs are already baked into context collected during offline data work (1:00:17). This separates slow, flaky tool execution from gradient update cycles.
 

At inference time, the model behaves normally. It writes out reasoning and tagged tool‑call blocks. A runtime executes real‑world tools, feeds observations back to the model, and repeats until the model outputs its final‑answer tag (1:04:22).
 

Some of the experimental results stood out for their generalization power. When Swirl trained on GSM8K math problems using a calculator tool, the resulting model performed better on HotpotQA multi‑hop questions requiring search tools — even though search never appeared in its training data. The reverse transfer also held true. Training on HotpotQA with search improved GSM8K calculator‑task results (1:08:32). This suggests Swirl does not merely teach narrow tool syntax. It teaches general multi‑step problem‑solving procedures.
 

Training built from process‑filtered data delivered better downstream generalization than outcome‑filtered data alone. The lecturer’s reasoning made sense to me. Outcome‑filtered data preserves trajectories that stumbled through poor intermediate steps yet still landed on correct final answers. Training on those examples teaches the model to copy lucky paths. Process‑filtered data rewards logically sound intermediate moves, even when the full trajectory fails, building more robust general‑problem‑solving habits (1:07:12).
 

Swirl’s multi‑step reinforcement‑learning also outperformed plain supervised fine‑tuning. Supervised fine‑tuning worked best when trained on trajectories passing both process and outcome filters. RL could learn effectively using only process‑filtered trajectories, thanks to local reward signals assigned to every individual step (1:12:42).
 

Unresolved points raised in the lecture

 

The presenter pointed out several open research gaps, and I also walked away with a few lingering questions.
 

Swirl’s judge LLM judges how reasonable a tool query is without seeing tool outputs. This works well for assessing high‑level intent. Still, it cannot catch subtle bugs inside tool inputs. A syntactically broken Python snippet might read like a perfectly reasonable query to the judge, yet crash when it runs live during inference. The offline judge cannot catch that category of error.
 

All three papers rely heavily on LLMs acting as judges to score states or individual steps. The lecture never explored how noisy or inconsistent those judge scores can turn out. If the judge frequently makes mistakes, the whole tree‑search or training pipeline will degrade in quality.
 

None of these three frameworks eliminate heavy compute demands entirely. LATS burns lots of extra compute during inference. Sprint cuts sequential runtime but increases total parallel compute load. Swirl needs large volumes of high‑quality synthetic trajectories generated by much larger seed models. Better reasoning still costs resources in one shape or another.
 

The lecture closed by noting multi‑step planning for self‑improving AI agents remains an active, fast‑moving area of research.

Content Disclaimer: This article is for general reference only and does not constitute professional R&D guidance, production process advice or quality certification. All material performance data has specific test premises; readers should verify parameters against actual equipment and working conditions.