A Smooth Sea Never Made a Skilled SAILOR: Robust Imitation via Learning to Search
Authors: Arnav Kumar Jain
Vibhakar Mohta
Subin Kim
Atiksh Bhardwaj
Juntao Ren
Yunhai Feng
Sanjiban Choudhury
Gokul Swamy
Venue: NeurIPS 2025 spotlight ()
Keywords: Imitation learning
world models
inverse reinforcement learning
Primary area: reinforcement_learning
TL;DR: Rather than directly learning a policy from expert demonstrations, we instead learn world and reward models, allowing us to search at test-time and recover from mistakes.

Abstract: The fundamental limitation of the behavioral cloning (BC) approach to imitation learning is that it only teaches an agent what the expert did at states the expert visited. This means that when a BC agent makes a mistake which takes them out of the support of the demonstrations, they often don't know how to recover from it. In this sense, BC is akin to *giving the agent the fish* -- giving them dense supervision across a narrow set of states -- rather than teaching them *to fish*: to be able to reason independently about achieving the expert's outcome even when faced with unseen situations at test-time. In response, we explore *learning to search* (L2S) from expert demonstrations, i.e. learning the components required to, at test time, plan to match expert outcomes, even after making a mistake. These include *(1)* a world model and *(2)* a reward model. We carefully ablate the set of algorithmic and design decisions required to combine these and other components for stable and sample/interaction-efficient learning of recovery behavior without additional human corrections. Across a dozen visual manipulation tasks from three benchmarks, our approach SAILOR consistently out-performs state-of-the-art Diffusion Policies trained via BC on the same data. Furthermore, scaling up the amount of demonstrations used for BC by 5-10x still leaves a performance gap. We find that SAILOR can identify nuanced failures and is robust to reward hacking. Our code is available at [https://github.com/arnavkj1995/SAILOR](https://github.com/arnavkj1995/SAILOR).

A Smooth Sea Never Made a Skilled SAILOR: Robust Imitation via Learning to Search
Arnav Kumar Jain"
Vibhakar Mohta" Carnegie Mellon University
Subin Kim Cornell University
Mila- Quebec Al Institute Université de Montréal
A Atiksh Bhardwaj Cornell University
Juntao Ren Cornell University
Yunhai Feng Cornell University
Sanjiban Choudhury Cornell University
Gokul Swamy Carnegie Mellon University
Abstract
The fundamental limitation of the behavioral cloning (BC) approach to imitation learning is that it only teaches an agent what the expert did at states the expert visited. This means that when BC agent makes mistake which takes them out the support of the demonstrations, they often don't know how to recover from it. In this sense, BC is akin to giving the agent the fish giving them dense supervision across a narrow set of states rather than teaching them to fish: to be able reason independently about achieving the expert's outcome even when faced with unseen situations at test-time. In response, we explore learning to search (L2S) from expert demonstrations, i.c. learning the components required to, at test time, plan to match expert outcomes, even after making a mistake. These include (1) world model and (2) reward model. We carefully ablate the of algorithmic and design decisions required to combine these and other components for stable and sample/interaciion httpp///mmeeracttooffffiiiennt learning of recovery behavior without additional human corrections. Across dozen visual manipulation tasks from three benchmarks, our approach SAILOR consistently out-performs fate-of-theeart Diffusion Policies trained via BC on the same data. Furthermore, scaling up the amount of demonstrations used for BC by 1000x still leaves a performance gap. We find that SAILOR can identify nuanced failures and is robust to reward hacking. Our code is available at https: //gi thub. cow/ jrnavkj 11w.k/9999999/SAR
Introduction
The workhorse of modern imitation learning (IL.) is behavioral cloning (BC, Pomerleau [1988]). From training Diffusion Policies (DPs, Chi et al. [2023]) on per-task expert demonstrations collected via variety of teleoperation interfaces [Zhao et al., 2023, Chi et al., 2024, Wu al., 2024] to Visual-Language-Actoon models (VLAs, Team et al. [2024], Kim et al. [2024], Intelligence et al. [2025]) trained on wider, multi-task datasets [Khazatsky et al., 2024, ONeill al., 2024]. we see the same recipe applied: collecting more data to train more expressive policy models. The latent hope here is that scaling will eventually lead to "ChatGPT moment" for robotics [Vemmpala et al., 2023].
'Equal Contribution.
Correspondence to Arnay carav-kumarrain
wuweeaueeeaaaauaaaeeeaaaaaaaaeee m and Gokul geswamy@cmu.com
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

SAILOR Searching Across Imagined Latents Online for Recovery
DP Fails OOD State
2. Local Search in WM against RM
3. SAILOR Avoids Faiture
(M,)
WMC
(214) at a
+ +13)
1+++
base
Latent
Space
base G
base
WMC
base
base
base "ctk
baise baise Fot
+ ret-
Base DP
SAILOR
4. Distillation
Figure 1: We introduce SAILOR, method for learning search from expert demonstrations. By learning world and reward models on a mixture of expert and base-policy data, we endow the agent with the ability to, at test time, reason about how to recover from mistakes that the base policy makes.
However, even for the simpler problem of language modeling where one doesn't 't have deal with the complexities of embodiment (c.g. grounding, stochastic dynamics, partial observability, safety), simply scaling next token prediction (i.e. BC) was insufficient: we needed interactive learning in the form of Reinforcement Learning from Human Feedback (RLHF, Stiennon et al. [2020], Ouyang et al. (2022]) and more recently, Test-Time Scaling (TTS, Jaech al. [2024], Guo et al. [2025]) to build robust systems. If internet-scale offline pretraining was insufficient to solve language modeling, it stands reason that we'll need similar interactive learning algorithms to train (embodied) agents.
At heart, this is because even when they are trained on large amounts of data and over expressive policy classes, agents still sometimes make mistakes that take them out of the support of the offline data. This is fundamental property of sequential decision-making: one has to deal with the consequences of their own prior actions. When faced with the resulting unseen situation, we'd d like our agents to attempt to recover and match the expert's 's outcome (whenever it is possible to do so). The most direct approach teach an agent this recovery behavior is ask sumaan-n--the-ooee to correct mistakes [Ross et al., 2011, 1, Kelly et al., 2019, Spencer et al., 2020]. While simple, such an approach can be difficult to scale as fundamentally makes human time the bottleneck for robot learning.
In an ideal world, we'd like our robots to be able to learn to recover from their own mistakes without additional human feedback. There are two fundamental capabilities agent needs to reason test-time about recovering from mistakes. The first is prediction: to understand the consequences of their proposed actions. The second is evaluation: to know which outcomes are preferable others.
We propose an algorithmic paradigm that allows us to acquire both of these capabilities without requiring any sources of human data beyond the standard imitation learning pipeline. In other words, a better recipe with the same ingredients. In particular, rather than merely learning a policy from expert demonstrations, we propose learning local world model (WM) and reward model (RM) from demonstration and base policy data [Ren et al., 2024b]. By combining these components with a planning algorithm, we have the capability to, at test time, reason about how recover from mistakes that the base policy makes by planning against our learned RM inside our learned WM. Thus, rather than mere imitation, we're learning search (L.2S, Ratliff et al. [2009]) from expert demonstrations. Our key insight is that we can infer the latent search process required to recover from local mistakes without any more human feedback (e.g., corrections) than a standard behavioral cloning pipeline.
Put differently, in contrast to approaches like BC and DAgger that give the agent the fish - i.c. relying on human teacher to demonstrate desired or recovery behavior, we focus on teaching the agent to fish: to develop the reasoning process required to match the expert's outcomes, even when faced with situations unseen in the training dataset, often result of the agent's own earlier mistakes.
In our work, we focus on long-horizon visual manipulation tasks and therefore instantiate the L2S paradigm by using base Diffusion Policies [Chi et al., 2023]. Dreamer World Models [Hafner al., 2024], and the Model-Predictive Path Integral Control (MPPI, Williams et al. [2017]) Planner. Concretely, we learn residual planner [Silver et al., 2018] that performs local search at test time

---

Door
1.00
1.00
0.99
Rate
0.65
0.57
ROC (1)
0.09
3
50
100
150
200
50
100
150
200
[D]
Timestap
(-x)
Tinestop = (-))
Cereal
Success Rate
0.95
0.882
0.83
RMC
0.35
0.25
20
25
30
50
100
150
50 Timestep (-)
100
160
[2]
Tinestep
(-+)
Square
Success Bate(1)) (1)
0.83
0.78
BM(s:20 (1))
0.44
0.48
0.29
100
50
100
150
200
50
100
150
200
[D]
Timestep & (-+)
Timestep t (-+)
Figure 2: Left: we see SAILOR consistently out- out-perform diffusion policies trained on the same demos across various visual manipulation problems at multiple dataset scales (D), Right: SAILOR's learned reward model is able to detect shared prefixes (black dots and frames), base policy failures (purple dots, lines, and frames) and SAILOR'S successes (orange dots, lines, and frames).
correct mistakes the base policy makes. We call composite architecture SAILOR: Searching Across Imagined Latents Online for Recovery. More specifically, our contributions are three-fold:
1. We demonstrate that across a dozen visual manipulation problems at three different dataset scales, SAILOR outperforms Diffuusio Policies. Simply scaling up the number demonstrations used for DP by 5-10x still leaves performance gap. Furthermore, SAILOR is also significantly more interaction-efficiem than more traditional model-free inverse RL methods that use state the art Diffusion Policy RL algorithms like DPPO [Ren et al., 2024a] directly update policy parameters.
We carefully ablate the algorithmic and design decisions required learn and combine the above components for stable and sample-efficient learning. Specifically, we find that "warm starting" the hybrid world model training period, doing online hybrid world model fine-tuning [Ross and Bagnell, 2012, Vemula et al., 2023, Ren et al., 2024b], and periodically distilling the learned search algorithm into the base policy (akin to expert iteration [Anthony et al., 2017, Sun et al., 2018] or Guided Policy Search [Levine and Koltun, 2013]) are critical for performance.
We investigate the fidelity and test-time scaling properties of our learned search algorithm. We find that our learned reward model is able to identify nuanced failures that occur at different stages of our oon--horizon manipulation tasks and that our composite stack is robust to reward hacking.
Related Work
Imitation Learning. The simplest approach to imitation learning is behavioral cloning (BC, Pomer- Ieau [1988], Chi et al. [2023]), where one learns policy via maximizing the likelihood of expert

---

actions at expert states. However, whether is due to limited demonstrations [Swamy et al., 2022b], optimization error [Swamy al., 2021], misspecification [Espinosa- Dice et al., 2025]. or partial observability [Swamy et al., 2022a]. a policy trained via BC will often make a mistake that takes it out of the support of the expert demonstrations, where it can continue to make errors, an issue known as compounding errors [Ross et al., 2011], which is unavoidable in general [Swamy et al., 2021].
Under the hood, the reason a BC policy makes mistakes is due the covariate shift between the training input distribution (expert states) and the testing input distribution (learnr states). Interacting with the environment allows us to generate samples from this test distribution. If we 're able to further query the demonstrator in the loop, we can ask them for action labels at these states [Ross et al., 2011, 1, Kelly al., 2019, Spencer et al., 2020]. However, such approaches are often labor-intensive.
Learning to Search = RM + WM. In Learning to Search (1.2S, Ratliff et al. [2009]), one learns the components (i.e., RM and WM) required to, test time, search for actions, rather than directly learning a policy. there are no meaningful dynamics or stochasticity (e.g., append-only, auto- regressive language generation), one can eschew the WM and plan over the entire horizon via repeated sampling and scoring under the RM, an approach known as Best-of-N (BoN, Brown et al. [2024]). Otherwise, one samples some plans, performs stochastic rollouts inside the WM, calling the RM along the way to estimate performance before updating the sampling distribution. We use the MPPI algorithm of Williams et al. [2017] as our search procedure due to the strong performance it has demonstrated when deployed inside a learned WM on robotics problems [Hansen et al., 2022, 2024]. In contrast to prior L2S approaches. SAILOR learns from expert demonstrations alone (i.e., without test-time access to ground-truth simulator as in Silver et al. [2017], Brown and Sandholm [2019] or any information about the ground-truth reward function as in Hansen et al. [2022, 2024]).
After expending computation at test-time for a search procedure, it is often valuable to distill the search process back into the base policy, an approach known as expert iteration or dual policy iteration [Anthony et al., 2017, Sun et al., 2018J. While such approaches have long demonstrated strong performance in continuous control [Levine and Koltun, 2013, Wang et al., 2025], they have attracted renewed interest in the context of LLMs [Zelikman et al., 2022, Gandhi et al., 2024, Hosseini et al., 2024, Jain al., 2025a]. SAILOR can be seen as a generalization of these ideas to continuous control problems with stochastic dynamics, visual observations, and unknown reward functions, which none of the above can handle directly. We ablate the value of ExIt-like updates and find they provide nontrivial improvement in policy performance. Recent work by Wu et al. [2025b] applies similar ideas on real-world visual manipulation problems but focuses on mode selection rather than more general behavior correction. SAILOR can be thought of fusing the paradigms of residual RL [Silver et al., 2018, Yuan et al., 2025, Ankile et al., 2024] and L2S by learning to search for residuals.
3 Robust Imitation via Learning to Search
We adopt the framework of Partially Observed Markov Decision Process (POMDP, Kaelbling et al. [1998]) and use O, A. and 2 to denote the observation, action, and latent spaces, respectively. Also, we use 7 E [0, 1] to denote the discount factor and & to denote our policy's planning horizon.
3.1 What is Learning to Search?
In SAILOR, we learn a search algorithm that generates residual plans Ait k to correct a nominal plan base generated by the base policy. Concretely, after learning a world model WM, reward model RM, and critic V from combination of expert demonstrations and base policy rollouts, we perform repeated stochastic rollouts of potential corrected plans inside our WM, scoring the latent states 2 with R and V, before selecting the plan with the highest estimated score. We then execute the first step of the corrected plan in the real world before re- planning in the style of model predictive control
We note briefly that the term "learning search" is also used describe class of methods for structured prediction problems, which encompass many natural language processing tasks [Chang et al., 2015, 2023]. However, these methods usually assume access to a queryable expert policy in the vein of DAgger [Ross et al., 2011] I] and VaTe [Ross and Bagnell, 2014j. In contrast, we make no such assumptions, and instead focus on how best to give the learner the ability to search independently test time without additional expert guidance.

---

SAILOR
Latent Rollouts
Nx
Current Obs.
RM( (M(2111)
((
QINY
totk
MPPI
Ot
WM
enc
A Lit+t
base
base Cxt++
Base Plan
Residual Plans
Execute 0f A
Figure 3: At inference time, SAILOR performs a search for residual plans to correct mistakes in the base policy's nominal plan in the latent world model WM against the learned reward model RM and critic V. It then executes the first step of the best corrected plan before re-planning. MPC-style.
Algorithm SAILOR (Inference)
1: Input: Base Policy π base World Model WM, Reward Model RM, Critic Network V, Obs. or.
2: Sample nominal plan from base: (Et+h π base (ot).
3: for iteration j in 1 1... do
4: // Can perform
parallel
5: for residual plan n in L...N do
6: Sample A 6.44 ~N(p),
22 7: Execute plan in WM: "1:++ ~ WM(O,, a 12++
base
+ An --(4))
8: Compute Qs: Q" Σ (MMM( (((
9: end for
// Update mean std
pi+1, 1-11, MPPI_update(ji os QU.N A1N
12:
end for Return corrected plan di.- (a + pp,a") o)).
(MPC). More formally, at inference time, we attempt to solve the following local search problem:
A *
argmax E B
RM be+h)) + ( (((+)))))))) Ot, aut+ h =
base det++ + Acttk
Auceo
We then execute the first of these actions, abase + Aj, in the environment before re-planning on top of the fresh observation and base plan. We describe this process in Alg. I and visualize in Fig. 3. We now describe each of these components before describing the phases of training. We use D to denote the of expert demonstrations and B denote replay buffer of on-policy learner rollouts.
Base Policy. We assume access to a base policy base that can generate k-step plans given an observation o;; d y base (o+). This could be arbitrary policy pretrained on wide set of data like a VI.A ['Team et al., 2024, Intelligence et al., 2025, Kim et al., 2024] or a task-specific Diffusion Policy (DP, Chi al. [2023]) trained via behavioral cloning. We adopt the latter for our experiments for simplicity, but note that the general SAILOR framework does not require doing so. As is standard practice [Chi al., 2023], we use ResNet-18 [He et al., 2016] encoder for our DP that takes in both image observations and the proprioceptive state of the robot and generates 8-step plans.
'Explicitly, we make assumptions on the base policy, other than the ability to potentially fine-tune it via behavioral cloning (i.e., maximum likelihood estimation) for the expert iteration subroutine.
5

---

Algorithm 2 SAILOR (Training)
1: Input: Base Policy π base Expert Demos D
2: // Phase I: Warm Start
3: Collect on-policy rollouts: B + {(oe at) ~ base ).
4: // Co-train components
hybrid data
5: WM 4- argmin (D, B, WM), RM + argmingy ((D, B, RM), V + argminy ((,, B.V) V).
6: // Phase II: Online Tuning
7: for iteration j in ....JJ do
8: Collect on-policy rollouts: B B a.) SAILOR].
9: WM <- g L(D, B, RM).RM t- argmingy ((D, B, RM), V + argminy (D,, v).
10:
// Phase III: Expert Iteration
11:
if % m = 0 then
12:
Relabel gdistil <- {(or, (o) E B).
13:
Distill base argmin, ((Bdistil) R).
14:
end if'
end for Return base WM. RM, V.
16:
World Model. avoid the complexity of modeling the dynamics of and planning directly in the space of high-dimensional observations, adopt the Recurrent State-Space Model architecture (RSSM, Hafner et al. [2019. 2024]). This means our world model is composed of three key components, i.c. WM = (enc, f, dec]. The first, the encoder enc Z x O x + Z encodes the observation into latent space i.c. žt = enc(z 1t),, The second, the latent dynamics model f Z x A + AZ)) predicts the next latent after taking an action, i.e. 2t R f a1-1). The third, the decoder dec: 2- + O tries to reconstruct the observation from the latent state, i.e. dec(2z) = ot. In SAILOR, we train all three of these components on hybrid data, i.e. a mixture of D and B see Appendix B for all loss functions. As argued by Ross and Bagnell [2012], Vemula et al. [2023]. Ren et al. [2024b], while such a world model only locally accurate on the expert's and learner's 's state distributions, policy that looks good when evaluated in such model is guaranteed to do well in the real world. 4
0r).
Reward Model. Intuitively speaking, we train reward model RM: - IR to score the latent states of our world model by how "expert-like" they are, which allows us to plan to match expert outcomes in the imagined future. More formally, we train a discriminator between the latent embeddings of expert and learner rollouts using the moment-matching loss proposed by Swamy et al. [2021]:
E(D, B, RM) = (z (nc a, o))] EE,,,,0~
RM enc(2, a, o))]-
(1)
We also add in a gradient penalty Gulrajani et al. [2017] to the above to stabilize training. We iteratively update the RM over the course of training to ensure that we are able to detect mistakes made by the current iteration of the composite SAILOR stack, similar to inverse RL procedures [Ziebart et al., 2008, Ho and Ermon, 2016]. While we do not dwell on the theoretical implications thereof, we note in passing that minimizing the above across a set of potential reward functions corresponds to bounding the performance difference between the expert and the learner under any of these rewards, thereby avoiding compounding errors unlike purely offline behavioral cloning [Swamy et al., 2021].
Critic. To enable truncated horizon rollouts in our WM of l steps, we learn critic V that acts terminal cost estimate. The critic V is trained to predict bootstrapped X-returns (Sutton et al. [1998]):
= RM(x)) + + ((1 A))(2) + Av V++ = V((+*)).
(2)
In our experiments, we train an ensemble of 5 critic networks (Ball et al. [2023], Chen et al. [2021]). When computing terminal cost estimates, we take the mean of 2 randomly sampled critics and subtract an uncertainty penalty proportional to the standard deviation across the entire ensemble.
The Three Phases of Training a Seaworthy SAILOR. As outlined in Algorithm 2, training proceeds in three phases. In Phase Warm-Start (Lines 2-5), we perform rollouts with the base policy base
*While do not investigate this in our experiments, the above theory still holds if the world model is trained on a wide data distribution that covers both the expert and learner's 's visitation distribution. Thus, a particularly interesting future direction is to train a powerful foundation world model [Agarwal et al., 2025, Bruce et al., 2024, Parker-Holde et al., 2024] on internet-scale robotics datasets [Khazatsky et al., 2024, ONeill et al., 2024]. before plugging into the SAILOR framework to allow for wider-ranging recovery from mistakes.

