PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2605.13013v1cs.LG

TL;DR

Pixel diffusion world models are computationally expensive, while existing latent diffusion improves efficiency but lacks end-to-end sequential prediction and trails DIAMOND. JEDI learns latents end to end through conditional denoising in a JEPA framework. It reports strong Atari100k results, improved efficiency over DIAMOND, and a different across-game performance profile.

  • Problem

    Pixel diffusion world models are costly in memory, sampling time, and training time, while existing latent diffusion improves efficiency but lacks end-to-end sequential prediction and does not yet match DIAMOND.

    Source: §1 · p. 1

  • Method

    JEDI trains an online latent diffusion world model by learning its encoder end to end through conditional denoising, together with reward and end prediction.

    Source: §1 · p. 2

  • Results

    JEDI reports strong Atari100k results and, relative to DIAMOND, uses 43% less VRAM, samples the world model more than 3× faster, and trains more than 2.5× faster, with a markedly different across-game performance profile.

    Source: §1 · p. 2, §4.1 · p. 6

  • Takeaways & Limitations

    The results suggest predictive latent diffusion is a promising alternative to separately trained reconstruction- and perceptual-loss latents for online MBRL.

    Source: §1 · p. 2

    Evidence is limited to Atari100k (5 seeds) and four Craftium tasks (3 seeds); JEDI has not been tested on DMControl, Procgen, or higher-resolution domains.

    Source: §5.3 · p. 9

Abstract

from arXiv · show

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.

1 Introduction

JEDI addresses the gap between expensive pixel diffusion and more efficient latent diffusion that lacks end-to-end sequential prediction. It learns predictive latents through conditional denoising and is theoretically motivated, with strong Atari100k results and improved efficiency over DIAMOND.

1 Introduction

  • Pixel diffusion world models such as DIAMOND incur memory, sampling-time, and training-time costs, motivating a move to latent space.Latent approaches such as HI improve efficiency but use latents learned without end-to-end sequential prediction and do not yet match DIAMOND’s performance.Source: §1 · p. 1
  • Diffusion objectives are presented as relevant to end-to-end latent world models because iterative denoising can support expressive prediction for multimodal, stochastic, or sequential targets.The introduction connects this motivation with end-to-end predictive representation learning for MBRL.Source: §1 · p. 2
  • JEDI introduces online end-to-end latent diffusion: its encoder learns through conditional denoising, together with reward and end prediction.The authors describe it as the first online end-to-end latent diffusion world model built on a JEPA framework.Source: §1 · p. 2
  • Latent conditional diffusion admits a variational information-bottleneck decomposition, theoretically motivating denoising as a basis for end-to-end latent learning.The paper presents this as motivation for representation-learning properties, not a guarantee of optimization.Source: §1 · p. 2
  • 43% less VRAM, more than 3× faster world-model sampling, and more than 2.5× faster training than DIAMOND accompany JEDI’s strong Atari100k results.The paper also reports a markedly different performance profile across games relative to DIAMOND.Source: §1 · p. 2

2 Theoretical Motivation: Information Bottleneck in JEPA

JEPA’s predictive objective has a variational information-bottleneck interpretation: it rewards representations useful for predicting targets while penalizing retained input information. JEDI’s conditional denoising admits a related decomposition, but its implementation optimizes the denoising term rather than the full bottleneck objective.

2 Theoretical Motivation: Information Bottleneck in JEPA

  • JEPA’s input branches can represent sequential MDP observations, with actions omitted for clarity, or two views in a multiview self-supervised setting.For categorical targets, the KL term reduces to cross-entropy; for fixed-variance Gaussian targets, it reduces to mean-squared error.Source: §2 · p. 3
  • The JEPA predictive loss appears within a variational bottleneck objective, retaining information useful for predicting the target while penalizing input information in both branches.Recovering the full objective would additionally require explicit optimization of the bottleneck regularizer.Source: §2 · p. 3

2.1 Information Bottleneck in JEDI´s Latent Conditional Diffusion

  • JEDI replaces one-step latent prediction with conditional denoising of a clean future latent, conditioned on the current clean latent and a noised future latent.The denoising objective consists of an endpoint loss and denoising KL terms.Source: §2.1 · p. 3, Figure 3 · p. 3
  • Conditional latent denoising admits a JEPA-style variational bottleneck decomposition, making denoising a predictive learning signal for the encoder rather than an exact optimization of the full objective.The current implementation optimizes the denoising term without the full objective’s additional bottleneck regularizer.Source: §2.1 · p. 3

3 Method

