IMLE Policy: Fast and Sample Efficient Visuomotor Policy Learning via Implicit Maximum Likelihood Estimation

TL;DR

IMLE Policy addresses the challenge of learning expressive, multi-modal visuomotor policies from limited demonstrations while enabling fast, single-step inference. By extending conditional RS-IMLE to behaviour cloning, it preserves mode coverage without mode collapse and uses rejection sampling to maintain data efficiency. Across eight simulation benchmarks and two real-world tasks, it achieves up to $38\%$ data savings and up to $97.3\%$ faster inference compared with diffusion-based baselines, while maintaining competitive performance in rich multi-modal settings. The work demonstrates that a simple generator-based architecture coupled with temporal consistency can deliver practical, robust robotic policies suitable for real-time deployment and open avenues for RL fine-tuning and MPC applications.

Abstract

Recent advances in imitation learning, particularly using generative modelling techniques like diffusion, have enabled policies to capture complex multi-modal action distributions. However, these methods often require large datasets and multiple inference steps for action generation, posing challenges in robotics where the cost for data collection is high and computation resources are limited. To address this, we introduce IMLE Policy, a novel behaviour cloning approach based on Implicit Maximum Likelihood Estimation (IMLE). IMLE Policy excels in low-data regimes, effectively learning from minimal demonstrations and requiring 38\% less data on average to match the performance of baseline methods in learning complex multi-modal behaviours. Its simple generator-based architecture enables single-step action generation, improving inference speed by 97.3\% compared to Diffusion Policy, while outperforming single-step Flow Matching. We validate our approach across diverse manipulation tasks in simulated and real-world environments, showcasing its ability to capture complex behaviours under data constraints. Videos and code are provided on our project page: https://imle-policy.github.io/.

Figures (9)

• Figure 1: Capturing Multi-Modality from Limited Demonstrations. We present IMLE Policy, a novel behaviour cloning approach that is capable of learning complex multi-modal action distributions from less data than prior approaches, while only requiring a single network forward pass. Top: When trained on 20 samples for 4000 epochs, IMLE Policy learns to accurately model the underlying multi-modal function, while Diffusion Policy with 100 de-noising steps struggles. Flow Matching with a single step collapses modes. Bottom: When trained on only 17 demonstrations for a shoe racking task, IMLE Policy maintains multi-modality, predicting a variety of trajectories from the starting position where either shoe could be selected first while the other methods predict poor trajectories.
• Figure 2: Qualitative Analysis of Mode Capturing Performance. Illustrating the trajectories generated by different methods (rows) for varying initial end-effector positions (columns) in the Push-T task. The goal is to push the T-shaped block from its initial position to align with the target (green). At the top of the T, the dataset exhibits multi-modal behaviour: in the centre, demonstrations are evenly split between pushing left and right, while towards the edges, the majority of demonstrations push towards the closer side, with only a few moving in the opposite direction. IMLE Policy successfully captures all modes without collapsing or biasing, maintaining trajectory diversity even in underrepresented modes. In contrast, Diffusion Policy tends to bias towards one mode when close to an edge, while Flow Matching collapses modes, producing averaged and unimodal trajectories.
• Figure 3: IMLE Policy Overview a) Training: The policy takes in a sequence of past observations $\mathcal{O}$ and $m$ sampled latents $\textbf{z}$ for which the policy generates $m$ sequences of predicted actions $\mathcal{A}$. Generated trajectories that lie within the rejection sampling threshold $\epsilon$ are rejected. From the remaining trajectories, the nearest-neighbour to the ground truth trajectory is selected for training. We minimise the the distance between this trajectory and the ground truth trajectory to optimise the policy. As the loss focuses on each data sample, it ensures that all modes are captured even from small datasets. b) When compared to baselines with similar multi-modal capturing capabilities, IMLE can generate actions with a single inference step as opposed to multi-step de-noising processes. c) For highly multi-modal tasks, we enhance the performance of IMLE Policy by introducing a simple inference procedure to induce consistency upon mode choice based on a nearest-neighbour search over batch-generated action proposals with the previously executed action sequence.
• Figure 4: Simulation Environments. Visualisation of the environments used to evaluate IMLE Policy. The first 5 tasks are part of the Robomimic benchmark robomimic and the other environments include Push-T chi2023diffusion, UR3 BlockPush kim2022automating and Franka Kitchen gupta2019relay.
• Figure 5: Dataset Size Study. We evaluate how each of the methods performs when trained with different size subsets of the full dataset used in the Push-T benchmark task. The dashed line indicates that to achieve the same reward of 0.5 on the task, IMLE Policy requires less than 29% of the data, while Diffusion Policy requires approximately 43% while Flow Matching Policy requires over 80% of the full dataset.