Offline Multi-Step Tool-Use RL (SWiRL)

Sources#

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

Summary#

SWiRL — step-wise RL for multi-step reasoning and tool use — is the third paper in CS329A lecture 5, presented as a preview of a COLM 2025 talk. It attacks the engineering problem that makes agentic RL painful: tools in the training loop. Live tools are slow, they fail, they have bugs, and training is already expensive; every rollout that has to hit a search engine or a Python sandbox couples the optimizer's throughput to an external system's availability.

SWiRL's answer is to move every tool call out of the RL loop entirely, into an offline data-generation phase, and to make the reward something you can compute without the tool's response.

Its stated design goals are a compact spec for what "multi-step tool use" even means as a skill: know when to call a tool, write the right query for it, hold accuracy across steps, recover from errors, and know when to stop and answer. The lecture's complaint about the prior art is that RLHF, RLAIF and execution-feedback RL are all optimized for a single step — the model does whatever it wants and one terminal reward is smeared back over all of it — where this needs "more governance on the process of generating these steps."

Evidence. CS329A Self-Improving AI Agents — Part 5: Planning and Multi-Step Reasoning (Azalia Mirhoseini solo, delivered 2025-10-06, published 2026-08-03, practitioner-opinion). Slide walkthrough; the paper is not in raw/. COI: her own lab's work, narrated in the first person throughout ("we did not need to train an LLM as a judge for this") and never disclosed as hers by name. All figures ASR-read off slides.

Phase 1: synthetic multi-step data, generated by iterative prompting#

Give a model a prompt, a tool list, and permission to do exactly one of three things per turn — reason, call a tool, or emit a final answer. Show it the environment's response. Re-prompt with the original question plus everything that has happened, and let it take the next action. Repeat until it answers. The trajectory length falls out of the task rather than being fixed: "from one to three or five and so on for different queries."

Then label. An LLM judge scores each action — given the prior context and this action (a reasoning step followed by a tool call), how good is it? This is entirely offline and embarrassingly parallel, which is the point: many questions, many trajectories, all the tool latency paid once.

Reported setup, with ASR hedges: the generator is Gemma 2 27B (moderate transcription confidence), questions come from HotpotQA (multi-hop search) and GSM8K (grade-school math with a calculator), and the corpus is on the order of 50K trajectories. The judge is prompted, not trained.

Phase 2: RL where the tool never runs#

Here is the move. During RL, the model is shown the prompt, the previously-collected actions, and the previously-collected environment responses from the offline dataset. It proposes the next action. That action is not executed — it is scored by the judge and the reward goes straight into the policy gradient. The objective is the expected reward of a single action conditioned on all prior states and actions, summed over the trajectory's steps.

A student asks the right question: how can the judge score a tool call it never sees the output of? Mirhoseini's answer is the paper's actual thesis:

"We are asking the judge to judge the quality of the query that is generated by the model to call the tool, rather than the output of the tool… the model doesn't need to know the age of this person to know this question was good or bad."

In the worked example ("who is older, Glenn or Rossling?") the first action is search for the age of the first person, and that is assessable as a well-formed step toward the goal without running the search. Tool results are not absent from training — they are present as frozen context from the offline collection. What is absent is any dependency on executing the action being optimized.

This is process supervision built out of a prompted judge on action proposals, and it is exactly the design that RLEF refuses: RLEF's whole signal is what happens when you run the thing.

The filtering result, and why it inverts between RL and SFT#

Having a per-step judgment and a final answer gives two independent filters on the trajectory corpus — keep only trajectories where every step was judged good (process-filtered), or keep only trajectories that ended correct (outcome-filtered), or both, or a random subset. The reported finding, which Mirhoseini calls initially non-intuitive:

• For RL, process-filtered wins — beating both outcome-filtered and process-plus-outcome-filtered.
• For SFT, the ordering reverses: process-plus-outcome-filtered data works best.

Her explanation is the mechanism worth keeping. Outcome filtering keeps only problems the model already solves, so training on it cannot teach the model to solve what it currently cannot; process filtering keeps good reasoning on problems that ended wrong, which is where the headroom is. And the reason the ordering flips is that SFT is imitation — showing it a trajectory that ends in a wrong answer teaches it to reproduce a wrong answer — while RL only conditions on the prior steps and rewards a new action, so it can "break out" of a trajectory whose ending was bad.

That is a clean, transferable statement of when correct-outcome filtering helps and when it is actively self-limiting, and it applies well beyond this paper: the same argument says an RLVR pipeline that keeps only verified-correct rollouts is discarding its most informative data. Compare lecture 3's arc, where automating step labels via rollout success rates quietly reintroduces outcome filtering through the back door.

The generalization result, which is the headline#

Trained on GSM8K with a calculator, evaluated on HotpotQA with a search engine — different task, different tool, different domain:

Training	Evaluated on HotpotQA
GSM8K + calculator (transfer)	65 → 75.1
HotpotQA + search (in-domain)	65 → 73

The direction is reported to hold both ways: HotpotQA-with-search training also improves GSM8K-with-Python. And the effect scales with data — from 100 to 10,000 synthetic training examples, out-of-domain accuracy keeps rising.

Read the table with care. As transcribed, the transfer arm beats the in-domain arm on the same evaluation, which is a strong and slightly implausible claim; the alternative reading is that 65 → 73 is an in-domain gain on a different eval and the two rows are not measured on one axis. The ASR does not disambiguate, and the slide is not in raw/. What is safe: cross-tool, cross-domain transfer is large and of the same order as in-domain training. Mirhoseini's own reading — "the model is learning how to think in steps and how to invoke a tool," not how to use a specific tool — does not depend on which row is bigger, and it is the same conclusion SPRINT reaches independently in the same lecture.

