Diffusion world models have quietly become one of the more compelling ideas in model-based reinforcement learning: instead of hand-rolling a dynamics predictor, let a diffusion model learn to "dream" the next frame of a game or a robot's camera feed, then train a policy inside that dream. The catch has been cost. The strongest online diffusion world model, DIAMOND, does its denoising directly in pixel space, which is accurate but punishing on memory, sampling time, and training time. A new paper from the National University of Singapore proposes a way out that is less about shaving compute and more about rethinking where the latent space comes from. It is called JEDI, and its claim is that a diffusion world model can learn its own representation end to end, straight from the denoising loss.
The problem: latents you didn't learn for the task
The obvious fix for expensive pixel diffusion is to move the denoising into a compressed latent space. Recent work such as Horizon Imagination (HI) does exactly that and gets a real efficiency win. But those latents are trained separately, typically with reconstruction and perceptual objectives borrowed from image generation, and they do not yet match DIAMOND's performance. The authors frame this as a missed opportunity. Much of modern MBRL progress, they argue, came from learning representations end to end against a world-model objective rather than bolting a pretrained autoencoder onto a dynamics model.
JEDI, short for Joint Embedding DIffusion, is the authors' answer, and they describe it as the first online end-to-end latent diffusion world model. The idea fuses two threads running in parallel across machine learning. One is JEPA, the joint-embedding predictive architecture line of work, which learns representations that capture the predictable, abstract structure of the world while ignoring nuisance detail that reconstruction-based methods waste capacity on. The other is the broad shift toward diffusion-style objectives, whose iterative refinement handles multimodal and stochastic targets well. JEDI asks whether a diffusion world model can learn its latent representation end to end from the denoising objective, inside a JEPA framework. The answer is yes.
The method: denoising as the representation-learning signal
Mechanically, JEDI keeps three pieces. An encoder compresses each 64x64x3 observation into a compact latent (16x8x8, bounded to a fixed range). A latent diffusion dynamics model predicts the next latent through conditional denoising, conditioned on a window of past latents and past actions plus the diffusion-time embedding. And a small head predicts reward and episode termination. The denoiser follows the EDM preconditioning recipe familiar from modern diffusion work.
The pivotal detail is where the gradients go. Both the diffusion denoising loss and the reward/termination loss flow back through the latent and update the encoder directly, so the representation is shaped by predictive world-model training rather than by a reconstruction objective the dynamics model never asked for. To keep this from collapsing, JEDI applies a stop-gradient to the future-latent target and runs the encoder at 0.3 times the denoiser's learning rate, echoing stabilization tricks from JEPA and TD-MPC2, and clamps latent activations with a differentiable tanh transform.
The authors also give the recipe a theoretical spine. They show that conventional JEPA objectives induce a predictive information bottleneck, and that conditional diffusion denoising admits a closely related predictive-compression decomposition. Their claim is deliberately measured: diffusion denoising supplies the predictive term in a JEPA-style bottleneck and therefore provides a principled end-to-end learning signal for the encoder, not that the method exactly optimizes the full variational objective.
The results: competitive scores, much cheaper compute
The headline numbers are about efficiency, and they are concrete. Relative to the pixel-diffusion DIAMOND baseline, JEDI uses 43 percent less VRAM, samples from its world model more than three times faster, and trains more than 2.5 times faster. Those are the margins that decide whether an online world model is practical to iterate on at all.
On capability, the paper reports that JEDI is competitive on the full Atari100k benchmark, the standard sample-efficiency suite where an agent gets only 100,000 environment steps, and that it outperforms the baseline with separately trained latents where the two are directly comparable, meaning it beats the HI-style approach it set out to improve on. The authors also run supporting experiments in Craftium, a richer 3D environment, to test the method beyond Atari's 2D games. The paper does not position JEDI as a clean sweep; it acknowledges DIAMOND remains a strong baseline. What it flags instead is more interesting: JEDI shows a markedly different performance profile across games, winning and losing on different titles than the pixel model. The authors read that as evidence that end-to-end predictive latents change what the model learns, not just how fast.
Why it matters
The significance here is architectural rather than a leaderboard win. If a diffusion world model can grow its own latent space from the prediction task, the field does not have to choose between DIAMOND's expensive pixel fidelity and the weaker performance of latents trained offline for a different objective. For anyone eyeing world models for robotics or control, where the useful information in a camera frame is a small fraction of the pixels, a representation that deliberately discards nuisance detail is exactly the right inductive bias, and the compute savings matter when training has to happen online.
What to watch
The open questions are the honest ones. The different per-game profile cuts both ways: it suggests the learned latents are doing something real, but it also means JEDI is not strictly dominant, and understanding when predictive latents help versus hurt is unresolved. The theory is framed as motivation rather than a tight optimality guarantee. And the demonstrations live on Atari100k and Craftium; the leap to continuous-control and real robotic domains, where diffusion world models are most tantalizing, is the test that will decide whether JEDI is a benchmark curiosity or a template. The paper's authors are Jing Yu Lim, Rushi Shah, Zarif Ikram, Samson Yu, Haozhe Ma, Tze-Yun Leong, and Dianbo Liu at the National University of Singapore.