Mastering Atari Games with Limited Data

1 Introduction . Reinforcement learning has achieved great success on many challenging problems . Notable work includes DQN [ 24 ] , AlphaGo [ 33 ] and OpenAI Five [ 5 ] . However , most of these works come at the cost of a large number of environmental interactions . For example , AlphaZero [ 34 ] needs to play 21 million games at training time . On the contrary , a professional human player can only play around 5 games per day , meaning it would take a human player 11,500 years to achieve the same amount of experience . The sample complexity might be less of an issue when applying RL algorithms in simulation and games . However , when it comes to real-world problems , such as robotic ∗ { ywr20 , liush20 } @ mails.tsinghua.edu.cn , gaoyangiiis @ tsinghua.edu.cn † { thanard.kurutach , pabbeel } @ berkeley.edu 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . manipulation , healthcare , and advertisement recommendation systems , achieving high performance while maintaining low sample complexity is the key to viability . People have made a lot of progress in sample efficient RL in the past years [ 8 , 10 , 35 , 22 , 21 , 32 , 18 ] . Among them , model-based methods have attracted a lot of attention , since both the data from real environments and the “ imagined data ” from the model can be used to train the policy , making these methods particularly sample-efficient [ 8 , 10 ] . However , most of the successes are in state-based environments . In image-based environments , some model-based methods such as MuZero [ 27 ] and Dreamer V2 [ 14 ] achieve super-human performance , but they are not sample efficient ; other methods such as SimPLe [ 18 ] is quite efficient but achieve inferior performance ( 0.144 human normalized median scores ) . Recently , data-augmented and self-supervised methods applied to modelfree methods have achieved more success in the data-efficient regime [ 32 ] . However , they still fail to achieve the levels which can be expected of a human . Therefore , for improving the sample efficiency as well as keeping superior performance , we find the following three components are essential to the model-based visual RL agent : a self-supervised environment model , a mechanism to alleviate the model compounding error , and a method to correct the off-policy issue . In this work , we propose EfficientZero , a model-based RL algorithm that achieves high performance with limited data . Our proposed method is built on MuZero . We make three critical changes : ( 1 ) use self-supervised learning to learn a temporally consistent environment model , ( 2 ) learn the value prefix in an end-to-end manner , thus helping to alleviate the compounding error in the model , ( 3 ) use the learned model to correct off-policy value targets . As illustrated as Figure 1 , our model achieves state-of-the-art performance on the widely used Atari [ 4 ] 100k benchmark and it achieves super-human performance with only 2 hours of real-time gameplay . More specifically , our model achieves 190.4 % mean human normalized performance and 116.0 % median human normalized performance . As a reference , DQN [ 24 ] achieves 220 % mean human normalized performance , and 96 % median human normalized performance , at the cost of 500 times more data ( 200 million frames ) . To further verify the effectiveness of EfficientZero , we conduct experiments on some simulated robotics environments of the DeepMind Control ( DMControl ) suite . It achieves state-of-the-art performance and outperforms the state SAC which directly learns from the ground truth states . Our sample efficient and high-performance algorithm opens the possibility of having more impact on many real-world problems . 2 Related Work . 2.1 Sample Efficient Reinforcement Learning . Sample efficiency has attracted significant work in the past . In RL with image inputs , model-based approaches [ 13 , 12 ] which model the world with both a stochastic and a deterministic component , have achieved promising results for simulated robotic control . Kaiser et al . [ 18 ] propose to use an action-conditioned video prediction model , along with a policy learning algorithm . It achieves the first strong performance on Atari games with as little as 400k frames . However , Kielak [ 19 ] and van Hasselt et al . [ 39 ] argue that this is not necessary to achieve strong results with model-based methods , and they show that when tuned appropriately , Rainbow [ 16 ] can achieve comparable results . Recent advances in self-supervised learning , such as SimCLR [ 6 ] , MoCo [ 15 ] , SimSiam [ 7 ] and BYOL [ 11 ] have inspired representation learning in image-based RL . Srinivas et al . [ 35 ] propose to use contrastive learning in RL algorithms and their work achieves strong performance on image-based continuous and discrete control tasks . Later , Laskin et al . [ 22 ] and Kostrikov et al . [ 21 ] find that contrastive learning is not necessary , but with data augmentations alone , they can achieve better performance . Schwarzer et al . [ 32 ] propose a temporal consistency loss , which is combined with data augmentations and achieves state-of-the-art performance . Notably , our self-supervised consistency loss is quite similar to Schwarzer et al . [ 32 ] , except we use SimSiam [ 6 ] while they use BYOL [ 11 ] as the base self-supervised learning framework . However , Schwarzer et al . [ 32 ] only apply the learned representations in a model-free manner , while we combine the learned model with model-based exploration and policy improvement , thus leading to more efficient use of the environment model . Despite the recent progress in the sample-efficient RL , today ’ s RL algorithms are still well behind human performance when the amount of data is limited . Although traditional model-based RL is considered more sample efficient than model-free ones , current model-free methods dominate in terms of performance for image-input settings . In this paper , we propose a model-based RL algorithm that for the first time , achieves super-human performance on Atari games with limited data . 2.2 Reinforcement Learning with MCTS . Temporal difference learning [ 24 , 38 , 40 , 16 ] and policy gradient based methods [ 25 , 23 , 29 , 31 ] are two types of popular reinforcement learning algorithms . Recently , Silver et al . [ 33 ] propose to use MCTS as a policy improvement operator and has achieved great success in many board games , such as Go , Chess , and Shogi [ 34 ] . Later , the algorithm is adapted to learn the world model at the same time [ 27 ] . It has also been extended to deal with continuous action spaces [ 17 ] and offline data [ 28 ] . These MCTS RL algorithms are a hybrid of model-based learning and model-free learning . However , most of them are trained with a lot of environmental samples . Our method is built on top of MuZero [ 27 ] , and we demonstrate that our method can achieve higher sample efficiency while still achieving competitive performance on the Atari 100k benchmark . de Vries et al . [ 9 ] have studied the potential of using auxiliary loss similar to our self-supervised consistency loss . However , they only test on two low dimensional state-based environments and find the auxiliary loss has mixed effects on the performance . On the contrary , we find that the consistency loss is critical in most environments with high dimensional observations and limited data . 2.3 Multi-Step Value Estimation . In Q-learning [ 41 ] , the target Q value is computed by one step backup . In practice , people find that incorporating multiple steps of rewards at once , i.e . zt = ∑k−1 i=0 γ iut+i + γ kvt+k , where ut+i is the reward from the replay buffer , vt+k is the value estimation from the target network , to compute the value target zt leads to faster convergence [ 24 , 16 ] . However , the use of multi-step value has off-policy issues , since ut+i are not generated by the current policy . In practice , this issue is usually ignored when there is a large amount of data since the data can be thought as approximately on-policy . TD ( λ ) [ 36 ] and GAE [ 30 ] improve the value estimation by better trading off the bias and the variance , but they do not deal with the off-policy issue . Recently , image input model-based algorithms such as Kaiser et al . [ 18 ] and Hafner et al . [ 12 ] use model imaginary rollouts to avoid the off-policy issue . However , this approach has the risk of model exploitation . Asadi et al . [ 2 ] proposed a multi-step model to combat the compounding error . Our proposed model-based off-policy correction method starts from the rewards in the real-world experience and uses model-based value estimate to bootstrap . Our approach balances between the off-policy issue and model exploitation . 3 Background . 3.1 MuZero . Our method is built on top of the MuZero Reanalyze [ 27 ] algorithm . For brevity , we refer to it as MuZero throughout the paper . MuZero is a policy learning method based on the Monte-Carlo Tree Search ( MCTS ) algorithm . The MCTS algorithm operates with an environment model , a prior policy function , and a value function . The environment model is represented as the reward functionR and the dynamic function G : rt = R ( st , at ) , ŝt+1 = G ( st , at ) , which are needed when MCTS expands a new node . In MuZero , the environment model is learned . Thus the reward and the next state are approximated . Besides , the predicted policy pt = acts as a search prior over actions of a node . It helps the MCTS focus on more promising actions when expanding the node . MCTS also needs a value function V ( st ) that measures the expected return of the node st , which provides a long-term evaluation of the tree ’ s leaf node without further search . MCTS will output an action visit distribution πt over the root node , which is potentially a better policy , compared to the current neural network . Thus , the MCTS algorithm can be thought of as a policy improvement operator . In practice , the environment model , policy function , and value function operate on a hidden abstract state st , both for computational efficiency and ease of environment modeling . The abstract state is extracted by a representation function H on observations ot : st = H ( ot ) . All of the mentioned models above are usually represented as neural networks . During training , the algorithm collects roll-out data in the environment using MCTS , resulting in potentially higher quality data than the current neural network policy . The data is stored in a replay buffer . The optimizer minimizes the following loss on the data sampled from the replay buffer : L ( ut , rt ) + λ1L ( πt , pt ) + λ2L ( zt , vt ) ( 1 ) Here , ut is the reward from the environment , rt = R ( st , at ) is the predicted reward , πt is the output visit count distribution of the MCTS , pt = P ( st ) is the predicted policy , zt = ∑k−1 i=0 γ iut+i+γ kvt+k is the bootstrapped value target and vt = V ( st ) is the predicted value . Specifically , the reward function R , policy function P , value function V , the representation functionH and the dynamics function G are trainable neural networks . It is worth noting that MuZero does not explicitly learn the environment model . Instead , it solely relies on the reward , value , and policy prediction to learn the model .

Summary: The authors consider the problem of sample-efficient reinforcement learning focused on image-based data, particularly in the Atari domain. Building on the MuZero model, they introduce three key changes to improve the learned model and thereby the sample efficiency: (a) a loss to enforce consistency between rolled-out and direct state abstractions learned by the model; (b) an end-to-end approach for predicting the sum of discounted rewards ("value prefix") used to generate a MCTS Q-value estimate, and (c) the use of an additional MCTS rollout for bootstrapping the state-value estimate with replay buffer data, to limit the off-policy bias of multi-step estimates from old data. They show that the proposed changes result in a significant performance boost in the low-data (100k) Atari domain, over both the original muZero and a variety of sample-efficient baselines. The also provide an ablation study that demonstrates improvement from the combination of the three interventions over each individually. The authors also release the source code.