Stanford CS329A: Planning and Multi-Step Reasoning in AI Agents

Overview

Multi-step reasoning in AI agents requires a cohesive integration of reasoning (planning what to do), acting (executing steps via tools or search), and searching (refining trajectories based on feedback). The core challenge is moving beyond simple sequential generation to encourage diversification, exploration, and efficiency in how models reach a correct solution.

Language Agent Tree Search (LATS)

LATS unifies reasoning, acting, and planning by applying Monte Carlo Tree Search (MCTS) to language model trajectories. It transforms the linear chain-of-thought process into a tree-based search to optimize the path toward a final answer.

The LATS Framework

LATS operates through six distinct stages:

  1. Selection: A node is selected for expansion using the Upper Confidence Bound applied to Trees (UCT) to balance exploration of new paths with exploitation of known high-value paths.
  2. Expansion: The model samples multiple potential actions from the selected node.
  3. Evaluation: Each resulting state is scored using a weighted average of an LLM-as-a-Judge (prompting the model for a 0-1 score) and a self-consistency score (the frequency with which a specific action is sampled).
  4. Simulation: The model greedily expands the highest-scoring state until it reaches a terminal state (success or failure) or hits a compute budget.
  5. Backpropagation: The outcome of the trajectory is used to update the values of all states leading to that result.
  6. Reflection: The model analyzes the trajectory to determine why it succeeded or failed, which is then appended to the context to improve future search.

Performance and Trade-offs

LATS demonstrated strong results on HotPotQA (multi-hop reasoning) and WebShop (practical e-commerce tasks), achieving performance close to human experts without fine-tuning. However, the approach introduces significant inference costs due to the repeated expansion and backpropagation of the tree. Additionally, LATS does not currently address irreversible actions (e.g., financial transactions) where trial-and-error search is not feasible.

SPRINT: Parallelizing Reasoning Traces

SPRINT addresses the inefficiency of sequential "thinking" in frontier models (like DeepSeek-R1 or Gemini), where longer reasoning chains often correlate with higher accuracy but increase latency and cost.

Parallel Planning and Execution

SPRINT is based on the observation that many reasoning steps are independent and can be executed in parallel. The framework uses a post-training/fine-tuning approach to teach models to act as both a planner and an executor:

  • Data Creation: Reasoning traces from models like DeepSeek-R1 are annotated by GPT-4o to identify planning steps and execution steps. These are then organized into a Directed Acyclic Graph (DAG) to identify which steps can run concurrently.
  • Fine-tuning: The model is supervised-fine-tuned (SFT) on these reformatted trajectories, teaching it to output specific tags for parallel plans (e.g., "Plan 1" and "Plan 2") and their corresponding executions.
  • Inference: At test time, the model generates multiple independent plans; these plans are then executed in parallel (e.g., via multiple Python tool calls), and the results are synced back into the context before the model proceeds to the next planning stage.

Key Findings

  • Efficiency: SPRINT significantly reduces the number of sequential tokens required to reach an answer, reducing sequential token count by approximately 40% on certain tasks.
  • Accuracy: Contrary to the goal of mere efficiency, the structured parallel thinking process actually improved accuracy on benchmarks like MATH, Countdown, and GPQA Diamond.
  • Generalization: Models trained on MATH data generalized their parallel-thinking capabilities to out-of-domain tasks without further training.

SWiRL: Synthetic Multi-Step RL

SWiRL focuses on teaching models to use tools and reason across multiple steps without the overhead and instability of executing tools live during the reinforcement learning (RL) process.

The SWiRL Methodology

SWiRL separates data collection from model optimization:

  1. Offline Synthetic Generation: The model generates multi-step trajectories (reasoning $\rightarrow$ tool call $\rightarrow$ environment response) iteratively.
  2. Process Labeling: An LLM-as-a-Judge assigns a reward to each individual step based on the quality of the proposed action/query, regardless of the tool's actual output.
  3. Multi-Step RL: The model is trained using RL to optimize the expected reward of a single action given the prior context. Crucially, the tool is not executed during training; the model is simply rewarded for proposing a "good" action based on the pre-collected offline data.

Data Filtering and Generalization

  • Process vs. Outcome Filtering: The researchers found that process-filtered data (keeping steps judged as "good" by the LLM) was more effective for training than outcome-filtered data (keeping only trajectories with correct final answers). This suggests that learning correct intermediate reasoning is more valuable than simply imitating a correct final result.
  • Cross-Domain Transfer: SWiRL showed remarkable generalization. Models trained on GSM8K (using a calculator tool) improved their performance on HotPotQA (using a search tool), and vice versa. This indicates the model is learning the general meta-skill of multi-step tool invocation and reasoning rather than just learning a specific tool's API.
  • RL vs. SFT: Multi-step RL significantly outperformed Supervised Fine-Tuning (SFT), as RL allows the model to explore and break away from incorrect trajectories present in the imitation data.

Sources