Title: MODEL BASED REINFORCEMENT LEARNING FOR ATARI

Abstract: Model-free reinforcement learning (RL) can be used to learn effective policies for complex tasks, such as Atari games, even from image observations. However, this typically requires very large amounts of interaction -substantially more, in fact, than a human would need to learn the same games. How can people learn so quickly? Part of the answer may be that people can learn how the game works and predict which actions will lead to desirable outcomes. In this paper, we explore how video prediction models can similarly enable agents to solve Atari games with fewer interactions than model-free methods. We describe Simulated Policy Learning (SimPLe), a complete model-based deep RL algorithm based on video prediction models and present a comparison of several model architectures, including a novel architecture that yields the best results in our setting. Our experiments evaluate SimPLe on a range of Atari games in low data regime of 100k interactions between the agent and the environment, which corresponds to two hours of real-time play. In most games SimPLe outperforms state-of-the-art model-free algorithms, in some games by over an order of magnitude.

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, Figure 1 : Main loop of SimPLe. 1) the agent starts interacting with the real environment following the latest policy (initialized to random). 2) the collected observations will be used to train (update) the current world model. 3) the agent updates the policy by acting inside the world model. The new policy will be evaluated to measure the performance of the agent as well as collecting more data (back to 1). Note that world model training is self-supervised for the observed states and supervised for the reward.

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. 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 L 2 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. 1Outside 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) . The input to the model is four stacked frames (as well as the action selected by the agent) while the output is the next predicted frame and expected reward. Input pixels and action are embedded using fully connected layers, and there is per-pixel softmax (256 colors) in the output. This model has two main components. First, the bottom part of the network which consists of a skip-connected convolutional encoder and decoder. To condition the output on the actions of the agent, the output of each layer in the decoder is multiplied with the (learned) embedded action. Second part of the model is a convolutional inference network which approximates the posterior given the next frame, similarly to Babaeizadeh et al. (2017a) . At training time, the sampled latent values from the approximated posterior will be discretized into bits. To keep the model differentiable, the backpropagation bypasses the discretization following Kaiser & Bengio (2018) . A third LSTM based network is trained to approximate each bit given the previous ones. At inference time, the latent bits are predicted auto-regressively using this network.

The deterministic model has the same architecture as this figure but without the inference network.

3. SIMULATED POLICY LEARNING (SIMPLE):

