arXiv:2605.13013cs.LG2026-05被引 1

JEDI首次实现端到端潜空间扩散世界模型,兼顾高效与性能。

JEDI: Joint Embedding Diffusion World Model for Online Model-Based Reinforcement Learning

论文配图:JEDI: Joint Embedding Diffusion World Model for Online Model-Based Reinforcement Learning
图 1 · 摘自论文原文
  • 用扩散去噪损失联合学习潜变量,无需预训练模型
  • 在Atari100k上表现优于独立训练潜变量的基线
  • 相比像素级扩散模型,显存减少43%,采样快3倍以上

扩散世界模型在在线模型基于强化学习中表现日益优异,但现有方法存在权衡:像素级扩散虽有效却计算开销大,而最新潜空间扩散虽提升效率但性能不足,且依赖独立训练的潜变量,而非端到端世界模型目标。本文提出联合嵌入扩散(JEDI),首个面向在线模型基于强化学习的端到端潜空间扩散世界模型。JEDI通过JEPA框架,直接从扩散去噪损失中学习潜空间,以去噪方式预测未来潜变量,而非依赖重建或预训练模型。理论分析表明,传统JEPA目标诱导预测信息瓶颈,而条件扩散去噪具备相似的预测-压缩分解结构。实验显示,JEDI在Atari100k上表现竞争力,优于直接可比的独立训练潜变量基线;相较像素级扩散基线,节省43%显存,世界模型采样速度超3倍,训练速度提升2.5倍。同时,任务层面性能特征显著不同于像素基线,表明端到端预测潜变量带来的改变远超单纯算力提升。

原文摘要 · Abstract (English)

Diffusion world models have recently become competitive for online model-based reinforcement learning, but current approaches expose a tension: pixel diffusion is effective but computationally expensive while the latest latent diffusion approach improves efficiency yet performs subpar. The latter also relies on separately trained latents rather than the end-to-end world-model objectives that have driven much of modern MBRL progress. In particular, JEPA-style predictive representation learning has emerged as an especially promising direction for world modeling and MBRL. Concurrently, diffusion-style objectives have gained traction across multiple domains, with iterative refinement as a promising approach for multimodal and stochastic targets. Taken together, these trends motivate Joint Embedding DIffusion (JEDI), the first online end-to-end latent diffusion world model. JEDI learns its latent space directly from the diffusion denoising loss with a JEPA framework, using denoising to learn and predict future latents rather than relying on reconstruction and pretrained models. We provide a theoretical motivation showing that conventional JEPA objectives induce a predictive information bottleneck, and that conditional diffusion denoising admits a closely related predictive-compression decomposition. Empirically, JEDI is competitive on Atari100k and outperforms the baseline with seperately trained latents where directly comparable. Relative to the pixel diffusion baseline, JEDI uses 43% less VRAM, over 3$\times$ faster world-model sampling, and 2.5$\times$ faster training. JEDI also exhibits a markedly different task-level performance profile from the pixel baseline, suggesting that end-to-end predictive latents change more than compute alone.

世界模型扩散模型强化学习潜空间

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。