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

• URL: http://arxiv.org/abs/2605.13013v1
• Date: Wed, 13 May 2026 05:07:32 GMT
• Title: JEDI: Joint Embedding Diffusion World Model for Online Model-Based Reinforcement Learning
• Abstract summary: Joint Embedding DIffusion (JEDI) is the first online end-to-end latent diffusion world model.<n>We show that JEDI is competitive on Atari100k and outperforms the baseline with seperately trained latents where directly comparable.
• Score: 6.0332038912905555
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Valdi: Value Diffusion World Models [9.342187139336017]
  World models can enable Model Predictive Control (MPC), but this requires dynamics prediction that is fast enough for online use and expressive enough to represent uncertain futures.<n>Diffusion models offer a natural mechanism for modeling uncertain dynamics, yet their iterative inference procedure makes them difficult to use for low-latency latent planning.<n>We bridge this gap with Value Diffusion World Models (Valdi), combining end-to-end online training for MPC with a latent diffusion dynamics model.
  arXiv  Detail & Related papers  (2026-07-01T13:22:18Z)
• Consistent Diffusion Language Models [22.253153649144476]
  Diffusion language models (DLMs) promise sub-time, parallel generation, yet practical gains remain elusive.<n>In continuous domains, consistency training along the probability-flow ODE is a popular recipe to accelerate diffusion.<n>We introduce Multi-Path Discrete Consistency (DCMP), a new principle that trains a denoiser to be path-in expectation across these bridges.<n>A single CDLM objective unifies masked diffusion, continuous consistency models, and progressive/discrete distillation as analytic limits or empirical approximations of one common view.
  arXiv  Detail & Related papers  (2026-04-30T19:31:02Z)
• Deep Leakage with Generative Flow Matching Denoiser [54.05993847488204]
  We introduce a new deep leakage (DL) attack that integrates a generative Flow Matching (FM) prior into the reconstruction process.<n>Our approach consistently outperforms state-of-the-art attacks across pixel-level, perceptual, and feature-based similarity metrics.
  arXiv  Detail & Related papers  (2026-01-21T14:51:01Z)
• Plug-and-Play Prompt Refinement via Latent Feedback for Diffusion Model Alignment [54.17386822940477]
  We introduce PromptLoop, a plug-and-play reinforcement learning framework that incorporates latent feedback into step-wise prompt refinement.<n>This design achieves a structural analogy to the Diffusion RL approach, while retaining the flexibility and generality of prompt-based alignment.
  arXiv  Detail & Related papers  (2025-10-01T02:18:58Z)
• Consistent World Models via Foresight Diffusion [56.45012929930605]
  We argue that a key bottleneck in learning consistent diffusion-based world models lies in the suboptimal predictive ability.<n>We propose Foresight Diffusion (ForeDiff), a diffusion-based world modeling framework that enhances consistency by decoupling condition understanding from target denoising.
  arXiv  Detail & Related papers  (2025-05-22T10:01:59Z)
• Generalized Interpolating Discrete Diffusion [65.74168524007484]
  Masked diffusion is a popular choice due to its simplicity and effectiveness.<n>We generalize a new family of general interpolating discrete diffusion (GIDD) which offers greater flexibility in the design of the noising processes.<n>Exploiting GIDD's flexibility, we explore a hybrid approach combining masking and uniform noise, leading to improved sample quality.
  arXiv  Detail & Related papers  (2025-03-06T14:30:55Z)
• Improving Vector-Quantized Image Modeling with Latent Consistency-Matching Diffusion [55.185588994883226]
  We introduce VQ-LCMD, a continuous-space latent diffusion framework within the embedding space that stabilizes training.<n>VQ-LCMD uses a novel training objective combining the joint embedding-diffusion variational lower bound with a consistency-matching (CM) loss.<n>Experiments show that the proposed VQ-LCMD yields superior results on FFHQ, LSUN Churches, and LSUN Bedrooms compared to discrete-state latent diffusion models.
  arXiv  Detail & Related papers  (2024-10-18T09:12:33Z)
• Denoising with a Joint-Embedding Predictive Architecture [21.42513407755273]
  We introduce Denoising with a Joint-Embedding Predictive Architecture (D-JEPA)<n>By recognizing JEPA as a form of masked image modeling, we reinterpret it as a generalized next-token prediction strategy.<n>We also incorporate diffusion loss to model the per-token probability distribution, enabling data generation in a continuous space.
  arXiv  Detail & Related papers  (2024-10-02T05:57:10Z)
• Self-Play Fine-Tuning of Diffusion Models for Text-to-Image Generation [59.184980778643464]
  Fine-tuning Diffusion Models remains an underexplored frontier in generative artificial intelligence (GenAI) In this paper, we introduce an innovative technique called self-play fine-tuning for diffusion models (SPIN-Diffusion) Our approach offers an alternative to conventional supervised fine-tuning and RL strategies, significantly improving both model performance and alignment.
  arXiv  Detail & Related papers  (2024-02-15T18:59:18Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.