PaperlayerView exact revision on arXiv

Source-linked AI summary

Diffusion for World Modeling: Visual Details Matter in Atari

Eloi Alonso, Adam Jelley, Vincent Micheli, Anssi Kanervisto, Amos Storkey, Tim Pearce, François Fleuret

arXiv:2405.12399v2cs.LGcs.AIcs.CV

TL;DR

Discrete latent world models may omit visual details important for reinforcement learning, motivating DIAMOND, an agent trained in a diffusion world model. The paper analyzes diffusion design choices for world modeling and reports strong Atari performance, while also demonstrating an interactive Counter-Strike neural game engine. Its scope remains limited by discrete-control evaluation, minimal memory from frame stacking, and deferred reward/termination integration.

  • Problem

    Discrete latent representations in recent world models may ignore visual details important for reinforcement learning, motivating investigation of diffusion-based world modeling.

  • Method

    DIAMOND trains a reinforcement-learning agent inside a diffusion world model and analyzes design choices needed to make diffusion stable and suitable for environment modeling.

  • Results

    DIAMOND achieves a mean human normalized score of 1.46 on Atari 100k, a new best among agents trained entirely within a world model, and models static Counter-Strike gameplay as an interactive neural game engine.

  • Takeaways & Limitations

    Improved visual detail is associated with stronger agent performance, while the diffusion world model also supports interactive modeling of a 3D environment.

  • Takeaways & Limitations

    The main evaluation focuses on discrete control, frame stacking provides minimal memory, and reward or termination prediction is left for future work.

Abstract

from arXiv · show

World models constitute a promising approach for training reinforcement learning agents in a safe and sample-efficient manner. Recent world models predominantly operate on sequences of discrete latent variables to model environment dynamics. However, this compression into a compact discrete representation may ignore visual details that are important for reinforcement learning. Concurrently, diffusion models have become a dominant approach for image generation, challenging well-established methods modeling discrete latents. Motivated by this paradigm shift, we introduce DIAMOND (DIffusion As a Model Of eNvironment Dreams), a reinforcement learning agent trained in a diffusion world model. We analyze the key design choices that are required to make diffusion suitable for world modeling, and demonstrate how improved visual details can lead to improved agent performance. DIAMOND achieves a mean human normalized score of 1.46 on the competitive Atari 100k benchmark; a new best for agents trained entirely within a world model. We further demonstrate that DIAMOND's diffusion world model can stand alone as an interactive neural game engine by training on static Counter-Strike: Global Offensive gameplay. To foster future research on diffusion for world modeling, we release our code, agents, videos and playable world models at https://diamond-wm.github.io.

1 Introduction

World models can improve reinforcement-learning sample efficiency, but discrete latent representations may discard visual information important for learning. DIAMOND instead uses diffusion for environment modeling, with visual fidelity as a central motivation.

  • Motivation: World models are generative environment models used to train reinforcement-learning agents with improved sample efficiency.This can support learning from real-world experience while addressing reinforcement learning’s sample inefficiency.
  • Motivation: Recent methods commonly represent environment dynamics as discrete latent-variable sequences to reduce compounding error over long horizons.The discrete representation helps stabilize multi-step modeling but compresses observations.
  • Motivation: Discrete compression may lose information, reducing generality and reconstruction quality when visual details matter.The passage identifies this as potentially problematic for more real-world scenarios.
  • Motivation: Diffusion offers an alternative generative modeling paradigm motivated by the goal of preserving visual details in world models.The supplied introduction frames diffusion’s use as a response to limitations of compact discrete representations.

2 Preliminaries

The paper formulates environments as partially observable decision processes and explains how diffusion models reverse a noising process to generate observations. A conditional score model can therefore predict future observations from history and actions for imagined rollouts.

  • Reinforcement learning and world models: The environment is modeled as a POMDP where agents observe images rather than latent states and choose actions to maximize discounted return.The transition, reward, and observation functions define the environment and its partial observability.
  • Reinforcement learning and world models: World models estimate environment dynamics and can serve as simulated environments in which reinforcement-learning agents train in imagination.Training cycles between collecting real-environment data, updating the world model, and training the agent within it.
  • Score-based diffusion models: Diffusion models generate data by reversing a process that transforms data distributions into tractable prior noise.The diffusion process uses clean data at τ=0 and an unstructured prior such as a Gaussian at τ=T.
  • Score-based diffusion models: A time-dependent score model estimates the score functions needed to reverse the forward noising process.Score matching trains this model from data without access to the true score, using noised samples from a known perturbation kernel.
  • Diffusion for world modeling: For world modeling, diffusion is conditioned on past observations and actions to generate the next observation directly in a partially observable environment.Sampling the next observation requires iteratively solving the reverse SDE, creating a trade-off between sampling quality and inference cost measured by NFE.

