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

TL;DR

JEDI introduces an end-to-end latent diffusion world model for reinforcement learning that improves efficiency and performance by directly learning predictive latents through diffusion denoising.

Contribution

It is the first online end-to-end latent diffusion world model that integrates diffusion objectives with JEPA-style predictive learning.

Findings

JEDI outperforms baseline with separately trained latents on Atari100k.

JEDI uses 43% less VRAM compared to pixel diffusion baseline.

JEDI achieves over 3× faster world-model sampling and 2.5× faster training.

Abstract

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.