mlsanityData to 14 Oct 2026

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

arXiv 2602.14351 ↗First posted 15 Feb 2026cs.LG cs.AIcited by 3

On the map: World models → Latent dynamics models and model-based RL

also in 1 more
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.

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.

Builds on

49 of 59 references.

Closest ancestors · 11

Same direction, newest first

  1. 2023TD-MPC2: Scalable, Robust World Models for Continuous Control25 Oct
  2. 2023STORM: Efficient Stochastic Transformer based World Models for Reinforcement Learning14 Oct
  3. 2023Transformer-based World Models Are Happy With 100k Interactions13 Mar
  4. 2023Mastering Diverse Domains through World Models10 Jan
  5. 2022Transformers are Sample-Efficient World Models1 Sep
  6. 2021Mastering Atari Games with Limited Data30 Oct
  7. 2020Mastering Atari with Discrete World Models5 Oct
  8. 2019Dream to Control: Learning Behaviors by Latent Imagination3 Dec
  9. 2019When to Trust Your Model: Model-Based Policy Optimization19 Jun
  10. 2018Recurrent World Models Facilitate Policy Evolution4 Sep
  11. 2018Model-Ensemble Trust-Region Policy Optimization28 Feb

Other references · 38

  1. 2017Proximal Policy Optimization Algorithmscited by 15,215
  2. 2022Generative Adversarial Networkscited by 5,627
  3. 2018Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actorcited by 4,885
  4. 2022Elucidating the Design Space of Diffusion-Based Generative Modelscited by 3,045
  5. 2018Learning Latent Dynamics for Planning from Pixelscited by 1,264
  6. 2022Planning with Diffusion for Flexible Behavior Synthesiscited by 1,061
  7. 2018DeepMind Control Suitecited by 956
  8. 2021Deep Reinforcement Learning at the Edge of the Statistical Precipicecited by 807
  9. 2017Rainbow: Combining Improvements in Deep Reinforcement Learningcited by 786
  10. 2017A Distributional Perspective on Reinforcement Learningcited by 765
  11. 2020CURL: Contrastive Unsupervised Representations for Reinforcement Learningcited by 760
  12. 2022Is Conditional Generative Modeling all you need for Decision-Making?cited by 516
26 more
  1. 2017Distributional Reinforcement Learning with Quantile Regressioncited by 474
  2. 2022Temporal Difference Learning for Model Predictive Controlcited by 397
  3. 2018Implicit Quantile Networks for Distributional Reinforcement Learningcited by 369
  4. 2020Model-Based Offline Planningcited by 129
  5. 2023Diffusion Model is an Effective Planner and Data Synthesizer for Multi-Task Reinforcement Learningcited by 115
  6. 2018Plan Online, Learn Offline: Efficient Learning and Exploration via Model-Based Controlcited by 113
  7. 2024HumanoidBench: Simulated Humanoid Benchmark for Whole-Body Locomotion and Manipulationcited by 112
  8. 2024Bigger, Regularized, Optimistic: scaling for compute and sample-efficient continuous controlcited by 108
  9. 2018Lipschitz Continuity in Model-based Reinforcement Learningcited by 91
  10. 2022MyoSuite -- A contact-rich simulation suite for musculoskeletal motor controlcited by 80
  11. 2024SimBa: Simplicity Bias for Scaling Up Parameters in Deep Reinforcement Learningcited by 76
  12. 2020Model-Augmented Actor-Critic: Backpropagating through Pathscited by 63
  13. 2019Combating the Compounding-Error Problem with a Multi-step Modelcited by 54
  14. 2018Implicit Maximum Likelihood Estimationcited by 52
  15. 2016Self-Correcting Models for Model-Based Reinforcement Learningcited by 43
  16. 2024Rethinking Aleatoric and Epistemic Uncertaintycited by 42
  17. 2025Towards General-Purpose Model-Free Reinforcement Learningcited by 41
  18. 2019Learning to Combat Compounding-Error in Model-Based Reinforcement Learningcited by 39
  19. 2021Survey on reinforcement learning for language processingcited by 31
  20. 2016Playing Atari Games with Deep Reinforcement Learning and Human Checkpoint Replaycited by 22
  21. 2020Learning Accurate Long-term Dynamics for Model-based Reinforcement Learningcited by 21
  22. 2025On Rollouts in Model-Based Reinforcement Learningcited by 10
  23. 2018Towards a Simple Approach to Multi-step Model-based Reinforcement Learningcited by 9
  24. 2024Diffusion Models as Optimizers for Efficient Planning in Offline RLcited by 8
  25. 2024Rejection Sampling IMLE: Designing Priors for Better Few-Shot Image Synthesiscited by 8
  26. 2025Learning with Imperfect Models: When Multi-step Prediction Mitigates Compounding Errorcited by 5

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2026: 3.

3 papers.

  1. 2026Transferable Delay-Aware Reinforcement Learning via Implicit Causal Graph Modelingcited by 0
  2. 2026Beyond One-Size-Fits-All: Diagnosis-Driven Online Reinforcement Learning with Offline Priorscited by 0
  3. 2026IMLE-VLA: Fast Single-Step Action Generation for Vision-Language-Action Policiescited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2025Uncertainty-Aware Robotic World Model Makes Offline Model-Based Reinforcement Learning Work on Real Robots
  2. 2026Scaling Tasks, Not Samples: Mastering Humanoid Control through Multi-Task Model-Based Reinforcement Learning
  3. 2026Optimistic World Models: Efficient Exploration in Model-Based Deep Reinforcement Learning
  4. 2022On Uncertainty in Deep State Space Models for Model-Based Reinforcement Learning
  5. 2026EfficientTDMPC: Improved MPC Objectives for Sample-Efficient Continuous Control
  6. 2018Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Models
  7. 2025Simulus: Combining Improvements in Sample-Efficient World Model Agents
  8. 2025Accelerating Model-Based Reinforcement Learning with State-Space World Models