3 Method

DIAMOND adapts the EDM diffusion formulation into a conditional world model that predicts observations from past context and actions. Its design combines noise-level-dependent preconditioning, a U-Net, efficient sampling, and separate reward and termination prediction.

  • Practical choice of diffusion paradigm: DIAMOND uses the EDM formulation rather than DDPM as its practical diffusion paradigm for world modeling.The paper treats this as a key design choice and discusses its implications experimentally.
  • Practical choice of diffusion paradigm: Network preconditioning parameterizes the denoiser as a weighted combination of the noisy observation and a neural-network prediction.The coefficients are designed to keep network inputs and outputs at unit variance across noise levels.
  • Practical choice of diffusion paradigm: The training target adaptively mixes clean signal and added noise according to the degradation level, preventing a trivial low-noise objective.Noise levels are sampled from a log-normal distribution to concentrate training around medium-noise regions.
  • World-model architecture: A U-Net conditions next-observation prediction on a buffer of past observations and actions by concatenating observations and injecting actions through adaptive normalization.This supplies the history needed to model partially observable environment dynamics.
  • Sampling: Euler’s method generates observations efficiently without the additional NFE or complexity of higher-order or stochastic samplers.The paper selected it as effective for sampling from the trained diffusion model.
  • World-model architecture: A separate CNN-LSTM reward and termination model complements the diffusion model for reinforcement-learning training in imagination.The recurrent component handles partial observability while predicting scalar rewards and termination signals.

4 Experiments

DIAMOND is evaluated on Atari 100k against world-model agents using standardized aggregate metrics and confidence intervals. It performs strongly across the benchmark, achieving a new best mean human-normalized score among agents trained entirely within a world model.

  • Benchmark and setup: Atari 100k evaluates 26 games after 100k real-environment actions, roughly two hours of human gameplay.Unconstrained Atari agents are usually trained for 50 million steps, or 500 times more experience.
  • Benchmark and setup: DIAMOND was trained from scratch with 5 random seeds per game, requiring approximately 2.9 days per run on one RTX 4090.Each run used around 12GB of VRAM, for 1.03 GPU years in total.
  • Baselines and metrics: The comparison includes recent agents trained entirely within world models, while broader model-free and search-based comparisons are reported separately because their techniques are not directly comparable.The broader comparison includes BBF and EfficientZero.
  • Baselines and metrics: Mean and interquartile mean human-normalized scores are reported with stratified bootstrap confidence intervals, alongside performance profiles and additional metrics.The aggregate metrics follow recommendations addressing limitations of point estimates.
  • Results: 1.46 mean HNS makes DIAMOND superhuman and establishes a new best among agents trained entirely within a world model.DIAMOND outperforms human players on 11 games and achieves an IQM of 0.64, on par with STORM and above the other listed baselines.

5 Analysis

The analysis examines diffusion-framework choice, denoising-step count, and visual consistency against IRIS. EDM is more stable than DDPM at low step counts, while richer visual details and consistency align with stronger agent performance.

  • 5.1 Choice of diffusion framework: DDPM trajectories quickly drift out of distribution through severe compounding error, whereas EDM remains stable over long horizons even with one denoising step.The comparison uses autoregressive trajectories up to t = 1000 with n ≤10 denoising steps.
  • 5.1 Choice of diffusion framework: EDM’s improved training objective and adaptive signal-noise mixing produce better score estimates when noise dominates, enabling higher-quality generations with fewer denoising steps.The DDPM objective can learn the identity function under dominant noise, degrading generation quality and causing compounding error.
  • 5.2 Choice of the number of denoising steps: Partial observability can create multimodal observation distributions, making single-step prediction an out-of-distribution expectation; DIAMOND therefore uses n = 3.Some deterministic games can be accurately modeled with one denoising step, but the paper reports a broader limitation for other games.
  • 5.2 Choice of the number of denoising steps: Single-step sampling can blur unpredictable outcomes by interpolating across modes, while multi-step sampling drives generation toward a particular mode and produces crisp images.In Boxing, both methods correctly predict the policy-controlled white player because its actions are known to the world model.
  • 5.3 Qualitative visual comparison with IRIS: Compared with IRIS, DIAMOND generates higher-quality, more faithful trajectories without the observed cross-frame inconsistencies in Asterix, Breakout, and Road Runner.In Breakout, DIAMOND reliably updates the score by +7 when a red brick is broken.
  • 5.3 Qualitative visual comparison with IRIS: Improved visual-detail consistency is generally reflected by greater agent performance, while DIAMOND uses 3 NFE per frame versus IRIS’s 16 NFE at the same 64 × 64 resolution.The paper notes that DIAMOND also has significantly fewer parameters and takes less time to train than IRIS.

