Title: A Smooth Sea Never Made a Skilled SAILOR: Robust Imitation via Learning to Search
Venue: NeurIPS 2025
Accepted: yes
Keywords: Imitation learning

world models

inverse reinforcement learning
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).


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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”).\n\nSTRENGTHS AND WEAKNESSES:\n**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.\n\nQUESTIONS:\nI 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.\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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.\n\nQUESTIONS:\nPlease address limitation above\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-01"}

comments: 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).

STRENGTHS AND WEAKNESSES:
### Strengths
+ The paper is well-motivated and presented
+ Detailed ablation analysis and I appreciate the clarity in disclosing the key findings in tuning the proposed method (section 4.2)

### Weaknesses
- 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.
- 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).
- 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?
- Missing an explicit limitation discussion section. 
- If possible, I'd like to see a comparison on data/compute-efficiency against GPS (guided policy search).

QUESTIONS:
- I’d be interested to learn a variant of figure 6 where the x-axis is the total amount of compute used (EFLOPs).
- 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.
TITLE: Official Review
CLARITY: 3
IMPACT: 2
ORIGINALITY: 3
DATE: 2025-07-01

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\n- No real robot experiments.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nSAILOR 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.\n\nSTRENGTHS AND WEAKNESSES:\n# 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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nNo 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-07-03"}