JEDI encodes observations into clean latents, predicts future latents through conditional denoising, and predicts reward and termination from latent states. Both world-model losses update the encoder directly, while practical choices stabilize representation learning and iterative denoising.

3.1 JEDI: Joint Embedding DIffusion World Model

  • JEDI compresses observations into clean latents, predicts the next latent through conditional denoising, and predicts reward and termination from the latent state.The encoder is learned end to end through conditional denoising rather than a reconstruction objective.Source: §3.1 · p. 4
  • Gradients from both world-model losses update the encoder through the clean latent, so predictive training learns the representation without reconstruction.Training perturbs the future latent with diffusion noise and trains the denoiser toward the clean next latent; inference initializes the future from noise and iteratively denoises it.Source: §3.1 · p. 4

3.2 Practical Design Choices

  • Stop-gradient on future targets and a slower encoder learning rate help prevent collapse, while differentiable latent clamping stabilizes iterative denoising.The encoder learning rate is 0.3 times the denoiser learning rate; latent clamping uses s = 3.Source: §3.2 · p. 4
  • Uniformly switching between denoiser and encoder outputs for next-step conditioning balances near-horizon consistency with direct encoder supervision.REINFORCE keeps policy learning lightweight by avoiding backpropagation through multi-step denoising.Source: §3.2 · p. 4

4 Experiments

Experiments evaluate JEDI on Atari100k and Craftium against pixel- and latent-diffusion baselines, finding strong benchmark performance and efficiency gains over DIAMOND. Ablations favor joint predictive latent learning with diffusion denoising, while stabilization choices provide additional support.

4 Experiments

  • Evaluation centers on Atari100k's 26 tasks with five seeds, plus Craftium with three seeds per task; DIAMOND and HI are the primary baselines.For Atari100k, the reported result for each task is the better of two diffusion-model batch sizes, 32 and 64.Source: §4 · p. 5

4.1 Main Atari100k Results

  • On Atari100k, JEDI matches or exceeds strong baselines on IQM, optimality gap, and state-of-the-art game scores, beats HI on four shared tasks, and consistently outperforms DreamerV3 with random frameskips.Figure 6 compares JEDI, HI, and DIAMOND; the random-frameskip comparison uses three Atari environments with skips from 2 to 6.Source: Figure 6 · p. 5, §4.1 · p. 6

4 Experiments

  • Har-HNS complements arithmetic mean HNS by emphasizing performance on tasks that remain very difficult, where arithmetic averaging can obscure relative gains.JEDI and DIAMOND have similar mean HNS but improve in different HNS regimes.Source: §4 · pp. 5–6

4.2 Craftium Results

  • On four Craftium tasks, JEDI performs better overall than HI, supporting results beyond 2D Atari without reconstruction-based latent training.Craftium is a 3D embodied Minecraft-like environment.Source: §4.2 · p. 6

4.3 Efficiency and Scalability

  • Relative to DIAMOND, JEDI uses 43% less VRAM, samples the world model more than 3× faster, and cuts A100 training time from 98 to 38 hours at the same parameter count.The efficiency gains come from moving diffusion from observations to a compact learned latent; JEDI remains competitive while HI trains faster.Source: §4.3 · pp. 6–7, Table 1 · p. 6, Figure 8 · p. 7

4.4 Ablation Study

  • Joint predictive training performs best: replacing denoising with next-state MSE reduces performance, removing diffusion-loss gradients hurts further, and separately trained VAE latents perform worst.Adding decoder gradients also slightly hurts performance.Source: §4.4 · p. 7, Figure 9 · p. 7
  • EMA target encoders, deterministic or absent random switching, and removing switching with latent clamping all underperform the default, though HNS remains high across these design ablations.The ablations characterize stabilization choices as helpful rather than critical; the predictive objective appears to be the main ingredient.Source: §4.4 · p. 7

5 Discussion

JEDI’s aggregate and task-level results differ from DIAMOND’s, with its strongest gains concentrated in games where DIAMOND has low HNS. The discussion relates this profile to game complexity and illustrates distinct learned policies, while noting limits on the evidence.

5.1 Performance Profile and Har-HNS

  • JEDI’s top relative-gain games occur where DIAMOND’s mean HNS is ≈0.06, while DIAMOND’s top-gain games have mean HNS ≈0.81.The task profiles show JEDI excelling against DIAMOND on low-HNS games and DIAMOND tending to outperform on high-HNS games; this contrast is less apparent against TWISTER.Source: §5.1 · p. 7, Figure 11 · p. 8
  • 0.377 Har-HNS puts JEDI ahead of DIAMOND (0.319), although JEDI’s arithmetic mean HNS (1.450) is slightly below DIAMOND’s (1.459).The authors interpret the Har-HNS result as JEDI failing less severely on games that remain hardest.Source: §5.1 · pp. 7–8

