# WIMLE: Uncertainty-Aware World Models with IMLE for Sample-Efficient Continuous Control

arXiv: https://arxiv.org/abs/2602.14351 · HTML: https://mlsanity.com/paper/2602.14351/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Placed in Latent dynamics models and model-based RL, one of 11 directions in World models.
- Builds on 49 of its 59 references that are in the map.
- First posted on arXiv 15 Feb 2026.

## Placed in

- World models / [Latent dynamics models and model-based RL](https://mlsanity.com/wm/latent-dynamics-models-and-model-based-rl/)
- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

## Builds on (most cited first, 30 of 49)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Generative Adversarial Networks](https://mlsanity.com/paper/2203.00667/) (2022)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Elucidating the Design Space of Diffusion-Based Generative Models](https://mlsanity.com/paper/2206.00364/) (2022)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [Planning with Diffusion for Flexible Behavior Synthesis](https://mlsanity.com/paper/2205.09991/) (2022)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [Deep Reinforcement Learning at the Edge of the Statistical Precipice](https://mlsanity.com/paper/2108.13264/) (2021)
- [Rainbow: Combining Improvements in Deep Reinforcement Learning](https://mlsanity.com/paper/1710.02298/) (2017)
- [Recurrent World Models Facilitate Policy Evolution](https://mlsanity.com/paper/1809.01999/) (2018)
- [When to Trust Your Model: Model-Based Policy Optimization](https://mlsanity.com/paper/1906.08253/) (2019)
- [A Distributional Perspective on Reinforcement Learning](https://mlsanity.com/paper/1707.06887/) (2017)
- [CURL: Contrastive Unsupervised Representations for Reinforcement Learning](https://mlsanity.com/paper/2004.04136/) (2020)
- [Is Conditional Generative Modeling all you need for Decision-Making?](https://mlsanity.com/paper/2211.15657/) (2022)
- [TD-MPC2: Scalable, Robust World Models for Continuous Control](https://mlsanity.com/paper/2310.16828/) (2023)
- [Distributional Reinforcement Learning with Quantile Regression](https://mlsanity.com/paper/1710.10044/) (2017)
- [Temporal Difference Learning for Model Predictive Control](https://mlsanity.com/paper/2203.04955/) (2022)
- [Implicit Quantile Networks for Distributional Reinforcement Learning](https://mlsanity.com/paper/1806.06923/) (2018)
- [Transformers are Sample-Efficient World Models](https://mlsanity.com/paper/2209.00588/) (2022)
- [Mastering Atari Games with Limited Data](https://mlsanity.com/paper/2111.00210/) (2021)
- [Model-Ensemble Trust-Region Policy Optimization](https://mlsanity.com/paper/1802.10592/) (2018)
- [Transformer-based World Models Are Happy With 100k Interactions](https://mlsanity.com/paper/2303.07109/) (2023)
- [Model-Based Offline Planning](https://mlsanity.com/paper/2008.05556/) (2020)
- [Diffusion Model is an Effective Planner and Data Synthesizer for Multi-Task Reinforcement Learning](https://mlsanity.com/paper/2305.18459/) (2023)
- [Plan Online, Learn Offline: Efficient Learning and Exploration via Model-Based Control](https://mlsanity.com/paper/1811.01848/) (2018)
- [HumanoidBench: Simulated Humanoid Benchmark for Whole-Body Locomotion and Manipulation](https://mlsanity.com/paper/2403.10506/) (2024)
- [Bigger, Regularized, Optimistic: scaling for compute and sample-efficient continuous control](https://mlsanity.com/paper/2405.16158/) (2024)

## Built on by (most cited first, 3 of 3)

- [Transferable Delay-Aware Reinforcement Learning via Implicit Causal Graph Modeling](https://mlsanity.com/paper/2605.12312/) (2026)
- [Beyond One-Size-Fits-All: Diagnosis-Driven Online Reinforcement Learning with Offline Priors](https://mlsanity.com/paper/2606.25527/) (2026)
- [IMLE-VLA: Fast Single-Step Action Generation for Vision-Language-Action Policies](https://mlsanity.com/paper/2609.10915/) (2026)

## Abstract (arXiv)

Model-based reinforcement learning promises strong sample efficiency but often underperforms in practice due to compounding model error, unimodal world models that average over multi-modal dynamics, and overconfident predictions that bias learning. We introduce WIMLE, a model-based method that extends Implicit Maximum Likelihood Estimation (IMLE) to the model-based RL framework to learn stochastic, multi-modal world models without iterative sampling and to estimate predictive uncertainty via ensembles and latent sampling. During training, WIMLE weights each synthetic transition by its predicted confidence, preserving useful model rollouts while attenuating bias from uncertain predictions and enabling stable learning. Across 40 continuous-control tasks spanning DeepMind Control, MyoSuite, and HumanoidBench, WIMLE achieves superior sample efficiency and competitive or better asymptotic performance than strong model-free and model-based baselines. Notably, on the challenging Humanoid-run task, WIMLE improves sample efficiency by over 50% relative to the strongest competitor, and on HumanoidBench it solves 8 of 14 tasks (versus 4 for BRO and 5 for SimbaV2). These results highlight the value of IMLE-based multi-modality and uncertainty-aware weighting for stable model-based RL.