---

LiftPeg
PokeCube
Pul Cube
BreadCanRandon
Success Rate
0.93
0.99
0.99
0.91
0.87
0.88
0.89
0.91
78
0.89
0.86
0.71
0.63
59
0.61
0.68
0.40
0.43
0.41
0.
0.23
15
20
25
5
10
15
10
15
20
50
75
100
[D)
[D]
[D]
[D]
Door
Stack
Bread
Round
Saccess Rate(
1.00
1.00
0.99
0.
0.87
0.86
0.85
0.89
0.93
0.87
0.95
0.95
0.657
0.63
0.55
0.62
68
0.71
0.
0.51
0.45
0.49
0.220
0.:09
20
25
15
20
25
[D)
[D]
[D]
[D]
SAILOR
DP
Cereal
Lift 1.00
Can
Square
Soccess Rate f
1.00
0.999
0.95
0.93
0.99
0.82
0.83
0.85
0.83
0.71
0.75
0.58
0.44
0.48
0.25
0.
0.35
0.33
0.29
20
5
10
15
50
100
(D)
[2]]
[D]
[2]
Figure 4: Across visual manipulation problems from 3 benchmarks, SAILOR consistently outper- forms diffusion policy (DP) trained on the same demos, where [D] denotes the number of demos.
to pre-fill the buffer B, before co-training the world model WM, reward model RM, and critic V mixture of data from B and the expert demonstrations D. In Phase II: Online Fine- Timinn (Lines 6-15), instead perform rollouts with the entire SAILOR stack, periodically performing hybrid training of the WM, RM, and V. In Phase III: Experi Iteration (Lines 1 1-14), we distill the outputs of test-time search into the base policy avoid the need expend compute if the learner ends up in similar situation in the future. More formally, we take some of the most recent (o, a) pairs in buffer B, post-hoc relabel them by calling the SAILOR stack the observation
generate new action labels. and fine-tune the base policy x base via behavioral cloning. This can be seen as using SA SAILOR as an expert for DAgger [Ross al., 2011] update of the base policy, recycling prior test-time compute.
4 Experiments
Benchmarks. We evaluate the efficacy of SAILOR on 12 challenging visual manipulation problems. This includes tasks from Robomimic [Mandlekar et al., 2021b] (Lift, Can, Square), 6 from Ro- boSuite [Zhu et al., 2020] {Door, Stack, Bread, Cereal, Round, BreadCanRandom) and from ManiSkill [Tao et al., 2025] (PullCube, PokeCube, LiftPeg). These include diverse tasks: multi-step pick-and-place (BreadCanRandom)) articulated object manipulation (Door), and tool use (Square, Round, PokeCube). For Robomimic tasks, we use the provided demonstrations, while collect
demonstrations using SpaceMouse controller for the other suites. To solve each task, SAILOR is provided with a set of expert demonstrations D, along with a budget specifying the maximum number of environment interactions it can perform. The observation space consists of RGB images from wrist-mounted camera and third-person camera mounted in front of the agent, and the proprioceptive states. More details are provided in App. and E.
"One can consider training on buffer B as a variant of Follow the Regularized Leader [McMahan, 201 1]. stable no-regret algorithm.. More formally, by plugging in the bounds of Ross and Bagnell [2012] into the reduction of Ren al. [2024b], derive performance bounds for SAILOR under standard assumptions.

---

BreadCanRandon
Round
Cereal
1.00
1.00
1.00-
Succesa Bate
0.75
0.75
0 0.75-
0.50
SAILOR
0.50
0.28
26 -0-00
0.25
150
200
[D]
[D]
(D)
Figure 5: We see that simply scaling up the amount of demos [D] used for training DP via behavioral cloning by 5-10x often plateaus in performance and unable to match the performance SAILOR.
Door [D] =5 5
Can [P] 10
Bread [D]
Succeas Rate(
1.0
1.0
000000
0.75
0.75
Pooooo
0.5
0.5
0. 0.5
0.25
0.
0.25
0.0
0.0
0.0
1E5
3E5
5E5
185
3E5
6E5
1E5
3E5
5E5
Eav Steps (-))
Eav Steps (-))
Eny Scepa (+))
SAILOR
DRPO-IRL
Figure 6: We find that SAILOR is significantly more interaction-efficient than a model-free inverse RL baseline that directly updates DP policy parameters via DPPO [Ren al., 2024a].
Algorithms. We compare three imitation learning methods: Diffusion Policies (DP, Chi et al. [2023]) trained via BC, a model-free inverse RL method (DPPO-IRL) that directly updates DP parameters via RL (DPPO, Ren et al. [2024a]) against a learned RM with a separate encoder, and SAILOR. For evaluation, we measure the Success Rate (SR) across 50 rollouts, and report the mean and standard error obtained with 3 seeds. More details on the implementation of the methods are provided in App. B, All methods were trained on 1 NVIDIA 6000 Ada GPU with 48 GB of memory.
4.1 Results
Can SAILOR outperform DP trained on the same D? Fig. 4 compares SAILOR and DP across various tasks and size of demonstration datasets. We observe that SAILOR significantly outperforms DP across all environments and dataset scales considered, indicating that the learning search paradigm allows us to squeeze out significantly more from the same expert data. We find particularly large gaps in the low-data regime, reflecting how I.2S inherits sample-efficiency benefits of inverse RI. approaches that learn verifiers from human data rather than just learning policies generators [Swamy et al., 2025].
Does DP catch up with more data? A natural question after viewing the preceding results might be as to whether more demonstrations could close the gap between DP and SAILOR, reflecting robot learning's current emphasis large-scale data collection [ONeill et al., 2024, Khazatsky al., 2024]. In Fig. 5, we observe that while DP improves with more expert data, the performance plateaus after 100 demonstrations. This means that when we scale up the number of provided demonstrations to 200, R 10x the amount provided to SAILOR, DP still unable to match our method's s performance.
Note further that these are expert demonstrations collected directly for the target environment. One can imagine that as practitioners choose to scale up offline data with sources such as human-video demonstrations and Internet-scale pretraining, data which is by design even less in-distribution to the task at hand, we expect a pure BC model to exhibit diminishing performance gains. SAILOR, by contrast, might be able to absorb that same large, noisy dataset into its WM and RM. Pre-training these components on broad "notions of success" may yield robust dynamics priors and value estimates that drive on-policy action distillation and keep improving the policy long after BC has plateaued.
How much real-world interaction does the WM save us? Model-based approaches are well- well-known to be more interaction-efficin than their model-free analogs, benefit SAILOR inherits. In particular, each observation isn't treated as a single data point. Instead, SAILOR uses as seed to spawn an entire tree of counter- factual trajectories, not only slashing the number of costly real trials needed
8

---

