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

arXiv: https://arxiv.org/abs/2605.13013 · HTML: https://mlsanity.com/paper/2605.13013/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in Latent dynamics models and model-based RL, one of 11 directions in World models.
- Builds on 36 of its 74 references that are in the map.
- First posted on arXiv 13 May 2026.

## Placed in

- World models / [Latent dynamics models and model-based RL](https://mlsanity.com/wm/latent-dynamics-models-and-model-based-rl/)
- Diffusion models / [Diffusion and flow policies for robotics](https://mlsanity.com/diffusion/diffusion-and-flow-policies-for-robotics/)
- Deep reinforcement learning methods / [RL with diffusion and flow models](https://mlsanity.com/deeprl/rl-with-diffusion-and-flow-models/)

## Builds on (most cited first, 30 of 36)

- [Stable Video Diffusion: Scaling Latent Video Diffusion Models to Large Datasets](https://mlsanity.com/paper/2311.15127/) (2023)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [Planning with Diffusion for Flexible Behavior Synthesis](https://mlsanity.com/paper/2205.09991/) (2022)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [Deep Variational Information Bottleneck](https://mlsanity.com/paper/1612.00410/) (2016)
- [V-JEPA 2: Self-Supervised Video Models Enable Understanding, Prediction and Planning](https://mlsanity.com/paper/2506.09985/) (2025)
- [Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution](https://mlsanity.com/paper/2310.16834/) (2023)
- [Mastering Chess and Shogi by Self-Play with a General Reinforcement Learning Algorithm](https://mlsanity.com/paper/1712.01815/) (2017)
- [Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning](https://mlsanity.com/paper/2208.06193/) (2022)
- [TD-MPC2: Scalable, Robust World Models for Continuous Control](https://mlsanity.com/paper/2310.16828/) (2023)
- [Opening the Black Box of Deep Neural Networks via Information](https://mlsanity.com/paper/1703.00810/) (2017)
- [Model-Based Reinforcement Learning for Atari](https://mlsanity.com/paper/1903.00374/) (2019)
- [Label-Efficient Semantic Segmentation with Diffusion Models](https://mlsanity.com/paper/2112.03126/) (2021)
- [Sora: A Review on Background, Technology, Limitations, and Opportunities of Large Vision Models](https://mlsanity.com/paper/2402.17177/) (2024)
- [Temporal Difference Learning for Model Predictive Control](https://mlsanity.com/paper/2203.04955/) (2022)
- [Transformers are Sample-Efficient World Models](https://mlsanity.com/paper/2209.00588/) (2022)
- [LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels](https://mlsanity.com/paper/2603.19312/) (2026)
- [LeJEPA: Provable and Scalable Self-Supervised Learning Without the Heuristics](https://mlsanity.com/paper/2511.08544/) (2025)
- [Transformer-based World Models Are Happy With 100k Interactions](https://mlsanity.com/paper/2303.07109/) (2023)
- [DiffuseVAE: Efficient, Controllable and High-Fidelity Generation from Low-Dimensional Latents](https://mlsanity.com/paper/2201.00308/) (2022)
- [Diffusion Beats Autoregressive in Data-Constrained Settings](https://mlsanity.com/paper/2507.15857/) (2025)
- [Diffusion Guidance Is a Controllable Policy Improvement Operator](https://mlsanity.com/paper/2505.23458/) (2025)
- [Policy-Guided Diffusion](https://mlsanity.com/paper/2404.06356/) (2024)
- [Temporal Straightening for Latent Planning](https://mlsanity.com/paper/2603.12231/) (2026)
- [Exploring Diffusion Time-steps for Unsupervised Representation Learning](https://mlsanity.com/paper/2401.11430/) (2024)
- [Diffusion World Model: Future Modeling Beyond Step-by-Step Rollout for Offline Reinforcement Learning](https://mlsanity.com/paper/2402.03570/) (2024)
- [Learning Transformer-based World Models with Contrastive Predictive Coding](https://mlsanity.com/paper/2503.04416/) (2025)
- [DiffEnc: Variational Diffusion with a Learned Encoder](https://mlsanity.com/paper/2310.19789/) (2023)

## Built on by (most cited first, 2 of 2)

- [Code as Worlds: Agentic Discovery of Executable World Representations for Physical Reasoning](https://mlsanity.com/paper/2608.27549/) (2026)
- [Branch-JEPA: Finite-Support Predictive Distributions for JEPA World Models](https://mlsanity.com/paper/2607.05238/) (2026)

## Abstract (arXiv)

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× faster world-model sampling, and 2.5× 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.