JEDI: Latent End-to-end Diffusion Mitigates Agent-Human Performance Asymmetry in Model-Based Reinforcement Learning

📄 arXiv: 2505.19698v2 📥 PDF

作者: Jing Yu Lim, Zarif Ikram, Samson Yu, Haozhe Ma, Tze-Yun Leong, Dianbo Liu

分类: cs.LG, cs.AI, cs.RO

发布日期: 2025-05-26 (更新: 2025-05-28)

备注: Preprint

💡 一句话要点

提出JEDI：通过端到端隐空间扩散模型缓解基于模型的强化学习中Agent-Human性能不对称问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 基于模型的强化学习 扩散模型 隐空间表示 Agent-Human性能不对称 Atari100k 端到端学习 时间结构建模

📋 核心要点

  1. 现有基于模型的强化学习方法在Atari100k上表现出Agent-Human性能不对称，即部分任务远超人类，部分任务远低于人类。
  2. 论文提出Joint Embedding DIffusion (JEDI)，一种端到端训练的隐空间扩散世界模型，旨在学习具有时间结构的隐空间。
  3. JEDI在human-optimal任务上超越现有模型，同时保持Atari100k整体竞争力，并显著提升了计算效率和内存占用。

📝 摘要（中文）

基于模型的强化学习(MBRL)的最新进展在Atari100k基准测试中取得了超人类水平的性能，这得益于在强大的扩散世界模型上训练的强化学习智能体。然而，我们发现当前的聚合掩盖了一个主要的性能不对称性：MBRL智能体在某些任务中显著优于人类，但在其他任务中却表现不佳，前者夸大了聚合指标。在基于像素的、使用扩散世界模型训练的智能体中，这一点尤其明显。在这项工作中，我们致力于解决基于像素的智能体中观察到的显著不对称性，作为扭转这种令人担忧的上升趋势的初步尝试。我们通过将所有任务划分为Agent-Optimal或Human-Optimal来解决有问题的聚合，并主张同等重视来自两组的指标。接下来，我们假设这种显著的不对称性是由于基于像素的方法中缺乏使用世界模型目标训练的、具有时间结构的隐空间。最后，为了解决这个问题，我们提出了一种新的联合嵌入扩散(JEDI)模型，这是一种与自洽目标端到端训练的隐空间扩散世界模型。JEDI在human-optimal任务中优于SOTA模型，同时在Atari100k基准测试中保持竞争力，并且比最新的基于像素的扩散基线快3倍，内存降低43%。总的来说，我们的工作重新思考了在Atari100k中真正超越人类水平的性能意味着什么。

🔬 方法详解

问题定义：现有基于模型的强化学习方法，尤其是在像素空间上训练的智能体，在Atari100k基准测试中表现出显著的Agent-Human性能不对称性。这意味着智能体在某些任务上远超人类水平，而在另一些任务上却远低于人类水平。这种不对称性导致整体性能指标被高估，无法真实反映智能体的泛化能力和鲁棒性。现有方法缺乏对时间结构信息的有效建模，导致智能体难以学习到符合人类直觉的策略。

核心思路：论文的核心思路是通过学习一个具有时间结构的隐空间来缓解Agent-Human性能不对称性。具体来说，论文提出Joint Embedding DIffusion (JEDI)模型，该模型通过端到端的方式，将视觉信息编码到隐空间，并利用扩散模型学习隐空间的动态特性。这种方法旨在使智能体能够更好地理解环境的时间演化规律，从而制定更符合人类直觉的策略。

技术框架：JEDI模型是一个端到端的隐空间扩散世界模型，其主要包含以下模块：1) 编码器：将像素级别的观测信息编码到隐空间中。2) 扩散模型：学习隐空间的动态特性，用于预测未来的状态。3) 解码器：将隐空间的状态解码回像素空间，用于重构观测信息。整个框架通过自洽性目标进行端到端训练，即通过预测未来的状态并重构观测信息，来学习隐空间的有效表示。

关键创新：JEDI模型的关键创新在于其端到端的训练方式和对时间结构的建模。与传统的基于像素的扩散模型不同，JEDI模型直接在隐空间中进行动态建模，从而避免了像素空间的高维性和噪声。此外，JEDI模型通过自洽性目标，鼓励隐空间学习到具有时间结构的表示，从而更好地捕捉环境的动态特性。

关键设计：JEDI模型的关键设计包括：1) 使用VAE结构进行隐空间编码和解码。2) 使用扩散模型学习隐空间的动态特性，扩散模型的具体结构未知。3) 使用自洽性损失函数进行端到端训练，损失函数包括重构损失和动态预测损失。4) 论文中未明确给出关键参数设置和网络结构的具体细节，这部分信息未知。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

JEDI模型在human-optimal任务上超越了现有SOTA模型，同时在Atari100k基准测试中保持了竞争力。此外，JEDI模型比最新的基于像素的扩散基线快3倍，内存降低43%。这些结果表明，JEDI模型在性能和效率方面都具有显著优势。

🎯 应用场景

该研究成果可应用于机器人控制、游戏AI、自动驾驶等领域。通过学习更符合人类直觉的策略，智能体可以更好地与人类协作，提高任务完成效率和安全性。此外，该方法还可以用于生成更逼真的虚拟环境，为训练智能体提供更丰富的训练数据。

📄 摘要（原文）

Recent advances in model-based reinforcement learning (MBRL) have achieved super-human level performance on the Atari100k benchmark, driven by reinforcement learning agents trained on powerful diffusion world models. However, we identify that the current aggregates mask a major performance asymmetry: MBRL agents dramatically outperform humans in some tasks despite drastically underperforming in others, with the former inflating the aggregate metrics. This is especially pronounced in pixel-based agents trained with diffusion world models. In this work, we address the pronounced asymmetry observed in pixel-based agents as an initial attempt to reverse the worrying upward trend observed in them. We address the problematic aggregates by delineating all tasks as Agent-Optimal or Human-Optimal and advocate for equal importance on metrics from both sets. Next, we hypothesize this pronounced asymmetry is due to the lack of temporally-structured latent space trained with the World Model objective in pixel-based methods. Lastly, to address this issue, we propose Joint Embedding DIffusion (JEDI), a novel latent diffusion world model trained end-to-end with the self-consistency objective. JEDI outperforms SOTA models in human-optimal tasks while staying competitive across the Atari100k benchmark, and runs 3 times faster with 43% lower memory than the latest pixel-based diffusion baseline. Overall, our work rethinks what it truly means to cross human-level performance in Atari100k.