The mechanism check is the last result: measure the judge's average per-step process reward before and after RL, and it rises both in-distribution (HotpotQA) and out-of-distribution (GSM8K). The model does not just answer better; it takes better individual steps on tasks it was never trained on. Finally, RL beats SFT on the same data "by a good amount" — this wiki's cleanest instance of the standard SFT-memorizes / RL-generalizes split, from the lecturer's own lab.

What the design costs#

Three things the lecture does not price, recorded here as the compile's reading:

• A well-formed query that returns nothing is rewarded. The judge scores the query, so "search for the age of Glenn" scores well whether or not the search engine has the page. That is precisely ReAct's documented failure mode — the hallucination-to-retrieval-failure swap — and SWiRL's reward is blind to it by construction. The one thing that could catch it, the query's actual result, is the thing removed from the loop.
• The judge is the ceiling. Nothing is trained, nothing is calibrated, nothing is ensembled. Everything the course's own lectures 2–3 establish about judge reliability (LLM-as-a-Judge, Weak-Verifier Ensembling) applies to a signal that here is the only signal at every step.
• The frozen context is off-policy by design. The model is scored on an action taken in a world where its earlier actions were someone else's. Whether that mismatch bites at longer horizons is not addressed, and it is the precise boundary between this and on-policy agentic RL (Single-Rollout Optimization, Group Relative Policy Optimization (GRPO)).

Connections#

• CS329A: Self-Improving AI Agents (Stanford) — lecture 5's third paper, and the course's training-time answer to the planning problem
• RL from Execution Feedback (RLEF) — the direct opposite. RLEF puts the interpreter inside the loop and makes execution the reward; SWiRL removes execution from the loop and rewards the proposed query instead. Both are taught in this course, three days apart, as answers to the same question about multi-step tool use
• Turn-Level Credit Assignment — the same dense-per-turn credit, bought the expensive way. TRACE's whole framing is getting turn-level signal without an LLM judge, step labels, or Monte-Carlo continuations; SWiRL is the judge branch it defines itself against, and the honest score is that SWiRL's judge is cheap offline and TRACE's probe needs the gold answer
• Process vs Outcome Reward Models — SWiRL's per-step judge is process supervision with a prompt in place of a trained PRM, and its process-versus-outcome filtering result is a new datapoint in that page's arc
• Reasoning–Acting Interleaving (ReAct) — the trajectory format SWiRL generates and trains on is ReAct's, one action–observation pair at a time; and ReAct's retrieval-failure mode is the blind spot SWiRL's query-only reward cannot see
• Intra-Trace Parallel Planning (SPRINT) — the sibling paper from the same lab in the same lecture: same trick of building training data by annotating one model's trajectories with another model, same finding that what generalizes is a manner of reasoning rather than a task skill
• Tree Search over Agent Trajectories (LATS) — lecture 5's first paper and the frozen-model alternative: search harder at inference instead of moving the weights
• LLM-as-a-Judge — the entire reward signal, unensembled and untrained
• The Verifiability Thesis — a domain with no cheap verifier gets a judge, and the judge's reliability becomes the ceiling on how far the flywheel turns
• Single-Rollout Optimization — the on-policy agentic-RL alternative and the systems problem SWiRL sidesteps by never touching an environment during training
• Rationale Bootstrapping (STaR) — the 2022 statement of the same hole, attacked from the other side: STaR keeps only correct-outcome generations too, and its answer is to manufacture a positive for a failed problem by supplying the answer as a hint, where SWiRL's is to keep the good steps of a trajectory that ended wrong
• Azalia Mirhoseini — her lab's paper, taught in the first person

Open Questions#

• The judge scores the query, never the result, so a fluent query that retrieves nothing is indistinguishable from one that retrieves the answer. Does a query-only reward measurably degrade retrieval quality relative to a result-aware reward — and would a hybrid (judge the query, verify the result offline once) recover it without putting tools back in the loop?
• Every action is scored against frozen prior context collected from a different policy. Does the off-policy mismatch grow with trajectory length, capping the method at the 3–5-step tasks demonstrated, or does step-level reward make it insensitive to horizon?

Sources#

• CS329A Self-Improving AI Agents — Part 5: Planning and Multi-Step Reasoning — CS329A Self-Improving AI Agents — Part 5: Planning and Multi-Step Reasoning, Azalia Mirhoseini solo, Stanford Online. Delivered 2025-10-06, published to YouTube 2026-08-03 (practitioner-opinion, YouTube auto-caption transcript, ~11.3k words). The SWiRL third of the lecture: the five design goals, the iterative-prompting synthetic-data pipeline and its LLM-judge step labels, the no-tools-during-RL objective and the "judge the query, not the output" exchange with a student, the four filtering arms and the RL/SFT inversion, the inference-time prompting protocol with tool and answer tags, the Gemma-2-27B / HotpotQA / GSM8K / ~50K-trajectory setup, the 65→75.1 transfer and 65→73 in-domain figures, the 100→10,000 data-scaling curve, the per-step process-reward measurement before and after RL, and the RL-beats-SFT comparison. The paper is not in raw/; her lab's, presented as a COLM 2025 preview. All figures ASR-read off slides; the transfer/in-domain table is transcribed ambiguously and is flagged inline rather than resolved