Mastering Atari with Discrete World Models

Problem

Framing

Atari world models had been sample-efficient but not competitive at 200M-frame scale. DreamerV2 closes that gap with a discrete latent RSSM and latent-space actor-critic training, reaching human-level Atari performance and beating single-GPU Rainbow and IQN on aggregate scores.

Currently Used Methods

Foundational

• @haWorldModels2018 — latent world models for policy learning from compact imagined trajectories.
  • Limitation in context: Atari-scale control remained far below strong model-free baselines.
• @hafnerDreamerV1_2019 — latent imagination with continuous RSSM and actor-critic learning.
  • Limitation in context: Gaussian latents underfit visually complex Atari dynamics.
• @mnihDQN2015 — value-based Atari control from pixels with model-free Q-learning.
  • Limitation in context: no learned world model for planning or imagination.
• @mnihA3C_2016 — asynchronous model-free policy optimization for Atari.
  • Limitation in context: weaker data efficiency than latent model-based training.
• "Rainbow: Combining Improvements in Deep Reinforcement Learning" — strong model-free Atari baseline combining several DQN extensions.
  • Limitation in context: needs environment interaction instead of imagined rollouts.

Proposed Method

Architecture

DreamerV2 keeps the Dreamer RSSM factorization but replaces Gaussian stochastic state with 32 categorical variables of 32 classes each. The world model predicts image, reward, and discount from latent state; actor and critic are MLPs trained on imagined latent trajectories.

Evaluation

Datasets

• Atari 55-game benchmark.
• 200M environment steps.
• Action repeat 4.

Metrics

• Gamer Median.
• Gamer Mean.
• Record Mean.
• Clipped Record Mean.

Headline results

• Atari-55: Gamer Median 2.15, Gamer Mean 11.33.
• Atari-55: Record Mean 0.44, Clipped Record Mean 0.28.
• Atari-55 with schedules: Gamer Median 2.64, Gamer Mean 10.45.
• IQN: Gamer Median 1.29, Gamer Mean 8.85.
• Rainbow: Gamer Median 0.93, Gamer Mean 7.03.

Ablations

• Remove discrete latents: Gamer Mean drops from 11.33 to 3.71.
• Remove KL balancing: Gamer Mean drops to 1.13.
• Remove image gradients: Gamer Mean drops to 1.74.
• Remove reward gradients: Gamer Mean drops to 6.18.

Method Strengths and Weaknesses

Strengths

• Discrete latents sharply improve Atari aggregate performance over Gaussian Dreamer-style latents.
• Latent-space imagination reaches human-level Atari scores with one GPU and one environment.
• KL balancing materially stabilizes representation learning; removing it collapses gamer mean.
• Same agent works across 55 games with shared hyperparameters.

Weaknesses

• Performance still depends on several coupled tricks: discrete latents, KL balancing, layer norm, reward gradients.
• Actor learning still needs mixed Reinforce and straight-through gradients.
• Main benchmark uses 200M steps, so sample efficiency is not the only claim.
• Sparse-reward games still need tuning; Montezuma improves with lower discount.

Suggestions from the authors

• Improve exploration for sparse-reward Atari tasks.
• Increase data efficiency by raising training frequency.
• Study better discount settings for sparse-reward environments.
• Extend discrete world models beyond Atari to broader domains.

Links

Prior Papers

• @haWorldModels2018 — early latent world-model agent that motivates learning compact predictive dynamics for control.
• @hafnerDreamerV1_2019 — direct predecessor introducing latent imagination and actor-critic learning that DreamerV2 upgrades with discrete latents.

Further Papers

• @hafnerDreamerV3_2023 — follow-on Dreamer paper that generalizes the world-model recipe across domains and scales.