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

📄 arXiv: 2602.14351 📥 PDF

作者: Mehran Aghabozorgi, Alireza Moazeni, Yanshu Zhang, Ke Li

分类: cs.LG, cs.AI

发布日期: 2026-04-07

💡 一句话要点

WIMLE：基于不确定性感知世界模型与IMLE的样本高效连续控制

🎯 匹配领域: 支柱一：机器人控制 (Robot Control) 支柱二：RL算法与架构 (RL & Architecture)

关键词: 模型强化学习 隐式最大似然估计 不确定性估计 多模态学习 连续控制

📋 核心要点

  1. 模型强化学习面临模型误差累积、世界模型无法捕捉多模态动态以及过度自信预测导致学习偏差等挑战。
  2. WIMLE通过将隐式最大似然估计（IMLE）引入模型强化学习，学习多模态世界模型并估计预测不确定性。
  3. 实验表明，WIMLE在多个连续控制任务中实现了优越的样本效率和性能，尤其在复杂任务中提升显著。

📝 摘要（中文）

本文提出了一种名为WIMLE的模型强化学习方法，旨在解决模型误差累积、世界模型单峰性以及过度自信预测等问题。WIMLE将隐式最大似然估计（IMLE）扩展到模型强化学习框架中，无需迭代采样即可学习随机、多模态的世界模型，并通过集成和潜在采样来估计预测不确定性。在训练过程中，WIMLE根据预测置信度对每个合成转移进行加权，保留有用的模型rollout，同时减轻不确定预测带来的偏差，从而实现稳定的学习。在DeepMind Control、MyoSuite和HumanoidBench等40个连续控制任务中，WIMLE实现了优于现有模型无关和基于模型的基线的样本效率和具有竞争力的渐近性能。尤其是在具有挑战性的Humanoid-run任务上，WIMLE相对于最强的竞争对手提高了50%以上的样本效率，并且在HumanoidBench上解决了14个任务中的8个（BRO为4个，SimbaV2为5个）。这些结果突出了基于IMLE的多模态和不确定性感知加权对于稳定模型强化学习的价值。

🔬 方法详解

问题定义：模型强化学习旨在通过学习环境模型来提高样本效率，但实际应用中常因模型误差累积、无法捕捉环境动态的多模态特性以及模型预测的过度自信而表现不佳。现有的世界模型往往是单峰的，无法准确表示复杂环境中的多种可能性，导致策略学习受限。

核心思路：WIMLE的核心思路是利用隐式最大似然估计（IMLE）来学习随机、多模态的世界模型，并结合不确定性估计来指导模型训练。通过IMLE，模型能够更好地拟合真实环境的复杂动态，而无需进行迭代采样。同时，利用集成方法和潜在采样来估计模型预测的不确定性，从而避免过度自信的预测对策略学习产生负面影响。

技术框架：WIMLE的整体框架包括以下几个主要模块：1）环境交互模块，用于收集真实环境中的数据；2）世界模型学习模块，使用IMLE学习环境的动态模型；3）不确定性估计模块，通过集成和潜在采样来估计模型预测的不确定性；4）策略优化模块，利用学习到的世界模型和不确定性估计来优化策略。在训练过程中，每个合成转移样本会根据其预测置信度进行加权，置信度高的样本权重较高，置信度低的样本权重较低，从而保留有用的模型rollout，并减轻不确定预测带来的偏差。

关键创新：WIMLE的关键创新在于将IMLE引入模型强化学习框架，并结合不确定性感知的加权方法。与传统的基于最大似然估计的方法相比，IMLE能够更好地学习多模态的世界模型，避免了对真实环境动态的过度简化。同时，不确定性感知的加权方法能够有效地抑制模型误差的累积，提高策略学习的稳定性。

关键设计：WIMLE的关键设计包括：1）使用集成方法（如Bootstrap）来估计模型预测的不确定性；2）使用潜在变量模型来表示环境的随机性；3）设计合适的损失函数，鼓励模型学习准确的动态模型和不确定性估计；4）根据预测置信度对合成转移样本进行加权，可以使用高斯分布或其他概率分布来表示置信度。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

WIMLE在40个连续控制任务中表现出色，尤其是在Humanoid-run任务上，相对于最强的竞争对手提高了50%以上的样本效率。在HumanoidBench上，WIMLE解决了14个任务中的8个，而BRO和SimbaV2分别只解决了4个和5个。这些结果表明，WIMLE在样本效率和性能方面均优于现有的模型无关和基于模型的强化学习方法。

🎯 应用场景

WIMLE具有广泛的应用前景，可应用于机器人控制、自动驾驶、游戏AI等领域。通过学习准确的环境模型和估计预测的不确定性，WIMLE能够帮助智能体更好地理解和适应复杂环境，从而实现更高效、更安全的决策。此外，WIMLE还可以用于模型预测控制等领域，提高控制系统的鲁棒性和适应性。

📄 摘要（原文）

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.