🤖 machine learning

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

JEDI introduces the first online end-to-end latent diffusion world model that unifies JEPA-style predictive representation learning with diffusion objectives to achieve superior efficiency and competitive performance in model-based reinforcement learning compared to both pixel-based and separately trained latent approaches.

Original authors: Jing Yu Lim, Rushi Shah, Zarif Ikram, Samson Yu, Haozhe Ma, Tze-Yun Leong, Dianbo Liu

Published 2026-05-14

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Jing Yu Lim, Rushi Shah, Zarif Ikram, Samson Yu, Haozhe Ma, Tze-Yun Leong, Dianbo Liu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to play video games. To do this efficiently, the robot needs a "world model"—a mental simulation of how the game works so it can practice in its head before actually playing.

For a long time, the best way to build this mental simulation was to have the robot try to reconstruct the game screen pixel-by-pixel, like a painter trying to copy a photo exactly. This is accurate but very slow and requires a massive amount of computer memory (like trying to carry a library in your backpack).

Recently, a new technique called Diffusion became popular. Think of diffusion like a sculptor who starts with a block of noisy, static-filled clay and slowly chisels away the noise to reveal a clear statue. This is great for understanding complex scenes, but doing it pixel-by-pixel is still too heavy for a robot to run in real-time.

Another approach tried to compress the game into a tiny, abstract "summary" (a latent space) first, but these summaries were often trained separately and didn't learn the game rules well enough on their own.

Enter JEDI (Joint Embedding Diffusion).

The authors of this paper created a new method that combines the best of both worlds. Here is how JEDI works, using simple analogies:

1. The "Mental Snapshot" vs. The "Painting"

Old Way (Pixel Diffusion): Imagine the robot trying to memorize every single grain of sand on a beach to understand the ocean. It's accurate, but it takes forever and uses up all its brainpower.
JEDI's Way: Instead of memorizing the sand, JEDI teaches the robot to take a mental snapshot of the ocean's essence (the waves, the color, the movement). It doesn't care about individual grains of sand. It compresses the game screen into a tiny, efficient "summary" that captures only what matters for winning.

2. The "Guessing Game" Training

How does the robot learn to take these snapshots?

The Old Way: The robot was often trained by being shown a picture and asked to re-paint it exactly. If it missed a detail, it got a penalty. This is like a student being graded on handwriting rather than understanding the story.
The JEDI Way: JEDI uses a predictive guessing game.
1. The robot sees the current game state.
2. It is asked to guess what the next "mental snapshot" will look like.
3. To make it harder (and smarter), the computer takes the actual next snapshot, adds some "static noise" to it (like blurring it), and asks the robot to denoise it back to clarity.
4. Crucially, the robot learns this while playing. It doesn't need a separate teacher to show it how to paint; it learns the rules of the game by trying to predict the future.

3. Why This is a Big Deal (The Results)

The paper claims JEDI is a major upgrade for three main reasons:

It's Leaner: Because JEDI works with tiny "mental snapshots" instead of full video frames, it uses 43% less computer memory. It's like switching from carrying a heavy backpack full of books to carrying a single, smart notebook.
It's Faster: The robot can think (sample) 3 times faster and train 2.5 times faster than the previous best pixel-based methods.
It Plays Differently: This is the most surprising part. The paper found that JEDI doesn't just play faster; it plays differently.
- On games that were already easy for other robots, JEDI was good but not always the absolute best.
- However, on the hardest, most chaotic games (like shooters with many moving targets), JEDI shined. It seemed to handle the chaos better because its "mental snapshot" focused on the strategy rather than getting distracted by visual noise.

The Bottom Line

The paper argues that you don't need to be a perfect painter (reconstructing every pixel) to be a great strategist. By teaching the robot to predict the "gist" of the future using a noise-removal game (diffusion), you get a system that is faster, cheaper to run, and surprisingly better at handling difficult, chaotic situations.

The authors call this JEDI, and they claim it is the first time this specific "predictive guessing" method has been successfully combined with "noise-removal" (diffusion) to teach a robot to play games online, without needing any pre-trained teachers.

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

Problem Statement

Online Model-Based Reinforcement Learning (MBRL) relies on learned world models to predict future outcomes and improve sample efficiency. While diffusion models have emerged as powerful backbones for world modeling in visually complex domains, current approaches face a significant tension:

Pixel Diffusion: Models like DIAMOND operate directly in pixel space. While effective, they are computationally expensive regarding memory (VRAM), sampling time, and training duration.
Latent Diffusion: Moving diffusion to latent space improves efficiency but often relies on separately trained latents (e.g., via reconstruction objectives) rather than end-to-end world-model objectives. These approaches, such as Horizon Imagination (HI), have yet to match the performance of pixel-based diffusion models and lack the end-to-end predictive representation learning that has driven modern MBRL progress.

The central question addressed is whether a diffusion world model can learn its latent representation end-to-end from the denoising objective using a Joint Embedding Predictive Architecture (JEPA) framework, thereby combining the efficiency of latent diffusion with the performance of end-to-end predictive learning.

Methodology: JEDI

The authors introduce JEDI (Joint Embedding Diffusion), the first online end-to-end latent diffusion world model built on a JEPA framework.

