M ASTERING A TARI WITH D ISCRETE W ORLD M ODELS
Anonymous authors
Paper under double-blind review
A BSTRACT
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-efﬁciency.
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 ﬁrst agent that achieves human-level performance on the Atari benchmark oftasks by learning behaviors inside a separately trained world model.
With the same computational budget and wall-clock time, DreamerV2 reaches 200M frames and exceeds the ﬁnal performance of the top single-GPU agents IQN and Rainbow.
I NTRODUCTION
To successfully operate in unknown environments, reinforcement learning agents need to learn about their environments over time.
World models are an explicit way to represent an agent’s knowledge about its environment.
Compared to model-free reinforcement learning that learn through trial and error, world models facilitate generalization and can predict the outcomes of potential actions to enable planning (Sutton, 1991).
Capturing general aspects of the environment, world models have been shown to be effective for transfer to novel tasks (Byravan et al., 2019), directed exploration (Sekar et al., 2020), and generalization from ofﬂine datasets (Yu et al., 2020).
When the inputs are high-dimensional images, latent dynamics models predict ahead in a abstract latent space (Watter et al., 2015;
Ha and Schmidhuber, 2018; Hafner et al., 2018; Zhang et al., 2019).
Predicting compact representations instead of images has been hypothesized to reduce the effect of accumulating errors and their small memory footprint enables thousands of parallel predictions on a single GPU (Hafner et al., 2018; 2019).
Leveraging this approach, the recent Dreamer agent (Hafner et al., 2019) has solved a wide range of continuous control tasks from image inputs.
Despite their intriguing properties, world models have so far not been accurate enough to compete with the state-of-the-art model-free algorithms on the most competitive benchmarks.
The well-established
[Figure]
J a m e s
B o n d
U p N
D o w n
G o p h
e r
A t l a n t i s Q b e r t
T i m e
P il o t A ss a u l t B r e a k o u t K r u ll
R o a d
R unn e r
I c e H o c k e y F r o s t b i t e K u n g
F u M a s t e r S k ii n g T e nn i s
W i z a r d
O f W o r
N a m e
T h
i s G a m e
A s t e r i x
G r a v i t a r P h
o e n i x Z a xx o n
B e a m
R i d e r
F i s h
i n g D e r b y M s P a c m a n
M o n t e z u m a
R e v .
E n d u r o B o w li n g
A li e n
A s t e r o i d s F r ee w a y
P i t f a ll K a n g a r oo
C e n t i p e d e
T u t a n k h a m
B e r z e r k
P o n g
S o l a r i s
B a n k H e i s t S e a q u e s t
D o u b l e D un k R o b o t a n k C r a z y
C li m b e r B a tt l e Z o n e R i v e rr a i d
P r i v a t e
E y e
Y a r s R e v e n g e
A m i d a r H e r
o
Sp a c e I n v a d e r s B o x i n g
V e n t
u
r e
C h
o pp e r C o m .
S t a r G unn e r D e m o n A tt a c k
V i
d e o
P i n b a ll
[Figure]
DreamerV2 is the ﬁrst agent that learns within a separately trained world model to achieve humanlevel performance on the Atari suite, demonstrating the high accuracy of its learned world model.
DreamerV2 further outperforms the state-of-the-art single-GPU agents Rainbow (Hessel et al., 2018) and IQN (Dabney et al., 2018), achieving comparable or higher performance on almost all games.
The baseline scores are takes from the Dopamine framework (Castro et al., 2018).
Atari benchmark (Bellemare et al., 2013) historically required model-free algorithms to achieve human-level performance, such as DQN (Mnih et al., 2015), A3C (Mnih et al., 2016), or Rainbow (Hessel et al., 2018).
Several attempts at learning accurate world models of Atari games have been made, without achieving competitive performance (Oh et al., 2015; Chiappa et al., 2017; Kaiser et al., 2019).
On the other hand, the recently proposed MuZero agent (Schrittwieser et al., 2019) shows that planning can achieve impressive performance on board games and deterministic Atari games given extensive engineering effort and a vast computational budget.
However, its implementation is not available to the public and it would require over 2 months of computation to train even one agent on a GPU, rendering it impractical for most research groups.
In this paper, we introduce DreamerV2, the ﬁrst reinforcement learning agent that achieves humanlevel performance on the Atari benchmark by learning behaviors purely within a separately trained world model, as shown in Figure 1.
Learning successful behaviors purely within the world model demonstrates that the world model learns to accurately represent the environment.
To achieve this, we apply small modiﬁcations to the Dreamer agent (Hafner et al., 2019), such as using discrete latents and balancing terms within the KL loss.
Using a single GPU and a single environment instance,
DreamerV2 outperforms top single-GPU Atari agents Rainbow (Hessel et al., 2018) and IQN (Dabney et al., 2018), which rest upon years of model-free reinforcement learning research (Van Hasselt et al., 2015; Schaul et al., 2015; Wang et al., 2016; Bellemare et al., 2017; Fortunato et al., 2017).
Moreover, aspects of these algorithms are complimentary to our world model and could be integrated into the
Dreamer framework in the future.
To rigorously compare the algorithms, we report scores normalized by both a human gamer (Mnih et al., 2015) and the human world record (Toromanoff et al., 2019) and make a suggestion for reporting scores going forward.
D REAMER V
We present DreamerV2, an evolution of the Dreamer agent (Hafner et al., 2019).
We refer to the original Dreamer agent as DreamerV1 throughout this paper.
This section describes the complete
DreamerV2 algorithm, consisting of the three typical components of a model-based agent (Sutton, 1991).
We learn the world model from a dataset of past experience, learn an actor and critic from imagined sequences of compact model states, and execute the actor in the environment to grow the experience dataset.
In Appendix A, we include a list of changes that we applied to the open sourced
DreamerV1 implementation 1 and which of them we found to increase empirical performance.
2. W ORLD M ODEL L EARNING
World models summarize an agent’s experience into a predictive model that can be used in place of the environment to learn behaviors.
When inputs are high-dimensional images, it is beneﬁcial to learn compact state representations of the inputs to predict ahead in this learned latent space (Watter et al., 2015; Karl et al., 2016;
Ha and Schmidhuber, 2018).
These models are called latent dynamics models.
Predicting ahead in latent space not only facilitates long-term predictions, it also allows to efﬁciently predict thousands of compact state sequences in parallel in a single batch, without having to generate images.
DreamerV2 builds upon the world model that was introduced by PlaNet (Hafner et al., 2018) and used in DreamerV1, by replacing its Gaussian latents with categorical variables.
Experience dataset The world model is trained from the agent’s growing dataset of past experience that contains sequences of of images x 1: T , actions a 1: T , rewards r 1: T , and discount factors γ 1: T .
The discount factors equal a ﬁxed hyper parameter γ = 0 .
995 for time steps within an episode and are set to zero for terminal time steps.
For training, we use batches of B = 50 sequences of ﬁxed length
L = 50 that are sampled randomly within the stored episodes.
To observe enough episode ends during training, we sample the start index of each training sequence uniformly within the episode and then clip it to not exceed the episode length minus the training sequence length.
Model components The world model consists of ﬁve components.
Given a batch of sequences from the dataset, the representation model summarizes the history of images x 1: t and actions a t − into a compact Markovian model state z t .
From this model state, we predict the current image x t , current reward r t , discount factor γ t , and next model state z t +1 .
The model components are:
Representation model:
[Equation]
All components are implemented as neural networks and φ describes their combined parameter vector.
The transition predictor guesses the next model state only from the current model state and the action but without using the next image, so that we can later learn behaviors by predicting sequences of model states without having to observe or generate images.
The discount predictor lets us estimate the probability of an episode end when learning behaviors from model predictions.
Neural networks The representation model is implemented as a Convolutional Neural Network (CNN; LeCun et al., 1989) followed by a Recurrent State Space Model (RSSM; Hafner et al., 2018).
The RSSM uses a Gated Recurrent Unit (GRU; Cho et al., 2014) that outputs a stochastic variable at each time step and feeds it back for the next time step.
The model state z t is the concatenation of deterministic GRU state and a sample of the stochastic variable.
The transition predictor is an RSSM that shares most of its parameters with the representation model as in PlaNet.
The image predictor is a transposed CNN and the reward and discount predictors are Multi-Layer Perceptrons (MLPs).
We down-scale the 84 × 84 grayscale images to 64 × 64 pixels so that we can apply the convolutional architecture of DreamerV1.
We use the ELU activation function for all components of the model (Clevert et al., 2015).
The world model uses a total of 20M trainable parameters.
Distributions
The image predictor outputs the mean of a diagonal Gaussian likelihood with unit variance, the reward predictor outputs a univariate Gaussian with unit variance, and the discount predictor outputs a Bernoulli likelihood.
In prior work, the latent variable in the model state was a diagonal Gaussian that used reparamterization gradients during backpropagation (Kingma and
Welling, 2013; Rezende et al., 2014).
In DreamerV2, we instead use a vector of several categorical variables and leverage straight-through gradients during backpropagation (Bengio et al., 2013), which are also easy to implement using automatic differentiation.
All components of the world model are optimized jointly.
The distributions produced by the image predictor, reward predictor, discount predictor, and transition predictor are trained to maximize the log-likelihood of their corresponding targets.
The representation model is trained to produce model states that facilitates these prediction tasks, through the expectation below.
Moreover, it is regularized to produce model states with high entropy, such that the model becomes robust to many different model states during training.
The loss function for learning the world model is:
[Equation]
[Table]
[Equation]
We jointly minimize the loss function with respect to the vector φ that contains all parameters of the world model using the Adam optimizer (Kingma and Ba, 2014).
We use the loss scales η i = 1 / (64 · 64 · 3) for the image, η r = 1 for the reward, η γ = 1 for the discount, η t = 0 .
08 for the transition, and η q = 0 .
02 for the entropy regularizer.
Scaling the transition loss up compared to the entropy regularizer allows us to encourage learning an accurate transition function.
Probabilistic interpretation The world model loss function in Equation 2 is the ELBO or variational free energy of an hidden Markov model that is conditioned on the action sequence.
The world model can thus be interpreted as a sequential VAE, where the representation model is the approximate posterior and the transition predictor is the temporal prior.
In the ELBO objective, the transition loss and entropy regularizer together are known as the KL regularizer.
Scaling the image loss relative to the KL regularizer is known as beta-VAE (Higgins et al., 2016).
Additionally, we scale the two terms within the KL term, the prior cross entropy and the posterior entropy, relative to each other.
We refer to this technique as KL balancing.
KL balancing encourages learning an accurate prior over increasing posterior entropy, so that the prior better approximates the aggregate posterior.
2. B EHAVIOR L EARNING
DreamerV2 learns long-horizon behaviors purely within its world model using an actor and a critic.
The actor chooses actions for predicting imagined sequences of compact model states.
The critic accumulates the future predicted rewards to take into account rewards beyond the planning horizon.
Both the actor and critic operate on top of the learned model states and thus beneﬁt from the representations learned by the world model.
The world model is ﬁxed during behavior learning, so the actor and value gradients do not affect its representations.
Not predicting images during behavior learning lets us efﬁciently simulate 2500 latent trajectories in parallel on a single GPU.
Imagination MDP To learn behaviors within the latent space of world model, we deﬁne the imagination MPD as follows.
The distribution of initial states ˆ z 0 in the imagination MDP is the distribution of compact model states encountered during world model training.
From there, the transition predictor p φ (ˆ z t | ˆ z t − 1 , ˆ a t − 1 ) outputs sequences ˆ z 1: H of compact model states up to the imagination horizon H = 15 .
The mean of the reward predictor p φ (ˆ r t | ˆ z t ) is used as reward sequence ˆ r 1: H .
The discount predictor p φ (ˆ γ t | ˆ z t ) outputs the discount sequence ˆ γ 1: H that is used to down-weight rewards.
Moreover, we weight the loss terms of the actor and critic by the cumulative predicted discount factors to softly account for the possibility of episode ends.
Model components To learn long-horizon behaviors in the imagination MDP, we leverage a stochastic actor that chooses actions and a deterministic critic.
The actor and critic are trained coorperatively, where the actor aims to output actions that lead to states that maximize the critic output, while the critic aims to accurately estimate the sum of future rewards achieved by the actor from each imagined state.
The actor and critic use the parameter vectors ψ and ξ , respectively: Actor:
[Equation]
[Equation]
In contrast to the actual environment, the latent state sequence is Markovian, so that there is no need for the actor and critic to condition on more than the current model state.
The actor and critic are both MLPs with ELU activations (Clevert et al., 2015) and use 1M trainable parameters each.
The actor outputs a categorical distribution over actions and the critic has a deterministic output.
The two components are trained from the same imagined trajectories but optimize separate loss functions.
Critic loss function The critic aims to predict the discounted sum of future rewards that the actor achieves in a given model state, known as the state value.
For this, we leverage temporal-difference learning, where the critic is trained toward a value target that is constructed from intermediate rewards and critic outputs for later states.
A common choice is the 1-step target that sums the current reward and the critic output for the following state.
However, the imagination MDP lets us generate on-policy trajectories of multiple steps, suggesting the use of n-step targets that incorporate reward information into the critic more quickly.
We follow DreamerV1 in using the more general λ -target (Sutton and
Barto, 2018; Schulman et al., 2015) that is deﬁned recursively as follows:
[Equation]
Intuitively, the λ -target is a weighted average of n-step returns for different horizons, where longer horizons are weighted exponentially less.
We set λ = 0 .
95 in practice, to focus more on long horizon targets than on short horizon targets.
Given a trajectory of model states, rewards, and discount factors, we train the critic to regress the λ -return using a squared loss:
[Equation]
We optimize the critic loss with respect to the critic parameters ξ using the Adam optimizer.
There is no loss term for the last time step because the target equals the critic at that step.
We stop the gradients around the targets, denoted by the sg( · ) function, as typical in the literature.
We stabilize value learning using a target network (Mnih et al., 2015), namely, we compute the targets using a copy of the critic that is updated every 100 gradient steps.
Actor loss function The actor aims to output actions that maximize the prediction of long-term future rewards made by the critic.
To incorporate intermediate rewards more directly, we train the actor to maximize the same λ -return that was computed for training the critic.
There are different gradient estimators for maximizing the targets with respect the the actor parameters.
DreamerV combines unbiased but high-variance Reinforce gradients with biased but low-variance straightthrough gradients.
Moreover, we regularize the entropy of the actor to encourage exploration where feasible while allowing the actor to choose precise actions when necessary.
Learning by Reinforce (Williams, 1992) maximizes the actor’s
probability of its own sampled actions weighted by the values of those actions.
The variance of this estimator can be reduced by subtracting the state value as baseline, which does not depend on the current action.
Intuitively, subtracting the baseline centers the weights and leads to faster learning.
The beneﬁt of Reinforce is that it produced unbiased gradients and the downside is that it can have high variance, even with baseline.
DreamerV1 relied entirely on reparameterization gradients (Kingma and Welling, 2013; Rezende et al., 2014) to train the actor directly by backpropagating value gradients through the sequence of sampled model states and actions.
DreamerV2 uses both discrete latents and discrete actions.
To backpropagate through the sampled actions and state sequences, we leverage straight-through gradients (Bengio et al., 2013).
This results in a biased gradient estimate with low variance.
The combined actor loss function is:
[Equation]
reinforce
[Equation]
We optimize the actor loss with respect to the actor parameters ψ using the Adam optimizer.
We use the loss scale η s = 0 .
9 for reinforce and linearly anneal the loss scale for dynamics backpropagation η d = 0 . 1 {+→ 0 . 
+}
and the entropy regularizer η e = 3 · 10 − 3 → 3 · 10 − 4 over the ﬁrst 10M environment frames.
We hypothesize that combining the two gradient estimators is beneﬁcial because the lowvariance straight-through gradients can accelerate early learning, while the unbiased Reinforce gradients can help ﬁnd a better ﬁnal solution.
E XPERIMENTS
We evaluate DreamerV2 on the well-established Atari benchmark with sticky actions, comparing to four strong model-free algorithms.
DreamerV2 outperforms the four model-free algorithms in all scenarios.
For an extensive comparison, we report four scores according to four aggregation protocols and give a recommendation for meaningfully aggregating scores across games going forward.
We also ablate the importance of discrete representations in the world model.
Our implementation of
DreamerV2 reaches 200M environment steps in under 10 days, while using only a single NVIDIA
V100 GPU and a single environment instance.
During the 200M environment steps, DreamerV learns its policy from 468B compact states imagined under the model, which is 10,000 × more than the 50M inputs received from the real environment after action repeat.
To help others build upon this work, we will release the implementation and training curves with the ﬁnal version of the paper.
Experimental setup We select the 55 games that prior works in the literature from different research labs tend to agree on (Mnih et al., 2016;
Brockman et al., 2016; Hessel et al., 2018; Castro et al., 2018; Badia et al., 2020) and recommend this set of games for evaluation going forward.
We follow the evaluation protocol of Machado et al.
(2018) with 200M environment steps, action repeat of 4, a time limit of 108,000 steps per episode that correspond to 30 minutes of game play, no access to life information, full action space, and sticky actions.
Because the world model integrates information over time, DreamerV2 does not use frame stacking.
The experiments use a single-task setup where a separate agent is trained for each game.
Moreover, each agent uses only a single environment instance.
We compare the algorithms based on both human gamer and human world record normalization (Toromanoff et al., 2019).
Model-free baselines We compare the learning curves and ﬁnal scores of DreamerV2 to four model-free algorithms, IQN (Dabney et al., 2018), Rainbow (Hessel et al., 2018), C51 (Bellemare et al., 2017), and DQN (Mnih et al., 2015).
We use the scores of these agents provided by the
Dopamine framework (Castro et al., 2018) that use sticky actions.
These may differ from the reported results in the papers that introduce these algorithms in the deterministic Atari setup.
The training time of Rainbow was reported at 10 days on a single GPU and using one environment instance.
[Figure]
The performance curves of DreamerV2 and four standard model-free algorithms are visualized in Figure 2.
The ﬁnal scores at 200M environment steps are shown in Table 1 and the scores on individual games are included in Table G.1.
There are different approaches for aggregating the scores across the 55 games and we show that this choice can have a substantial impact on the relative performance between algorithms.
To extensively compare DreamerV2 to the model-free algorithms, we consider the following four aggregation approaches:
• Task Median, Gamer Normalized Atari scores are commonly normalized based on a random policy and a professional gamer, and the median over tasks is reported (Mnih et al., 2015; 2016).
However, if almost half of the scores would be zero, the median would not be affected.
Thus, we argue that median scores are not reﬂective of the robustness of an algorithm.
• Task Mean, Gamer Normalized Compared to the task median, the task mean considers all tasks.
However, the gamer performed poorly on a small number of games, such as Crazy Climber, James Bond, and Video Pinball.
This makes it easy for algorithms to achieve a high normalized score on these few games, which then dominate the task mean.
• Task Mean, Record Normalized Instead of normalizing based on the professional gamer, Toromanoff et al.
(2019) suggest to normalize based on the registered human world record of each game.
This partially addresses the outlier problem but the mean is still dominated by a small number of games, where the algorithms achieve super-human performance.
• Task Mean, Clipped Record Norm
To overcome these limitations, we recommend normalizing by the human world record and then clipping the scores to not exceed a value of 1, so that performance above the record does not further increase the score.
The result is a robust measure of algorithm performance on the Atari suite that considers performance across all games [Table] Latent Variables
CategoricalGaussian 0.0 0.5 1.0 1.5 2. KL Balancing EnabledDisabled 0.0 0.5 1.0 1.5 2. Image Gradients
EnabledDisabled 0.0 0.5 1.0 1.5 2.0 1e 0. Reward Gradients EnabledDisabled are shown in Figure F.1. The ablations highlight the of over Gaussian latent variables and of
KL balancing.
Moreover, they show that the world model relies on image gradients for learning its representations. Stopping reward gradients even improves performance on some tasks, suggesting that representations that are not speciﬁcally trained to predict previously experienced rewards may generalize better to new situations. From Figure 2 and Table 1, we see that the different aggregation approaches let us examine agent performance from different angles.
Interestingly, Rainbow clearly outperforms IQN in the ﬁrst aggregation method but IQN clearly outperforms Rainbow in the remaining setups.
DreamerV outperforms the model-free agents in all four metrics, with the largest margin in record normalized mean performance.
Despite this, we recommend clipped record normalized as the most meaningful aggregation method, as it considers all tasks to a similar degree without being dominated by a small number of outlier scores.
Individual games The scores on individual Atari games at 200M environment steps are included in
Table G.1, alongside the model-free algorithms and the baselines of random play, human gamer, and human world record.
We ﬁlled in reasonable values for the 2 out of 55 games that have no registered world record.
Figure C.1 compares the score differences between DreamerV2 and each model-free algorithm for the individual games.
DreamerV2 achieves comparable or higher performance on most games except for Video Pinball.
We hypothesize that the reconstruction loss of the world model does not encourage learning a meaningful latent representation because the most important object in the game, the ball, occupies only a single pixel.
One the other hand, DreamerV2 achieves the strongest improvements over the model-free agents on the games James Bond, Up N Down, and Assault. 3. A BLATION S TUDY
To understand which ingredients of DreamerV2 are responsible for its success, we conduct an extensive ablation study.
We compare equipping the world model with categorical latents, as in
DreamerV2, to Gaussian latents, as in DreamerV1.
Moreover, we study the importance of KL balancing. Finally, we investigate the importance of gradients from image reconstruction and reward prediction for learning the model representations, by stopping one of the two gradient signals before entering the model states. The results of the ablation study are summarized in Figure 3 and Table 2. Refer to the appendix for the score curves of the individual tasks.
[Table]
Categorical latents Categorical latent variables outperform than Gaussian latent variables ontasks, achieve lower performance on 9 tasks, and are tied on 5 tasks.
We deﬁne a tie as being within 5% of another.
While we do not know the reason why the categorical variables are beneﬁcial, we state several hypotheses that can be investigated in future work.
First, categorical variables could be beneﬁcial because they better capture the discrete objects present in many Atari games.
Second, they may be easier to optimize, possibly by reducing the amount of exploding and vanishing gradients.
This, a sample from the 32 categorical with 32 classes each results in a spare binary vector of length 1024 with 32 active bits.
The enforced level of sparsity could be beneﬁcial for generalization.
KL balancing KL balancing outperforms the standard KL regularizer on 44 tasks, achieves lower performance on 9 tasks, and is tied on 2 tasks.
Learning accurate prior dynamics of the world model is critical because it is used for imagining latent state trajectories using policy optimization.
By scaling up the prior cross entropy relative to the posterior entropy, the world model is encouraged to minimize the KL by improving its prior dynamics toward the more informed posteriors, as opposed to reducing the KL by increasing the posterior entropy.
KL balancing may also be beneﬁcial for probabilistic models with learned priors beyond world models.
Model gradients Stopping the image gradients increases performance on 2 tasks, decreases performance on 51 tasks, and is tied on 2 tasks.
The world model of DreamerV2 thus heavily relies on the learning signal provided by the high-dimensional images.
Stopping the reward gradients increases performance on 13 tasks, decreases performance on 19 tasks, and is tied on 23 tasks.
Figure F. further shows that the difference in scores is small.
In contrast to MuZero, DreamerV2 thus learns general representations of the environment state from image information alone.
Stopping reward gradients improved performance on a number of tasks, suggesting that the representations that are not speciﬁc to previously experienced rewards may generalize better to unseen situations.
R ELATED W ORK
Model-free Atari
The majority of agents applied to the Atari benchmark have been trained using model-free algorithms.
DQN (Mnih et al., 2015) showed that deep neural network policies can be trained using Q-learning by incorporating experience replay and target networks.
Several works have extended DQN to incorporate bias correction as in DDQN (Van Hasselt et al., 2015), prioritized experience replay (Schaul et al., 2015), architectural improvements (Wang et al., 2016), and distributional value learning (Bellemare et al., 2017; Dabney et al., 2017; 2018).
Besides value learning, agents based on policy gradients have targeted the Atari benchmark, such as ACER (Schulman et al., 2017a), PPO (Schulman et al., 2017a), ACKTR (Wu et al., 2017), and Reactor (Gruslys et al., 2017).
Another line of work has focused on improving performance by distributing data collection, often while increasing the budget of environment steps beyond 200M (Mnih et al., 2016; Schulman et al., 2017b; Horgan et al., 2018;
Kapturowski et al., 2018; Badia et al., 2020).
World models Several model-based agents focus on proprioceptive inputs (Watter et al., 2015; Gal et al., 2016; Higuera et al., 2018;
Henaff et al., 2018;
Chua et al., 2018; Wang et al., 2019; Wang and
Ba, 2019), model images without using them for planning (Oh et al., 2015; Krishnan et al., 2015;
Karl et al., 2016; Chiappa et al., 2017; Babaeizadeh et al., 2017; Gemici et al., 2017; Denton and
Fergus, 2018; Buesing et al., 2018; Doerr et al., 2018; Gregor and Besse, 2018), or combine the beneﬁts of model-based and model-free approaches (Kalweit and Boedecker, 2017; Nagabandi et al., 2017;
Weber et al., 2017; Kurutach et al., 2018; Buckman et al., 2018; Ha and Schmidhuber, 2018;
Wayne et al., 2018; Igl et al., 2018; Srinivas et al., 2018; Lee et al., 2019).
Most world model agents with image inputs have thus far been limited to relatively simple control tasks (Watter et al., 2015;
Ebert et al., 2017; Ha and Schmidhuber, 2018; Hafner et al., 2018; Zhang et al., 2019; Hafner et al., 2019).
We explain the two model-based approaches that were applied to Atari in detail below.
SimPLe The SimPLe agent (Kaiser et al., 2019) learns a video prediction model in pixel-space and uses its predictions to train a PPO agent (Schulman et al., 2017a), as shown in Table 3.
The model directly predicts each frame from the previous four frames and receives an additional discrete latent variable as input.
The authors evaluate SimPLe on a subset of Atari games for 400k and 2M environment steps, after which they report diminishing returns.
Some recent model-free methods have followed the comparison at 400k steps (Srinivas et al., 2020; Kostrikov et al., 2020).
However, the highest performance achieved in this data-efﬁcient regime is a gamer normalized median score of 0.28 (Kostrikov et al., 2020) that is far from human-level performance. Instead, we focus on the [Table]
DreamerV2 and SimPLe learn complete models of the environment by leveraging the learning signal provided by the image inputs, while MuZero learns its model through value gradients that are speciﬁc to an individual task.
The Monte-Carlo tree search used by MuZero is effective but adds complexity and is challenging to parallelize.
This component is orthogonal to the world model proposed here. well-established and competitive evaluation after 200M frames, where many successful model-free algorithms are available for comparison.
MuZero The MuZero agent (Schrittwieser et al., 2019) learns a sequence model of rewards and values (Oh et al., 2017) to solve reinforcement learning tasks via Monte-Carlo Tree Search (MCTS;
Coulom, 2006; Silver et al., 2017).
The sequence model is trained purely by predicting task-speciﬁc information and does not incorporate explicit representation learning using the images, as shown in Table 3.
MuZero shows that with signiﬁcant engineering effort and a vast computational budget, planning can achieve impressive performance on several board games and deterministic Atari games.
However, MuZero is not publicly available, and it would require over 2 months to train an Atari agent on one GPU.
By comparison, DreamerV2 is a simple algorithm that achieves human-level performance on Atari on a single GPU in 10 days, making it reproducible for many researchers.
Moreover, the advanced planning components of MuZero are complimentary and could be applied to the accurate world models learned by DreamerV2.
DreamerV2 leverages the additional learning signal provided by the input images, analogous to recent successes by semi-supervised image classiﬁcation (Chen et al., 2020;
He et al., 2020; Grill et al., 2020).
D ISCUSSION
We present DreamerV2, a model-based agent that achieves human-level performance on the Atari 200M benchmark by learning behaviors purely from the latent-space predictions of a separately trained world model.
Using a single GPU and a single environment instance, DreamerV2 outperforms top model-free single-GPU agents Rainbow and IQN using the same computational budget and training time.
To develop DreamerV2, we apply several small modiﬁcations to the Dreamer agent (Hafner et al., 2019). We conﬁrm
experimentally that learning a categorical latent space and using KL balancing improves the performance of agent.
Moreover, we ﬁnd the DreamerV2 relies on image information for learning generally useful representations — its performance is not impacted by whether the representations are especially learned for predicting rewards.
DreamerV2 serves as proof of concept, showing that model-based RL can outperform top model-free algorithms on the most competitive RL benchmarks, despite the years of research and engineering effort that modern model-free agents rest upon.
Beyond achieving strong performance on individual tasks, world models open avenues for efﬁcient transfer and multi-task learning, sample-efﬁcient learning on physical robots, and global exploration based on uncertainty estimates.
A S UMMARY OF M ODIFICATIONS
To develop DreamerV2, we used the Dreamer agent (Hafner et al., 2019) as a starting point.
This subsection describes the changes that we applied to the agent to achieve high performance on the
Atari benchmark, as well as the changes that were tried but not found to increase performance and thus were not not included in DreamerV2.
Summary of changes that were tried and were found to help:
• Categorical latents Using categorical latent states using straight-through gradients in the world model instead of Gaussian latents with reparameterized gradients.
• Mixed actor gradients Combining Reinforce and dynamics backpropagation gradients for learning the actor instead of dynamics backpropagation only.
• Policy entropy Regularizing the policy entropy for exploration both in imagination and during data collection, instead of using external action noise during data collection.
• KL balancing Separately scaling the prior cross entropy and the posterior entropy in the KL loss to encourage learning an accurate temporal prior, instead of using free nats.
• Model size Increasing the number of units or feature maps per layer of all model components, resulting in a change from 13M parameters to 22M parameters.
• Layer norm Using layer normalization in the GRU that is used as part of the RSSM latent transition model, instead of no normalization.
Summary of changes that were tried but were not shown to help:
• Binary latents Using a larger number of binary latents for the world model instead of using categorical latents, which could have encouraged a more disentangled representation.
• Long-term entropy Including the policy entropy into temporal-difference loss of the value function, so that the actor seeks out states with high action entropy beyond the planning horizon.
• Scheduling Scheduling the learning rate, KL scale, free bits.
Only scheduling the entropy regularizer and the amount of value backpropagation for the policy was beneﬁcial.
• Reinforce only Using only Reinforce gradients for the actor worked for many games but fails to start learning on some games, possibly because of the high variance of reinforce gradients.
Due to the large computational requirements, a comprehensive ablation study on this list of changes is unfortunately infeasible for us.
This would require 55 tasks times 5 seeds for 10 days per change to run, resulting in over 60,000 GPU hours per change.
However, we aim to include ablations for the most critical changes in the ﬁnal version of this paper.
[Table]
J a m e
[Table]
o n d
U p N
D
[Table]
n
[Table]
p h e r [Table] T i m e P il
[Table]
t
[Table]
k
[Table]
A
[Table]
a
[Table]
i s
[Table]
t
[Table]
unn e
[Table]
I
[Table]
F
[Table]
i
[Table]
r
[Table]
s
[Table]
r
[Table]
F
[Table]
y
[Table]
i
[Table]
P
[Table]
P r i
[Table]
k
[Table]
m
[Table]
k
[Table]
S
[Table]
e
[Table]
m
[Table]
pp e
[Table]
m .
[Table]
e Sp [Table] I n
[Table]
r s
B
[Table]
n g Y a r s R e v e n g e S t a r G unn e r D e m o n A tt a c k
V i
d e o
P i n b a ll
DreamerV2 vs IQN
J a m e s
B o n d
U p N
D o w n
G o p h
e r
A t l a n t i s Q b e r t
T i m e
P il o t A ss a u l t B r e a k o u t K r u ll
R o a d
R unn e r
I c e H o c k e y F r o s t b i t e K un g
F u M a s t e r
S k ii n g T e nn i s
W i z a r d
O f W o r
N a m e
T h
i s G a m e
A s t e r i x
G r a v i t a r P h
o e n i x Z a xx o n
B e a m
R i d e r
F i s h
i n g D e r b y M s P a c m a n
M o n t e z u m a
R e v .
E n d u r o B o w li n g
A li e n
A s t e r o i d s F r ee w a y
P i t f a ll K a n g a r oo
C e n t i p e d e
T u t a n k h a m
B e r z e r k
P o n g
S o l a r i s B a n k
H e i s t S e a q u e s t
D o u b l e
D un k R o b o t a n k C r a z y
C li m b e r
B a tt l e
Z o n e
R i v e rr a i d
P r i v a t e E y e
Y a r s R e v e n g e A m i d a r H e r
o Sp a c e I n v a d e r s B o x i n g
V e n t
u
r e
C h
o pp e r C o m .
S t a r G unn e r D e m o n A tt a c k
V i
d e o
P i n b a ll
DreamerV2 vs Rainbow
J a m e s
B o n d
U p N
D o w n
G o p h
e r
T i m e
P il o t Q b e r t A ss a u l t A t l a n t i s
D o u b l e
D un k
B r e a k o u t
R o a d
R unn e r
F r o s t b i t e
I c e H o c k e y K r u ll
W i z a r d
O f W o r
K un g F u M a s t e r
K a n g a r oo
A s t e r i x Z a xx o n E n d u r o S k ii n g
P h
o e n i x
F i s h
i n g D e r b y G r a v i t a r
B e a m
R i d e r
B a n k H e i s t R o b o t a n k
N a m e T h i s G a m e M s P a c m a n
R i v e rr a i d C r a z y C li m b e r
B a tt l e
Z o n e
B o x i n g A li e n A m i d a r
Y a r s R e v e n g e
M o n t e z u m a R e v .
B o w li n g B e r z e r k
T e nn i s
A s t e r o i d s
F r ee w a y
P i t
f a ll P o n g T u t a n k h a m
P r i v a t e
E y e
S o l a r i s H
e r o
C e n t
i p e d e
S t a r G unn e r
C h
o pp e r C o m .
S e a q u e s t
Sp a c e I n v a d e r s
V e n t
u
r e
D e m o n A tt a c k
V i
d e o
P i n b a ll
[Table]
U
[Table]
D
[Table]
n
[Table]
e r
T i m e
P
[Table]
D
[Table]
u b l e
D un k Q b e r t A ss a u l t A t l a n t i s B r e a k o u t
R o a d
R unn e r F r o s t b
i t
e K r u ll
I c e H o c k e y A s t e r i x
W i z a r d O f W o r
K un g F u M a s t e r
E n d u r o T e nn i s Z a xx o n C r a z y C li m b e r
N a m e
T h
i s G a m e
F i s h i n g D e r b y
P h
o e n i x
G r a v i t a r
T u t a n k h a m Sp a c e I n v a d e r s
B a n k
H e i s t S k ii n g
B o x i n g H e r o R o b o t a n k K a n g a r oo
B a tt l e
Z o n e
B e a m
R i d e r
C e n t
i p e d e
M o n t e z u m a
R e v .
M s P a c m a n R i v e rr a i d F r ee w a y
A m i d a r A li e n S e a q u e s t B o w li n g
B e r z e r k
P o n g
A s t e r o i d s
P i t
f a ll
P r i v a t e
E y e
C h
o pp e r C o m .
S o l a r i s
V e n t
u r e
Y a r s R e v e n g e
S t a r G unn e r D e m o n A tt a c k
V i
d e o
P i n b a ll
[Figure]
D
M ODEL -F REE C
[Figure]
Montezuma Rev
[Figure]
Yars
[Figure]
E L ATENTS AND KL W EIGHTING A BLATIONS
[Figure]
Rev [Figure] [Figure] [Figure]
Montezuma Rev. [Figure] G A TARI T ASK S CORES [Table]
Algorithms
[Table]
Record
Rainbow
[Table]
DreamerV
[Table]
[Table]
[Table]
[Table] [Table] 2195 Atlantis     999395 Bank [-Heist     1063 Battle Zone     27483 Beam Rider     10193 Berzerk     660 Bowling     
-][Table]
[Table]
[Table]
[Table] 786 Bowling [Table] 84 Breakout     345 Centipede[-4898 Chopper Command     1630 Crazy Climber     134512 Demon Attack     
-]
[Table]
Command     2676 Crazy
[Table] [Table] [Table] - -    22 Enduro     2297 Fishing Derby - -    61 Freeway     34 Frostbite     17088 Gopher     53891 Gravitar     3507 Hero     30714 Ice Hockey -    - 30 James Bond     9835 Kangaroo     12742 Krull     9641 Kung Fu Master
Montezuma Revenge
Ms Pacman     5576 Name This Game     14653 Phoenix     12015 Pitfall -   - - -1 Pong19 Private Eye     154 Qbert     172022 Riverraid     16792 Road Runner     85288 Robotank     64 Seaquest     6795 Skiing - - - - - -8444 Solaris     1073 Space Invaders     2707 Star Gunner     {+
+} Tennis - -    23 Time Pilot
Tutankham
[Table]