WIMLE: Uncertainty-Aware World Models with IMLE for Sample-Efficient Continuous Control
This work addresses sample efficiency and stability issues in continuous control tasks for reinforcement learning practitioners, representing a strong specific gain rather than a foundational breakthrough.
The paper tackled the problem of sample inefficiency and performance limitations in model-based reinforcement learning due to compounding model error and overconfident predictions by introducing WIMLE, which uses IMLE to learn stochastic, multi-modal world models with uncertainty estimation, achieving over 50% sample efficiency improvement on Humanoid-run and solving 8 out of 14 tasks on HumanoidBench.
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.