NVIDIA teaches AI agents to bounce back after one bad move

NVIDIA researchers, working with Princeton University and the University of Maryland, have introduced PivotOPD, a training method that helps AI agents recover after making a critical early mistake during a multi-step task ¹. The approach targets what the team calls a "pivotal mistake" — an action that pushes the agent farther from its goal or makes the task unsolvable.

In tests across three Qwen3 model sizes (8B to 235B), more than half of failed task rollouts contained such a mistake, and it typically appeared early, around turn 8 to 12 out of 30 ¹. Fixing just that one turn in replay experiments boosted success from 8% to 59% on the ALFWorld benchmark.

How standard methods miss the problem

Standard on-policy distillation (OPD) — a technique where a smaller student model learns from a larger teacher's outputs on trajectories the student itself generates — reduced overall failure rates from 79% to 56% in early experiments ¹. But failures that followed a pivotal turn barely improved, dropping only from 51% to 49% .

The reason is straightforward: the correct action after a pivot often has less than a 1% probability, so in 8 rollout attempts the student almost never samples it. Outcome-based reinforcement learning shares the same blind spot — if every attempt fails, the group-relative advantage signal is zero ¹.

What PivotOPD does differently

PivotOPD adds three components to a group-based RL update ¹. First, a larger teacher model reviews each rollout after the fact, identifies candidate turns where the student deviated from an optimal path, and names a gold correction action at each point. On ALFWorld, these detected pivots landed within one turn of the oracle's ground truth in 77.8% of failed rollouts .

Second, preventive distillation uses reverse KL divergence to push the student away from the committed mistake, given the teacher's gold action as a hint ¹. Third, recovery distillation uses forward KL divergence to transfer recovery behaviors that the student rarely tries on its own — naming a corrective action for up to K turns after the pivot .

Results

With a Qwen3-1.7B student, PivotOPD averaged 73.7% on ALFWorld (up 5.5 points over the strongest baseline), 44.5% on Search-based QA, and 76.6% success rate on WebShop ¹. With the 8B student, it reached 93.0% on ALFWorld and 81.9% on WebShop . Even when the 8B model taught itself, it still beat all 13 baselines by at least 1.5 points across benchmarks .

On SWE-Bench Verified, a Nemotron-3.5 student rose from 62.8% to 66.0%, compared with 63.0% for standard OPD and 73.0% for the teacher ¹. In replay tests, PivotOPD recovered from 72.7% of pivotal mistakes versus 20.3% for standard OPD .

The method adds no inference cost and requires only K=1 to K=2 extra recovery turns, with training overhead ranging from about 12% to 94% depending on K ¹. Code will be released on GitHub soon .

Key facts

More than half of failed agent rollouts contain one early, recoverable pivotal mistake
PivotOPD recovers from 72.7% of replayed mistakes vs 20.3% for standard OPD
Best average across 13 baselines on ALFWorld, WebShop and Search-based QA
Nemotron-3.5 student improved from 62.8% to 66.0% on SWE-Bench Verified
Zero inference cost; training runs on NVIDIA H100 GPUs with code coming soon