Core Architecture

JEDI learns a latent space encoder ( $E_\phi$ ) and a dynamics model simultaneously. The system consists of:

Encoder: Compresses observations $x_t$ into clean latents $z^0_t$ .
Latent Diffusion Dynamics: Predicts the next latent state through conditional denoising. Instead of a single-step prediction, it uses a reverse diffusion chain conditioned on the current latent $z^0_t$ and past actions to predict the clean future latent $z^0_{t+1}$ .
Reward and Termination Heads: Predict rewards and termination signals directly from the latent state.

Training Mechanism

Unlike traditional latent diffusion models that rely on reconstruction losses (e.g., VAEs) to train the encoder, JEDI trains the encoder end-to-end using the diffusion denoising loss.

Joint Embedding Objective: The encoder is updated via gradients from the conditional denoising objective. The future latent target is treated as a "label" derived from the environment, and the model learns to denoise from a noisy version of this future latent back to the clean state, conditioned on the current state.
Stop-Gradient: To prevent representation collapse and stabilize training, a stop-gradient is applied to the future latent target during encoder updates, following JEPA and TDMPC2 principles.
Stabilization Techniques: The model employs differentiable tanh clamping on latent activations and a random switching mechanism between the denoiser output and encoder output for next-step conditioning to balance consistency and supervision.

Theoretical Motivation

The paper provides a theoretical justification linking JEDI to the Information Bottleneck principle.

JEPA as a Bottleneck: The standard JEPA objective is shown to induce a variational information bottleneck, rewarding representations that retain information useful for predicting the future while compressing irrelevant details.
Diffusion-JEPA Decomposition: The authors demonstrate that latent conditional diffusion admits a similar variational information-bottleneck decomposition. The conditional denoising loss acts as the predictive term within this bottleneck structure. This suggests that diffusion denoising provides a principled, end-to-end predictive learning signal for the latent encoder, rather than merely serving as a generative surrogate.

Key Contributions

First End-to-End Latent Diffusion World Model: JEDI is the first online MBRL model to train its latent encoder end-to-end using a diffusion denoising objective within a JEPA framework, eliminating the need for separate pretraining or reconstruction losses.
Theoretical Connection: The paper establishes a mathematical link between conditional diffusion objectives and the variational information bottleneck, explaining why denoising can effectively support predictive representation learning.
Efficiency Gains: By moving from pixel space to a compact learned latent space ( $16 \times 8 \times 8$ ) and training end-to-end, JEDI significantly reduces computational overhead compared to pixel-based diffusion.

Experimental Results

JEDI was evaluated on the Atari100k benchmark (26 tasks, 5 seeds) and the Craftium environment (3D embodied decision-making).

Performance

Atari100k: JEDI achieves competitive performance, matching or exceeding strong baselines (DIAMOND, DreamerV3, STORM) on Interquartile Mean (IQM) and Optimality Gap metrics. It achieves State-of-the-Art (SOTA) results in 7 games and runner-up performance in the number of superhuman tasks.
Craftium: JEDI outperforms Horizon Imagination (HI) on all four tested tasks, demonstrating that end-to-end predictive latents can support strong performance in 3D environments without reconstruction-based latent training.
Stochasticity: In experiments with stochastic frame-skipping, JEDI consistently outperforms DreamerV3, suggesting diffusion models may be more robust to high aleatoric uncertainty.

Efficiency

Relative to the pixel-based baseline DIAMOND, JEDI demonstrates substantial efficiency improvements:

VRAM Usage: Uses 43% less VRAM.
Sampling Speed: Achieves >3× faster world-model sampling.
Training Speed: Trains >2.5× faster (38 hours vs. 98 hours on A100 for Atari100k).
Model Size: Maintains a similar parameter count (13.5M) to DIAMOND but significantly outperforms HI (97M parameters) in efficiency while retaining competitive performance.

Performance Profile Shift

A distinct finding is that JEDI exhibits a different task-level performance profile compared to DIAMOND. While DIAMOND excels in high-HNS (Human Normalized Score) regimes, JEDI shows marked improvements in low-HNS regimes (difficult games). This is captured by the Har-HNS metric, where JEDI outperforms DIAMOND, indicating it fails less severely on the most difficult tasks. The authors attribute this to the end-to-end latent interface reducing the dimensionality of the input for the actor-critic, allowing better optimization in complex, high-action-space games.

Significance and Claims

The paper claims that JEDI demonstrates that predictive latent diffusion is a viable and efficient alternative to both pixel-space diffusion and separately trained latent models for online MBRL.

Beyond Compute: The results suggest that end-to-end predictive latents change the behavior of the world model itself, not just the computational cost. The shift in performance profiles indicates that the learned representation fundamentally alters how the agent interacts with complex environments.
Principled Approach: By grounding the method in the information bottleneck theory, the paper positions predictive latent diffusion as a principled direction for scaling online MBRL.
Limitations: The authors modestly note that evidence is currently limited to Atari100k and Craftium, and the method has not yet been tested on domains like DMControl or Procgen. They also acknowledge that performance depends on specific design choices (e.g., clamping, random switching) and that the theoretical derivation motivates representation properties rather than guaranteeing optimization of the full variational objective.