5.2 Why Does the Performance Profile Change?

  • JEDI’s largest gains cluster in higher-complexity games, especially those with larger action spaces and shooter-style dynamics.JEDI supplies the actor-critic a latent z0 in R16×8×8 instead of a 64×64×3 observation, making its input about 12× smaller; the authors present the proposed explanation as interpretation, not formal causal proof.Source: §5.2 · p. 8, Figure 12 · p. 8

5.3 Qualitative Analysis

  • On BankHeist, DemonAttack, and Hero, JEDI learns different policies from DIAMOND, including collecting bank spawns, destroying obstacles, and tracking enemies.The paper describes these tasks as belonging to JEDI’s top relative-performance quantile and presents the trajectories as consistent with its broader performance-profile account.Source: §5.3 · p. 8, Figure 13 · p. 9
  • The evidence covers Atari100k and four Craftium tasks, and does not test JEDI on DMControl, Procgen, or higher-resolution domains.The reported results use five seeds for Atari100k and three for the Craftium tasks.Source: §5.3 · p. 9

6 Related Works

Related work spans predictive world models, JEPA representations, and diffusion approaches, while JEDI combines end-to-end predictive latents with diffusion denoising. The appendix frames conventional JEPA prediction as one component of a variational decomposition of the two-view likelihood.

6 Related Works

  • Earlier world-model work spans predictive control [9], imagined-trajectory actor-critic learning [3] [10] [11], Atari100k sequence models [46] [13] [12] [47], pixel diffusion [7], and separately learned latent diffusion [8].The paper places search-based methods [51] [52] [53] and methods adding annotations, event supervision, or explicit exploration [45] [54] [55] outside its scope.Source: §6 · p. 9
  • JEDI differs from prior diffusion uses by learning and predicting compressed latents end to end through a JEPA-style diffusion world-model objective.The related-work discussion contrasts this with separately learned latent spaces [58] [59] [60] [4], reconstruction-based end-to-end latent generation [61], and prior diffusion representation-learning approaches [66] [67] [68] [70] [71] [72] [73].Source: §6 · p. 9

A.1 JEPA as a Variational Information Bottleneck

  • Under the two-view model, the likelihood decomposition separates view-likelihood terms, a JEPA prediction term, a bottleneck regularizer on z1, and a posterior approximation gap.The predictor pθ(z2 | z1) predicts the second-view latent from the first-view latent.Source: §A.1 · p. 14
  • The resulting identity expresses the negative JEPA loss through mutual information between views, latent information terms, a bottleneck regularizer, and the posterior gap.The decomposition is stated in the appendix’s equations for I(X1; X2) and −LJEPA.Source: §A.1 · p. 15

A.2 Conditional Diffusion-JEPA as a Variational Information Bottleneck

The appendix extends the variational JEPA analysis by replacing one-step latent prediction with a conditional reverse diffusion path. Its decomposition connects conditional denoising to the same information-bottleneck structure, with diffusion-specific terms and a posterior gap.

A.2 Conditional Diffusion-JEPA as a Variational Information Bottleneck

  • Conditional diffusion-JEPA predicts the clean target latent z0_2 from noise while conditioning each reverse transition on the clean context latent z0_1.The reverse diffusion path acts as a multi-step stochastic JEPA predictor from z0_1 to z0_2.Source: §A.2 · p. 15
  • The likelihood decomposition separates view-likelihood terms, a conditional diffusion-JEPA prediction term, a bottleneck regularizer on z0_1, and a posterior approximation gap.The path loss is the direct analogue of one-step JEPA prediction loss, but predicts the full reverse diffusion path.Source: §A.2 · pp. 16–17
  • For a Markov forward diffusion process, the path loss decomposes into target-encoder entropy, terminal prior matching, endpoint reconstruction, and conditional denoising KL terms.These terms distinguish endpoint and prior contributions from the intermediate denoising terms.Source: §A.2 · p. 17
  • Even when training uses only conditional denoising, its loss has the variational information-bottleneck structure of one-step JEPA, up to terminal regularization, endpoint constants, and the posterior gap.The appendix gives this equivalence with the stated remainder terms.Source: §A.2 · p. 18

A.2.1 Remark on optimizing the bottleneck regularizer.

The appendix qualifies the information-bottleneck account: JEDI trains the conditional prediction term rather than the exact regularizer. It also describes remainder-term conditions, the training procedure, and implementation changes relative to DIAMOND.

