Model Based Reinforcement Learning for Atari

1 INTRODUCTION . Human players can learn to play Atari games in minutes ( Tsividis et al. , 2017 ) . However , some of the best model-free reinforcement learning algorithms require tens or hundreds of millions of time steps – the equivalent of several weeks of training in real time . How is it that humans can learn these games so much faster ? Perhaps part of the puzzle is that humans possess an intuitive understanding of the physical processes that are represented in the game : we know that planes can fly , balls can roll , and bullets can destroy aliens . We can therefore predict the outcomes of our actions . In this paper , we explore how learned video models can enable learning in the Atari Learning Environment ( ALE ) benchmark Bellemare et al . ( 2015 ) ; Machado et al . ( 2018 ) with a budget restricted to 100K time steps – roughly to two hours of a play time . Although prior works have proposed training predictive models for next-frame , future-frame , as well as combined future-frame and reward predictions in Atari games ( Oh et al . ( 2015 ) ; Chiappa et al . ( 2017 ) ; Leibfried et al . ( 2016 ) ) , no prior work has successfully demonstrated model-based control via predictive models that achieve competitive results with model-free RL . Indeed , in a recent survey ( Section 7.2 in Machado et al . ( 2018 ) ) this was formulated as the following challenge : “ So far , there has been no clear demonstration of successful planning with a learned model in the ALE ” . Using models of environments , or informally giving the agent ability to predict its future , has a fundamental appeal for reinforcement learning . The spectrum of possible applications is vast , including learning policies from the model ( Watter et al. , 2015 ; Finn et al. , 2016 ; Finn & Levine , 2017 ; Ebert et al. , 2017 ; Hafner et al. , 2019 ; Piergiovanni et al. , 2018 ; Rybkin et al. , 2018 ; Sutton & Barto , ∗Equal contribution , authors listed in random order . BO performed the work partially during an internship at Google Brain . Correspondence to : b.osinski @ mimuw.edu.pl 1 Published as a conference paper at ICLR 2020 2017 , Chapter 8 ) , capturing important details of the scene ( Ha & Schmidhuber , 2018 ) , encouraging exploration ( Oh et al. , 2015 ) , creating intrinsic motivation ( Schmidhuber , 2010 ) or counterfactual reasoning ( Buesing et al. , 2019 ) . One of the exciting benefits of model-based learning is the promise to substantially improve sample efficiency of deep reinforcement learning ( see Chapter 8 in Sutton & Barto ( 2017 ) ) . Our work advances the state-of-the-art in model-based reinforcement learning by introducing a system that , to our knowledge , is the first to successfully handle a variety of challenging games in the ALE benchmark . To that end , we experiment with several stochastic video prediction techniques , including a novel model based on discrete latent variables . We present an approach , called Simulated Policy Learning ( SimPLe ) , that utilizes these video prediction techniques and trains a policy to play the game within the learned model . With several iterations of dataset aggregation , where the policy is deployed to collect more data in the original game , we learn a policy that , for many games , successfully plays the game in the real environment ( see videos on the project webpage https : //goo.gl/itykP8 ) . In our empirical evaluation , we find that SimPLe is significantly more sample-efficient than a highly tuned version of the state-of-the-art Rainbow algorithm ( Hessel et al. , 2018 ) on almost all games . In particular , in low data regime of 100k samples , on more than half of the games , our method achieves a score which Rainbow requires at least twice as many samples . In the best case of Freeway , our method is more than 10x more sample-efficient , see Figure 3 . Since the publication of the first preprint of this work , it has been shown in van Hasselt et al . ( 2019 ) ; Kielak ( 2020 ) that Rainbow can be tuned to have better results in low data regime . The results are on a par with SimPLe – both of the model-free methods are better in 13 games , while SimPLe is better in the other 13 out of the total 26 games tested ( note that in Section 4.2 van Hasselt et al . ( 2019 ) compares with the results of our first preprint , later improved ) . 2 RELATED WORK . Atari games gained prominence as a benchmark for reinforcement learning with the introduction of the Arcade Learning Environment ( ALE ) Bellemare et al . ( 2015 ) . The combination of reinforcement learning and deep models then enabled RL algorithms to learn to play Atari games directly from images of the game screen , using variants of the DQN algorithm ( Mnih et al. , 2013 ; 2015 ; Hessel et al. , 2018 ) and actor-critic algorithms ( Mnih et al. , 2016 ; Schulman et al. , 2017 ; Babaeizadeh et al. , 2017b ; Wu et al. , 2017 ; Espeholt et al. , 2018 ) . The most successful methods in this domain remain model-free algorithms ( Hessel et al. , 2018 ; Espeholt et al. , 2018 ) . Although the sample complexity of these methods has substantially improved recently , it remains far higher than the amount of experience required for human players to learn each game ( Tsividis et al. , 2017 ) . In this work , we aim to learn Atari games with a budget of just 100K agent steps ( 400K frames ) , corresponding to about two hours of play time . Prior methods are generally not evaluated in this regime , and we therefore optimized Rainbow ( Hessel et al. , 2018 ) for optimal performance on 1M steps , see Appendix E for details . 2 Published as a conference paper at ICLR 2020 Oh et al . ( 2015 ) and Chiappa et al . ( 2017 ) show that learning predictive models of Atari 2600 environments is possible using appropriately chosen deep learning architectures . Impressively , in some cases the predictions maintain low L2 error over timespans of hundreds of steps . As learned simulators of Atari environments are core ingredients of our approach , in many aspects our work is motivated by Oh et al . ( 2015 ) and Chiappa et al . ( 2017 ) , however we focus on using video prediction in the context of learning how to play the game well and positively verify that learned simulators can be used to train a policy useful in original environments . An important step in this direction was made by Leibfried et al . ( 2016 ) , which extends the work of Oh et al . ( 2015 ) by including reward prediction , but does not use the model to learn policies that play the games . Most of these approaches , including ours , encode knowledge of the game in implicit way . Unlike this , there are works in which modeling is more explicit , for example Ersen & Sariel ( 2014 ) uses testbed of the Incredible Machines to learn objects behaviors and their interactions . Similarly Guzdial et al . ( 2017 ) learns an engine predicting interactions of predefined set of sprites in the domain of Super Mario Bros. Perhaps surprisingly , there is virtually no work on model-based RL in video games from images . Notable exceptions are the works of Oh et al . ( 2017 ) , Sodhani et al . ( 2019 ) , Ha & Schmidhuber ( 2018 ) , Holland et al . ( 2018 ) , Leibfried et al . ( 2018 ) and Azizzadenesheli et al . ( 2018 ) . Oh et al . ( 2017 ) use a model of rewards to augment model-free learning with good results on a number of Atari games . However , this method does not actually aim to model or predict future frames , and achieves clear but relatively modest gains in efficiency . Sodhani et al . ( 2019 ) proposes learning a model consistent with RNN policy which helps to train policies that are more powerful than their model-free baseline . Ha & Schmidhuber ( 2018 ) present a way to compose a variational autoencoder with a recurrent neural network into an architecture that is successfully evaluated in the VizDoom environment and on a 2D racing game . The training procedure is similar to Algorithm 1 , but only one iteration of the loop is needed as the environments are simple enough to be fully explored with random exploration . Similarly , Alaniz ( 2018 ) utilizes a transition model with Monte Carlo tree search to solve a block-placing task in Minecraft . Holland et al . ( 2018 ) use a variant of Dyna ( Sutton , 1991 ) to learn a model of the environment and generate experience for policy training in the context of Atari games . Using six Atari games as a benchmark Holland et al . ( 2018 ) measure the impact of planning shapes on performance of the Dyna-DQN algorithm and include ablations comparing scores obtained with perfect and imperfect models . Our method achieves around 330 % of the Dyna-DQN score on Asterix , 120 % on Q-Bert , 150 % on Seaquest and 80 % on Ms. Pac-Man . Azizzadenesheli et al . ( 2018 ) propose an algorithm called Generative Adversarial Tree Search ( GATS ) and for five Atari games train a GAN-based world model along with a Q-function . Azizzadenesheli et al . ( 2018 ) primarily discuss various failure modes of the GATS algorithm . Our method achieves around 64 times the score of GATS on Pong and 10 times on Breakout . 1 Outside of games , model-based reinforcement learning has been investigated at length for applications such as robotics ( Deisenroth et al. , 2013 ) . Though most of such works do not use image observations , several recent works have incorporated images into real-world ( Finn et al. , 2016 ; Finn & Levine , 2017 ; Babaeizadeh et al. , 2017a ; Ebert et al. , 2017 ; Piergiovanni et al. , 2018 ; Paxton et al. , 2019 ; Rybkin et al. , 2018 ; Ebert et al. , 2018 ) and simulated ( Watter et al. , 2015 ; Hafner et al. , 2019 ) robotic control . Our video models of Atari environments described in Section 4 are motivated by models developed in the context of robotics . Another source of inspiration are discrete autoencoders proposed by van den Oord et al . ( 2017 ) and Kaiser & Bengio ( 2018 ) . The structure of the model-based RL algorithm that we employ consists of alternating between learning a model , and then using this model to optimize a policy with model-free reinforcement learning . Variants of this basic algorithm have been proposed in a number of prior works , starting from Dyna Q Sutton ( 1991 ) to more recent methods that incorporate deep networks Heess et al . ( 2015 ) ; Feinberg et al . ( 2018 ) ; Kalweit & Boedecker ( 2017 ) ; Kurutach et al . ( 2018 ) . 1Comparison with Dyna-DQN and GATS is based on random-normalized scores achieved at 100K interactions . Those are approximate , as the authors Dyna-DQN and GATS have not provided tabular results . Authors of Dyna-DQN also report scores on two games which we do not consider : Beam Rider and Space Invaders . For both games the reported scores are close to random scores , as are GATS scores on Asterix . 3 Published as a conference paper at ICLR 2020

Summary: This paper proposes a model-based reinforcement learning algorithm suitable for high-dimensional visual environments like Atari. The algorithmic loop is conceptually simple and comprises 1) collecting real data with the current policy 2) updating an environment model with old and newly acquired data and 3) updating the policy "virtually" inside of the environment model using PPO. The approach is evaluated on 26 games from the Atari benchmark and compared against the model-free baselines Rainbow DQN and PPO. The newly proposed model-based method clearly outperforms both model-free baselines in low training regimes (100,000 steps). Further ablation studies are provided, e.g. similar results are obtained in a more stochastic setting of ALE with sticky actions.