Mastering Atari with Discrete World Models

1 INTRODUCTION 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 learns 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 offline datasets ( Yu et al. , 2020 ) . When the inputs are high-dimensional images , latent dynamics models predict ahead in an 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 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 stateof-the-art model-free algorithms on the most competitive benchmarks . The well-established Atari benchmark ∗Correspondence to : Danijar Hafner < mail @ danijar.com > . ( 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 first 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 modifications 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 complementary 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 . 2 DREAMERV2 . 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 C , we include a list of changes that we applied to DreamerV1 and which of them we found to increase empirical performance . 2.1 WORLD MODEL LEARNING . 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 beneficial 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 efficiently 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 images x1 : T , actions a1 : T , rewards r1 : T , and discount factors γ1 : T . The discount factors equal a fixed hyper parameter γ = 0.999 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 fixed 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 an image encoder , a Recurrent State-Space Model ( RSSM ; Hafner et al. , 2018 ) to learn the dynamics , and predictors for the image , reward , and discount factor . The world model is summarized in Figure 2 . The RSSM uses a sequence of deterministic recurrent states ht , from which it computes two distributions over stochastic states at each step . The posterior state zt incorporates information about the current image xt , while the prior state ẑt aims to predict the posterior without access to the current image . The concatenation of deterministic and stochastic states forms the compact model state . From the posterior model state , we reconstruct the current image xt and predict the reward rt and discount factor γt . The model components are : RSSM Recurrent model : ht = fφ ( ht−1 , zt−1 , at−1 ) Representation model : zt ∼ qφ ( zt | ht , xt ) Transition predictor : ẑt ∼ pφ ( ẑt | ht ) Image predictor : x̂t ∼ pφ ( x̂t | ht , zt ) Reward predictor : r̂t ∼ pφ ( r̂t | ht , zt ) Discount predictor : γ̂t ∼ pφ ( γ̂t | ht , zt ) . ( 1 ) 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 ending 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 Multi-Layer Perceptron ( MLP ) that receives the image embedding and the deterministic recurrent state . The RSSM uses a Gated Recurrent Unit ( GRU ; Cho et al. , 2014 ) to compute the deterministic recurrent states . The model state is the concatenation of deterministic GRU state and a sample of the stochastic state . The image predictor is a transposed CNN and the transition , reward , and discount predictors are MLPs . We down-scale the 84 × 84 grayscale images to 64× 64 pixels so that we can apply the convolutional architecture of DreamerV1 . Algorithm 1 : Straight-Through Gradients with Automatic Differentiation sample = one_hot ( draw ( logits ) ) # sample has no gradient probs = softmax ( logits ) # want gradient of this sample = sample + probs - stop_grad ( probs ) # has gradient of probs 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 reparameterization gradients during backpropagation ( Kingma and Welling , 2013 ; Rezende et al. , 2014 ) . In DreamerV2 , we instead use a vector of several categorical variables and optimize them using straight-through gradients ( Bengio et al. , 2013 ) , which are easy to implement using automatic differentiation as shown in Algorithm 1 . We discuss possible benefits of categorical over Gaussian latents in the experiments section . Loss function 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 : L ( φ ) .= Eqφ ( z1 : T | a1 : T , x1 : T ) [ ∑T t=1 − ln pφ ( xt | ht , zt ) image log loss − ln pφ ( rt | ht , zt ) reward log loss − ln pφ ( γt | ht , zt ) discount log loss +βKL [ qφ ( zt | ht , xt ) ∥∥ pφ ( zt | ht ) ] KL loss ] . ( 2 ) 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 scale the KL loss by β = 0.1 for Atari and by β = 1.0 for continuous control ( Higgins et al. , 2016 ) . KL balancing The world model loss function in Equation 2 is the ELBO or variational free energy of a 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 KL loss serves two purposes : it trains the prior toward the representations , and it regularizes the representations toward the prior . However , learning the transition function is difficult and we want to avoid regularizing the representations toward a poorly trained prior . To solve this problem , we minimize the KL loss faster with respect to the prior than the representations by using different learning rates , α = 0.8 for the prior and 1 − α for the approximate posterior . We implement this technique as shown in Algorithm 2 and refer to it as KL balancing . KL balancing encourages learning an accurate prior over increasing posterior entropy , so that the prior better approximates the aggregate posterior . KL balancing is different from and orthogonal to beta-VAEs ( Higgins et al. , 2016 ) .

Summary: The authors introduce DreamerV2, a modification of the influential Dreamer RL agent (hereafter refered to as DreamerV1). The primary changes from DreamerV1 are a discrete latent space and a modified loss function (and with it, a modified optimization scheme). As in DreamerV1, the agent trains a world model with environment experience, and the policy is learned by "imagining" within the learned latent space using the world model to simulate transitions and rewards. They demonstrate superior performance over a variety of successful benchmarks that use similar compute (1 GPU, 1 environment -- e.g. MuZero, which requires vastly more, is not considered) on Atari. Further, they analyze several ways of aggregating Atari scores, and (while their algorithm performs best of those tried in each aggregation), they recommend one aggregation method (along with several other choices made for benchmarking) going forward.