6 Scaling the diffusion world model to Counter-Strike: Global Offensive4

DIAMOND’s diffusion world model was scaled to static CS:GO gameplay to create an interactive neural game engine for Dust II. The model generated stable long trajectories but retained important limitations in rarely visited areas and state memory.

  • 87 hours of static CS:GO gameplay trained DIAMOND’s diffusion world model as an interactive neural game engine for Dust II.
  • The CS:GO model reduced observations to 56 × 30 resolution and used a 381M-parameter diffusion system with a 51M-parameter upsampler.The combined model was trained for 12 days on an RTX 4090.
  • Stochastic sampling and 10 upsampler denoising steps improved visual quality while preserving a three-step dynamics model running at 10Hz on an RTX 3090.
  • The model generated stable trajectories over hundreds of timesteps but was more likely to drift out-of-distribution in less frequently visited map areas.
  • Limited memory could cause state loss near walls or when visibility ended, while sparse multi-jump examples led to repeated mid-air jumps.

7 Related work

Related work spans world models for reinforcement learning, generative vision and video models, diffusion methods for reinforcement learning, and neural game engines. DIAMOND connects these lines by using diffusion for environment modeling.

  • World models: World models support reinforcement learning in imagination, with prior Atari and latent-state approaches addressing sample efficiency and compounding error.
  • Generative vision models: Generative vision research includes discrete-token autoregressive models and diffusion models for image generation.
  • Generative vision models: Video-generation methods extend these ideas to sequential visual prediction through discrete autoencoders, transformers, prompt conditioning, and improved autoregressive modeling.
  • Diffusion for reinforcement learning: Diffusion has been applied to reinforcement-learning policies, planners, reward models, and offline trajectory augmentation.
  • Generative game engines: Neural game engines use generative architectures such as GANs, Neural GPUs, and diffusion models to produce playable environments.

8 Limitations

The paper identifies three limitations: evaluation is focused on discrete control, memory relies on minimal frame stacking, and reward or termination prediction is separate from diffusion.

  • The main evaluation focuses on discrete control, leaving application of DIAMOND to continuous environments for future study.
  • Frame stacking provides only minimal memory; an autoregressive transformer could support longer-term memory and better scalability.
  • Reward and termination prediction remain separate because integrating these objectives into diffusion is nontrivial and could make the world model unnecessarily complex.

9 Conclusion and Broader Impact

DIAMOND is a reinforcement learning agent trained in a diffusion world model, designed to model environment dynamics while retaining visual detail. It achieved strong Atari performance and also demonstrated interactive 3D game modeling, with sampling cost and visual quality requiring a trade-off.

  • Conclusion: DIAMOND is a reinforcement learning agent trained entirely within a diffusion world model.
  • Conclusion: 1.46 mean human normalized score on Atari 100k was a new best among agents trained entirely within a world model.
  • Conclusion: The paper attributes improved performance in some games to better modeling of critical visual details.
  • Conclusion: The diffusion world model modeled 3D CS:GO environments from static gameplay and served as a real-time neural game engine.
  • Broader Impact: World models may improve sample efficiency and safety by letting agents train in simulated environments, although model imperfections can produce suboptimal or unexpected behavior.
  • Sampling: Reverse-SDE sampling generates the next observation from past observations and actions, with solver choice creating a visual-quality versus inference-cost trade-off.
  • Architecture: DIAMOND conditions a U-Net diffusion model on the last four frames, actions, and diffusion time, while separate CNN-LSTM models predict reward and termination.
  • Reinforcement learning objectives: The policy uses REINFORCE with a value baseline and entropy maximization, while the value network learns from λ-returns over imagined trajectories.

