JEDI：基于在线模型的强化学习的联合嵌入扩散世界模型

扩散世界模型最近在基于在线模型的强化学习中变得具有竞争力，但当前的方法暴露了一种张力：像素扩散是有效的，但计算成本昂贵，而最新的潜在扩散方法提高了效率，但性能不佳。后者还依赖于单独训练的潜在模型，而不是推动现代 MBRL 进步的端到端世界模型目标。特别是，JEPA 式的预测表示学习已成为世界建模和 MBRL 的一个特别有前途的方向。与此同时，扩散式目标已在多个领域获得关注，迭代细化作为多模式和随机目标的一种有前途的方法。总而言之，这些趋势激发了联合嵌入扩散（JEDI），这是第一个在线端到端潜在扩散世界模型。 JEDI 使用 JEPA 框架直接从扩散去噪损失中学习其潜在空间，使用去噪来学习和预测未来的潜在空间，而不是依赖于重建和预训练模型。我们提供的理论动机表明，传统的 JEPA 目标会导致预测信息瓶颈，并且条件扩散去噪允许密切相关的预测压缩分解。根据经验，JEDI 在 Atari100k 上具有竞争力，并且在可直接比较的单独训练的潜在变量中优于基线。相对于像素扩散基线，JEDI 使用的 VRAM 减少了 43%，世界模型采样速度快了 3 美元\倍，训练速度快了 2.5 美元\倍。 JEDI 还表现出与像素基线明显不同的任务级性能概况，这表明端到端预测潜伏的变化比单独计算的变化更大。