Surrece Place
Alland Greup
BoCo (t)
Mosil Pull
F0M (t)
Wond Orey
Nittral Haang
Timeatep I (1+)
Cizeetop t (++)
Timestop (-+)
Sequential Pick
Tool Grasp
Tool Hang
Figure 7: At different stages of multi-step task execution, we see our learned reward model RM able to identify a variety of nuanced failures and SAILOR able counterfactually avoid them. We generate these plots by rolling out the base policy until it fails at the given task (purple line), resetting the agent state where failure is not yet guaranteed, and letting SAILOR complete the rest of the episode (orange line). We use black dots to mark shared prefixes, purple dots to mark the base policy's failure suffixes, and orange dots mark SAILOR's succesful suffixes.
Can D
Stack D
15
Round [
Success (ate(())
0.9-
0.9p
1.0
0.75
0.76
0.75-
SAILOR
SAILOR
SAILOR
23
2
2
2
2
23
2
2
2
2
2
2
2
2
2
MPPI Samples (logz scale)
MPPI Samples (logz Scale)
MPPI Samples (logz Scale)
Figure 8: We use 256 samples (analogous to in BoN) for the MPPI planning process at train-time (black cross). We see that more compute leads to higher performance up to 256 and do not see degradation performance even with 8x as many samples, indicating robustness to reward hacking.
to reach given performance, but also pruning potentially high-variance or dangerous paths before executing them on the physical hardware. To quantify the size this benefit, we compare SAILOR to model-free IRL algorithm DPPO-IR IRL that directly updates policy parameters after performing rollouts in the real world. In Fig. 6, we consistently see that even with 5x the interaction budget, DPPO-IRA IRL is unable match the performance of SAILOR, reflecting the importance of the WM.
Can the RM detect nuanced failures? We now explore qualitatively what our learned RM is detecting. We do this by sampling trajectories that fail to ultimately complete the task from the base DP, truncate at a prefix where failure is not yet guaranteed, and then roll out SAILOR counterfactually from these states to recover from mistakes. In Fig. 7, we see that our RM is able to detect nuanced failures that occur at various stages complex task (e.g. a narrowly missed tool hang after a successful grasp).
How robust is SAILOR to reward hacking? A common concern with test-time scaling approaches that with enough compute, they may "hack" learned proxy reward model and perform worse on the ground-truv by scaling up the number of samples generated in the MPPI planning procedure, which analogous to the N in BoN. We see in Fig. 12 that up to the number of samples used for training (256), more compute consistently leads to better performance. Furthermore, we do not see degradation in performance even with 8x as many samples, reflecting high degree of robustness to reward hacking.
metric [Gao et al., 2023]. We explore the test-time scaling properties of SAILOR
4.2 What Matters in Learning To Search?
We now ablate the importance several parts of the overall SAILOR pipeline.
Warm-Start. We found that allocating 20% of the interaction budget to a "warm-start phase" significantly boosts final performance, as shown in the leftmost part of Fig. 9. We attribute this to having the WM. RM, and V be accurate on x base s distribution, reducing exploration Phase II. We find performance plateaus when using 20% of the interaction budget for warm- warm-starting (App. A.1).
Hybrid WM Training. Another component that has a significant effect on model performance using a mix of expert and learner data to update the WM. We observe in Fig. 9, center, that this "hybrid" fitting of the WM not only has the potential to improve initial model performance (as for PokeCube), but can also improve final performance (as for Cereal), reflecting the insights of Ren et al. [2024b].

---

Round (P) 25 Equare [M] =
PakeClibe (2)
Cereal (2) 25
PakeCiba (D) = 10 Cereal [P] [
0.78
Succesa Bazelf)
1.0-
1.0
0.76
76
0.34
0.50
0.5
0.6
0,80
0.25
0.91
1E5
285
3G3
3E5
6E5
Stepa (-+)
Env Stepa (-+)
SKILCR
BAILOR
Hybrid WM
BAILOR
ExdT
Figure 9: We ablate the impact of three components of our overall training pipeline: warm starting, hybrid world model training, and expert iteration, and find that all three improve performance.
Expert Iteration. In Fig. 9, right, we see that the inclusion of the expert iteration subroutine boosts model performance. This can be attributed to the base policy knowing how to recover from mistakes its prior iterations would have made, effectively recycling test- tis-timm compute into a new base policy.
Residual Planning. We ablate the importance of using test- time planning search for residuals online instead of precomputed residual policy (RP). Similar to the residual planner, the RP takes input both the latent state and base policy action. is trained via standard actor-critic techniques as in Dreamer [Hafner et al., 2024]. In Fig. 10, we see that across the board, explicit test-time planning out-performs also include
Lift [D]
Cas D) 10 Squaze [D] 50
0.99
0.03
0.78
0.58
residual policy. As additional statisties, numbers for the wall-clock time overhead of running MPPI, which we find to be somewhat limited. We also ablate the amount of MPPI iterations performed in the loop in App. A.2. There, we see fast convergence to an improved plan.
SAT1OOR
Figure 10: We find that explicit resid- ual planning out-performs Residual Policy (RP) across the board.
5 Discussion
In summary, we see that across dozen visual manipulation tasks, SAILOR able to train robust agents that recover from from the failures of the base policy via test-time planning and match the expert's outcomes, all without on-policy human feedback. Furthermore, scaling up the amount of data used to train diffusion policy via behavioral cloning by 2 5-10x is unable to close this performance gap. We also see that our learned RM is able to detect nuanced failures, while the SAILOR agent is able to recover from them. As mentioned above, while for our experiments policy, MPPI planner, and Dreamer world model, the general SAILOR architecture is fully compatible with foundation behavior (e.g. VLAs) and world models trained on internet-scale data, which might unlock further levels of generalization and robustness we leave this as a promising future direction.
use base diffusion
Before we close, it is worth discussing two reasons why L2S is particularly scalable algorithmic paradigm. The first is data-efficiency: as with any inverse RL approach that learns a verifier (i.e., reward model) from demonstrations rather than just a generator (i.e., policy), L2S requires less data than direct policy learning approaches on problems with senerution-verifiaatao gap [Swamy et al., 2025]. Given we have far less data in robotics than in domains like language modeling, squeezing as much as we can out of every demonstration is of paramount importance. Furthermore, as we scale up task horizons, the size of the generation- veriftiitiie gap grows. This means that L2S techniques could exhibit better "data scaling laws" with task complexity than behavioral cloning techniques.
The second reason is compute-efficiency in L22,, we only need to plan at the states actually encoun- tered at test-time, rather than all states we could potentially encounter as in a train-time procedure. This provides provable computational benefits [Kearns et al., 2002]. Furthermore, expert iteration- style distillation lets us recycle this compute, so we never need to solve the same problem twice. Thus, L2S techniques could also exhibit better "compute scaling laws" than vanilla policy learning.
To close, we observe that in his oft-quoted essay The Bitter Lesson, Richard Sutton writes that:
"One thing that should be learned from the bitter lesson is the great power of general purpose methods... The two methods that seem to scale arbitrarily in this way are search and learning."
In SAILOR, rather than pick one, we fuse learning and search to chart course for imitation.
10

---

Acknowledgments
Calcul We thank Kensuke Nakamura, Abigail DeFranco, and Andrea Bajesy for their help in the early stages of this project and thoughtful advice throughout its completion. We thank Mohan Kumar for his help with setting up our demonstration collection interface and Yilin Wu for her advice on world models and demo collection. We thank Kensuke Nakamura, Harley Wiltzer, Jesse Farebrother, and Andrea Bajesy for feedback on our initial draft of this paper. GKS is supported by a STTR grant. SC is supported in part by Google Faculty Research Award, an OpenAl SuperAlignment Grant, an ONR Young Investigator Award, NSF RI #23 12956, and NSF FRR #2327973, AJ is supported by Fonds de Recherche du Quebec (FRQ) (DOI assigned: http:://doi.org/10.6977/7 Quebec, and Canada Excellence Research Chairs (CERC) program.
Contribution Statements
AJ co-lead the implementation of all components, supervised the work of other students and helped write the paper. VM co-lead the implementation of all components and helped write the paper.
SK implemented the first-pass version of our planner. SK, JR, and YF and implemented and ran experiments for model-free inverse RL baseline. JR also helped write the paper.
AB collected the demonstrations used for training all methods.
SC helped advise the project, shaping its direction and presentation, suggested experiments and ablations, and provided computational and personnel resources.
GS initially conceived of the project and algorithmic paradigm, gathered the team of students, managed the project over several phases, made all the figures, and wrote most of the paper.
References
Niket Agarwal, Arslan Ali, Maciej Bala, Yogesh Balaji, Erik Barker, Tiffany Cai, Prithvijit Chat- topadhyay, Yongxin Chen, Yin Cui, Yifan Ding, et al. Cosmos world foundation model platform for physical ai. arXiv preprint Ttt.:2500.033755 2025.
Lars Ankile, Anthony Simeonov, Idan Shenfeld, Marcel Torne, and Pulkit Agrawal. From imitation to refinemeet--esiddual rl for precise visual assembly. arXiv preprini arXiv:2407. 166777777.677, 2024.
Thomas Anthony, Zheng Tian, and David Barber. Thinking fast and slow with deep learning and tree search. Advances in neural information processing systems, 30, 2017.
James Bagnell, Sham M Kakade, Jeff Schneider, and Andrew Ng. Policy search by dynamic programming. Advances neural information processing systems, 16, 2003.
Philip J Ball, Laura Smith, Ilya Kostrikov, and Sergey Levine. Efficient online reinforcement learning with offline data. In International Conference on Machine Learning, pages 1577-1594, PMLR, 2023.
Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating propagating gradients through stochastic neurons for conditional computation. arXiv preprint / 2013.
Bradley Brown, Jordan Juravsky, Ryan Ehrlich, Ronald Clark, Quoc V Le, Christopher Ré, and Azalia Mirhoseini. Large language monkeys: Scaling inference compute with repeated sampling. arXiv preprint arXi:2240.221787 2024.
Noam Brown and Tuomas Sandholm. Superhuman ai for multiplayer poker. Science, 365(6456): 885-890, 2019.
Jake Bruce, Michael D Dennis, Ashley Edwards, Jack Parker-Holder, Yuge Shi, Edward Hughes, Matthew Lai, Aditi Mavalankar, Richie Steigerwald, Chris Apps, et al. Genie: Generative interactive environments. In Forty-first International Conference Machine Learning, 2024.
Jonathan D. Chang, Kiante Brantley, Rajkumar Ramamurthy, Dipendra Misra, and Wen Sun. Learning to generate better than your 1lm, 2023, URL hattps://artn org/ ab//23/6....////////// 11816.
11

---

Kai- Wei Chang, Akshay Krishnamurthy. Alekh Agarwal, Hal Daumé III, and John Langford. Learning to search better than your teacher. In International Conference on Machine Learning, pages 2058- 2066. PMLR, 2015.
Xinyue Chen, Che Wang, Zijian Zhou, and Keith Ross. Randomized ensembled double q-learning: Learning fast without a model. arXiv preprint arXiv:2101.05962 2021.
Cheng Chi, Zhenjia Xu, Siyuan Feng. Eric Cousineau, Yilun Du, Benjamin Burchfiel, Russ Tedrake, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. The International Journal of Robotics Research, page 02783649241273666, 2023.
Cheng Chi, Zhenjia Xu, Chuer Pan, Eric Cousineau, Benjamin Burchfiel, Siyuan Feng, Russ Tedrake, and Shuran Song. Universal manipulation interface: In-the-wild robot teaching without in-the-wild robots. arXiv preprint UrXiv:2402.10329 2024.
Sudeep Dasari, Oier Mees, Sebastian Zhao, Mohan Kumar Srirama, and Sergey Levine The ingredients for robotic diffusion transformers. arXiv preprint 0rXiv:2410.10085 2024.
Nicolas Espinosa- Dice, Sanjiban Choudhury, Wen Sun, and Gokul Swamy, Efficient imitation under misspecification. arXiv preprint arXiv:2503. /3/62, 2025.
Kanishk Gandhi, Denise Lee, Gabriel Grand, Muxin Liu, Winson Cheng, Archit Sharma, and Noah D Goodman. Stream of search (sos): Learning to search in language. arXiv preprint 000::2444.06833. 2024.
Leo Gao, John Schulman, and Jacob Hilton. Scaling laws for reward model overoptimization. In International Conference Machine Learning, pages 10835-10866. PMLR, 2023.
Ishaan Gulrajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron C Courville. Improved training of wasserstein gans. Advances in neural information processing systems, 30, 2017.
Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. Deepseck-rl: Incentivizing reasoning capability in Ilms via reinforcement learning. arXiv preprint arXiv:2501.12943 2025.
David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv: 1803. /0/22, 2018.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning. volume 97 of Proceedings of Machine Learning Research, pages 2555-2565. PMLR, 09-15 Jun 2019. URL https: http///pceeee....e.ee/// mlr wwwppeeeeerreee http9///aaafeerr/ html.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference Learning Representations, 2020.
URL attps:// http:::///penreviii .nt wttprreiiweeeeeeeeoerereeeeeeeee htte//orummmi---1100e1aa
Danijar Hafner, Timothy P Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. International Conference on Learning Representations, 2021. URL
h/tpp//////pereviex / wiwpmmmddddddooobbll00
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models, 2024. URL https: //arxiva.vvvv////////// org/ httg/bbs/23000 04104.
Siddhant Haldar, Vaibhav Mathur, Denis Yarats, and Lerrel Pinto. Watch and match: Supercharging imitation with regularized optimal transport. In Conference on Robot Learning, pages 32-43.
PMLR, 2023a.
Siddhant Haldar, Jyothish Pari, Anant Rai, and Lerrel Pinto. Teach robot to fish: Versatile imitation from one minute of demonstrations. arXiv preprint arXiv:2303.0142 2023b.
Nicklas Hansen, Xiaolong Wang, and Hao Su. Temporal difference learning for model predictive control. In ICML, 2022.
12

---

Nicklas Hansen, Hao Su, and Xiaolong Wang. Td-mpe2: Scalable, robust world models for continuous control, 2024.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference computer vision and pattern recognition, pages 770-778, 2016.
Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. Advances in neural information processing systems, 29, 2016.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851, 2020.
Arian Hosseini, Xingdi Yuan, Nikolay Malkin, Aaron Courville, Alessandro Sordoni, and Rishabh Agarwal. V-star: Training verifiers for self-taught reasoners. arXiv preprint arXiv:2402.066 4444576457 2024.
Physical Intelligence, Kevin Black, Noah Brown, James Darpinian, Karan Dhabalia, Danny Driess, Adnan Esmail, Michael Equi, Chelsea Finn, Niccolo Fusai, Manuel Y. Galliker, Dibya Ghosh, Lachy Groom, Karol Hausman, Brian Ichter, Szymon Jakubezak, Tim Jones, Liyiming Ke, Devin LeBlanc, Sergey Levine, Adrian Li-Bell, Mohith Mothukuri, Suraj Nair, Karl Pertsch, Allen Z. Ren, Lucy Xiaoyang Shi, Laura Smith, Jost Tobias Springenberg, Kyle Stachowicz, James Tanner, Quan Vuong, Homer Walke, Anna Walling, Haohuan Wang, Lili Yu, and Ury Zhilinsky, To.5 vision-language-action model with open-world generalization, 2025. URL https: /ttp:///aaaiviaeee org/ abs/2504. 16054.
Aaron Jaech, Adam Kalai, Adam Lerer, Adam Richardson, Ahmed El-Kishky, Aiden Low, Alec Helyar, Aleksander Madry, Alex Beutel, Alex Cameey, et al. Openai ol system card. arXiv preprint 1rX10:2412.16720, 2024.
Arnav Kumar Jain, Shiva Kanth Sujit, Shruti Joshi, Vincent Michalski, Danijar Hafner, and Samira Ebrahimi Kahou. Learning robust dynamics through variational sparse gating. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Informa- lion Processing Systems, 2022. URL https: http///pppeeeeeeeee///ee/e wttpiiiwwwwetteeee www.orrmmmiiiiiissooeoeee
Arnav Kumar Jain, Gonzalo
Wayne Chen, Alexander M Rush, Wenting Zhao, and Sanjiban Choudhury. Multi-tum code generation through single-step rewards. In Workshop on Reasoning and Planning for Large Language Models, 2025a. URL //tttps://ppprree forum?id= BBkkbdddBbkbb
net/
Arny Kumar Jain, Harley Wiltzer, Jesse Farebrother, Irina Rish, Glen Berseth, and Sanjiban Choudhury. Non-adversarial inverse reinforcement learning via successor feature matching. In The Thirteenth International Conference on Learning Representations, 2025b. URL https:
//openreview.o wttpr/iwwwweeeetutieeeeeeee htt.roormmiiddiiimm id=..RdgsvLLL....
Leslie Pack Kaelbling, Michael L Littman, and Anthony R Cassandra. Planning and acting in partially observable stochastic domains. Artificial intelligence, 101(1-2):99-134 1998.
Michael Kearns, Yishay Mansour, and Andrew Y Ng. A sparse sampling algorithm for near-optimal planning in large markov decision processes. Machine learning, 49:193-208, 2002.
Michael Kelly, Chelsea Sidrane, Katherine Driggs-Campbell, and Mykel J Kochenderfer. Hg-dagger: Interactive imitation learning with human experts. 2019 International Conference on Robotics and Automation (ICRA), pages 8077-8083. IEEE, 2019.
Alexander Khazatsky, Karl Pertsch, Suraj Nair, Ashwin Balakrishna, Sudeep Dasari, Siddharth Karamcheti, Soroush Nasiriany, Mohan Kumar Srirama, Lawrence Yunliang Chen, Kirsty Ellis, al. Droid: large-scale in-the-wild robot manipulation dataset. arXiv preprint 4rX1v:2403.12945. 2024.
Moo Jin Kim, Karl Pertsch, Siddharth Karamcheti, Ted Xiao, Ashwin Balakrishna, Suraj Nair, Rafael Rafailov, Ethan Foster, Grace Lam, Pannag Sanketi, et al. Openvla: An open-source vision- Assionlaaggaage-attion model. arXiv preprint 0801:224060024666 2024.
13

---

Sergey Levine and Vladlen Koltun. Guided policy search. In International conference on machine learning, pages 1-9. PMLR, 2013.
Ajay Mandlekar, Danfei Xu, Josiah Wong, Soroush Nasiriany, Chen Wang, Rohun Kulkarni, Fei- Fei, Silvio Savarese, Yuke Zhu, and Roberto Martin-Martin. What matters in learning from offline human demonstrations for robot manipulation. arXiv preprint arXiv:2108.03298, 2021a.
Ajay Mandlekar, Danfei Xu, Josiah Wong, Soroush Nasiriany, Chen Wang. Rohun Kulkarni, Li Fei- Fei, Silvio Savarese, Yuke Zhu, and Roberto Martin-Martin. What matters in learning from offline human demonstrations for robot manipulation. In Conference on Robot Learning (CoRL), 2021
Brendan McMahan. Follow-the-regulee reguurrizeddeeeaaeeeeee and mirror descent: Equivalence theorems and 11 regularization. In Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics, pages 525 533. JMLR Workshop and Conference Proceedings, 2011.
Russell Mendonca, Oleh Rybkin, Kostas Danilidis, Danijar Hafner, and Deepak Pathak. Discovering and achieving goals via world models, 2021.
Diganta Misra. Mish: A self regularized non- non-monotonic activation function. arXiv preprint arXiv: /908. 08681, 2019.
Andrew Y Ng. Stuart Russell, et al. Algorithms for inverse reinforcement learning. In loml, volume 1, page 2, 2000.
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International conference on machine learning, pages 8162-811 2171. PMLR, 2021.
Abby ONeill, Abdul Rehman, Abhiram Maddukuri, Abhishek Gupta, Abhishek Padalkar, Abraham Lee, Acorn Pooley, Agrim Gupta, Ajay Mandlekar, Ajinkya Jain, et al. Open x-embodiment: Robotic learning datasets and rt-x models: Open x-embodiment collaboration 0. In 2024 IEEE International Conference Robotics and Automation (ICRA), pages
IEEE, 2024.
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems, 35: 27730-27744, 2022.
J Parker-Holder, P Ball, J Bruce, V Dasagi, K Holsheimer, C Kaplanis, A Moufarek, G Scully, Shar, Shi, et al. Genie 2: A large-scale foundation world model. URL: https://deepmind 2024.
Dean A Pomerleau. Alvinn: An autonomous land vehicle in neural network. Advances in neural information processing systems, 1, 1988.
Alexander Popov, Alperen Degirmenci, David Wehr, Shashank Hegde, Ryan Oldja, Alexey Kamenev, Bertrand Douillard, David Nistér, Urs Muller, Ruchi Bhargava, al. Mitigating covariate shift in imitation learning for autonomous vehicles using latent space generative world models. arXiv preprint 4rXiv:2409.1666 16663, 2024.
Nathan D Ratliff, David Silver, and Andrew Bagnell. Learning search: Functional gradient
techniques for imitation learning. Autonomous Robots, 27:25-53, 2009.
Allen Z Ren, Justin Lidard, Lars L Ankile, Anthony Simeonov, Pulkit Agrawal, Anirudha Majumdar, Benjamin Burchfiel, Hongkai Dai, and Max Simchowitz. Diffusion policy policy optimization. arXiv preprint 000x:2409900588, 2024a.
Juntao Ren, Gokul Swamy, Zhiwei Steven Wu, J Andrew Bagnell, and Sanjiban Choudhury. Hybrid inverse reinforcement learning. arXiv preprint arXiv:2402.08888, 2024b.
Stephane Ross and J Andrew Bagnell. Agnostic system identification for model-based reinforcement learning. arXiv preprint / 1007, 2012.
Stephane Ross and J Andrew Bagnell. Reinforcement and imitation learning via interactive no-regret learning. arXiv preprint arXiv:1406.5973 2014.
14

---

Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. A reduction of imitation learning and structured prediction no-regret online learning. In Proceedings of the fourteenth international conference on artificial intelligence and stafistics, pages 627-635. JMLR Workshop and Conference Proceedings, 2011.
David Silver, Julian Schrittwieser Karen Simonyan, loannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, el. Mastering the game of go without human knowledge. nature, 350(7676):354-355 2017.
Tom Silver, Kelsey Allen, Josh Tenenbaum, and Leslie Kaelbling. Residual policy learning. arXiv preprint arXiv:/8/2.06220 2018.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502, 2020.
Jonathan Spencer, Sanjiban Choudbury, Matthew Barnes, Matthew Schmittle, Mung Chiang, Peter Ramadge, and Siddhartha Srinivasa. Learning from interventions: Human-robot interaction as both explicit and implicit feedback. In 16th roboties: science and systems, RSS 2020. MIT Press Journals, 2020.
Nisan Stiennon, Long Ouyang, Jeffrey Wu, Daniel Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul F Christiano. Learning to summarize with human feedback. Advances in Neural Information Processing Systems, 33:3008-3022 -0008,3-21-200 2020.
Wen Sun, Geoffrey J Gordon, Byron Boots, and J Bagnell. Dual policy iteration. Advances in Neural Information Processing Systems, 31, 2018.
Richard S Sutton, Andrew G Barto, et al. Reinforcement learning: An introduction, volume 1. MIT press Cambridge, 1998.
Gokul Swamy, Sanjiban Choudhury. J Andrew Bagnell, and Steven Wu. Of moments and matching: game-theoretic framework for closing the imitation gap. In International Conference on Machine Learning, pages 10022-10032. PMLR, 2021.
Gokul Swamy, Sanjiban Choudhury, J Bagnell, and Steven Z Wu. Sequence model imitation learning with unobserved contexts. Advances in Neural Information Processing Systems, 35:17665-17676,
2022a.
Gokul Swamy, Nived Rajaraman, Matt Peng, Sanjiban Choudhury, J Bagnell, Steven Z Wu, Jiantao Jiao, and Kannan Ramchandran. Minimax optimal online imitation learning via replay estimation. Advances in Neural Information Processing Systems, 35:7077-7088, 2022b.
Gokul Swamy, David Wu, Sanjiban Choudhury, Drew Bagnell, and Steven Wu. Inverse reinforcement learning without reinforcement learning. In International Conference on Machine Learning, pages 09299-33318. PMLR, 2023.
Gokul Swamy, Sanjiban Choudhury, Wen Sun, Zhiwei Steven Wu, and J Andrew Bagnell. All roads lead to likelihood: The value of reinforcement learning in fine-tuning. arXiv preprint arXiv:2503.01067, 2025.
Umar Syed and Robert Schapire. A game-theoretic approach apprenticeship learning. Advances in nearal information processing systems, 20, 2007.
Stone Tao, Fanbo Xiang, Arth Shukla, Yuzhe Qin, Xander Hinrichsen, Xiaodi Yuan, Chen Bao, Xinsong Lin, Yulin Liu, Tse kai Chan, Yuan Gao, Xuanlin Li, Tongzhou Mu, Nan Xiao, Amay Gurha, Viswesh Nagaswamy Rajesh, Yong Woo Choi, Yen-Ru Chen, Zhiao Huang, Roberto Calandra, Rui Chen, Shan Luo, and Hao Su. Maniskill3: Gpu parallelized robotics simulation and rendering for generalizable embodied ai. Robotics: Science and Systems, 2025.
Octo Model Team, Dibya Ghosh, Homer Walke, Karl Pertsch, Kevin Black, Oier Mees, Sudeep Dasari, Joey Hejna, Tobias Kreiman, Charles Xu, et Octo: An open-source generalist robot policy. arXiv preprint
2024.
15

---

Sai Vemprala, Rogerio Bonatti, Arthur Bucker, and Ashish Kapoor. Chatgpt for robotics: Design principles and model abilities, 2023, URL
orgpiii.orraaae httgg/b///306... 17582.
Anirudh Vemula, Yuda Song, Aarti Singh, Drew Bagnell, and Sanjiban Choudhury. The virtues of laziness in model-based rl: A unified objective and algorithms. In International Conference Machine Learning, pages 34978-35005. PMLR, 2023.
Yuhang Wang, Hanwei Guo, Sizhe Wang. Long Qian, and Xuguang Lan. Bootstrapped model predictive control. In The Thirteenth International Conference on Learning Representations, 2025.
URL httts::/ wwwpeeewwwwwwweeeeeeeeeeeeeeeeee met/forummid=i1 AYFYDcM.
Grady Williams, Andrew Aldrich, and Evangelos A Theodorou. Model predictive path integral control: From theory parallel computation. Journal of Guidance, Control, and Dynamics, 40(2): 2017.
Philipp Wu, Alejandro Escontrela, Danijar Hafner, Pieter Abbeel, and Ken Goldberg. Daydreamer: World models for physical robot learning. In Conference on robot learning, pages 2226-2240.
PMLR, 2023.
Philipp Wu, Yide Shentu, Zhongke Yi, Xingyu Lin, and Picter Abbeel. Gello: A general, low-cost, and intuitive teleoperation framework for robot manipulators. In 2024 IEEE/RSJ International Conference Intelligent Robots and Systems (IROS), pages 12156 12163. IEEE, 2024.
Runzhe Wu, Yiding Chen, Gokul Swamy, Kianté Brantley, and Wen Sun. Diffusing states and matching scores: A new framework for imitation learning. In The Thirteenth International Conference on Learning Representations, 2025a. URL https:///peerreeieet wwt. formmiiid KWRKNDUGUN.
Yilin Wu, Ran Tian, Gokul Swamy, and Andrea Bajesy. From foresight to forethought: VIm-in-the- loop policy steering via latent alignment. arXiv preprint arXiv:2502.01828. 2025b.
Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European conference on computer vision (ECCV), pages 3- 3-19, 2018.
Xiu Yuan, Tongzhou Mu, Stone Tao, Yunhao Fang, Mengke Zhang, and Hao Su. Policy decora- tor: Model-agnostic online refinement for large policy model. In The Thirteenth International Conference on Learning Representations, 2025.
Eric Zelikman, Yuhuai Wu, Jesse Mu, and Noah Goodman. Star: Bootstrapping reasoning with reasoning. Advances in Neural Information Processing Systems, 35:15476-15488, 2022.
Tony 7. Zhao, Vikash Kumar, Sergey Levine, and Chelsea Finn. Learning fine-grained bimanual manipulation with low-cost hardware. arXiv preprint arXiv:2304. 13705, 2023.
Gaoyue Zhou, Hengkai Pan, Yann LeCun, and Lerrel Pinto. DINO-WM. World models on pre-trained visual features enable zero-shot planning. In secty- International Conference on Machine Learning, 2025. URL /ttps://opptt //pprreveee/eeee http:/rriiiiieeeeeeeeeeerreeeeee /ttp://rrmmddd-dM../2E-
Yuke Zhu, Josiah Wong, Ajay Mandlekar, Roberto Martín-Martín, Abhishek Joshi, Soroush Nasiriany, Yifeng Zhu, and Kevin Lin. robosuite: A modular simulation framework and benchmark for robot learning. In arXiv preprint arXiv:2009. /2293, 2020.
Brian D Ziebart, Andrew Bagnell, Anind K Dey, et al. Maximum entropy inverse reinforcement learning. 2008
16

---

A
Additional Results
We now discuss two additional ablations of the SAILOR stack.
A.1 Warm-Start Fraction
The warm-start fraction refers to the portion of total interaction budget we use collect data with just the base policy (i.e., not performing any policy updates or using the planning stack). We then use this data to train the WM, RM and Critic. We performed ablation of this fraction on the Square task with [D] = 50 demonstrations.
Succeas Rata(")
1.0
Square [D]
50
0.78
0.5
0.25
0.0
0.1
0.2
0.3
Fraction
DP SALLOR
In Fig. 11. we observe that using warm- start fraction of 0.1 improves performance over not warm-starting. However, performance plateaus after using warm-start fraction of 0.2, which we used for all of our experiments.
Figure 11: We ablate the warm-start percentage and find performance plateaus after 0.2. We use total interaction budget of 500k environment steps and measure performance across three seeds.
A.2
Residual Search vs. Residual Policy
Succesa Bate(1)
1.0
Lift (D] =5
1.0
Can
D
Square [D]
50
0.75
0.75
0.75
0.5
0.5
0.5-
Iterations
MPPI Iterations
MPPI Iveratione
SAILOR
Distilled
We use 6 iterations of the MPPI planning process (black line). We see that more iterations lead higher performance up to 6 and do not see a degradation performance even with 9 iterations. Iteration 0 represents the performance of the base DP. We also explored a post-hoc distillation step (blue line) for reducing wall-clock time. Doing so matches SAILOR on Lift and Can but not Square
Figure 12:
In Fig, 12, compare the performance of SAILOR with different number of iterations of the MPPI planning process. We find that the performance converges after 2-3 iterations for most tasks and there only a marginal improvement in later iterations. suggesting we could truncate SAILOR 's search process early. Doing so adaptively on state-wise basis is an interesting future research direction.
Additionally, measured that SAILOR spends 0.014s time per iteration of MPPI, while it takes 0.066s for the base diffusion policy to generate actions, making test-time planning a somewhat limited overhead. This is likely because the MPPI planner operates top of lower-dimensional latent states from the world model, rather than high-dimensional image observations like the base DP.
Lastly, we explored performing final post-hoc distillation of the SAILOR agent into the base policy (Distilled DP) for fast inference at test-time. As seen in Fig. 12, the performance of Distilled DP is similar to the SAILOR agent on Lift and Can tasks. However, test-time search outperforms the Distilled DP. This is potentially due to the limited capacity of the diffusion policy for representing the full planning circuit for sufficiently complex tasks.
the harder Square task,
17

---

B Implementation Details
In this section we first describe the network architecture for each component of SAILOR in Sec. B.1 following details of training pipeline and hyperparameters in Sec. B.2
B.1 Network Architecture
Base Policy. SAILOR uses Diffusion Policy (DP, Chi al. [2023]) as the base policy. The base policy takes the stack of current and previous observations ((o-11,9)) as input and predicts the k-step action plan d The observation has proprioceptive states and RGB images from wrist camera and a front camera. To encode RGB images, DP uses a ResNet- 18 encoder initialized with the ImageNet1KK V1 weights, and the intermediate activations are passed through Spatial Softmax layer [Mandlekar et al., 2021a] to get the final image embedding. Here, each input image uses different copy of the encoder, and all encoded image inputs are concatenated together with the proprioceptive state and finally passed to the noise network
The noise network is conditional ID U-Net, where conditioning is incorporated via FiLM mod- ulation, following the approach of [Chi et al., 2023]. The network is conditioned on the diffusion timestep (represented as a 16 dimentional embedding generated by sinosuidal encoding layer followed by a small MLP), and the encoded observations. The UNet consists of multiple convolution and transposed convolution layers with channels [64, 128, 256]. kernel size 3, GroupNorm Wu and He [2018] for stable training and Mish [Misra, 2019] activation function. The noise network is trained to reverse the forward noising process of adding Gaussian noise at each step (DDPM, Ho al. [2020], Nichol and Dhariwal [2021]). Lastly, DP trained using behavior cloning objective, formulated as denoising task where the model learns to predict the noise added to a k-step action chunk, conditioned on the given observation and diffusion timestep.
During inference, DP uses the DDIM sampling [Song et al., 2020] with the noise network to generate the action chunk for a given observation. In this work, we use the first action generated pase to step in the environment. The agent uses a action blending mechanism that smoothens the action at current step using the predictions from earlier observations [Dasari al., 2024]. We found this to significantly boost the performance of DP in our experiments.
World Model. The world model used in SAILOR uses the architecture of DreamerV3 [Hafner et al., 2024]. The encoder (enc) encodes the pixel-based inputs with stride 2 convolutions to a resolution of 4 x 4 and state inputs are embedded with a 5-layer MLP. Note that the is using a separate encoder for observations as we found it to work well in our experiments. The decoder (dec) uses transposed convolutions with stride 2 to reconstruct image observations and a 5-layer MLP to reconstruct the proprioceptive states. Note that the RGB images are downscaled to size 64 x 64 for training. The latent representation 2 a combination of a deterministic recurrent state h2 and a stochastic state st. The deterministic state is has a GRU network with 512 dimensional hidden state. The latent state 2;-1 and action a from previous time step are used to estimate the deterministic component h; at current step. The deterministic state is fed through the dynamics model f to sample the prior stochastic state â, Moreover, the deterministic state combined with the observation is passed through the encoder enc to get the posterior state The stochastic representation of 1024 dimensions is sampled through a vector of multinomial distributions where the gradients are backpropagated through straight-through estimators [Bengio et al., 2013] while learning.
The world model is trained by hybrid learning where the half of the batch of sequences is obtained from the demonstration D and other half comes from the replay buffer B. More formally, consider sampled batch of observation subsequences Ot:t+a actions a.s+a continuation flag Cu:+u (to predict the end of episode), where t ~ ~411 "ill....,Ty, ,T), u denotes the length of the sequence, and T denotes the length of trajectory. We minimize combination of a prediction loss fpred. a dynamics loss layn and a representation loss frep over samples drawn from DD and B:
1
((D, B3, WM)
E Eoute Onees Ceee ( (-) Soyn (.) BropCrop()
2
+
E wew.ennnntsssstsseeee (apreeewwwsspppdppppeeee (-)d Poyn (-) + Prop (opp)rrep))))
(3)
2
18

---

where Spred = 1.0, Sdyn = .1 and Brep .5 represent the loss weights. The loss functions are:
(o:tt Ourtur Ctot+ta WM)
In dec =n) In dec(calzn),
h=t
(oyntOr: (ret--rry Ctat+a WM)
max(1, Dxx. (sg(enc(z) i 0h-1, dec F4 1n
k=t 1+u
Ge:t+r Ct:t+a WM)
max 1, DxL (ww.(ennlln-n,, ah on) sg(dec((n) 04-1))
The prediction loss optimizes for reconstructing the observations via Mean Squared Loss (MSE) criterion and the continuation predictor via logistic regression. The dynamics and representation losses optimize the same objective is optimized with different set of parameters: the former updates the dynamics model to predict the posterior state while the latter term ensures that the stochastic term is more predictable. Note that the dynamics and representation loss only differ in the coopggradient term sg and the loss weight terms.
Reward Model. The reward model RM takes the latent representation 2, as input and uses a 2-layer MLP to predict scalar value expressing the desirability of being in state. The RM optimized with the loss function defined in Eq. 1. The RM is updated with a gradient penalty term [Gulrajani et al., 2017] with a coefficient of 10. To stabilize learning, the RM updated less frequently than the world model. Moreover, we do not update the parameters of the world model with the loss of the RM.
Critic Network. Similar RM, the critic network uses 2-layer MLP and predicts the discounted average rewards of future states with the latent representation 2n as input. The critic is updated with the Mean Squared Loss (MSE) with target va defined in Eq. 2:
ElEertus V) --=
(V(z)) - v
(4)
Unlike RM, the critic is updated with the WM. Training the critic more frequently than RM is important to ensure it accurately estimates the average rewards and remains synchronized with the changing reward function. Similar to Hansen et al. [2024], SAILOR maintains and uses an ensemble of 5 value networks. Like Dreamer, SAILOR also maintains a slow value network compute a slow target for the critic network and uses this as a regularizer for critic loss. The slow networks are updated with EMA over the critic parameters.
MPPI Planner. SAILOR uses MPPI for planning withing the world model in this work. The planner maintains a gaussian distribution with mean μ and diagonal covariance o predict the residual action Ai. The planning procedure is described in Alg. 1 where the parameters are initialized with 0 mean and a fixed standard deviation. At each iteration, 256 action chunks are sampled and scored with the RM and by imagining future latent with WM. The top 64 sequences with highest n-step returns used to update the parameters μ and a. Unrolling in the latent space allows evaluating large batches in parallel on a single GPU, and thereby makes planning efficient. After 6 iterations, an action plan Ai.t is sampled using final parameters μ and a*.
B.2 Training details
With the demonstrations D, SAILOR first pretrains the base policy- DP. The training procedure of SAILOR was outlined in Alg. 2. The pretrained DP is used to collect rollouts in the environment and uses around 20% of total environment steps. For instance, when the agent is tasked with 100K environment steps, the pretrained DP is deployed to collect for 20K transitions. SAILOR adds a small noise sampled from (O to the action to promote exploration while collecting data from the environment. The agent maintains a uniform replay buffer with an online queue of size During the warmstart phase, the WM, RM and critic V are updated with hybrid batch sampled from demonstration and replay buffer B. Note that, the gradient steps of WM and was set to 1.5x the number of transitions collected. For training stability, the RM is updated slowly and once every 100 updates to the WM. After warmstart, SAILOR uses online finetuned with batch data collection and updates over multiple rounds. At each round, SAILOR deploys the planner to collect on-policy
10
19

---

trajectories for 3500 environment steps, The WM and V are then updated for 5000 gradient steps where the RM is updated once every 100 gradient steps of WM. After every 10 rounds. the last 64 collect trajectories are relabeled with the planner and the base policy is updated for 1000 iterations using a hybrid batch composed of relabeled data and expert demonstrations. In Table. 1, we provide the details of hyperparameters used in this work. Our agents where trained on a single NVIDIA 6000Ada GPU with 48 GB memory and takes hours for training end-to-end with 500K environment steps.
C Baselines
For baselines, compare SAILOR with the pretrained DP policy and a diffusion based IRL method (DPPO-IRL). Specifically, we apply DPPO [Ren et al., 2024a] a model-free RL algorithm that directly updates DP parameters with feedback from reward model learned in the same for SAILOR (Eq. 1). In contrast to learning a WM and learning to search in SAILOR, DPPO-IRL optimizes the policy directly using the outcomes of the learned reward model. To encode the observation, the reward model used the same encoding as the base DP of DPPO and 2-layer MLP of width 256. We tuned the hyperparameters
the reward model (update epochs and batch size), and observed that the best version used 2 as update epochs and a batch size of 100. As for SAILOR, we used gradient penalty term with coefficient of 10 stabilize learning of the reward model.
20

---

Name
Value
DP Pretraining
Batch Size
256 AdamW W 24,000 Cosine Annealing 100
Optimizer
Training iterations
LR scheduler
LR scheduler warmup steps
LR range
x 10 10-
World Model
Replay capacity
1 x 10
Batch size
Batch length
32 Adam
Optimizer
Reconstruction Loss Scale
1.0
LR
1 x 10-4
Demo Sampling Ratio
50%
Reward Model
Optimizer
Adam 3 10-5
LR
Gradient Penalty coefficient
10
Critic
Discount factor y
.997
Return lambda À
.95
EMA regularizer
EMA decay
.98
Optimizer
Adam
LR
3 x 10-5
Ensemble size
MPPI Planner
Iterations
Samples
256
Top candidates
32
Temperature
General
Warmstart env step ratio
20%
Env steps per round
3500
Update steps per round
5000
Distillation frequency
10
Trajectories relabeled for distillation
64
Distillation steps
1000
Table Hyperparameters. For training, recommend tuning training parameters in the General section that includes the update steps per round E [1000, 2000, 3500, 5000, 10000), distillation fre- quency of the base policy [1, 5, 10, 50] and distillation steps of the base policy E [500, 1000, 2000). LR, Env and EMA denotes learning rate, environment and exponential moving average.
21

---

D Benchmarks
In this section, we describe the environments used in this work. The experiments are conducted on multiple robotic manipulation environments from RoboMimic [Mandlekar et al., 2021b], Robo- Suite [Zhu al., 2020] and ManiSkilI3 [Tao al., 2025]. Fig. 13 presents a visual image for each of the task used in this work. For each task, the agent is provided with an RGB image from the wrist camera and an RGB image from camera front of the agent. The agent is also given proprioceptive states composed of position, orientation of the end effector and the position of gripper. In Table 2, we provide the episode horizon, the environment steps used for IRI. training and brief description of the task. The action space is 7-dimensional vector with values between [-1, 1]. The first 6 dimensions of the action control the change in position and orientation of the end-effector and the last value opens closes the gripper.
Domain
Task
Horizon
Env Steps
Description
Lift
100
x 105
Lift Block above the desk. Lift can and place correct bin.
RoboMimic
Can
200
10a
Square
200
5 x 10a
Pick square tool and insert in slot.
Door
200
10
Pull down handle and open door. Lift block and place above other block. Lift bread and place in correct bin.
Stack
150
x 10
Bread
200
x 105
RoboMimic
Cereal
150
x 105
Lift cereal and place in correct bin. Pick round tool and place in slot. Place both objects in respective bins.
Round
300
x 105
BreadCan
400
X 10$
PullCube
50
x 10a
Pull cube to the marked area Use tool to poke cube to marked area Lift peg to make stand upright
ManiSkill
PokeCube
100
x 10$
LiftPeg
150
x 10
Table 2: The maximum episode length (Horizon), environment steps and description of the tasks.
22

---

RoboMimic- Lift
RoboMimic - Can
RoboMimic Square
RoboSuite Door
RoboSuite- Stack
RoboSuite Bread
RoboSuite- Cereal
RoboSuite Round
RoboSuite BreadCan
ManiSkill - Lift Peg
ManiSkill Pull Cube
ManiSkill Poke Cubc
Figure 13: Visual description of all tasks used from RoboMimic (top row), RoboSuite (middle rows) and ManiSkill (bottom row).
23

---

E Demonstrations
In this section, we describe the demonstrations used for learning. For RoboMimic tasks, used the demonstrations provided in the dataset. For tasks in RoboSuite and ManiSKill, we collected upto 200 demonstrations using 3D SpaceMouse. The camera angles were adjusted for each task to make the process intuitive to the human teleoperator. Below we provide a table of the number of demonstrations used per task for Fig. 4. For BreadCan, Cereal and Round we collected upto 200 demonstration for our result in Fig. 5. We plan to release the demonstrations and the codebase.
Domain
Task
Number of Demonstrations
Lift
5, 10, 15 10, 15, 20 50, 75, 100
RoboMimic
Can
Square
Door
1,3, 5
Stack
15, 20, 25
Bread
15, 20, 25
RoboMimic
Cereal
20, 25, 30
Round
25, 50, 75
BreadCan
50, 75, 100
PullCube
10, 15, 20
ManiSkill
PokeCube
5, 10, 15
LiftPeg
15, 20, 25
Table 3: The number of demonstrations used for each task.
24

---

F Related Work (Extended)
Reward Models. Another approach to interactive imitation learning is inverse reinforcement learning (IRL, INg et al., 2000, Syed and Schapire, 2007, Ziebart et al., 2008, Ho and Ermon, 2016, Swamy al., 2021J). which does not require that maximally differentiates learner from expert behavior and uses it as a reward model (RM) for RL-based policy updates. Different forms of reward models include successor features [Jain et al., 2025b], score matching [Wu et al., 2025a], and optimal transport metrics Haldar et al. [2023b,a]. Unfortunately, the RL step of IRL is often rather interaction-ineffcientt. Recent theoretical work has argued that rather than global RL procedure, local search procedure" is sufficient for performant imitation [Swamy ct al., 2023, Espinosa- -iieooss--eeeeeeeo et al., 2025]. SAILOR fits within the overarching algorithmic paradigm of local search IRL but learns search rather than learning policy directly.
-n---------e-oop queries. In IRL. one learns a classifier
World Models. World models (WMs, Ha and Schmidhuber [2018]) have been an integral component of impressive RL results in variety of domains [Hafner et al., 2020, 2021, Jain et al., 2022, Hansen al., 2022, Hafner et al., 2024, Hansen et al., 2024, Zhou et al., 2025, Bruce et al., 2024, Parker- Holder et al., 2024, Agarwal et al., 2025], including real-world robotics [Mendonca et al., 2021, Wu et al., 2023]. While our overall algorithmic framework is agnostic to the choice of WM architecture, we use Dreamer-style world models [Hafner et al., 2020] in our experiments due to their ubiquity. Popov et al. [2024] use rollouts in a learned world model to minimize a trajectory-leve divergence between the expert and the learner at train time autonomous driving problems. In contrast, the world model at rest-time to enable recovery from just the mistakes the learner actually makes, which might be more computationally efficient than global search procedure [Kearns et al., 2002].
use
"By local search, we mean merely competing with the expert rather than the optimal policy for an adversari- ally chosen reward. For the former, there exist algorithms that avoid the corst-case, exponential-in the-horizzo exploration complexity of RL [Bagnell al., 2003, Ross and Bagnell, 2014, Swamy et al., 2023].
25

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We show that our method outperforms DP in Figure present DP cannot catch-up with 5-10x data in Figure 5, and show robustness of learned RM in Figure 12 and Figure
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model cell-specificatinn. asymptotic approximations only holding locally). ). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
26

---

Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We describe the architecture details of our method in Appendix B. The baselines are described in C. For the tasks where we collected demonstrations, we have added details in Appendix E.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
27

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We have released the data and code and have provided details in main paper and appendix for reproducibility.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.. cc/
onPolicy) for more details.
While encourage the release of code and data, we understand that this might not be possible, so âÄlJNoâÄ is acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed run reproduce the results. See the NeurlPS code and data submission guidelines (https:
ctaass.ccpublicc// htpppulicgguuideeeeeeee gwwdeu/de///e// /CodeSubmias ssi onPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We have added the training and test details in Sec, 4.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: seeds.
We report all our results with mean score and standard error obtained with
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
28

---

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out of range (c.g. negative emror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes)
Justification:
We add details on compute details under Implementation Details in Sec
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: : //wurrrrrs.p.. cc http/pplic/tthcccciiieeeeoooeede
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
29

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We have added citations for code and benchmarks in the main paper.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
30

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
The code and dataset used in this work are released.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/cooe/mmoel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question:
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
31

---

Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components.
Please refer to our LLM policy (https for what should or should not be described.
32


Reviews:
[{"review_id": "GDxzWrlh0d", "reviewer": "Reviewer_yCh2", "summary": "SAILOR augments an off-the-shelf behaviour-cloned diffusion policy with a learned world model, reward discriminator and latent-space MPPI planner. At test time it treats the base policy’s $k$-step plan as a proposal, samples residual trajectories in the world model, scores them with the learned reward model and critic, and executes the first action of the best-scoring plan. Across 12 simulated manipulation tasks and three data scales, SAILOR consistently beats the underlying DP, even when DP is given 5–10 × more demonstrations, and does so with far fewer real-environment interactions. The paper backs the design with three key ablations, demonstrates that SAILOR can directly correct its base policy, and shows that planning remains stable when the search budget is pushed.", "questions": "1. In Figure 2, the reward model (RM) is also used to evaluate the trajectories generated by DP. However, DP does not incorporate such a component. Where does the RM used to evaluate DP originate? Is it the same RM trained for SAILOR?\n2. Why are different tasks used in Figure 9 when ablating warm-start compared to the other components? Would all tasks confirm your findings?\n3. In the hybrid WM ablation (Fig. 9, center), does “no hybrid” mean using only expert demos or only on-policy rollouts?", "limitations": "No potential negative societal impact. I commend the authors for boldly marking “N/A” in the checklist regarding the limitations of their work. However, some practical constraints and assumptions, such as reliance on a competent base policy and computational demands of test-time planning, as discussed above, to name a few, deserve acknowledgment.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 4, "strengths_and_weaknesses": "# Strengths\n\n1. With the same number of expert demonstrations, SAILOR outperforms DP on a wide array of *12* tasks from different robotic arm suites, and DP can’t catch up even with 10× data. SAILOR also requires far fewer interactions than a model-free IRL baseline (DPPO-IRL). Its benefits extend beyond the typical few-shot imitation learning use case.\n2. Figures 2 and 7 effectively illustrate SAILOR’s ability to recover from base policy errors. The visualizations combine scenario snapshots with reward model curves, showing clear course-correction behavior during execution. Moreover, consistent color coding across all figures enhances readability and makes comparisons between SAILOR and DP intuitive.\n3. The method only requires a nominal plan from any pretrained policy along with an image observation, making it agnostic to the underlying architecture. Nothing ties it specifically to diffusion models, suggesting a modular, plug-and-play design that could generalize to a broader range of base policies.\n4. The paper is very well written and easy to follow. Design choices are consistently justified with logical reasoning, clear explanations, and relevant citations. It is well-situated within the literature on imitation learning, planning, L2S, and reward modeling, with strong connections drawn to prior work throughout.\n5. The ablation study sharpens the case for the chosen model/method components. By isolating warm-starting, hybrid world model updates, and expert iteration, the authors show that each component significantly affects performance.\n6. As demonstrated in the ablation study, the expert iteration loop enables SAILOR to progressively refine the base policy, reducing reliance on costly test-time planning in recurring scenarios.\n\n# Weaknesses\n\n1. The paper treats L2S as a universally appealing paradigm due to its training-time planning efficiency. However, this ignores scenarios where inference-time compute is limited, yet ample compute is available during training. In such settings, SAILOR's reliance on expensive MPPI rollouts at every timestep may be impractical, and more traditional offline methods or distilled policies would be preferable. While the authors briefly mention expert iteration for distillation, they don’t evaluate or emphasize it as a viable path for deployment, leaving a gap between their method's design and real-world utility.\n2. According to the ablation study, warm-starting plays a critical role in overall performance. To facilitate this, SAILOR assumes the base policy is sufficiently competent to provide useful warm-starts for learning a good prior. If the base policy performs poorly, the initial rollouts used to train the world model and reward model will be low-quality, leading to inaccurate dynamics and reward predictions. This, in turn, undermines the effectiveness of MPPI planning.\n3. While SAILOR clearly outperforms DP, this isn’t the fairest comparison. DP is a purely offline, feedforward policy with no planning, reward signal, or test-time adaptation, whereas SAILOR leverages a learned world model, reward model, and expensive online planning. The absence of any measurement of inference latency, MPPI rollout time, or runtime comparisons with DP or DPPO-IRL is a major omission, especially given SAILOR’s heavy test-time compute. DPPO-IRL serves as a more appropriate baseline, as it also learns a reward model and improves the policy beyond BC. However, it shifts the compute burden to training via model-free RL, whereas SAILOR achieves higher performance with significantly fewer environment interactions. A more meaningful comparison would account for total compute, both in terms of sample efficiency and wall-clock time.\n4. While the framework is presented as modular and capable of operating with any arbitrary base policy capable of generating a sequence of future actions, in practice, the paper only demonstrates integration with DP. To substantiate this claim, the authors could empirically evaluate SAILOR with alternative base policies. \n5. The authors state that as a key insight “*we can infer the latent search process required to recover from local mistakes from the same source of human data required for the standard behavioral cloning pipeline.*” While SAILOR does leverage expert demonstrations, this claim overstates their centrality. The core training of the world model, reward model, and critic relies heavily on **on-policy rollouts** from the base policy and the SAILOR stack. Expert data is used, but not exclusively, and much of the method’s performance stems from its ability to adapt online. The implication that SAILOR learns solely from offline BC-style data is misleading.\n\n# Minor Suggestions\n\n1. It took me a while to figure out that the dots on the line plots in Figure 2 (Right) represent when the image snapshots above were taken. Perhaps mention this in the caption for clarity.\n2. If I understand correctly, the x-axis values in Figure 8 should be powers of two, instead of values represented with E notation.", "was_revised": true, "final_justification": "In light of the clarifications to my comments (W2: on-policy training, W3: compute time, and W5: claim about offline data), the extra distillation results (W4), and responses to other reviewers, I have decided to **increase the score**.", "rebuttal": "Thank you for your valuable review of our paper. We are pleased to hear that you found the paper easy to understand and appreciated SAILOR’s ability to recover from base policy errors with fewer environment interactions. We respond to questions below.\n\n> W1: While the authors briefly mention expert iteration for distillation, they don’t evaluate or emphasize it as a viable path for deployment.\n\nIn terms of total compute used for policy updates across training and test times (not just at training time), we believe that L2S methods are likely more efficient than traditional policies. Intuitively, this is because for a policy to be able to respond well to a scenario, it likely needs to have been trained on something similar at training time. **In the worst case, we need to handle *all* such scenarios at train time, while an L2S method only needs to handle what is seen at test time.** A more formal version of this argument can be found in the classic paper https://www.cis.upenn.edu/~mkearns/papers/sparsesampling-journal.pdf. Furthermore, because we iteratively distill the output of the search process back into the base policy via the expert iteration subroutine, SAILOR avoids re-solving the same residual search problem in the future, essentially recycling the test-time compute.\n\nThat said, we agree with the reviewer that in situations where computation is highly limited at inference time, it may not be possible to successfully employ an L2S strategy. However, we found that for our SAILOR implementation, **we only spent ~0.014s time per iteration on the MPPI residual search, compared to ~0.066s on the DP forward pass** (we used 6 iterations of MPPI in reported results). This is likely because the MPPI planner operates on top of lower-dimensional *latent states*. In addition, we saw that performance converges after 2-3 iterations, and there is only a marginal improvement in later iterations. \n\nInspired by the reviewer's question, **we explored an additional option: performing a post-hoc final distillation of the SAILOR agent to base policy**. This helps the base policy to close the gap in performance with the search policy (shown in the table below). As this base policy will be as fast as DP at test time, it is suitable for deployment. In addition, the test-time planner can provide additional gains with more inference budget, and we can see that the planner adds marginal gains after 3 iterations. Lastly, we would like to highlight that MPPI planners for test-time planning have been extensively deployed on real robots, including the work of [1,2,3,4], implying that it is not a barrier to deployment. We thank the reviewer for suggesting this experiment and have added it to the paper. \n\n|        | Lift (5 Demos) | Can (10 Demos) | Square (50 Demos) | Inference Time (s)  | \n|--------|----------------|----------------|-------------------|----------------|\n| DP (pre-training)   | 0.58           | 0.33           |  0.28             | 0.066             |\n| DP (post-trainng)  | 0.94 |0.81 | 0.56 | 0.066 |\n| + MPPI 1 Iter | 0.98 | 0.84           |  0.71             | 0.08          |\n| + MPPI 3 Iter | 0.99 | 0.84          |  0.8             | 0.109          |\n| + MPPI 6 Iter | 0.99           | 0.85           |  0.8             | 0.151          |\n| + MPPI 9 Iter | 0.98           | 0.85           |  0.81             | 0.195          |\n\n> W2: If the base policy performs poorly, the initial rollouts used to train the world model and reward model will be low-quality, leading to inaccurate dynamics and reward predictions. \n\nWe agree with the reviewer that when the base policy is weaker, the learner is fundamentally faced with solving a harder search problem. However, **because both the world and reward model are trained on on-policy data, we should still be able to improve upon the base policy via search**. In some sense, the \"quality\" of a world/reward model is perhaps better measured by how useful they are for improving the current policy, rather than how accurate they are on a good policy's distribution -- see [5] for more information. That said, we do see at least some empirical evidence that SAILOR is able to achieve a high success rate even when the base policy does not. For example, for Door (1 demos) and PokeCube (5 demos), the pretrained DP has 0.09 and 0.2 SR, and SAILOR achieves 1.0 and .91 SR after training.\n\n> W3: The absence of any measurement of inference latency, MPPI rollout time, or runtime comparisons with DP or DPPO-IRL is a major omission.\n\nWe thank the reviewer for the suggestions. First, we would like to emphasize that no method in our paper (including SAILOR) gets access to any sort of \"ground-truth\" reward signal -- both DPPO-IRL and SAILOR infer a reward model from demonstrations rather than, say, regressing ground-truth reward labels. That said, we agree with the reviewer that including more information about the wall-clock time each of these methods takes at test-time would significantly strengthen our paper. As mentioned above, we only spent ~0.014s per iteration on the MPPI residual search, compared to ~0.066s on the DP forward pass (which all methods pay for).\n\nIn terms of wall-clock training time, SAILOR takes around 40 hours and DPPO-IRL takes around 11 hours for 500K environment steps on a single GPU (NVIDIA 6000Ada with 48 GB memory). We attribute this higher train time to the computation required to train a reconstruction-based world model like DREAMER. However, in terms of environment interactions required to reach a particular performance, we robustly find SAILOR requires significantly fewer than DPPO-IRL, as our results in the paper provide evidence for. Furthermore, if instead we measure the model performance given a fixed wall-clock time, we often find SAILOR to be significantly better than DPPO-IRL. For example, **on the Door task, SAILOR converges to a success rate of 1.0 in ~11 hrs with about 100k environment steps, while DPPO-IRL reaches a success rate of just 0.75 (with 500k interactions and ~11 hrs of training)**. We have included these numbers in the paper as they paint a more nuanced picture of the compute requirements for different methods across training and testing. Please let us know if there are any other comparisons you think would strengthen our story.\n\n> W4:  To substantiate this claim, the authors could empirically evaluate SAILOR with alternative base policies.\n\nWe agree that this is an interesting and promising direction for future work! In particular, we'd be interested in seeing how SAILOR could be applied to base VLA policies.\n\n> W5: The implication that SAILOR learns solely from offline BC-style data is misleading.\n\nWe thank the reviewer for highlighting this. In our original statement, we wanted to emphasize that SAILOR does not require any additional interaction/supervision from the *human* expert, just like behavioral cloning. This is in contrast to a method like DAgger [6]. We agree that saying we do not need any more \"human data\" is not as clear as we could have been, and have rephrased this in the paper.\n\n> Q1: Where does the RM used to evaluate DP originate? Is it the same RM trained for SAILOR?\n\nYes, the RM trained with SAILOR was used for the evaluation of all policies. In particular, we first generate rollouts using the DP and score them with the reward model, then reset to a state on this trajectory to rollout from with SAILOR. This shows that the SAILOR RM is able to distinguish between successful and unsuccessful rollouts and identify nuanced failures in the base policy. We have added a note on this to the caption of Figure 2.\n\n> Q2: Why are different tasks used in Figure 9 when ablating warm-start compared to the other components? Would all tasks confirm your findings?\n\nIn our experiments, different tasks showed different amounts of dependence on each component of the overall SAILOR stack. That said, all components show a significant impact on at least a small subset of tasks, as the plots confirm. Thus, to have a method that performs uniformly well across all tasks, we believe we need the full SAILOR stack. We have added this point to the caption in the paper.\n\n> Q3: In the hybrid WM ablation (Fig. 9, center), does “no hybrid” mean using only expert demos or only on-policy rollouts?\n\nIn the Fig.9, no hybrid means using only on-policy rollouts. We have updated the description of Fig. 9 to reflect this point.\n\n> Minor suggestions\n\nThis is a great catch. We have revised the caption of Figure 2 and updated Figure 9 to change E notation to powers of two.\n\n> Limitations \n\nWe have added a discussion and limitations section to the revised version. We discuss the lack of real-robot experiments and that learning from scratch often requires too many interactions to be feasible in practice, potentially incentivizing the use of foundation models for base policies and world models. Due to space constraints, we refer the reviewer to our response to Reviewer kaSH (W3) in the rebuttal for further details.\n\nThank you for the thoughtful comments. We hope we addressed most of your concerns, and hope you consider increasing your score. \n\nBest regards,\\\nThe Authors\n\n### References\n[1] Minarik et al., Model Predictive Path Integral Control for Agile Unmanned Aerial Vehicles, 2024\\\n[2] Testouri et al., Towards a Safe Real-Time Motion Planning Framework for Autonomous Driving Systems: A Model Predictive Path Integral Approach, 2024.\\\n[3] Han et al., Model Predictive Control for Aggressive Driving Over Uneven Terrain, RSS 2024.\\\n[4] Lancaster et al., MoDem-V2: Visuo-Motor World Models for Real-World Robot Manipulation, 2024.\\\n[5] Ross et al., Agnostic System Identification for Model-Based Reinforcement Learning, ICML 2012.\\\n[6] Ross et al, A Reduction of Imitation Learning and Structured Prediction to No-Regret Online Learning, AISTATS 2011."}, {"review_id": "ArWe3U1QB9", "reviewer": "Reviewer_dq9G", "summary": "This paper introduces SAILOR, an online framework for learning an imitation policy with robust recovery. SAILOR learns a world model, reward models, and a policy on mixed online interaction and offline expert demonstration data. SAILOR first trains a base BC policy on expert data, and populate the buffer with online rollouts. Then, it pretrains the world model and reward models on the mixed buffer. The reward models (RM, V) are trained to discriminate between expert and base policy rollouts. Then, during inference, SAILOR searches in world model latent space and optimizes the trajectory to match the expert demonstrations by maximizing the predicted rewards, and rolls out with receding horizon MPC. Results in various environments show that SAILOR is more robust than vanilla DP and is more efficient than inverse RL. Ablations validate the necessity of each component in SAILOR.", "questions": "- The paper mentioned that 20\\% of the budget is allocated for warm-start. For each task, how many frames of online vs. offline trajectories do you have? How many frames minimum does it require to yield improved WM/RM/policy learning afterwards? It would be great to see a sweep of warm-start budget from 0\\% to 20\\% and its impact on downstream performance.\n- For the distillation step, is it doing straightforward BC training with the new data? Do you have any data quality assumption on the collected online trajectories (e.g. suboptimal trajectories, success/failure assumptions)?\n- In Figure 4, why do you think vanilla DP should plateau, given the same interaction budget? How do the SAILOR rollouts differ from the expert demos? Is it a data coverage/diversity problem?\n- Reference links didn't seem to work.", "limitations": "- No real robot experiments.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "- Robust recovery from out-of-distribution states for imitation learning is an important problem.\n- The proposed method is well-motivated and novel. The authors introduce a model-based learning/planning approach for learning a robust imitation policy from expert demos and additional online interactions.\n- Experiment results in various sim visual manipulation environments show that SAILOR is more robust than vanilla DP, and more efficient than inverse RL baselines. The reward model learns interpretable rewards for the task.\n- Extensive ablations validate the necessity of each component in SAILOR.\n- No real robot experiments.", "was_revised": true, "final_justification": "I think SAILOR proposes a valid and novel idea for robust imitation learning, given the extensive evaluations on simulated environments, as well as related works on real robots exploring similar directions. I maintain my overall rating of accept.", "rebuttal": "Thank you for your valuable feedback and questions. We are glad that you appreciate the robust recovery at OOD states, and the Interpretability of learned rewards of SAILOR. We respond to your questions below: \n\n> Q1: The paper mentioned that 20% of the budget is allocated for warm-start. For each task, how many frames of online vs. offline trajectories do you have? How many frames minimum does it require to yield improved WM/RM/policy learning afterwards? It would be great to see a sweep of warm-start budget from 0% to 20% and its impact on downstream performance.\n\nThe number of frames of offline/expert data depends both on the number of expert demonstrations provided, as well as the length of each episode (100-250, depending on the task). The number of frames of online data is equal to the number of environment interactions -- we provide these numbers in Table 2 in the Appendix in the supplementary material. The warm-start percentage of 20% refers to the percentage of total interaction budget we use to collect data with the base policy to warm-start the WM and Critic. As requested, we performed an ablation of this percentage on the Square task with 50 demonstrations. We train agents with {0, 10, 20, 30}% of interaction dedicated to the warm-start. In the table below, **we observe that using 10% the interaction budget for a warm-start improves performance over no warm-start (performance measured across 3 seeds). However, performance plateaus after using 20% for the warm-start**. Ideally, we would want to allocate most of the interaction budget to online collection and policy update; hence, we used 20% for our experiments. We have added this ablation to the paper.\n\n| | Square (50 demos) |\n|----------|----------|\n|DP | .28 |\n| No Warmstart | .54    |\n| Warmstart (10%) |  .72  |\n| Warmstart (20%) |  .78  |\n| Warmstart (30%) |  .79  |\n\n> Q2: For the distillation step, is it doing straightforward BC training with the new data? Do you have any data quality assumptions on the collected online trajectories (e.g., suboptimal trajectories, success/failure assumptions)?\n\nYes! We are literally doing behavioral cloning (BC) for the policy update, which makes our expert iteration procedure simple and scalable. More explicitly, as outlined in Algorithm 2, after sampling some observations from the agent's replay buffer, we relabel these with actions generated by the current SAILOR stack, and use this as the dataset to update the base policy. **There are no data quality assumptions as we do not have access to any environment rewards or success flags**.\n\n> Q3: In Figure 4, why do you think vanilla DP should plateau, given the same interaction budget? How do the SAILOR rollouts differ from the expert demos? Is it a data coverage/diversity problem?\n\nTo clarify, in Figure 4, the x-axis represents the number of expert demonstrations provided to all methods, rather than the interaction budget -- the base DP does not have any interaction budget as it is trained via purely offline behavioral cloning. That said, the reviewer asks two interesting questions. The first is **whether BC can plateau at suboptimal performance even when provided with ample expert demonstrations**. The answer to this question is a **resounding yes** when the learner cannot perfectly imitate the expert -- i.e., the \"misspecified setting\" [1]. In such settings, mistakes are inevitable, and the goal is to avoid particularly bad mistakes that compound over the horizon (e.g., falling off a cliff) [2]. However, without the ability to observe the consequences of one's actions, a purely offline algorithm is unable to distinguish between mistakes that don't matter (e.g., an unnecessary lane change) and those that do (e.g., hitting another car) [3]. In such settings, the hope is that an interactive imitation learning algorithm like SAILOR is able to learn to correct from recoverable mistakes back to the expert's path. The second question the reviewer asks is also quite insightful: **If the expert demonstrations contained the desired recovery behavior, is there a need for interaction?** If, for example, one is able to pre-hoc predict *all* mistakes *any* policy could make, one could include this in the demonstration set, but this is, of course, a tall order. In general, this can require an amount of recovery demonstrations that scales *exponentially* with the task horizon. Instead, however, one can just observe what mistakes a particular chosen policy makes, and **can either have a human show how to recover from these mistakes (as in DAgger) or have the learner figure this out without further human guidance (as in SAILOR)**.\n\n> Q4: Reference links didn't seem to work.\n\nAh, thank you for the catch! We had a Chrome extension to cut the PDF for submission that hid this fact from us. The links are working in the paper submitted with the Appendix in the supplementary material, and we will fix this in the revised version of the paper. \n\n> Q5: No real robot experiments.\n\nWe agree with the reviewer that real robot experiments would strengthen our story and leave it as a promising direction for future work (and one that we are currently working on!). However, the reviewer might find it interesting that the just-published FOREWARN paper [4] showed that a similar setup to SAILOR can be scaled to real robots.\n\nWe hope we addressed most of your concerns, and hope you consider increasing your score. \n\nBest regards,\\\nThe Authors\n\n#### References\n[1] Espinosa-Dice et al., Efficient Imitation under Misspecification, ICLR 2025.\\\n[2] Ross et al., A Reduction of Imitation Learning and Structured Prediction to No-Regret Online Learning, AISTATS 2011.\\\n[3] Swamy et al., Of Moments and Matching: A Game-Theoretic Framework for Closing the Imitation Gap, ICML 2021.\\\n[4] Wu et al., From Foresight to Forethought: VLM-In-the-Loop Policy Steering via Latent Alignment, RSS 2025."}, {"review_id": "iBpHn8S7E2", "reviewer": "Reviewer_AkFH", "summary": "The paper presents SAILOR. Given nominal trajectories predicted from a diffusion policy, SAILOR tries to find a residual trajectory by searching under a world (RSSM) and a reward model (discriminator on “expert-like” behaviors).", "questions": "- I’d be interested to learn a variant of figure 6 where the x-axis is the total amount of compute used (EFLOPs).\n- How does the proposed approach work on a real robot?", "limitations": "Missing an explicit limitation discussion section. It's unclear how the proposed approach would work on a real robots.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "### Strengths\n+ The paper is well-motivated and presented\n+ Detailed ablation analysis and I appreciate the clarity in disclosing the key findings in tuning the proposed method (section 4.2)\n\n### Weaknesses\n- My main concern is the evaluation setup. The main metrics measured is success rate yet the core idea is to improve BC. Success rate, especially in simulators, can be easily improved using RL without using world model or reward model.\n- Related to above, please also consider comparing to other closed-loop IL methods, such as GAIL (that also trains the reward model on-the-fly).\n- The proposed method on a high-level looks very much like Dreamer, especially given the RSSM world model, yet did not ablate the core difference against Dreamer which is the local search. How important is the search compared to backprop through model in Dreamer?\n- Missing an explicit limitation discussion section. \n- If possible, I'd like to see a comparison on data/compute-efficiency against GPS (guided policy search).", "was_revised": true, "final_justification": "The responses to my questions regarding comparison to GPS (guided policy search) and Dreamer are fair. While I agree on that, it does not change the fact that despite not requiring ground-truth reward / dynamics, the proposed method only evaluates on simulation environments that do provide them, rather than on real-robots where the method is motivated on (not requiring ground-truth reward and dynamics). Therefore I am not fully convinced as the paper does not closed the \"motivation-method-experiment\" loop. However, given the paper's clarity and thorough experiments, and that the authors have addressed my other concerns, I want to raise my score to 4 at this point. I still strongly advise the authors to include real-robot experiments in the final version of the paper, if the paper were accepted.", "rebuttal": "Thank you for your valuable comments that helped improve the paper. We are delighted you found our analysis detailed and appreciated the key findings of SAILOR in our ablations. We respond to your questions below:\n\n> W1: My main concern is the evaluation setup. The main metrics measured is success rate yet the core idea is to improve BC. Success rate, especially in simulators, can be easily improved using RL without using world model or reward model.\n\nWe would like to clarify that we do not assume access to the ground truth reward function (e.g., success rate), which means we cannot directly use RL, as there is no reward function to run it with. More generally, **specifying reward functions (i.e., solving a *reward design problem*, [1]) is often quite challenging in practice**. For example, in domains like self-driving, it is hard to hand-engineer a reward function that captures what good driving is -- one would need to write down the precise trade-offs between getting an inch closer to another car and arriving at the destination a few minutes earlier. Instead, in inverse RL techniques like SAILOR, one *infers* such a reward function from demonstrations of good behavior, side-stepping this problem. \n\n> W2: Related to above, please also consider comparing to other closed-loop IL methods, such as GAIL (that also trains the reward model on-the-fly).\n\nOur DPPO-IRL baseline is a stronger version of traditional GAIL that optimizes over diffusion policies. The policy is optimized via (D)PPO as in GAIL. The discriminator/reward model is trained iteratively as in GAIL. We show that across multiple tasks, SAILOR is significantly more interaction-efficient than DPPO-IRL.\n\n> W3: The proposed method on a high-level looks very much like Dreamer, especially given the RSSM world model, yet did not ablate the core difference against Dreamer which is the local search. How important is the search compared to backprop through model in Dreamer?\n\nWe would like to again clarify that we do not assume access to a ground-truth reward function and therefore cannot run vanilla Dreamer, as there is no reward signal. That said, once we have inferred a reward model via training a classifier between expert demonstrations and learner rollouts as in inverse RL (e.g. GAIL), it is an interesting question whether residual search using the learned reward/world model performs better than learning a *residual policy* via backpropping through the learned world model -- we thank the reviewer for suggesting this ablation! More explicitly, the residual actor network is tasked with predicting a corrective action, taking as input the latent state and base policy action. We trained this policy in standard DREAMER fashion (i.e., via backprop through the world model). Our experiments in the table below (averaged with 3 seeds per task) show that **SAILOR with test-time planning is better than learning a residual actor**. We have added these results to our paper. We have added these results to the paper, where this study demonstrates the importance of test-time planning.\n\n| | Lift (5 Demos) | Can (10 Demos) | Square (50 Demos) |\n|----------|----------|----------|----------|\n| DP |.58    | .33    | .28    |\n| SAILOR (Residual Policy) |.85    | .61    | .54    |\n| SAILOR|.99    | .82    | .78    |\n\n> W4: Missing an explicit limitation discussion section.\n\nWe agree with the reviewer that adding an explicit discussion and limitation section would further strengthen our manuscript. We have added the following and are happy to move it to the additional page for camera-ready or the appendix:\n\n### Limitations\nAt a high level, we believe there are two key limitations in our work that we believe offer promising opportunities for future work. The first is to perform experiments on real robots -- recent work FOREWARN [2] has shown that similar ideas to SAILOR can be scaled to real robots via the combination of diffusion policies, the Dreamer WM architecture, using VLMs as reward functions, and test-time planning. We believe that adding in other components of the SAILOR architecture (e.g., the use of a residual search rather than mode selection as in FOREWARN) could further improve performance. Second, we currently train the world model from scratch. Instead, leveraging a pre-trained foundation world model like Dino WM or V-JEPA could further improve sample efficiency and potentially eliminate the need for a warm start.\n\n### Discussions\nIn summary, we see that across a dozen visual manipulation tasks, SAILOR is able to train robust agents that recover from the failures of the base policy via test-time planning and match the expert’s outcomes, all without any additional human data. Furthermore, scaling up the amount of data used to train a diffusion policy via behavioral cloning by ≈ 5-10× is unable to close this performance gap. We also see that our learned RM is able to detect nuanced failures, while the SAILOR agent is able to recover from them. As mentioned above, while for our experiments we use a base diffusion policy, MPPI planner, and Dreamer world model, the general SAILOR architecture is fully compatible with foundation behavior (e.g., VLAs) and world models trained on internet-scale data, which might unlock further levels of generalization and robustness – we leave this as a promising future direction.\n\n> If possible, I'd like to see a comparison on data/compute-efficiency against GPS (guided policy search).\n\nIf by GPS, the reviewer is referring to [3] (an excellent paper we are huge fan of), we would like to note both that (a) we cannot directly apply the original GPS from image observations and (b) SAILOR can already be thought of as a generalization of many of the ideas in GPS. Ignoring some of the lower-level details, GPS essentially boils down to using an iLQR/DDP planner as a local expert that gets distilled into a policy via expert iteration. This requires access to a low-dimensional state to apply iLQR/DDP (we operate from images) and assumes dynamics are well-approximated by piecewise linear functions (we make no such assumptions, which likely do not hold in image-space). However, via its expert iteration subroutine, SAILOR periodically distills the output of the search process into the base policy. In this sense, **SAILOR can be seen as a modern generalization of GPS that uses Dreamer to handle image observations, MPPI rather than iLQR / DDP to handle non-smooth dynamics, and does not require access to a ground-truth reward function**. We have updated our related work section to make this point more explicit.\n\n> Q1: I’d be interested to learn a variant of figure 6 where the x-axis is the total amount of compute used (EFLOPs).\n\nWe thank the reviewer for the suggestion. We have computed the total amount of compute for our model in GFLOPs and have included these numbers in the revised plot in our paper.\n\n**Base Policy:** 14.554 GFLOPs\\\n**MPPI Planner:** 0.1637 x `num_itrs` x `num_samples` GFLOPs\n\nWe would like to point out that since the MPPI planner can process all samples in parallel, the inference time only depends on `num_itrs` and is 0.014s per MPPI iteration (and 0.066s for the base diffusion policy). In the paper, we compute numbers with 256 samples and 6 iterations, but we observed that there are marginal gains after 2-3 iterations. \n\n> Q2: How does the proposed approach work on a real robot?\n\nIn this work, we conduct experiments in simulation. FOREWARN [2] deployed DP with test-time planning with Dreamer world model and VLM on a real robot, and we believe similar setups can be utilized to deploy SAILOR on real hardware. We thank the reviewer for pointing this out and have added this in the limitations of our work.\n\nWe hope our responses and proposed changes have resonated with you, and you consider increasing your score. \n\nBest regards,\\\nThe Authors\n\n#### References\n[1] Hadfield-Menell, Inverse Reward Design, NeurIPS 2025.\\\n[2]  Wu et al., From Foresight to Forethought: VLM-In-the-Loop Policy Steering via Latent Alignment, RSS 2025.\\\n[3] Levine et al., Guided Policy Search, ICML 2013."}, {"review_id": "uqD4i0ld2p", "reviewer": "Reviewer_FQJB", "summary": "The paper proposes SAILOR, an imitation learning approach that aims at overcoming the usual out-of-distribution errors of behavior cloning by \"learning to search\" on top of learned world and reward models. The idea is to train these models on a mixture of expert data and data generated from a base policy (trained via BC). At inference time, SAILOR performs residual planning via MPPI to refine the plans proposed by the base policy, hence enabling eg recovering from failures. The method is tested on top of diffusion policies (DPs) and shown to perform better than plain BC (on top of DPs) and then a model-free IRL method (DPPO-IRL) on 12 complex visual manipulation tasks.", "questions": "Please address limitation above", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. The paper considers a relevant and well-motivated problem: BC algorithms are ubiquitous in practice (especially now in the era of LLMs/VLMs/VLAs) and their out-of-distribution issues are well known for sequential decision-making problems\n2. The proposed algorithm is simple yet quite powerful. I particularly like the idea of including world modeling and planning to recover from mistakes and improve performance\n3. The paper is extremely well written and easy to read. Pseudo-codes are clear and self-contained, while the figures do a great job in supporting the presentation\n4. Existing literature is extensively discussed\n5. Experiments are done on sufficiently complex tasks (12 visual manipulation problems from different suites). I particularly appreciate that the models take as input both proprioceptive and visual observations, which reflects real-world usecases\n\nWeaknesses:\n1. My main concern is about the experimental evaluation, as I found the results not fully convincing to support the main claims. In particular:\n   -  \"Can SAILOR outperform DP trained on the same D?\" Unless I am missing something, SAILOR and DP are not trained on the same data: both use the expert demonstrations D, but SAILOR is additionally generating online rollouts in the environment. If this is the case, this makes the comparison pretty unfair, and especially makes it hard to conclude whether the ability of SAILOR to recover from failures really stems from world modeling and planning (as claimed) or whether it is simply because those situations are generated by rollouts at train time, so SAILOR has evidence to recover that DP does not have\n   - \"Does DP catch up with more data?\" May have the same issue as above, as SAILOR always uses more data\n   - \"How much real-world interaction does the WM save us?\" Given the many differences between SAILOR and DPPO-IRL, it is hard to conclude if this is due to the WM or other algorithmic choices. For instance, DPPO-IRL uses PPO as the base optimizer, which is known to be very sample inefficient, while SAILOR is closer to an off-policy method as it reuses efficiently previous data. I wonder if other baselines can shed more light on this (e.g., can an off-policy diffusion-based like Diffusion Q-learning be adapted as an IRL variant?)\n2. The claim \"Our key insight is that we can infer the latent search process required to recover from local mistakes from the same source of human data required for the standard behavioral cloning pipeline\" in the introduction also seems to contradict a bit the functionality of the algorithm: as SAILOR can rollout data in the environment, how can we conclude that it can recover from mistakes from expert data only?\n3. Section 3.1 mentions that \"Beyond these benefits, the unique advantage of the L2S paradigm is a computational one: the ability to only need to plan at the states encountered at test-time\". I was actually thinking the opposite, as planning methods are known to be computationally much more inefficient at test-time than pre-trained policies (of course if you ignore pre-training time, but that should be eventually much smaller than the cost of doing multiple test-time evaluations). I think one advantage is instead that planning implicitly expands the space of possible behaviors that can be generated (as opposed to a pre-trained policy which is limited to the capacity of a fixed neural net)\n4. The paper mentions data SAILOR can be applied on top of any base policy, but then the method is tested only on top of DPs. To really show this generality, I'd be better to add experiments on other base policies (eg recent VLAs as mentioned in the text)\n\nOverall I think this is a very promising paper which I really enjoyed reading, though I have to vote for rejection given the limitations above. I am more than happy to increase my score if I am missing something or if additional experiments are provided to reinforce the main conclusions.", "was_revised": true, "final_justification": "Reasons why I have increased my score from 3 to 4:\n- The authors provided additional experiments that make the main conclusions convincing\n- The authors also provided experiments to assess the scalability of the approach in terms of planning times\n\nReasons not to increase further:\n- Testing only on top of diffusion policies is not sufficient to assess the generality and broad applicability of the proposed method\n- I believe the fact that SAILOR has online access to the environment is the key factor to enable eg error recovery and makes the contribution of the main algorithmic components (eg WM + planning) milder. For instance, I don't expect those components to work that well if trained entirely on expert data (no online samples)", "rebuttal": "Thank you for your feedback, which helped strengthen our work. We are glad you enjoyed the simplicity of our approach and the ability of SAILOR to achieve good performance on a complex set of tasks having visual and proprioceptive inputs, a setting that reflects real-world use cases. We respond to your questions below.\n\n> W1 (a): \"Can SAILOR outperform DP trained on the same D?\" Unless I am missing something, SAILOR and DP are not trained on the same data: both use the expert demonstrations D, but SAILOR is additionally generating online rollouts in the environment. If this is the case, this makes the comparison pretty unfair, and especially makes it hard to conclude whether the ability of SAILOR to recover from failures really stems from world modeling and planning (as claimed) or whether it is simply because those situations are generated by rollouts at train time, so SAILOR has evidence to recover that DP does not have\n\n> W1 (b): \"Does DP catch up with more data?\" May have the same issue as above, as SAILOR always uses more data\n\nWe believe that these two questions essentially boil down to asking whether it is fair to compare SAILOR to DP, given that the former gets access to online interaction while the latter does not. There are two points worth mentioning here. The first is that all methods (interactive and not) get access to the same set of expert demonstrations – it is standard practice in imitation learning to compare methods based on how they perform based on the same set of expert demonstrations, even when some are interactive (e.g., in the GAIL paper [1]). More generally, perhaps, the entire field of inverse RL (which includes SAILOR) focuses on interactively learning better policies from the same set of expert demonstrations that an offline algorithm like behavioral cloning uses. The second is that, given we do not assume access to any ground-truth reward information, online data provides no intrinsic information about how to recover. More explicitly, if we simply see a rollout and have no way to judge how good it is (i.e., whether it “recovers” or whether it even made a mistake in the first place), it is not clear how to use it to improve the performance of the policy. For example, if we know that we dropped an object but not that dropping an object is bad, it is unclear what to do with this information. It is only because we infer a reward model from demonstrations in SAILOR (and, more generally, in inverse RL) that we are able to use this on-policy data to learn how to recover. In short, (a) the comparison we make in this paper is extremely standard in the literature, and (b) merely providing an algorithm with access to online samples provides no intrinsic value.\n\nThat said, there is a related question we find to be quite interesting: **How much does planning corrective actions against a learned RM + WM improve performance over using a residual policy that does not perform a test-time search** (and therefore may only handle situations seen at train-time). Both of these options get access to interaction, but only the former is able to respond to new situations unseen at training time. To compare these two options, we added an additional baseline in which we replaced the residual search process in SAILOR with a *residual policy*. More explicitly, the residual actor network is tasked to predict a corrective action, taking the latent state and base policy action as input. We train this residual actor via Dreamer against the learned reward model. Experiments were conducted on Robomimic tasks, and the numbers are reported in the table below (averaged with three seeds per task). We observe that SAILOR -- with test-time planning -- is better than learning a residual policy with backprop. We have added these results to the paper, demonstrating the importance of test-time planning.\n\n| | Lift (5 Demos) | Can (10 Demos) | Square (50 Demos) |\n|----------|----------|----------|----------|\n| DP |.58    | .33    | .28    |\n| WM + RM + Residual Policy|.85    | .61    | .54    |\n| SAILOR|.99    | .82    | .78    |\n\n\n> W1 (c): \"How much real-world interaction does the WM save us?\" Given the many differences between SAILOR and DPPO-IRL, it is hard to conclude if this is due to the WM or other algorithmic choices. For instance, DPPO-IRL uses PPO as the base optimizer, which is known to be very sample inefficient, while SAILOR is closer to an off-policy method as it reuses efficiently previous data. I wonder if other baselines can shed more light on this (e.g., can an off-policy diffusion-based approach like Diffusion Q-learning be adapted as an IRL variant?)\n\nWe agree that adding an additional baseline where we use an off-policy RL algorithm as the policy optimizer for our model-free IRL baseline would strengthen our paper. However, we were not able to find a method that (a) trains diffusion policies, (b) handles image observations, and (c) has publicly released code, which made it a challenge to implement such a baseline within the rebuttal time frame. For example, none of the other algorithm implementations in the DPPO codebase work from images other than DPPO, and the off-policy RLPD codebase cannot directly train a diffusion model. We also extensively searched for an implementation of DQL and IDQL that worked from image observations and could not find one. We are happy to add this as a point to our limitations section.\n\n> W2: The claim \"Our key insight is that we can infer the latent search process required to recover from local mistakes from the same source of human data required for the standard behavioral cloning pipeline\" in the introduction also seems to contradict a bit the functionality of the algorithm: as SAILOR can rollout data in the environment, how can we conclude that it can recover from mistakes from expert data only?\n\nWe thank the reviewer for highlighting this. In our original statement, we wanted to emphasize that SAILOR does not require any additional interaction/supervision from the *human* expert, just like behavioral cloning. This is in contrast to a method like DAgger [2]. We agree that saying we do not need any more \"human data\" is not as clear as we could have been, and have rephrased this in the paper.\n\n> W3: Section 3.1 mentions that \"Beyond these benefits, the unique advantage of the L2S paradigm is a computational one: the ability to only need to plan at the states encountered at test-time\". I was actually thinking the opposite, as planning methods are known to be computationally much more inefficient at test-time than pre-trained policies. I think one advantage is instead that planning implicitly expands the space of possible behaviors that can be generated (as opposed to a pre-trained policy which is limited to the capacity of a fixed neural net)\n\nWe agree with the reviewer that in situations where computation is highly limited at inference time, it may not be possible to do multiple test-time evaluations. However, we found that for our SAILOR implementation, we only spent ~0.014s time per iteration on the MPPI residual search, compared to ~0.066s on the DP forward pass, a somewhat minimal overhead. This is likely because the MPPI planner operates on top of lower-dimensional *latent states*, rather than high-dimensional image observations like the base DP. We used the default parameters, which include 6 iterations of planning, and we empirically observed that the performance converges after 2-3 iterations with a marginal improvement in later iterations (as seen from the table below). We thank the reviewer for bringing this up, and have added this to the paper\n\n| Method                 | Lift (5 Demos) | Can (10 Demos) | Square (50 Demos)  |\n|-----|-----|----|----|\n|BaseDP|0.58|0.33|0.28|\n|1 Itr|0.99|0.78|0.61|\n|3 Itrs|0.99|0.83|0.72|\n|6 Itrs (SAILOR)|0.99|0.83|0.78|\n|9 Itrs|0.98|0.87|0.77|\n\nThat said, in terms of total compute used for policy updates across training and test times (not just at training time), we believe that L2S methods are likely more efficient than traditional policies. Intuitively, this is because for a policy to be able to respond well to a scenario, it likely needs to have been trained on something similar at training time. In the worst case, we need to handle *all* such scenarios at train time, while an L2S method only needs to handle what is seen at test time. A more formal version of this argument can be found in the classic paper https://www.cis.upenn.edu/~mkearns/papers/sparsesampling-journal.pdf. Furthermore, **because we iteratively distill the output of the search process back into the base policy via the expert iteration subroutine, SAILOR avoids re-solving the same residual search problem in the future, essentially recycling the test-time compute**.\nAlso, our ablation on learning a residual actor to predict actions shows that test-time planning outperforms learning with a residual actor-- validating the limited capacity with a fixed neural net as policy. \n\n> W4: The paper mentions data SAILOR can be applied on top of any base policy, but then the method is tested only on top of DPs. To really show this generality, I'd be better to add experiments on other base policies (eg recent VLAs as mentioned in the text)\n\nWe agree that this is an interesting and promising direction for future work! In particular, we'd be interested in seeing how SAILOR could be applied to base VLA policies, and, similar to VLAs, can leverage text information as input while planning.  \n\nWe hope our response has addressed your concerns and resonated with you. We're happy to answer any further questions you may have and hope you consider increasing your score.\n\n\nBest Regards,\\\nThe Authors\n\n#### References\n[1] Ho et al., Generative Adversarial Imitation Learning, NeurIPS 2016.\\\n[2] Ross et al., A Reduction of Imitation Learning and Structured Prediction to No-Regret Online Learning, AISTATS 2011."}, {"review_id": "WXVA2aJ7yT", "reviewer": "Reviewer_kASh", "summary": "The paper proposes a novel architecture and training scheme for learning robust policies from expert demonstrations. In particular, the paper learns a latent-space (“world”) model to simulate rollouts, a reward function to quantify the quality of latent states, and a critic to estimate the residual value of truncated latent-state trajectories. These three components can be used to plan the next action (via MPPI in the paper). Additionally the paper uses a base policy (Diffusion based) for more efficient search, and to distill spent test-time compute spent on searching good actions back into the base policy. The paper demonstrates that this leads to good performance across a range of simulated robotic manipulation tasks with relatively few expert demonstrations. The method performs significantly better than naive behavioral cloning, even if the available number of demonstrations for the latter is increased by a factor of 5 to 10. The main premise is that the method addresses one of the main failure cases of behavioral cloning, which is going off distribution, by having the ability to recover by planning via the learned components (“learning to search”).", "questions": "I have no major concerns or questions (please see my comments on “main weaknesses” above, which I do not consider to be major). Here are two very minor comments:\n\n1. The second fundamental problem with expert imitation on partially observable decision-making tasks (besides having no off-distribution coverage) is that models need to learn to update their belief state differently when it is the result of their own action, compared to observing an expert action. Technically, the model’s own actions must be treated as interventions, whereas expert actions can be treated via standard probabilistic conditioning. In BC, models cannot learn this (as all actions are conditioning), and are thus likely to suffer from self-delusions (see: Shaking the foundations: delusions in sequence models for interaction and control, Ortega et al. 2021). When actively interacting (whether model based or not) models naturally learn the correct update mechanisms. The proposed method in the paper thus also seems to address this second, and quite subtle, fundamental problem with BC in partially observed interactive environments (where experts may hold unobserved “privileged” belief states).\n2. Though the paper does not introduce the term “Learning to Search”, I think it is not ideal, as it suggests imitating an expert policy where search emerges purely from learning (e.g., BC). Rather, the method in the paper is explicitly programmed for “Searching with Learned Models”. No strict need to change this in the paper, but maybe it’s still time to stop the proliferation of a misleading term.", "limitations": "The paper checklist for the discussion of limitations is answered with “NA” (which corresponds to “The answer NA means that the paper has no limitation”). I don’t think that this is what the authors meant, and I would encourage them to add a brief discussion of limitations, particularly w.r.t. the generality of the current findings and scaling up the method. Alternatively, please add a justification for the answer to the paper checklist.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Main contributions**\n1. A non-trivial combination of previously proposed and used parts, techniques, losses, etc. The choice of many of these components is discussed and contrasted against alternatives, and making all these components play well together is far from being guaranteed.\n2. Good empirical results, particularly against naive BC (even at increased amounts of data) and a version of inverse RL (DPPO).\n3. Ablations for the main parts of the method, and some behavioral analysis of the ability to recover from off-distribution situations.\n4. The paper is generally well written and easy to follow (though, given the number of parts of the model certain sections must be dense to stay within the page limit).\n\n**Main weaknesses**\n1. Part of the writing (particularly early in the paper) sounds as if the combination of a learned world model, reward model, and a planner (with base policy) had never been used before. This is clarified in the related work section, which makes it more clear that this particular combination of these parts is novel (and I agree it is non-trivial), along with a number of losses / training schemes for parts to make everything work together. Since I am not an expert in robotics, I cannot confidently judge the novelty and originality (and have lowered my confidence accordingly), but I would be surprised if there are no comparable roughly similar approaches in the literature.\n2. Related to 1. I am wondering whether there are stronger baselines that one could compare against, particularly methods that have a similar aim (planning via learned model / reward function). I may be wrong, and leave it to the expert reviewers to point out any such methods. I do appreciate the naive BC baseline to support the main claims and the inverse RL baseline is also a good addition.\n3. The paper is dense and is missing a discussion and conclusion and a limitations section. While I understand that the page limitations require a trade-off, and am happy that the many parts and pieces are explained well and important ablations are shown in the main paper, not having these additional sections is a small weakness (I do not expect a change to the main paper; I am merely pointing out a small weak point).\n\n**Verdict:**\n\nBased on the empirical results, the non-trivial combination of a number of components and the protocol to make these components play well together, the ablations, and the good writing, the paper looks solid and ready for publication to me. As stated above, I am not an expert in robotics, so I cannot confidently judge the novelty, originality, and potential impact. I am therefore currently between a weak accept and an accept, leaning towards the latter. I am particularly uncertain whether the method can relatively simply be put together from the current literature and whether strong baselines are missing, but hopefully other reviewers will address this. Currently, I have no reasons to recommend rejecting the paper.", "was_revised": true, "final_justification": "I had no major issues and only minor clarifying questions/comments after the initial round of reviews, and the authors have nicely addressed all of them. I therefore continue to argue in favor of accepting the paper and have raised my confidence (still relatively low as I am no expert in robotics and do not precisely know the SOTA methodologies).", "rebuttal": "Thank you for your thoughtful review of our work. We are glad that you found the integration of multiple learned components compelling and appreciated the thoroughness of our experiments. We address your questions below: \n\n> W1: Since I am not an expert in robotics, I cannot confidently judge the novelty and originality (and have lowered my confidence accordingly), but I would be surprised if there are no comparable, roughly similar approaches in the literature.\n\nWe thank the reviewer for their honesty here. We agree that while subsets of using a learned latent-space world model, reward model, planner, and base policy have been explored in prior work (as referenced in our related works section), we struggled to find work in which *all* of these components were used together. Furthermore, in much of the prior work, some of these components are learned via auxiliary sources of information (e.g., ground-truth reward labels) or need access to lower-dimensional state representations. In contrast, **we believe one of the key novelties of our framework is that all of these components can be learned just from visual demonstrations / on-policy rollouts, with no reward labels / further human supervision required**.\n\n> W2: Related to 1. I am wondering whether there are stronger baselines that one could compare against, particularly methods that have a similar aim (planning via learned model/reward function). I may be wrong, and leave it to the expert reviewers to point out any such methods. I do appreciate the naive BC baseline to support the main claims, and the inverse RL baseline is also a good addition.\n\nWe agree that more baselines could only make our work stronger. As suggested by reviewer AkFH, we added in a baseline where, rather than learning a residual planner, we instead learned a residual *policy*, allowing us to ablate the benefit that explicit planning at test time provides. We trained this policy in standard DREAMER fashion (i.e., via backprop through the world model). Our experiments in the table below (averaged with 3 seeds per task) show that **SAILOR with test-time planning is better than learning a residual actor**. We have added these results to our paper.\n\nSecond, we'd like to make a note on the strength of our inverse RL baseline. The reason we used DPPO as the policy optimizer for our inverse RL baseline is that it has been shown to significantly outperform other RL algorithms for training diffusion policies (e.g., DQL, IDQL, QSM) and demonstrated strong performance with ground-truth rewards on the Robomimic tasks we consider [1].\n\n| | Lift | Can | Square |\n|----------|----------|----------|----------|\n| DP |.58    | .33    | .28    |\n| WM + RM + Residual Policy|.85    | .61    | .54    |\n| SAILOR|.99    | .82    | .78    |\n\n> W3: The paper is dense and is missing a discussion and conclusion, and a limitations section.\n\nWe agree with the reviewer that adding an explicit discussion and limitation section would further strengthen our manuscript. We are happy to move this to the additional page for camera-ready or to the appendix. In particular, we have added the following:\n\n### Limitations\nAt a high level, we believe there are two key limitations in our work that we believe offer promising opportunities for future work. The first is to perform experiments on real robots -- recent work FOREWARN [2] has shown that similar ideas to SAILOR can be scaled to real robots via the combination of diffusion policies, the Dreamer WM architecture, using VLMs as reward functions, and test-time planning. We believe that adding in other components of the SAILOR architecture (e.g., the use of a residual search rather than mode selection as in FOREWARN) could further improve performance. Second, we currently train the world model from scratch. Instead, leveraging a pre-trained foundation world model like Dino WM or V-JEPA could further improve sample efficiency and potentially eliminate the need for a warm start.\n\n### Discussions\nIn summary, we see that across a dozen visual manipulation tasks, SAILOR is able to train robust agents that recover from the failures of the base policy via test-time planning and match the expert’s outcomes, all without any additional human data. Furthermore, scaling up the amount of data used to train a diffusion policy via behavior cloning by ≈ 5-10× is unable to close this performance gap. We also see that our learned RM is able to detect nuanced failures, while the SAILOR agent is able to recover from them. As mentioned above, while for our experiments we use a base diffusion policy, MPPI planner, and Dreamer world model, the general SAILOR architecture is fully compatible with foundation behavior (e.g., VLAs) and world models trained on internet-scale data, which might unlock further levels of generalization and robustness – we leave this as a promising future direction.\n\n> Q1:  The second fundamental problem with expert imitation on partially observable decision-making tasks (besides having no off-distribution coverage) is that models need to learn to update their belief state differently when it is the result of their own action, compared to observing an expert action. Technically, the model’s own actions must be treated as interventions, whereas expert actions can be treated via standard probabilistic conditioning. In BC, models cannot learn this (as all actions are conditioning), and are thus likely to suffer from self-delusions (see: Shaking the foundations: delusions in sequence models for interaction and control, Ortega et al. 2021). When actively interacting (whether model based or not) models naturally learn the correct update mechanisms. The proposed method in the paper thus also seems to address this second, and quite subtle, fundamental problem with BC in partially observed interactive environments (where experts may hold unobserved “privileged” belief states).\n\nYes -- in fact, it *provably* does under some assumptions (see [3] for more)! In greater detail, **interactive approaches to imitation -- both DAgger and inverse RL -- are able to handle unobserved confounders (e.g., context variables that only the expert can see) as the learner actually *intervenes* in the environment by taking actions and observing their consequences**, rather than learning from off-policy / \"observational\" data. In contrast, at test time, a BC policy treats the learner's past actions as though they were from the *expert*, which leads to the \"self-delusions\" / compounding errors the reviewer mentions. We mentioned partial observability as one possible cause of compounding errors in our related works and have added Pedro's lovely paper to the corresponding section. \n\n> Q2: Though the paper does not introduce the term “Learning to Search”, I think it is not ideal, as it suggests imitating an expert policy where search emerges purely from learning (e.g., BC). Rather, the method in the paper is explicitly programmed for “Searching with Learned Models”. No strict need to change this in the paper, but maybe it’s still time to stop the proliferation of a misleading term.\n\nWe thank the reviewer for their comment. In some sense, we believe that the term \"learning to search\" is not entirely inaccurate, in the sense that we're learning the components required to execute a search algorithm at test-time. Furthermore, we believe that an interesting direction for future work could be to replace the fixed MPPI algorithm with a learned operator to combine the learned world and reward models, and that SAILOR serves as an important first step towards doing so. That said, we do agree that we could add more clarification about what we mean by the term in our work (i.e., using learned models as part of a fixed search algorithm) and would be happy to do so, say, in the introduction.\n\n> Limitations\n\nAs described above, we have added discussion and limitations sections, as well as updated the paper checklist accordingly.\n\nWe hope our responses and proposed changes have resonated with you and hope you consider increasing your score.\n\nBest Regards,\\\nThe Authors\n\n#### References\n[1] Ren et al., Diffusion Policy Policy Optimization, ICLR 2025.\\\n[2] Wu et al., From Foresight to Forethought: VLM-In-the-Loop Policy Steering via Latent Alignment, RSS 2025.\\\n[3] Swamy et al., Sequence Model Imitation Learning with Unobserved Contexts, NeurIPS 2022."}]

Author rebuttal:

Meta-review: The paper addresses the the key limitation of Behaviour Cloning: lack of ability to recover from unseen states by extending BC with three blocks (namely a world model, a reward model and a latent-space planner) and a training regime (including ideas from data aggregation and policy distillation) that shifts the paradigm from imitation learning towards learning to search. The architecture is also modular and agnostic to the policy architecture. Simulation results show that the proposed approach outperforms vanilla BC, even when BC is trained with 5-10 times the data, demonstrating that it can robustly recover from policy errors, and it is also more sample efficient than a model-free inverse-RL baseline. The paper is also clear and well written. 

Reviewers appreciated all these contributions and expressed concerns about the evaluation, applicability and assumptions. The authors addressed these comments with the rebuttal and during the discussion period, running experiments with a new baseline and clarifying multiple points satisfactorily. They also acknowledge the limitation of claiming the method agnostic on the policy while it has been trained on diffusion BC only, and agreed with the reviewers that trying with VLAs would be an interesting future work.

I recommend acceptance as a spotlight because the contributions are significant, the idea is novel and sound, and the paper is clear and well-written. Overall, it is a very promising direction in imitation learning that I believe will inspire further research.
Decision: Accept (spotlight)