A.2.1 Remark on optimizing the bottleneck regularizer.

  • JEDI does not directly optimize the exact KL regularizer R1, so its bottleneck interpretation is a theoretical design lens rather than the full objective optimized in experiments.Explicitly optimizing R1 is identified as a possible extension.Source: §A.2.1 · p. 18

A.2.2 Remark on C2 remainder term

  • The remainder C2 is constant with respect to reverse-process parameters ψ when the target encoder has fixed entropy, but may depend on φ when that encoder is learnable through the objective.The stated fixed-entropy examples are deterministic or fixed-variance target encoders with stop-gradient.Source: §A.2.2 · p. 19

A.3 JEDI Algorithm

  • JEDI training alternates experience collection with latent-diffusion, reward/termination, and actor-critic updates; imagined actor-critic rollouts sample next latents by reverse diffusion.The reward and termination model is trained with cross-entropy losses, while the actor-critic uses imagined trajectories.Source: §A.3 · p. 21

B Changes Compared to Baseline and Hyperparameters

  • Relative to DIAMOND, implementation changes include a clamped 16×8×8 encoder output, a single-layer diffusion U-Net without downsampling, and no downsampling in the reward/termination encoder.The stated hyperparameter changes were a 1000-step diffusion warm-up and input-data sigma of 1.Source: B Changes Compared to Baseline and… · p. 22

B.1 Benchmark and environment settings

Atari100k follows the standard 26-game, 100k-interaction protocol with five seeds per task, while Craftium evaluation covers four tasks. Atari performance is summarized with aggregate human-normalized metrics.

B.1 Benchmark and environment settings

  • Atari100k uses 26 games, 100k environment interactions, five seeds per task, 64 × 64 RGB observations, and frame skip 4.The setup follows the standard Atari100k protocol and matches the image interface used by the JEDI encoder and world model.Source: §B.1 · p. 22
  • Atari metrics include mean and IQM HNS, plus optimality gap, the average shortfall below HNS = 1, for which lower is better.IQM averages the middle 50% of HNS values across game-seed runs and is more robust to outliers than the arithmetic mean.Source: §B.1 · p. 22
  • Craftium evaluation covers four tasks with three seeds per task, 64 × 64 RGB observations, and training-return reporting.The evaluation uses 300 training steps and 200 collection steps per epoch; SmallRoom-v0 uses 30k environment steps and the other environments use 100k.Source: §B.1 · p. 22

B.2 Third-party assets and licenses

The work uses third-party code and environment assets under their stated licenses. The supplied appendix materials also identify architecture and hyperparameter tables and Atari training curves.

B.1 Third-party assets and licenses

  • The implementation builds on DIAMOND [7] under MIT, IRIS [46] under GPL v3, and Craftium [35] under MIT for code and CC-BY-SA 3.0 for game content.The authors state that the assets are cited and used according to their repository-stated licenses.Source: §B.2 · p. 22
  • Table 2 presents JEDI architecture details, including changes from the pixel-based diffusion baseline [7] and the factor-of-two spatial downsampling operation.The caption identifies the downsampling operation as Maxpool.Source: Table 2 · p. 23
  • Table 3 lists JEDI training hyperparameters and marks changes relative to the pixel-based diffusion world-model baseline [7].Source: Table 3 · p. 24
  • Figure 14 compares JEDI and HI training curves on Atari100k.Source: Figure 14 · p. 25

C Action Space and Shooter Tasks

The Atari100k task table records task-level action counts and shooter-task labels.

C Action Space and Shooter Tasks

  • Table 4 catalogs the Atari100k tasks included in the benchmark.Source: Table 4 · p. 25
  • The task table provides the number of actions for each listed Atari100k task.Source: Table 4 · p. 25
  • The task table also labels which Atari100k tasks are shooter tasks.Source: Table 4 · p. 25

D Supplementary Results

JEDI achieves state-of-the-art optimality gap and IQM on Atari100k, with state-of-the-art results in seven games. Its task-level performance profile differs across low- and high-HNS regimes relative to DIAMOND.

D Supplementary Results

  • JEDI wins disproportionately often in low-HNS regimes, while DIAMOND’s largest relative wins occur in higher-HNS regimes.Source: Figure 15 · p. 26
  • JEDI achieves state-of-the-art optimality gap and IQM, runner-up performance in number of super-human games, and state-of-the-art results in seven games.Table 5 reports batch sizes 32 and 64; TWISTER does not publish seed-level results, preventing exact IQM and optimality-gap computation.Source: Table 5 · p. 26

Related papers

Loading 2605.13013v1…