Source-linked AI summary
Mastering Atari with Discrete World Models
Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, Jimmy Ba
TL;DR
Competitive Atari benchmarks had remained difficult for world models, which had not matched state-of-the-art model-free agents. DreamerV2 learns behavior from latent-space predictions of a separately trained world model using discrete representations, and achieves human-level Atari performance while surpassing leading single-GPU model-free agents.
Problem
World models had not been accurate enough to achieve competitive performance on challenging Atari benchmarks.
Method
DreamerV2 learns a world model from past experience, then trains an actor and critic on imagined sequences of compact model states.
Results
DreamerV2 achieves human-level Atari performance and outperforms top single-GPU agents Rainbow and IQN using the same computational budget and training time.
Takeaways & Limitations
The results provide a proof of concept that model-based reinforcement learning can outperform top model-free algorithms on competitive reinforcement-learning benchmarks.
Abstract
from arXiv · showhide
Intelligent agents need to generalize from past experience to achieve goals in complex environments. World models facilitate such generalization and allow learning behaviors from imagined outcomes to increase sample-efficiency. While learning world models from image inputs has recently become feasible for some tasks, modeling Atari games accurately enough to derive successful behaviors has remained an open challenge for many years. We introduce DreamerV2, a reinforcement learning agent that learns behaviors purely from predictions in the compact latent space of a powerful world model. The world model uses discrete representations and is trained separately from the policy. DreamerV2 constitutes the first agent that achieves human-level performance on the Atari benchmark of 55 tasks by learning behaviors inside a separately trained world model. With the same computational budget and wall-clock time, Dreamer V2 reaches 200M frames and surpasses the final performance of the top single-GPU agents IQN and Rainbow. DreamerV2 is also applicable to tasks with continuous actions, where it learns an accurate world model of a complex humanoid robot and solves stand-up and walking from only pixel inputs.
1 INTRODUCTION
World models promise generalization and planning, but have not been accurate enough to match state-of-the-art model-free agents on competitive Atari benchmarks. DreamerV2 addresses this gap by learning behavior entirely within a separately trained world model and surpassing leading single-GPU model-free agents.
- World models represent environmental knowledge explicitly and can predict potential action outcomes to support planning and generalization.They can also support transfer to novel tasks and directed exploration.
- Competitive Atari benchmarks have historically required model-free algorithms to achieve human-level performance.Earlier attempts to learn accurate Atari world models did not achieve competitive performance.
- DreamerV2 learns successful behaviors purely within a separately trained world model and achieves human-level performance on Atari.The approach uses discrete latent representations and modifications to the Dreamer agent, including KL-loss balancing.
- DreamerV2 outperforms the top single-GPU Atari agents Rainbow and IQN using a single GPU and a single environment instance.The comparison uses the same computational budget and training time.
2 DREAMERV2
DreamerV2 learns a world model from past experience, then trains actor and critic behaviors from imagined sequences of compact latent states. Its discrete-state RSSM, KL balancing, and separate behavior learning support prediction without generating images during policy training.
- Algorithm: DreamerV2 learns the world model from past experience, trains actor and critic networks on imagined compact-state sequences, and executes the actor to expand the dataset.The three components are world-model learning, behavior learning from imagined sequences, and environment interaction.
- World model: The RSSM combines deterministic recurrent states with posterior states informed by images and prior states that predict them without current images.The concatenated deterministic and stochastic states form the compact model state used for reconstruction and reward and discount prediction.
- Behavior learning: The learned transition predictor enables imagination without observing or generating images, while the fixed world model supports actor and critic learning on imagined trajectories.The actor selects actions in latent rollouts, and the critic accumulates predicted future rewards.
- World model: DreamerV2 represents stochastic states with multiple categorical variables rather than the diagonal Gaussian latent variables used in earlier Dreamer models.The categorical representation is used in the world model's stochastic state.
- World-model optimization: KL balancing trains the prior faster than the approximate posterior, encouraging an accurate prior while avoiding regularization toward a poorly trained transition model.The method uses different learning rates for the prior and posterior contributions to the KL loss.
- Behavior learning: The critic uses λ-targets over multi-step imagined trajectories, incorporating reward information across different horizons for value learning.Imagined on-policy trajectories make n-step targets available, and λ-targets combine returns across horizons.
3 EXPERIMENTS
DreamerV2 is evaluated on 55 Atari games under multiple aggregation protocols and against four model-free baselines. It outperforms the model-free agents across metrics, while ablations identify discrete latents, KL balancing, and image-based representation learning as important components.
- Experimental setup: DreamerV2 reaches 200M environment steps in under 10 days using one NVIDIA V100 GPU and one environment instance.During training, it learns from 468B imagined compact states versus 50M real-environment inputs after action repeat.
- Experimental setup: The evaluation uses 55 Atari games, 200M environment steps, action repeat of 4, and sticky actions.The protocol also specifies a 108,000-step episode limit and no access to life information.
- Atari performance: DreamerV2 outperforms IQN, Rainbow, C51, and DQN across the reported Atari aggregation metrics.The comparison uses sticky-action scores from the Dopamine framework, which may differ from results reported in the original deterministic setups.
- Atari performance: The recommended metric is task mean clipped record-normalized score because it considers games similarly without domination by outlier scores.Scores are normalized by human world records and clipped so exceeding a record does not increase the score further.
- Atari performance: DreamerV2 achieves comparable or higher performance on most individual games except Video Pinball.Its strongest improvements over model-free agents occur on James Bond, Up N Down, and Assault.
- Ablation study: Categorical latents outperform Gaussian latents on 42 tasks, underperform on 8, and tie on 5.KL balancing outperforms the standard KL regularizer on 44 tasks, underperforms on 6, and ties on 5.
- Ablation study: Stopping image gradients decreases performance on 51 tasks, whereas stopping reward gradients improves performance on 15, decreases it on 22, and ties on 18.These results indicate strong reliance on image learning signals and suggest that reward-independent representations may generalize better to unseen situations.
- Ablation study: Using only Reinforce gradients improves performance on 18 tasks, decreases it on 24, and ties on 13; using only straight-through gradients improves 5, decreases 44, and ties 6.Mixing Reinforce and straight-through gradients substantially improves James Bond and Seaquest and raises gamer-normalized task mean.
4 RELATED WORK
Related work spans model-free Atari agents, video-prediction approaches, latent world models, and planning-based systems such as MuZero. DreamerV2 is evaluated in the established 200M-frame regime, where prior approaches had not combined competitive Atari performance with its approach.
- Model-free Atari: Model-free Atari research extends DQN with bias correction, prioritized replay, architectural improvements, and distributional value learning.
- SimPLe: SimPLe predicts video frames in pixel space and uses those predictions to train PPO on Atari.Its highest reported gamer-normalized median score in the data-efficient regime was 0.28, far from human-level performance.
- Evaluation setting: DreamerV2 is compared at 200M frames because this is the established competitive evaluation regime with many successful model-free algorithms.
- MuZero: MuZero uses a reward-and-value sequence model with Monte-Carlo Tree Search, achieving strong results with substantial engineering and computational resources.DreamerV2 is described as simpler and reproducible on a single GPU in 10 days.
5 DISCUSSION
DreamerV2 demonstrates that latent-space model-based reinforcement learning can reach human-level Atari performance and outperform leading model-free agents. The discussion also reports applicability to continuous-control humanoid tasks from pixel inputs and identifies broader possibilities for world models.
- Main findings: DreamerV2 achieves human-level performance on Atari by learning behaviors purely from latent-space predictions of a separately trained world model.
- Main findings: Using one GPU and one environment instance, DreamerV2 outperforms Rainbow and IQN under the same computational budget and training time.
- Representations: DreamerV2’s performance is not affected by whether its representations are specially trained to predict rewards, indicating reliance on image information for generally useful representations.
- Broader implications: World models may support transfer, multi-task learning, sample-efficient physical-robot learning, and exploration based on uncertainty estimates.
- Continuous control: DreamerV2 also applies to continuous-action control, including a humanoid environment using only image inputs.The humanoid task uses a 21-dimensional continuous action space.
B MONTEZUMA’S REVENGE
On the sparse-reward Montezuma’s Revenge task, DreamerV2 achieves strong performance without an explicit exploration mechanism. A lower discount factor is used for this game, and performance matches the ICM exploration method.
- Performance: DreamerV2 reaches about the same performance as ICM on Montezuma’s Revenge without an explicit exploration mechanism.
- Performance: DreamerV2 outperforms existing model-free approaches on Montezuma’s Revenge, a hard-exploration Atari game.
- Configuration: Using γ = 0.99 improves DreamerV2’s Montezuma’s Revenge performance, possibly by stabilizing value learning under sparse rewards.
- Interpretation: The authors suggest that world-model benefits for sparse rewards may involve improved generalization, compact-latent policy optimization, or reward-predictor generalization.
C SUMMARY OF MODIFICATIONS
DreamerV2 evolves Dreamer through targeted modifications to its latent representation, optimization, gradients, and model capacity. The reported ablations support categorical latents and KL balancing, while the evaluation remains computationally constrained and uses several Atari comparison views.
- Overview: DreamerV2 is developed by modifying the original Dreamer agent to improve Atari performance.
- Changes that helped: Helpful changes include categorical latents, KL balancing, Atari-specific Reinforce gradients, and increased model size.
- Limitations: A comprehensive ablation study was infeasible because each change would require over 60,000 GPU hours across 55 tasks, 5 seeds, and 10 days.
- Hyperparameters: The recommended Atari tuning ranges include β ∈ {0.1, 0.3, 1, 3} for KL loss scale and γ ∈ {0.99, 0.999} for discount factor.
- Atari comparison: At 200M steps, DreamerV2 outperforms four model-free algorithms on gamer-normalized scores except Video Pinball.The authors hypothesize that reconstruction loss overlooks the one-pixel ball in Video Pinball.
- Atari comparison: DreamerV2 outperforms IQN and Rainbow in all four aggregated scores, while exhibiting a different task-performance profile.
G LATENTS AND KL BALANCING ABLATIONS
The ablations compare DreamerV2 with Gaussian latent variables and without KL balancing. Categorical latents and KL balancing both substantially improve performance across many tasks and all four aggregated scores.
- Categorical latent variables substantially improve performance across many tasks.
- Removing KL balancing substantially reduces performance across many tasks.
- Both techniques improve all four aggregated scores.
H REPRESENTATION LEARNING ABLATIONS
The ablations compare image prediction, reward prediction, and both for learning model representations. Image gradients are crucial, whereas reward gradients are unnecessary and can be stopped; image-only representations outperform reward-specific ones on several tasks.
- Image gradients are crucial for successful world-model representation learning.
- Reward gradients are not necessary for the world model to succeed and can be stopped.
- Image-only representations outperform reward-specific representations on a number of tasks.They are not biased toward previously encountered rewards, suggesting better generalization to unseen situations.
I POLICY LEARNING ABLATIONS
The policy-learning ablations compare Reinforce gradients, straight-through gradients, and their combination. Reinforce gradients are crucial, while straight-through gradients are usually unimportant but substantially improve a small number of games when combined with Reinforce gradients.
- Reinforce gradients are crucial for policy learning across the evaluated tasks.
- Straight-through gradients are not important for most tasks.
- Combining straight-through and Reinforce gradients substantially improves performance on a small number of games, notably Seaquest.
- The authors conjecture that straight-through gradients have low variance, while Reinforce gradients are unbiased.The conjecture attributes early-learning help to low variance and reliable policy updates to unbiased gradients.
J ADDITIONAL ABLATIONS
Additional ablations examine layer normalization in the GRU and training from random-policy experience. Layer normalization has task-dependent effects, while random-data comparisons identify tasks requiring non-trivial exploration.
- Layer normalization in the GRU increases and decreases performance on roughly equal numbers of tasks.Its benefit therefore depends on the task at hand.
- Training from uniformly random-policy experience highlights which tasks require non-trivial exploration.
- The random-data comparison can guide future directed-exploration work using world models.
K ATARI TASK SCORES
Table K.1 compares DreamerV2, IQN, and Rainbow across 55 Atari games against random-action, human-gamer, and human-world-record baselines.
- The evaluation covers 55 Atari games selected because they are common among most papers in the literature.
- DreamerV2, IQN, and Rainbow are compared with random actions, DeepMind’s human gamer, and the human world record.
- Algorithm scores are bold when they fall within 5% of the best algorithm.
- Individual scores are averaged across seeds, while aggregated scores must be computed before averaging across seeds.