SAILOR

Robust imitation via learning to search

Sep 2024 to Jan 2025 · Cornell · with Sanjiban Choudhury and Gokul Swamy

NeurIPS 2025 · Spotlight

Robot learning Manipulation Planning & control

Where behavior cloning runs out

Door opening, base diffusion policy on the left and SAILOR on the right. The search recovers after the gripper slips off the handle, which is exactly the behavior that behavior cloning cannot produce.

Behavior cloning supervises a policy only at the states the expert visited. The first small error moves the robot off that support, where there is no supervision at all, and the errors compound from there. Adding demonstrations shrinks the off-support region but never removes it. The standard fixes (DAgger-style human corrections, ground-truth reward labels, a privileged simulator) all require exactly the supervision you do not have.

Learning to search from demonstrations alone

The loop in full. The observation is encoded, the base diffusion policy proposes a plan, MPPI samples N residual plans and scores them inside the world model with the learned reward plus a terminal value from the critic, and the best correction is added to the base action before execution.

Learn a latent world model and a reward model from a mixture of expert demonstrations and the base policy's own rollouts. The reward model is trained to tell expert states from base-policy states in latent space, so no ground-truth reward labels are needed.
Pre-train a diffusion policy on the expert data as the base policy, then alternate rounds of on-policy collection, world and reward model updates, and optional distillation back into the policy.
At test time, run MPPI inside the learned latent space to search for a correction to the base policy's nominal plan, scoring candidates with the learned reward plus a terminal value from the critic. Execute the first action, replan on the next observation, MPC style.
Evaluate on 12 visual manipulation tasks across RoboMimic, RoboSuite and ManiSkill at three dataset scales, against diffusion policies trained on identical data. The RoboSuite and ManiSkill demonstrations were teleoperated by hand with a 3D SpaceMouse, up to 200 per task.
Scale the search at test time to 8× the sample budget seen during training, to check whether harder optimization against a learned reward starts hacking it.

SAILOR reaches 2 to 3x higher success than state-of-the-art diffusion policies trained on exactly the same data, across all 12 tasks, and it gets there with no additional demonstrations, no DAgger-style corrections and no ground-truth reward labels. Scaling the baseline up by 5 to 10x still leaves a gap. The paper was a Spotlight at NeurIPS 2025, the top 3% of submissions.

Why a local search was enough

Door, Cereal and Square at three dataset sizes each. Orange is SAILOR, purple is the diffusion policy. The traces on the right show the learned reward over a rollout, with the frames where each policy diverges.

The search only has to be local. The base diffusion policy already proposes a reasonable multi-step plan, so MPPI only searches for a correction to it rather than for a plan from scratch. That keeps the optimization tractable and means an imperfect world model is still useful.
The failure data is the reward signal. Expert demonstrations versus the base policy's own rollouts is precisely the contrast that defines going wrong, so a discriminator over that pair yields a dense latent reward with no labels, no human in the loop, and no simulator reward function.
No reward hacking at 8× the training sample budget. The usual objection to test-time scaling is that more optimization finds the learned reward's blind spots. Performance improved up to the training-time budget and then held flat rather than collapsing, plausibly because the models are refit on the same on-policy data the planner drives the agent into.