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

TL;DR

WIMLE is a novel model-based reinforcement learning method that leverages IMLE for multi-modal world modeling and uncertainty estimation, leading to improved sample efficiency and performance in continuous control tasks.

Contribution

It introduces WIMLE, extending IMLE to learn stochastic, multi-modal world models with uncertainty-aware weighting, enhancing stability and efficiency in model-based RL.

Findings

WIMLE outperforms baselines in 40 continuous-control tasks.

On Humanoid-run, WIMLE improves sample efficiency by over 50%.

Solves 8 out of 14 tasks on HumanoidBench, surpassing competitors.

Abstract

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.