J Broader comparison to model-free and search-based methods

DIAMOND remains competitive against model-free and search-based Atari 100k agents, despite not using their additional search or advanced training techniques. Its compounding-error analysis quantitatively compares DDPM and EDM world models.

  • Broader comparison to model-free and search-based methods: DIAMOND outperforms EfficientZero and BBF on some Atari 100k games, although those methods achieve greater overall performance using lookahead search or advanced reinforcement learning techniques.EfficientZero uses computationally expensive lookahead Monte Carlo tree search, while BBF uses periodic network resets with hyperparameter scheduling.
  • Broader comparison to model-free and search-based methods: Figure 8 quantitatively measures the compounding error of DDPM- and EDM-based world models using average pixel drift from expert reference trajectories.The trajectories span 1000 timesteps and follow identical action sequences from the same starting frame.
  • Broader comparison to model-free and search-based methods: DDPM becomes more stable with more denoising steps but remains less stable than 1-step EDM, even with 10 denoising steps.The comparison uses 400 held-out reference trajectories in Breakout.

L Quantitative ablation on reducing the number of denoising steps

The paper examines denoising-step choices alongside environment and architecture design for diffusion world models. It uses static datasets for visual-model evaluation and does not perform reinforcement learning in the non-Atari environments.

  • Quantitative ablation on reducing the number of denoising steps: Table 7 compares DIAMOND with 3 denoising steps against a single-step EDM variant across its 10 highest-performing games.The 3-step configuration is the default used for Table 1, while the 1-step results use a single seed.
  • Quantitative ablation on reducing the number of denoising steps: Single-step agents generally perform worse than the default 3-step agents, particularly in Boxing, despite apparent similarity in Figure 8.Because the 1-step results use one seed, they have higher variance, but still provide supporting evidence for the qualitative analysis.
  • Environments: The CS:GO dataset contains 190k frames of high-skill human gameplay, with 150k frames for training and 40k for evaluation.Observations are resized to 64×64 pixels and actions include mouse and keyboard inputs captured at 16Hz.
  • Environments: The motorway dataset provides 4.4 hours of daylight human-driving data, split into five training trajectories and two testing trajectories.The data are downsampled to 10Hz and resized to 64×64 observations.
  • Environments: These static-dataset experiments evaluate DIAMOND’s diffusion model without reinforcement learning because no standard reinforcement learning protocol exists for the environments.The investigation focuses on training and evaluating the diffusion model itself.
  • Diffusion Model Architectures: DIAMOND evaluates frame-stacking and autoregressive cross-attention architectures, while U-Net 3D is shown as a comparison architecture that diffuses frame blocks jointly.Frame-stacking concatenates previous frames with the next noised frame; cross-attention uses a history encoder.

M.4 Analysis

DIAMOND produces visually strong trajectories in CS:GO and motorway driving, and its frame-stack model responds plausibly to many action inputs. Longer rollouts under unlikely action sequences can become unstable or degenerate.

  • Visual quality: DIAMOND outperforms the baselines across all visual-quality metrics, with frame-stacking ranking above cross-attention and IRIS variants.The reported qualitative ordering is DIAMOND frame-stack > DIAMOND cross-attention ≈ IRIS 64 > IRIS 16 > DreamerV3.
  • Sampling rate and visual quality: DIAMOND frame-stack is faster than IRIS with K = 16 tokens, while IRIS with K = 64 is slowed by an additional 2.8×.DreamerV3 is an order of magnitude faster, but its trajectories have lower visual quality.
  • Qualitative trajectories: Generated trajectories are often plausible over reasonable horizons, including correct CS:GO level geometry and a motorway car overtaking on the left.These examples are sampled every 25 timesteps from the generated trajectories.
  • Action-conditioned rollouts: In motorway driving, steering actions move the camera as expected, while slowing down causes the model to predict stationary traffic ahead.The action-conditioned rollouts share the same six initial real frames but use five different action sequences.
  • Action-conditioned rollouts: In CS:GO, common mouse and fire actions have intended immediate effects but become unstable after a few frames when the action sequence is unlikely in the demonstrations.Longer rollouts can degenerate, including implausible views directly into the ground.
Loading 2405.12399v2…