Reinforcement learning is formalized in Markov decision processes (MDP). An MDP is defined as a tuple (S, A, P, r, γ), where S is a state space, A is a set of actions available to an agent, P is the unknown transition kernel, r is the reward function and γ ∈ (0, 1) is the discount factor. In this work we refer to MDPs as environments and assume that environments do not provide direct access to the state (i.e., the RAM of Atari 2600 emulator In Atari 2600 games our goal is to find a policy which maximizes the value function from the beginning of the game. Crucially, apart from an Atari 2600 emulator environment env we will use a neural network simulated environment env which we call a world model and describe in detail in Section 4. The environment env shares the action space and reward space with env and produces visual observations in the same format, as it will be trained to mimic env. Our principal aim is to train a policy π using a simulated environment env so that π achieves good performance in the original environment env. In this training process we aim to use as few interactions with env as possible. The initial data to train env comes from random rollouts of env. As this is unlikely to capture all aspects of env, we use the iterative method presented in Algorithm 1.

4. WORLD MODELS:

In search for an effective world model we experimented with various architectures, both new and modified versions of existing ones. This search resulted in a novel stochastic video prediction model (visualized in Figure 2 ) which achieved superior results compared to other previously proposed models. In this section, we describe the details of this architecture and the rationale behind our design decisions. In Section 6 we compare the performance of these models.

Deterministic Model. Our basic architecture, presented as part of Figure 2 , resembles the convolutional feedforward network from Oh et al. (2015) . The input X consists of four consecutive game frames and an action a. Stacked convolution layers process the visual input. The actions are one-hot-encoded and embedded in a vector which is multiplied channel-wise with the output of the convolutional layers. The network outputs the next frame of the game and the value of the reward.

In our experiments, we varied details of the architecture above. In most cases, we use a stack of four convolutional layers with 64 filters followed by three dense layers (the first two have 1024 neurons). The dense layers are concatenated with 64 dimensional vector with a learnable action embedding. Next, three deconvolutional layers of 64 filters follow. An additional deconvolutional layer outputs an image of the original 105 × 80 size. The number of filters is either 3 or 3 × 256. In the first case, the output is a real-valued approximation of pixel's RGB value. In the second case, filters are followed by softmax producing a probability distribution on the color space. The reward is predicted by a softmax attached to the last fully connected layer. We used dropout equal to 0.2 and layer normalization.

Loss functions. The visual output of our networks is either one float per pixel/channel or the categorical 256-dimensional softmax. In both cases, we used the clipped loss max(Loss, C) for a constant C. We found that clipping was crucial for improving the models (measured with the correct reward predictions per sequence metric and successful training using Algorithm 1). We conjecture that clipping substantially decreases the magnitude of gradients stemming from fine-tuning of big areas of background consequently letting the optimization process concentrate on small but important areas (e.g. the ball in Pong). In our experiments, we set C = 10 for L 2 loss on pixel values and to C = 0.03 for softmax loss. Note that this means that when the level of confidence about the correct pixel value exceeds 97% (as -ln(0.97) ≈ 0.03) we get no gradients from that pixel any longer.

Scheduled sampling. The model env consumes its own predictions from previous steps and due to compounding errors, the model may drift out of the area of its applicability. Following Bengio et al. (2015) ; Venkatraman et al. (2016) , we mitigate this problem by randomly replacing in training some frames of the input X by the prediction from the previous step while linearly increasing the mixing probability to 100% around the middle of the first iteration of the training loop.

Stochastic Models. A stochastic model can be used to deal with limited horizon of past observed frames as well as sprites occlusion and flickering which results to higher quality predictions. Inspired by Babaeizadeh et al. (2017a) , we tried a variational autoencoder (Kingma & Welling, 2014) to model the stochasticity of the environment. In this model, an additional network receives the input frames as well as the future target frame as input and approximates the distribution of the posterior. At each timestep, a latent value z t is sampled from this distribution and passed as input to the original predictive model. At test time, the latent values are sampled from an assumed prior N (0, I). To match the assumed prior and the approximate, we use the Kullback-Leibler divergence term as an additional loss term (Babaeizadeh et al., 2017a) .

We noticed two major issues with the above model. First, the weight of the KL divergence loss term is game dependent, which is not practical if one wants to deal with a broad portfolio of Atari games. Second, this weight is usually a very small number in the range of [10 -3 , 10 -5 ] which means that the approximated posterior can diverge significantly from the assumed prior. This can result in previously unseen latent values at inference time that lead to poor predictions. We address these issues by utilizing a discrete latent variable similar to Kaiser & Bengio (2018) .

As visualized in Figure 2 , the proposed stochastic model with discrete latent variables discretizes the latent values into bits (zeros and ones) while training an auxiliary LSTM-based Hochreiter & Schmidhuber (1997) recurrent network to predict these bits autoregressively. At inference time, the latent bits will be generated by this auxiliary network in contrast to sampling from a prior. To make the predictive model more robust to unseen latent bits, we add uniform noise to approximated latent values before discretization and apply dropout (Srivastava et al., 2014) on bits after discretization. More details about the architecture is in Appendix C.

5. POLICY TRAINING:

We will now describe the details of SimPLe, outlined in Algorithm 1. In step 6 we use the proximal policy optimization (PPO) algorithm (Schulman et al., 2017) with γ = 0.95. The algorithm generates rollouts in the simulated environment env and uses them to improve policy π. The fundamental difficulty lays in imperfections of the model compounding over time. To mitigate this problem we use short rollouts of env . Typically every N = 50 steps we uniformly sample the starting state from the ground-truth buffer D and restart env (for experiments with the value of γ and N see Section 6.4). Using short rollouts may have a degrading effect as the PPO algorithm does not have a way to infer effects longer than the rollout length. To ease this problem, in the last step of a rollout we add to the reward the evaluation of the value function. Training with multiple iterations re-starting from trajectories gathered in the real environment is new to our knowledge. It was inspired by the classical Dyna-Q algorithm and, notably, in the Atari domain no comparable results have been achieved.

The main loop in Algorithm 1 is iterated 15 times (cf. Section 6.4). The world model is trained for 45K steps in the first iteration and for 15K steps in each of the following ones. Shorter training in later iterations does not degrade the performance because the world model after first iteration captures already part of the game dynamics and only needs to be extended to novel situations.

In each of the iterations, the agent is trained inside the latest world model using PPO. In every PPO epoch we used 16 parallel agents collecting 25, 50 or 100 steps from the simulated environment env (see Section 6.4 for ablations). The number of PPO epochs is z • 1000, where z equals to 1 in all passes except last one (where z = 3) and two passes number 8 and 12 (where z = 2). This gives 800K•z interactions with the simulated environment in each of the loop passes. In the process of training the agent performs 15.2M interactions with the simulated environment env .

6. EXPERIMENTS:

We evaluate SimPLe on a suite of Atari games from Atari Learning Environment (ALE) benchmark.

In our experiments, the training loop is repeated for 15 iterations, with 6400 interactions with the environment collected in each iteration. We apply a standard pre-processing for Atari games: a frame skip equal to 4, that is every action is repeated 4 times. The frames are down-scaled by a factor of 2.

Because some data is collected before the first iteration of the loop, altogether 6400 • 16 = 102, 400 interactions with the Atari environment are used during training. This is equivalent to 409, 600 frames from the Atari game (114 minutes at 60 FPS). At every iteration, the latest policy trained under the learned model is used to collect data in the real environment env. The data is also directly used to train the policy with PPO. Due to vast difference between number of training data from simulated environment and real environment (15M vs 100K) the impact of the latter on policy is negligible.

We evaluate our method on 26 games selected on the basis of being solvable with existing state-ofthe-art model-free deep RL algorithms2 , which in our comparisons are Rainbow Hessel et al. (2018) and PPO Schulman et al. (2017) . For Rainbow, we used the implementation from the Dopamine package and spent considerable time tuning it for sample efficiency (see Appendix E).

For visualization of all experiments see https://goo.gl/itykP8 and for a summary see Figure 3 . It can be seen that our method is more sample-efficient than a highly tuned Rainbow baseline on almost all games, requires less than half of the samples on more than half of the games and, on Freeway, is more than 10x more sample-efficient. Our method outperforms PPO by an even larger margin. We also compare our method with fixed score baselines (for different baselines) rather than counting how many steps are required to match our score, see Figure 4 for the results. For the qualitative analysis of performance on different games see Appendix B. The source code is available as part of the Tensor2Tensor library and it includes instructions on how to run the experiments3 .

6.1. SAMPLE EFFICIENCY:

The primary evaluation in our experiments studies the sample efficiency of SimPLe, in comparison with state-of-the-art model-free deep RL methods in the literature. To that end, we compare with Rainbow (Hessel et al., 2018; Castro et al., 2018) , which represents the state-of-the-art Q-learning method for Atari games, and PPO (Schulman et al., 2017) , a model-free policy gradient algorithm (see Appendix E for details of tuning of Rainbow and PPO). The results of the comparison are presented in Figure 3 . For each game, we plot the number of time steps needed for either Rainbow or PPO to reach the same score that our method reaches after 100K interaction steps. The red line indicates 100K steps: any bar larger than this indicates a game where the model-free method required more steps. SimPLe outperforms the model-free algorithms in terms of learning speed on nearly all of the games, and in the case of a few games, does so by over an order of magnitude. For some games, it reaches the same performance that our PPO implementation reaches at 10M steps. This indicates that model-based reinforcement learning provides an effective approach to learning Atari games, at a fraction of the sample complexity.

The results in these figures are generated by averaging 5 runs for each game. The model-based agent is better than a random policy for all the games except Bank Heist. Interestingly, we observed that the best of the 5 runs was often significantly better. For 6 of the games, it exceeds the average human score (as reported in Table 3 of Pohlen et al. (2018) ). This suggests that further stabilizing SimPLe should improve its performance, indicating an important direction for future work. In some cases during training we observed high variance of the results during each step of the loop. There are a number of possible reasons, such as mutual interactions of the policy training and the supervised training or domain mismatch between the model and the real environment. We present detailed numerical results, including best scores and standard deviations, in Appendix D. Figure 5 : Behaviour with respect to the number of used samples. We report number of frames required by PPO to reach the score of our models. Results are averaged over all games.

6.2. NUMBER OF FRAMES:

We focused our work on learning games with 100K interaction steps with the environment. In this section we present additional results for settings with 20K, 50K, 200K, 500K and 1M interactions; see Figure 5 (a) . Our results are poor with 20K interactions. For 50K they are already almost as good as with 100K interactions. From there the results improve until 500K samples -it is also the point at which they are on par with model-free PPO. Detailed per game results can be found in Appendix F.

This demonstrates that SimPLe excels in a low data regime, but its advantage disappears with a bigger amount of data. Such a behavior, with fast growth at the beginning of training, but lower asymptotic performance is commonly observed when comparing model-based and model-free methods (Wang et al. (2019) ). As observed in Section 6.4 assigning bigger computational budget helps in 100K setting. We suspect that gains would be even bigger for the settings with more samples.

Finally, we verified if a model obtained with SimPLe using 100K is a useful initialization for modelfree PPO training. Based on the results depicted in Figure 5 (b) we can positively answer this conjecture. Lower asymptotic performance is probably due to worse exploration. A policy pre-trained with SimPLe was meant to obtain the best performance on 100K, at which point its entropy is very low thus hindering further PPO training.

6.3. ENVIRONMENT STOCHASTICITY:

A crucial decision in the design of world models is the inclusion of stochasticity. Although Atari is known to be a deterministic environment, it is stochastic given only a limited horizon of past observed frames (in our case 4 frames). The level of stochasticity is game dependent; however, it can be observed in many Atari games. An example of such behavior can be observed in the game Kung Fu Master -after eliminating the current set of opponents, the game screen always looks the same (it contains only player's character and the background). The game dispatches diverse sets of new opponents, which cannot be inferred from the visual observation alone (without access to the game's internal state) and thus cannot be predicted by a deterministic model. Similar issues have been reported in Babaeizadeh et al. (2017a) , where the output of their baseline deterministic model was a blurred superposition of possible random object movements. As can be seen in Figure 11 in the Appendix, the stochastic model learns a reasonable behavior -samples potential opponents and renders them sharply. Given the stochasticity of the proposed model, Sim-PLe can be used with truly stochastic environments.

To demonstrate this, we ran an experiment where the full pipeline (both the world model and the policy) was trained in the presence of sticky actions, as recommended in (Machado et al., 2018, Section 5). Our world model learned to account for the stickiness of actions and in most cases the end results were very similar to the ones for the deterministic case even without any tuning, see Figure 6 .

6.4. ABLATIONS:

To evaluate the design of our method, we independently varied a number of the design decisions. Here we present an overview; see Appendix A for detailed results.

Model architecture and hyperparameters. We evaluated a few choices for the world model and our proposed stochastic discrete model performs best by a significant margin. The second most important parameter was the length of world model's training. We verified that a longer training would be beneficial, however we had to restrict it in all other ablation studies due to a high cost of training on all games. As for the length of rollouts from simulated env , we use N = 50 by default. We experimentally shown that N = 25 performs roughly on par, while N = 100 is slightly worse, likely due to compounding model errors. The discount factor was set to γ = 0.99 unless specified otherwise. We see that γ = 0.95 is slightly better than other values, and we hypothesize that it is due to better tolerance to model imperfections. But overall, all three values of γ perform comparably.

Model-based iterations. The iterative process of training the model, training the policy, and collecting data is crucial for non-trivial tasks where random data collection is insufficient. In a game-by-game analysis, we quantified the number of games where the best results were obtained in later iterations of training. In some games, good policies could be learned very early. While this might have been due to the high variability of training, it does suggest the possibility of much faster training (i.e. in fewer step than 100k) with more directed exploration policies. In Figure 9 in the Appendix we present the cumulative distribution plot for the (first) point during learning when the maximum score for the run was achieved in the main training loop of Algorithm 1. Random starts. Using short rollouts is crucial to mitigate the compounding errors in the model. To ensure exploration, SimPLe starts rollouts from randomly selected states taken from the real data buffer D. Figure 9 compares the baseline with an experiment without random starts and rollouts of length 1000 on Seaquest which shows much worse results without random starts.

7. CONCLUSIONS AND FUTURE WORK:

We presented SimPLe, a model-based reinforcement learning approach that operates directly on raw pixel observations and learns effective policies to play games in the Atari Learning Environment. Our experiments demonstrate that SimPLe learns to play many of the games with just 100K interactions with the environment, corresponding to 2 hours of play time. In many cases, the number of samples required for prior methods to learn to reach the same reward value is several times larger.

Our predictive model has stochastic latent variables so it can be applied in highly stochastic environments. Studying such environments is an exciting direction for future work, as is the study of other ways in which the predictive neural network model could be used. Our approach uses the model as a learned simulator and directly applies model-free policy learning to acquire the policy. However, we could use the model for planning. Also, since our model is differentiable, the additional information contained in its gradients could be incorporated into the reinforcement learning process. Finally, the representation learned by the predictive model is likely be more meaningful by itself than the raw pixel observations from the environment. Incorporating this representation into the policy could further accelerate and improve the reinforcement learning process.

While SimPLe is able to learn more quickly than model-free methods, it does have limitations. First, the final scores are on the whole lower than the best state-of-the-art model-free methods. This can be improved with better dynamics models and, while generally common with model-based RL algorithms, suggests an important direction for future work. Another, less obvious limitation is that the performance of our method generally varied substantially between different runs on the same game.

The complex interactions between the model, policy, and data collection were likely responsible for this. In future work, models that capture uncertainty via Bayesian parameter posteriors or ensembles (Kurutach et al., 2018; Chua et al., 2018) may improve robustness. Finally, the computational and time requirement of training inside world model are substantial (see Appendix C), which makes developing lighter models an important research direction.

In this paper our focus was to demonstrate the capability and generality of SimPLe only across a suite of Atari games, however, we believe similar methods can be applied to other environments and tasks which is one of our main directions for future work. As a long-term challenge, we believe that model-based reinforcement learning based on stochastic predictive models represents a promising and highly efficient alternative to model-free RL. Applications of such approaches to both high-fidelity simulated environments and real-world data represent an exciting direction for future work that can enable highly efficient learning of behaviors from raw sensory inputs in domains such as robotics and autonomous driving. 

A ABLATIONS:

To evaluate the design of our method, we independently varied a number of the design decisions: the choice of the model, the γ parameter and the length of PPO rollouts. The results for 7 experimental configurations are summarized in the Table 1 .

Models. To assess the model choice, we evaluated the following models: deterministic, deterministic recurrent, and stochastic discrete (see Section 4). Based on Table 1 it can be seen that our proposed stochastic discrete model performs best. Figures 7a and 7b show the role of stochasticity and recurrence.

Steps. See Figure 7d . As described in Section 5 every N steps we reinitialize the simulated environment with ground-truth data. By default we use N = 50, in some experiments we set N = 25 or N = 100. It is clear from the table above and Figure 7d that 100 is a bit worse than either 25 or 50, likely due to compounding model errors, but this effect is much smaller than the effect of model architecture.

Gamma. See Figure 8b . We used the discount factor γ = 0.99 unless specified otherwise. We see that γ = 0.95 is slightly better than other values, and we hypothesize that it is due to better tolerance to model imperfections. But overall, all three values of γ seem to perform comparably at the same number of steps.

Model-based iterations. The iterative process of training the model, training the policy, and collecting data is crucial for non-trivial tasks where simple random data collection is insufficient. In the game-by-game analysis, we quantified the number of games where the best results were obtained in later iterations of training. In some games, good policies could be learned very early. While this might have been due simply to the high variability of training, it does suggest the possibility that much faster training -in many fewer than 100k steps -could be obtained in future work with more directed exploration policies. We leave this question to future work.

In Figure 9 we present the cumulative distribution plot for the (first) point during learning when the maximum score for the run was achieved in the main training loop of Algorithm 1.

On Figure 7c we show results for experiments in which the number samples was fixed to be 100K but the number of training loop varied. We conclude that 15 is beneficial for training.

Long model training Our best results were obtained with much 5 times longer training of the world models, see Figure 8a for comparison with shorter training. Due to our resources constraints other ablations were made with the short model training setting.

Random starts. Using short rollouts is crucial to mitigate the compounding errors under the model.

To ensure exploration SimPLe starts rollouts from randomly selected states taken from the real data buffer D. In Figure 9 we present a comparison with an experiment without random starts and rollouts of length 1000 on Seaquest. These data strongly indicate that ablating random starts substantially deteriorate results.

B QUALITATIVE ANALYSIS:

This section provides a qualitative analysis and case studies of individual games. We emphasize that we did not adjust the method nor hyperparameters individually for each game, but we provide specific qualitative analysis to better understand the predictions from the model.4 

Solved games. The primary goal of our paper was to use model-based methods to achieve good performance within a modest budget of 100k For two games, Pong and Freeway, our method, SimPLe, was able to achieve the maximum score.

Exploration. Freeway is a particularly interesting game. Though simple, it presents a substantial exploration challenge. The chicken, controlled by the agents, is quite slow to ascend when exploring randomly as it constantly gets bumped down by the cars (see the left video https://goo.gl/YHbKZ6). This makes it very unlikely to fully cross the road and obtain a non-zero reward. Nevertheless, SimPLe is able to capture such rare events, internalize them into the predictive model and then successfully learn a successful policy.

However, this good performance did not happen on every run. We conjecture the following scenario in failing cases. If at early stages the entropy of the policy decayed too rapidly the collected experience stayed limited leading to a poor world model, which was not powerful enough to support exploration (e.g. the chicken disappears when moving to high). In one of our experiments, we observed that the final policy was that the chicken moved up only to the second lane and stayed waiting to be hit by the car and so on so forth.

Pixel-perfect games. In some cases (for Pong, Freeway, Breakout) our models were able to predict the future perfectly, down to every pixel. This property holds for rather short time intervals, we observed episodes lasting up to 50 time-steps. Extending it to long sequences would be a very exciting research direction. See videos https://goo.gl/uyfNnW.

Benign errors. Despite the aforementioned positive examples, accurate models are difficult to acquire for some games, especially at early stages of learning. However, model-based RL should be tolerant to modest model errors. Interestingly, in some cases our models differed from the original games in a way that was harmless or only mildly harmful for policy training.

For example, in Bowling and Pong, the ball sometimes splits into two. While nonphysical, seemingly these errors did not distort much the objective of the game, see Figure 10 and also https://goo.gl/JPi7rB.

In Kung Fu Master our model's predictions deviate from the real game by spawning a different number of opponents, see Figure 11 . In Crazy Climber we observed the bird appearing earlier in the game. These cases are probably to be attributed to the stochasticity in the model. Though not aligned with the true environment, the predicted behaviors are plausible, and the resulting policy can still play the original game.

Failures on hard games. On some of the games, our models simply failed to produce useful predictions. We believe that listing such errors may be helpful in designing better training protocols and building better models. The most common failure was due to the presence of very small but highly relevant objects. For example, in Atlantis and Battle Zone bullets are so small that they tend to disappear. Interestingly, Battle Zone has pseudo-3D graphics, which may have added to the difficulty. See videos https://goo.gl/uiccKU. Another interesting example comes from Private Eye in which the agent traverses different scenes, teleporting from one to the other. We found that our model generally struggled to capture such large global changes.

D NUMERICAL RESULTS:

Below we present numerical results of our experiments. We tested SimPLe on 7 configurations (see description in Section A). For each configuration we run 5 experiments. For the evaluation of the i-th experiments we used the policy given by softmax(logits(π i )/T ), where π i is the final learnt policy in the experiment and T is the temperature parameter. We found empirically that T = 0.5 worked best in most cases. A tentative explanation is that polices with temperatures smaller than 1 are less stochastic and thus more stable. However, going down to T = 0 proved to be detrimental in many cases as, possibly, it makes policies more prone to imperfections of models.

In Table 2 we present the mean and standard deviation of the 5 experiments. We observed that the median behaves rather similarly, which is reported it in Table 4 . In this table we also show maximal scores over 5 runs. Interestingly, in many cases they turned out to be much higher. This, we hope, indicates that our methods has a further potential of reaching these higher scores.

Human scores are "Avg. Human" from Table 3 in Pohlen et al. (2018) . 

E BASELINES OPTIMIZATION:

To assess the performance of SimPle we compare it with model-free algorithms. To make this comparison more reliable we tuned Rainbow in the low data regime. To this end we run an hyperparameter search over the following parameters from https://github.com/google/dopamine/ blob/master/dopamine/agents/rainbow/rainbow_agent.py:

• update_horizon in {1, 3}, best parameter = 3

• min_replay_history in {500, 5000, 20000}, best parameter = 20000

• update_period in {1, 4}, best parameter = 4

• target_update_period {50, 100, 1000, 4000}, best parameter = 8000

• replay_scheme in {uniform, prioritized}, best parameter = prioritized Each set of hyperparameters was used to train 5 Rainbow agents on the game of Pong until 1 million of interactions with the environment. Their average performance was used to pick the best hyperparameter set.

For PPO we used the standard set of hyperparameters from https://github.com/openai/ baselines. 

F RESULTS AT DIFFERENT NUMBERS OF INTERACTIONS:



:

Comparison 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.

Specifically, for the final evaluation we selected games which achieved non-random results using our method or the Rainbow algorithm using 100K interactions.

https://github.com/tensorflow/tensor2tensor/tree/master/tensor2tensor/ rl

We strongly encourage the reader to watch accompanying videos https://goo.gl/itykP8

Figures and Tables:

Figure2: Architecture of the proposed stochastic model with discrete latent. The input to the model is four stacked frames (as well as the action selected by the agent) while the output is the next predicted frame and expected reward. Input pixels and action are embedded using fully connected layers, and there is per-pixel softmax (256 colors) in the output. This model has two main components. First, the bottom part of the network which consists of a skip-connected convolutional encoder and decoder. To condition the output on the actions of the agent, the output of each layer in the decoder is multiplied with the (learned) embedded action. Second part of the model is a convolutional inference network which approximates the posterior given the next frame, similarly toBabaeizadeh et al. (2017a). At training time, the sampled latent values from the approximated posterior will be discretized into bits. To keep the model differentiable, the backpropagation bypasses the discretization followingKaiser & Bengio (2018). A third LSTM based network is trained to approximate each bit given the previous ones. At inference time, the latent bits are predicted auto-regressively using this network. The deterministic model has the same architecture as this figure but without the inference network.

Figure 3: Comparison with Rainbow and PPO. Each bar illustrates the number of interactions with environment required by Rainbow (left) or PPO (right) to achieve the same score as our method (SimPLe). The red line indicates the 100K interactions threshold which is used by the our method.

Figure 4: Fractions of Rainbow and PPO scores at different numbers of interactions calculated with the formula (SimP Le_score@100K -random_score)/(baseline_score -random_score); if denominator is smaller than 0, both nominator and denominator are increased by 1. From left to right, the baselines are: Rainbow at 100K, Rainbow at 200K, PPO at 100K, PPO at 200K. SimPLe outperforms Rainbow and PPO even when those are given twice as many interactions.

Figure 6: Impact of the environment stochasticity. The graphs are in the same format as Figure 3: each bar illustrates the number of interactions with environment required by Rainbow to achieve the same score as SimPLe (with stochastic discrete world model) using 100k steps in an environment with and without sticky actions.

Effect of adjusting of number of epochs. Effect of adjusting of number of steps.

Figure 7: Ablations part 1. The graphs are in the same format as Figure 3: each bar illustrates the number of interactions with environment required by Rainbow to achieve the same score as a particular variant of SimPLe. The red line indicates the 100K interactions threshold which is used by SimPLe.

Effect of adjusting γ in PPO training

Figure 8: Ablations part 2. The graphs are in the same format as Figure 3: each bar illustrates the number of interactions with environment required by Rainbow to achieve the same score as a particular variant of SimPLe. The red line indicates the 100K interactions threshold which is used by SimPLe.

Figure 9: (left) CDF of the number of iterations to acquire maximum score. The vertical axis represents the fraction of all games. (right) Comparison of random starts vs no random starts on Seaquest (for better readability we clip game rewards to {-1, 0, 1}). The vertical axis shows a mean reward and the horizontal axis the number of iterations of Algorithm 1.

Figure 10: Frames from the Pong environment.

Figure 11: Frames from the Kung Fu Master environment (left) and its model (right).

(a) Fraction at 100K clipped to 10. (b) Fraction at 200K (c) Fraction at 500K. (d) Fraction at 1M.

Figure12: Fractions of the rainbow scores at given number of samples. These were calculate with the formula (SimP Le_score -random_score)/(rainbow_score -random_score); if denominator is smaller than 0, both nominator and denominator are increased by 1.

(a) Fraction at 100K clipped to 10. (b) Fraction at 200K (c) Fraction at 500K. (d) Fraction at 1M.

Figure13: Fractions of the ppo scores at given number of samples. These were calculate with the formula (SimP Le_score -random_score)/(ppo_score -random_score); if denominator is smaller than 0, both nominator and denominator are increased by 1.

Summary of SimPLe ablations. For each game, a configuration was assigned a score being the mean over 5 experiments. The best and median scores were calculated per game. The table reports the number of games a given configuration achieved the best score or at least the median score, respectively.

Table content: <table><tr><td>model</td><td colspan="2">best at least median</td></tr><tr><td>deterministic</td><td>0</td><td>7</td></tr><tr><td>det. recurrent</td><td>3</td><td>13</td></tr><tr><td>SD</td><td>8</td><td>16</td></tr><tr><td>SD γ = 0.9</td><td>1</td><td>14</td></tr><tr><td>default</td><td>10</td><td>21</td></tr><tr><td>SD 100 steps</td><td>0</td><td>14</td></tr><tr><td>SD 25 steps</td><td>4</td><td>19</td></tr><tr><td colspan="3">All our code is available as part of the Tensor2Tensor library and it includes instructions on how</td></tr><tr><td colspan="3">to run our experiments: https://github.com/tensorflow/tensor2tensor/tree/</td></tr><tr><td>master/tensor2tensor/rl.</td><td/><td/></tr></table>

Models comparison. Mean scores and standard deviations over five training runs. Right most columns presents score for random agent and human.

Table content: <table><tr><td>Game</td><td>Ours, deterministic Ours, det. recurrent</td><td>Ours, SD long</td><td>Ours, SD</td><td>Ours, SD γ = 0.90 Ours, SD γ = 0.95 Ours, SD 100 steps</td><td>Ours, SD 25 steps</td><td>random</td><td>human</td></tr></table>

Comparison of our method (SimPLe) with model-free benchmarks -PPO and Rainbow, trained with 100 thousands/500 thousands/1 million steps. (1 step equals 4 frames)

Table content: <table><tr><td>Game</td><td>SimPLe</td><td>PPO_100k</td><td>PPO_500k</td><td>PPO_1m</td><td>Rainbow_100k</td><td>Rainbow_500k</td><td>Rainbow_1m</td><td>random</td><td>human</td></tr><tr><td>Alien</td><td>616.9 (252.2)</td><td>291.0 (40.3)</td><td>269.0 (203.4)</td><td>362.0 (102.0)</td><td>290.6 (14.8)</td><td>828.6 (54.2)</td><td>945.0 (85.0)</td><td>184.8</td><td>7128.0</td></tr><tr><td>Amidar</td><td>74.3 (28.3)</td><td>56.5 (20.8)</td><td>93.2 (36.7)</td><td>123.8 (19.7)</td><td>20.8 (2.3)</td><td>194.0 (34.9)</td><td>275.8 (66.7)</td><td>11.8</td><td>1720.0</td></tr><tr><td>Assault</td><td>527.2 (112.3)</td><td>424.2 (55.8)</td><td>552.3 (110.4)</td><td>1134.4 (798.8)</td><td>300.3 (14.6)</td><td>1041.5 (92.1)</td><td>1581.8 (207.8)</td><td>233.7</td><td>742.0</td></tr><tr><td>Asterix</td><td>1128.3 (211.8)</td><td>385.0 (104.4)</td><td>1085.0 (354.8)</td><td>2185.0 (931.6)</td><td>285.7 (9.3)</td><td>1702.7 (162.8)</td><td>2151.6 (202.6)</td><td>248.8</td><td>8503.0</td></tr><tr><td>Asteroids</td><td>793.6 (182.2)</td><td>1134.0 (326.9)</td><td>1053.0 (433.3)</td><td>1251.0 (377.9)</td><td>912.3 (62.7)</td><td>895.9 (82.0)</td><td>1071.5 (91.7)</td><td>649.0</td><td>47389.0</td></tr><tr><td>Atlantis</td><td colspan="3">20992.5 (11062.0) 34316.7 (5703.8) 4836416.7 (6218247.3)</td><td>-(-)</td><td colspan="5">17881.8 (617.6) 79541.0 (25393.4) 848800.0 (37533.1) 16492.0 29028.0</td></tr><tr><td>BankHeist</td><td>34.2 (29.2)</td><td>16.0 (12.4)</td><td>641.0 (352.8)</td><td>856.0 (376.7)</td><td>34.5 (2.0)</td><td>727.3 (198.3)</td><td>1053.3 (22.9)</td><td>15.0</td><td>753.0</td></tr><tr><td>BattleZone</td><td>4031.2 (1156.1)</td><td>5300.0 (3655.1)</td><td>14400.0 (6476.1)</td><td>19000.0 (4571.7)</td><td colspan="2">3363.5 (523.8) 19507.1 (3193.3)</td><td>22391.4 (7708.9)</td><td>2895.0</td><td>37188.0</td></tr><tr><td>BeamRider</td><td>621.6 (79.8)</td><td>563.6 (189.4)</td><td>497.6 (103.5)</td><td>684.0 (168.8)</td><td>365.6 (29.8)</td><td>5890.0 (525.6)</td><td>6945.3 (1390.8)</td><td>372.1</td><td>16926.0</td></tr><tr><td>Bowling</td><td>30.0 (5.8)</td><td>17.7 (11.2)</td><td>28.5 (3.4)</td><td>35.8 (6.2)</td><td>24.7 (0.8)</td><td>31.0 (1.9)</td><td>30.6 (6.2)</td><td>24.2</td><td>161.0</td></tr><tr><td>Boxing</td><td>7.8 (10.1)</td><td>-3.9 (6.4)</td><td>3.5 (3.5)</td><td>19.6 (20.9)</td><td>0.9 (1.7)</td><td>58.2 (16.5)</td><td>80.3 (5.6)</td><td>0.3</td><td>12.0</td></tr><tr><td>Breakout</td><td>16.4 (6.2)</td><td>5.9 (3.3)</td><td>66.1 (114.3)</td><td>128.0 (153.3)</td><td>3.3 (0.1)</td><td>26.7 (2.4)</td><td>38.7 (3.4)</td><td>0.9</td><td>30.0</td></tr><tr><td>ChopperCommand</td><td>979.4 (172.7)</td><td>730.0 (199.0)</td><td>860.0 (285.3)</td><td>970.0 (201.5)</td><td>776.6 (59.0)</td><td>1765.2 (280.7)</td><td>2474.0 (504.5)</td><td>671.0</td><td>7388.0</td></tr><tr><td>CrazyClimber</td><td colspan="2">62583.6 (16856.8) 18400.0 (5275.1)</td><td>33420.0 (3628.3)</td><td colspan="3">58000.0 (16994.6) 12558.3 (674.6) 75655.1 (9439.6)</td><td>97088.1 (9975.4)</td><td>7339.5</td><td>35829.0</td></tr><tr><td>DemonAttack</td><td>208.1 (56.8)</td><td>192.5 (83.1)</td><td>216.5 (96.2)</td><td>241.0 (135.0)</td><td>431.6 (79.5)</td><td>3642.1 (478.2)</td><td>5478.6 (297.9)</td><td>140.0</td><td>1971.0</td></tr><tr><td>FishingDerby</td><td>-90.7 (5.3)</td><td>-95.6 (4.3)</td><td>-87.2 (5.3)</td><td>-88.8 (4.0)</td><td>-91.1 (2.1)</td><td>-66.7 (6.0)</td><td>-23.2 (22.3)</td><td>-93.6</td><td>-39.0</td></tr><tr><td>Freeway</td><td>16.7 (15.7)</td><td>8.0 (9.8)</td><td>14.0 (11.5)</td><td>20.8 (11.1)</td><td>0.1 (0.1)</td><td>12.6 (15.4)</td><td>13.0 (15.9)</td><td>0.0</td><td>30.0</td></tr><tr><td>Frostbite</td><td>236.9 (31.5)</td><td>174.0 (40.7)</td><td>214.0 (10.2)</td><td>229.0 (20.6)</td><td>140.1 (2.7)</td><td>1386.1 (321.7)</td><td>2972.3 (284.9)</td><td>74.0</td><td>-</td></tr><tr><td>Gopher</td><td>596.8 (183.5)</td><td>246.0 (103.3)</td><td>560.0 (118.8)</td><td>696.0 (279.3)</td><td>748.3 (105.4)</td><td>1640.5 (105.6)</td><td>1905.0 (211.1)</td><td>245.9</td><td>2412.0</td></tr><tr><td>Gravitar</td><td>173.4 (54.7)</td><td>235.0 (197.2)</td><td>235.0 (134.7)</td><td>325.0 (85.1)</td><td>231.4 (50.7)</td><td>214.9 (27.6)</td><td>260.0 (22.7)</td><td>227.2</td><td>3351.0</td></tr><tr><td>Hero</td><td>2656.6 (483.1)</td><td>569.0 (1100.9)</td><td>1824.0 (1461.2)</td><td>3719.0 (1306.0)</td><td>2676.3 (93.7)</td><td>10664.3 (1060.5)</td><td>13295.5 (261.2)</td><td>224.6</td><td>30826.0</td></tr><tr><td>IceHockey</td><td>-11.6 (2.5)</td><td>-10.0 (2.1)</td><td>-6.6 (1.6)</td><td>-5.3 (1.7)</td><td>-9.5 (0.8)</td><td>-9.7 (0.8)</td><td>-6.5 (0.5)</td><td>-9.7</td><td>1.0</td></tr><tr><td>Jamesbond</td><td>100.5 (36.8)</td><td>65.0 (46.4)</td><td>255.0 (101.7)</td><td>310.0 (129.0)</td><td>61.7 (8.8)</td><td>429.7 (27.9)</td><td>692.6 (316.2)</td><td>29.2</td><td>303.0</td></tr><tr><td>Kangaroo</td><td>51.2 (17.8)</td><td>140.0 (102.0)</td><td>340.0 (407.9)</td><td>840.0 (806.5)</td><td>38.7 (9.3)</td><td>970.9 (501.9)</td><td>4084.6 (1954.1)</td><td>42.0</td><td>3035.0</td></tr><tr><td>Krull</td><td>2204.8 (776.5)</td><td>3750.4 (3071.9)</td><td>3056.1 (1155.5)</td><td>5061.8 (1333.4)</td><td>2978.8 (148.4)</td><td>4139.4 (336.2)</td><td>4971.1 (360.3)</td><td>1543.3</td><td>2666.0</td></tr><tr><td>KungFuMaster</td><td>14862.5 (4031.6)</td><td>4820.0 (983.2)</td><td>17370.0 (10707.6)</td><td>13780.0 (3971.6)</td><td colspan="2">1019.4 (149.6) 19346.1 (3274.4)</td><td>21258.6 (3210.2)</td><td>616.5</td><td>22736.0</td></tr><tr><td>MsPacman</td><td>1480.0 (288.2)</td><td>496.0 (379.8)</td><td>306.0 (70.2)</td><td>594.0 (247.9)</td><td>364.3 (20.4)</td><td>1558.0 (248.9)</td><td>1881.4 (112.0)</td><td>235.2</td><td>6952.0</td></tr><tr><td>NameThisGame</td><td>2420.7 (289.4)</td><td>2225.0 (423.7)</td><td>2106.0 (898.8)</td><td>2311.0 (547.6)</td><td>2368.2 (318.3)</td><td>4886.5 (583.1)</td><td>4454.2 (338.3)</td><td>2136.8</td><td>8049.0</td></tr><tr><td>Pong</td><td>12.8 (17.2)</td><td>-20.5 (0.6)</td><td>-8.6 (14.9)</td><td>14.7 (5.1)</td><td>-19.5 (0.2)</td><td>19.9 (0.4)</td><td>20.6 (0.2)</td><td>-20.4</td><td>15.0</td></tr><tr><td>PrivateEye</td><td>35.0 (60.2)</td><td>10.0 (20.0)</td><td>20.0 (40.0)</td><td>20.0 (40.0)</td><td>42.1 (53.8)</td><td>-6.2 (89.8)</td><td>2336.7 (4732.6)</td><td>26.6</td><td>69571.0</td></tr><tr><td>Qbert</td><td>1288.8 (1677.9)</td><td>362.5 (117.8)</td><td>757.5 (78.9)</td><td>2675.0 (1701.1)</td><td>235.6 (12.9)</td><td>4241.7 (193.1)</td><td>8885.2 (1690.9)</td><td>166.1</td><td>13455.0</td></tr><tr><td>Riverraid</td><td>1957.8 (758.1)</td><td>1398.0 (513.8)</td><td>2865.0 (327.1)</td><td>2887.0 (807.0)</td><td>1904.2 (44.2)</td><td>5068.6 (292.6)</td><td>7018.9 (334.2)</td><td>1451.0</td><td>17118.0</td></tr><tr><td>RoadRunner</td><td>5640.6 (3936.6)</td><td>1430.0 (760.0)</td><td>5750.0 (5259.9)</td><td>8930.0 (4304.0)</td><td colspan="2">524.1 (147.5) 18415.4 (5280.0)</td><td>31379.7 (3225.8)</td><td>0.0</td><td>7845.0</td></tr><tr><td>Seaquest</td><td>683.3 (171.2)</td><td>370.0 (103.3)</td><td>692.0 (48.3)</td><td>882.0 (122.7)</td><td>206.3 (17.1)</td><td>1558.7 (221.2)</td><td>3279.9 (683.9)</td><td>61.1</td><td>42055.0</td></tr><tr><td>UpNDown</td><td>3350.3 (3540.0)</td><td>2874.0 (1105.8)</td><td>12126.0 (1389.5)</td><td>13777.0 (6766.3)</td><td>1346.3 (95.1)</td><td>6120.7 (356.8)</td><td>8010.9 (907.0)</td><td>488.4</td><td>11693.0</td></tr><tr><td>YarsRevenge</td><td>5664.3 (1870.5)</td><td>5182.0 (1209.3)</td><td>8064.8 (2859.8)</td><td>9495.0 (2638.3)</td><td>3649.0 (168.6)</td><td>7005.7 (394.2)</td><td>8225.1 (957.9)</td><td>3121.2</td><td>54577.0</td></tr></table>

Models comparison. Scores of median (left) and best (right) models out of five training runs. Right most columns presents score for random agent and human.

Table content: <table><tr><td>Game</td><td>Ours, deterministic</td><td>Ours, det. recurrent</td><td>Ours, SD long</td><td>Ours, SD</td><td>Ours, SD γ = 0.90 Ours, SD γ = 0.95</td><td>SD 100 steps</td><td>Ours, SD 25 steps</td><td>random</td><td>human</td></tr></table>

Peer review (iclr2020):
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. Quality This paper has a strong applied focus and needs to be judged based on its experiments. The quality of those are high. The method is evaluated on a suite of 26 games, compared to strong model-free baselines and results are averaged over 5 seeds. One concern I have is that the method is only evaluated in low training regimes. While I do understand that increasing the training horizon is computationally demanding, results in the appendix (Figure 11a) indicate that the proposed model-based method has worse asymptotic performance compared to the model-free baselines. After 500,000 training steps the effect of sample efficiency vanishes and the final performance results are far away from the final performance results of the model-free baselines after 50,000,000 training steps. Also, a plot similar to Figure 11a) from the appendix for Rainbow DQN would be good (but I do understand this might be difficult to obtain in the course of the review period should this require more experiments). Clarity The paper is clearly written and easy to follow. However, the authors could state in the main paper more clearly that their method excels in low training regimes and that the sample efficiency effect seems to vanish when increasing training iterations from 100,000 to 500,000 steps. In fact, Figure 11a) from the appendix should go into the main paper, and it should be also mentioned that there is a huge discrepancy between the maximum performance achieved by the proposed model-based method and the maximum performance achieved by the model-free baselines when training for 50,000,000 steps. Based on the experiments, it is not clear at all if the new method will eventually catch up with best model-free results from the literature when training time is increased, or stall in low-performance regimes indefinitely. Originality The originality of this paper is not very high since the proposed algorithm and its components are not novel (there might be some minor novelty in the environment model architecture). However, this paper should not be judged based on its originality but based on its significance. Significance A working model-based RL algorithm for Atari is clearly a huge gap in the current literature and this paper takes an important step towards this direction. Demonstrating improved sample efficiency compared to strong model-free baselines in low training regimes is a significant result. The significance is however decreased by the fact that the paper does not answer the question how to obtain good asymptotic performance that matches (or comes close to) model-free state-of-the-art results. I therefore vote for weak accept at this stage. Minor details On a side note, there are two citations missing related to model-based RL in visual domains: - S. Alaniz. Deep Reinforcement Learning with Model Learning and Monte Carlo Tree Search in Minecraft. In the 3rd Multidisciplinary Conference on Reinforcement Learning and Decision Making (RLDM), 2017. - F. Leibfried and P. Vrancx. Model-Based Regularization for Deep Reinforcement Learning with Transcoder Networks. In NIPS Deep Reinforcement Learning Workshop, 2018. Update After the author response, my review remains the same. I think this paper is worthwhile publishing at ICLR.