Search papers, labs, and topics across Lattice.
The paper introduces WIMLE, a model-based RL method that leverages Implicit Maximum Likelihood Estimation (IMLE) to learn stochastic, multi-modal world models and estimate predictive uncertainty. WIMLE weights synthetic transitions by their predicted confidence, mitigating bias from uncertain predictions and stabilizing learning. Experiments across 40 continuous control tasks show WIMLE achieves superior sample efficiency and competitive asymptotic performance compared to strong model-free and model-based baselines, particularly on challenging humanoid locomotion tasks.
By weighting model rollouts based on predicted confidence, WIMLE achieves state-of-the-art sample efficiency in model-based RL, outperforming existing methods on complex continuous control tasks.
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.
