JEDI: Joint Embedding Diffusion World Model for Online Model-Based Reinforcement Learning
TLDR
JEDI combines JEPA-style predictive learning with diffusion denoising for an end-to-end latent diffusion world model, achieving competitive results on Atari100k.
Reasoning
The paper introduces a novel integration of JEPA and diffusion for online MBRL, with theoretical motivation and empirical results on Atari100k. Strengths include end-to-end learning and efficiency gains, but the evaluation is limited to a single simulated benchmark without real-world experiments.
Read-first score
Read-first score 55.1, weighted from topical fit, citation, graph, method, reproducibility, and recency signals. Original total remains 51.
Field roles
Rank sensitivity
Stability: volatile; rank range: 432.
Keyword Scores
Deep Analysis
Innovations
- First online end-to-end latent diffusion world model for model-based reinforcement learning
- Learns latent space directly from diffusion denoising loss within a JEPA framework
- Uses denoising to learn and predict future latents instead of reconstruction and pretrained models
- Theoretical motivation linking conventional JEPA objectives to a predictive information bottleneck and conditional diffusion denoising to a predictive-compression decomposition
Methodology
JEDI is a latent diffusion world model for online model-based reinforcement learning. It learns its latent space directly from the diffusion denoising loss using a JEPA framework, where denoising is used to learn and predict future latents rather than relying on reconstruction and pretrained models. The model is trained end-to-end.
Key Results
JEDI is competitive on Atari100k and outperforms the baseline with separately trained latents where directly comparable. Relative to the pixel diffusion baseline, JEDI uses 43% less VRAM, over 3× faster world-model sampling, and 2.5× faster training.