Mastering Atari Games with Limited Data
Year: 2022

Abstract:
Reinforcement learning has achieved great success in many applications. However, sample efficiency remains a key challenge, with prominent methods requiring millions (or even billions) of environment steps to train. Recently, there has been significant progress in sample efficient image-based RL algorithms; however, consistent human-level performance on the Atari game benchmark remains an elusive goal. We propose a sample efficient model-based visual RL algorithm built on MuZero, which we name EfficientZero. Our method achieves 190.4% mean human performance and 116.0% median performance on the Atari 100k benchmark with only two hours of real-time game experience and outperforms the state SAC in some tasks on the DMControl 100k benchmark. This is the first time an algorithm achieves super-human performance on Atari games with such little data. EfficientZero’s performance is also close to DQN’s performance at 200 million frames while we consume 500 times less data. EfficientZero’s low sample complexity and high performance can bring RL closer to real-world applicability. We implement our algorithm in an easy-to-understand manner and it is available at https://github.com/YeWR/EfficientZero. We hope it will accelerate the research of MCTS-based RL algorithms in the wider community. Figure 1: Our proposed method EfficientZero is 170% and 180% better than the previous SoTA performance in mean and median human normalized score and is the first to outperform the average human performance on the Atari 100k benchmark. The high sample efficiency and performance of EfficientZero can bring RL closer to the real-world applications.

heading: 1 Introduction
text: 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.

heading: 2 Related Work

heading: 2.1 Sample Efficient Reinforcement Learning
text: 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.

heading: 2.2 Reinforcement Learning with MCTS
text: 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.

heading: 2.3 Multi-Step Value Estimation
text: 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.

heading: 3 Background

heading: 3.1 MuZero
text: 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.

heading: 3.2 Monte-Carlo Tree Search
text: Monte-Carlo Tree Search [1, 33, 34, 14], or MCTS, is a heuristic search algorithm. In our setup, MCTS is used to find an action policy that is better than the current neural network policy.
More specifically, MCTS needs an environment model, including the reward function and the nextstate function. It also needs a value function and a policy function, which act as heuristics for the tree search. MCTS operates by expanding a search tree from the current node. It saves computation by selectively expanding a few nodes. In order to find a high-quality decision, the tree expansion process has to balance between exploration versus exploitation, i.e. balance between expanding a node that is promising with many visits versus expanding a node with lower performance but fewer visits. MCTS employs the UCT [26, 20] rule, i.e. UCB [3] on trees. At every node expansion step, UCT will select a node as follows [14]:
ak = argmax a
{ Q(s, a) + P (s, a) √∑ bN(s, b)
1 +N(s, a)
( c1 + log (∑ bN(s, b) + c2 + 1
c2
))} (2)
where, Q(s, a) is the current estimate of the Q-value, P (s, a) is the current neural network policy for selecting this action, helping the MCTS prioritize exploring promising part of the tree. During training time, P (s, a) is usually perturbed by noises to allow explorations. N(s, a) denotes how many times this state-action pair is visited in the tree search, and N(s, b) denote that of a’s siblings. Thus this term will encourage the search to visit the nodes whose siblings are visited often, but itself less visited. Finally, the last term gives a weights to the previous terms.
After expanding the nodes for a pre-defined number of times, the MCTS will return how many times each action under the root node is visited, as the improved policy to the root node. Thus, MCTS can be considered as a policy improvement operator in the RL setting.

heading: 4 EfficientZero
text: Model-based algorithms have achieved great success in sample-efficient learning from lowdimensional states. However, current visual model-based algorithms either require large amounts of training data or exhibit inferior performance to model-free algorithms in data-limited settings [32]. Many previous works even suspect whether model-based algorithms can really offer data efficiency when using image observations [39]. We provide a positive answer here. We propose the EfficientZero, a model-based algorithm built on the MCTS, that achieves super-human performance on the 100k Atari benchmark, outperforming the previous SoTA to a large degree.
When directly running MCTS-based RL algorithms such as MuZero, we find that they do not perform well on the limited-data benchmark. Through our ablations, we confirm the following three issues which pose challenges to algorithms like MuZero in data-limited settings.
Lack of supervision on environment model. First, the learned model in the environment dynamics is only trained through the reward, value and policy functions. However, the reward is only a scalar signal and in many scenarios, the reward will be sparse. Value functions are trained with bootstrapping, and thus are noisy. Policy functions are trained with the search process. None of the reward, value and policy losses can provide enough training signals to learn the environment model.
Hardness to deal with aleatoric uncertainty. Second, we find that even with enough data, the predicted rewards still have large prediction errors. This is caused by the aleatoric uncertainty of the underlying environment. For example, the environment is hard to model. The reward prediction
errors will accumulate when expanding the MCTS tree to a large depth, resulting in sub-optimal performance in exploration and evaluation.
Off-policy issues of multi-step value. Lastly, when computing the value target, MuZero uses the multi-step reward observed in the environment. Although this allows the reward to be propagated to the value function faster, we find that it suffers from severe off-policy issues and hinders convergence in the limited data scenario.
To address the above issues, we propose the following three critical modifications, which can greatly improve performance when samples are limited.
4.1 Self-Supervised Consistency Loss
In previous MCTS RL algorithms, the environment model is either given or only trained with rewards, values, and policies, which cannot provide sufficient training signals due to their scalar nature. The problem is more severe when the reward is sparse or the bootstrapped value is not accurate. The MCTS policy improvement operator heavily relies on the environment model. Thus, it is vital to have an accurate one.
We notice that the output ŝt+1 from the dynamic function G should be the same as st+1, i.e. the output of the representation functionH with input of the next observation ot+1 (Fig. 2). This can help to supervise the predicted next state ŝt+1 using the actual st+1, which is a tensor with at least a few hundred dimensions. This provides ŝt+1 with much more training signals than the default scalar reward and value.
More specifically, we adopt the recently proposed SimSiam [7] self-supervised framework. SimSiam [7] is a self-supervised method that takes two augmentation views
of the same image and pulls the output of the second branch close to that of the first branch, where the first branch is an encoder network without gradient, and the second branch is the same encoder network with the gradient and a predictor head. The predictor head can simply be a two-layer MLP.
Note that SimSiam only learns the representation of individual images, and is not aware of how different images are connected. The learned image representations of SimSiam might not be a good candidate for learning the environment transition function, since adjacent observations might be encoded to very different representation encodings. We propose a self-supervised method that learns the transition function, along with the image representation function in an end-to-end manner. Figure 2 shows our method. Since we aim to learn the transition between adjacent observations, we pull ot and ot+1 close to each other. The transition function is applied after the representation of ot, such that st is transformed to ŝt+1, which now represents the same entity as the other branch. Then both of st+1 and ŝt+1 go through a common projector network. Since st+1 is potentially a more accurate description of ot+1 compared to ŝt+1, we make the ot+1 branch as the target branch. It is common in self-supervised learning that the second or the third layer from the last is chosen as the features for some reason. Here, we choose the outputs from the representation network or the dynamics network as the hidden states rather than those from the projector or the predictor.
The two adjacent observations provide two views of the same entity. In practice, we find that applying augmentations to observations such as a random small shift of 0-4 pixels on the image helps to further improve the learned representation quality [35, 32]. We also unroll the dynamic function recurrently for 5 further steps and also pull ŝt+k close to st+k (k = 1, ..., 5). Please see the Appendix for more implementation details.

heading: 4.2 End-To-End Prediction of the Value Prefix
text: In model-based learning, the agent needs to predict the future states conditioned on the current state and a series of hypothetical actions. The longer the prediction, the harder to predict it accurately, due
to the compounding error in the recurrent rollouts. This is called the state aliasing problem. The environment model plays an important role in MCTS. The state aliasing problem harms the MCTS expansion, which will result in sub-optimal exploration as well as sub-optimal action search.
Predicting the reward from an aliased state is a hard problem. For example, as shown in Figure 3, the right agent loses the ball. If we only see the first observation, along with future actions, it is very hard both for an agent and a human to predict at which exact future timestep the player would lose a point. However, it is easy to predict the agent will miss the ball after a sufficient number of timesteps if he does not
move. In practice, a human will never try to predict the exact step that he loses the point but will imagine over a longer horizon and thus get a more confident prediction.
Inspired by this intuition, we propose an end-to-end method to predict the value prefix. We notice that the predicted reward is always used in the estimation of the Q-value Q(s, a) in UCT of Equation 2
Q(st, a) = k−1∑ i=0 γirt+i + γ kvt+k (3)
, where rt+i is the reward predicted from unrolled state ŝt+i. We name the sum of rewards∑k−1 i=0 γ irt+i as the value prefix, since it is used as a prefix in the later Q-value computation.
We propose to predict value prefix from the unrolled states (st, ŝt+1, · · · , ŝt+k−1) in an end-to-end manner, i.e. value-prefix = f(st, ŝt+1, · · · , ŝt+k−1). Here f is some neural network architecture that takes in a variable number of inputs and outputs a scalar. We choose the LSTM in our experiment. During the training time, the LSTM is supervised at every time step, since the value prefix can be computed whenever a new state comes in. This per-step rich supervision allows the LSTM can be trained well even with limited data. Compared with the naive per step reward prediction and summation approach, the end-to-end value prefix prediction is more accurate, because it can automatically handle the intermediate state aliasing problem. See Experiment Section 5.3 for empirical evaluations. As a result, it helps the MCTS to explore better, and thus increases the performance. See the Appendix for architectural details.

heading: 4.3 Model-Based Off-Policy Correction
text: In MCTS RL algorithms, the value function fits the value of the current neural network policy. However, in practice as MuZero Reanalyze does, the value target is computed by sampling a trajectory from the replay buffer and computing: zt = ∑k−1 i=0 γ iut+i + γ kvt+k. This value target suffers from off-policy issues, since the trajectory is rolled out using an older policy, and thus the value target is no longer accurate. When data is limited, we have to reuse the data sampled from a much older policy, thus exaggerating the inaccurate value target issue.
In previous model-free settings, there is no straightforward approach to fix this issue. On the contrary, since we have a model of the environment, we can use the model to imagine an "online experience". More specifically, we propose to use rewards of a dynamic horizon l from the old trajectory, where l < k and l should be smaller if the trajectory is older. This reduces the policy divergence by fewer rollout steps. Further, we redo an MCTS search with the current policy on the last state st+l and compute the empirical mean value at the root node. This effectively corrects the off policy issue using imagined rollouts with current policy and reduces the increased bias caused by setting l less than k. Formally, we propose to use the following value target:
zt = l−1∑ i=0 γiut+i + γ lνMCTSt+l (4)
where l <= k and the older the sampled trajectory, the smaller the l. νMCTS(st+l) is the root value of the MCTS tree expanded from st+l with the current policy, as MuZero non-Reanalyze does. See the Appendix for how to choose l. In practice, the computation cost of the correction is two times on the reanalyzed side. However, the training will not be affected due to the parallel implementation.

heading: 5 Experiments
text: In this section, we aim to evaluate the sample efficiency of the proposed algorithm. Here, the sample efficiency is measured by the performance of each algorithm at a common, small amount of environment transitions, i.e. the better the performance, the higher the sample efficiency. More specifically, we use the Atari 100k benchmark. Intuitively, this benchmark asks the agent to learn to play Atari games within two hours of real-world game time. Additionally, we conduct some ablation studies to investigate and analyze each component on Atari 100k. To further show the sample efficiency, we apply EfficientZero to some simulated robotics environments on the DMControl 100k benchmark, which contains the same 100k environment steps.

heading: 5.1 Environments
text: Atari 100k Atari 100k was first proposed by the SimPLe [18] method, and is now used by many sample-efficient RL works, such as Srinivas et al. [35], Laskin et al. [22], Kostrikov et al. [21], Schwarzer et al. [32]. The benchmark contains 26 Atari games, and the diverse set of games can effectively measure the performance of different algorithms. The benchmark allows the agent to interact with 100 thousand environment steps, i.e. 400 thousand frames due to a frameskip of 4, with each environment. 100k steps roughly correspond to 2 hours of real-time gameplay, which is far less than the usual RL settings. For example, DQN [24] uses 200 million frames, which is around 925 hours of real-time gameplay. Note that the human player’s performance is tested after allowing the human to get familiar with the game after 2 hours as well. We report the raw performance on each game, as well as the mean and median of the human normalized score. The human normalized score is defined as: (scoreagent − scorerandom)/(scorehuman − scorerandom). We compare our method to the following baselines. (1) SimPLe [18], a model-based RL algorithm that learns an action conditional video prediction model and trains PPO within the learned environment. (2) OTRainbow [19], which tunes the hyper-parameters of the Rainbow [16] method to achieve higher sample efficiency. (3) CURL [35], which uses contrastive learning as a side task to improve the image representation quality. (4) DrQ [21], which adds data augmentations to the input images while learning the original RL objective. (5) SPR [32], the previous SoTA in Atari 100k which proposes to augment the Rainbow [16] agent with data augmentations as well as a multi-step consistency loss using BYOL-style self-supervision. (6) MuZero [27] with our implementations and the same hyper-parameters as EfficientZero. (7) Random Agent (8) Human performance.
DeepMind Control 100k Tassa et al. [37] propose the DMControl suite, which includes some challenging visual robotics tasks with continuous action space. And some works [12, 35] have benchmarked for the sample efficiency on the DMControl 100k which contains 100k environment steps data. Since the MCTS-based methods cannot deal with tasks with continuous action space, we discretize each dimension into 5 discrete slots in MuZero [27] and EfficientZero. To avoid the dimension explosion, we evaluate EfficientZero in three low-dimensional tasks.
We compare our method to the following baselines. (1) Pixel SAC, which applies SAC directly to pixels. (2) SAC-AE [42], which combines the SAC and an auto-encoder to handle image-based inputs. (3) State SAC, which applies SAC directly to ground truth low dimensional states rather than the pixels. (4) Dreamer [12], which learns a world model and is trained in dreamed scenarios. (5) CURL [35], the previous SoTA in DMControl 100k. (6) MuZero [27] with action discretizations.

heading: 5.2 Results
text: Table 1 shows the results of EfficientZero on the Atari 100k benchmark. Normalizing our score with the score of human players, EfficientZero achieves a mean score of 1.904 and a median score of 1.160. As a reference, DQN [24] achieves a mean and median performance of 2.20 and 0.959 on these 26 games. However, it is trained with 500 times more data (200 million frames). For the first time, an agent trained with only 2 hours of game data can outperform the human player in terms of the mean and median performance. Among all games, our method outperforms the human in 14 out of 26 games. Compared with the previous state-of-the-art method (SPR [32]), we are 170% and 180% better in terms of mean and median score respectively.
Apart from the Atari games, EffcientZero achieves remarkable results in the simulated tasks with continuous action space. As shown in Table 2, EffcientZero outperforms CURL, the previous SoTA,
to a considerable degree and keeps a smaller variance but MuZero cannot work well here. Notably, EfficientZero achieves comparable results to the state SAC, which consumes the ground truth states as input and is considered as the oracles.

heading: 5.3 Ablations
text: In Section 4, we discuss three issues that prevent MuZero from achieving high performance when data is limited: (1) the lack of environment model supervision, (2) the state aliasing issue, and (3) the off-policy target value issue. We propose three corresponding approaches to fix those issues and demonstrate the usefulness of the combination of those approaches on a wide range of 26 Atari games. In this section, we will analyze each component individually.
Each Component Firstly, we do an ablation study by removing the three components from our full model one at a time. As shown in Table 3, we find that removing any one of the three components will lead to a performance drop compared to our full model. Furthermore, the richer learning signals are the aspect Muzero lacks most in the low-data regime as the largest performance drop is from the version without consistency supervision. As for the performance in the high-data regime, We find that the temporal consistency can significantly accelerate the training. The value prefix seems to be helpful during the early learning process, but not as much in the later stage. The off-policy correction is not necessary as it is specifically designed under limited data.
Table 3: Ablations of the self-supervised consistency, end-to-end value prefix and model-based off-policy correction. We remove one component at a time and evaluate the corresponding version on the 26 Atari games. Each component matters and the consistency one is the most significant. The detailed results are attached in the Appendix .
Game Full w.o. consistency w.o. value prefix w.o. off-policy correction
Normed Mean 1.904 0.881 1.482 1.475 Normed Median 1.160 0.340 0.552 0.836
Figure 4: Evaluations of image reconstructions based on latent states extracted from the model with or without self-supervised consistency. The predicted next states with consistency can basically be reconstructed into observations while the ones without consistency cannot.
Temporal Consistency As the version without self-supervised consistency cannot work well in most of the games, we attempt to dig into the reason for such phenomenon. We design a decoder D to reconstruct the original observations, taking the latent states as inputs. Specifically, the architecture of D and the H are symmetrical, which means that all the convolutional layers are replaced by deconvolutional layers in D and the order of the layers are reversed in D. Therefore,H is an encoder to obtain state st from observation ot and D tries to decode the ot from st. In this ablation, we freeze all parameters of the trained EfficientZero network with or without consistency respectively and the reconstructed results are shown in different columns of Figure 4. We regard the decoder as a tool to visualize the current states and unrolled states, shown in different rows of Figure 4. Here we note that Mcon is the trained EfficientZero model with consistency andMnon is the one without consistency. As shown in Figure 4, in terms of the current state st, the observation is reconstructed well enough in the two versions. However, it is remarkable that the the decoder givenMnon can not reconstruct images from the unrolled predicted states ŝt+k while the one givenMcon can reconstruct the basic observations.
To sum up, there are some distributional shifts between the latent states from the representation network and the states from the dynamics function without consistency. The consistency component can reduce the shift and provide more supervision for training the dynamics network.
Value Prefix We further validate our assumptions in the end-to-end learning of value prefix, i.e. the state aliasing problem will cause difficulty in predicting the reward, and end-to-end learning of value prefix can alleviate this phenomenon.
To fairly compare directly predicting the reward versus end-to-end learning of the value prefix, we need to control for the dataset that both methods are trained on. Since during the RL training, the dataset distribution is determined by the method, we opt to load a half-trained Pong model and rollout total 100k steps as the common static dataset. We split this dataset into a training set and a validation set. Then we run both the direct reward prediction and the value prefix method on the training split.
As shown in Figure 5, we find that the direct reward prediction method has lower losses on the training set. However, the value prefix’s validation error is much smaller when unrolled for 5 steps. This shows that the value prefix method avoids overfitting the hard reward prediction problem, and thus it can reduce the state aliasing problem, reaching a better generalization performance.
Off-Policy Correction To prove the effectiveness of the off-policy correction component, we compare the error between the target values and the ground truth values with or without off-policy correction. Specifically, the ground truth values are estimated by Monte Carlo sampling.
We train a model for the game UpNDown with total 100k training steps, and collect the trajectories at different training stages respectively (20k, 40k, ..., 100k steps). Then we calculate the ground truth values with the final model. We choose the trajectories at the same stage (20k) and use the final model to evaluate the target values with or without off-policy correction, following the Equation 4. We evaluate the L1 error of the target values and the ground truth, as shown in Table 4. The error of unrolled next 5 states means the average error of the unrolled 1-5 states with dynamics network from current states. The error is smaller in both current states and the unrolled states with off-policy correction. Thus, the correction component does reduce the bias caused by the off-policy issue.
Furthermore, we also ablate the value error of the trajectories at distinct stages in Table 5. We can find that the value error becomes smaller as the trajectories are fresher. This indicates that the off-policy issue is severe due to the staleness of the data. More significantly, the off-policy correction can provide more accurate target value estimation for the trajectories at distinct time-steps as all the errors with correction shown in the table are smaller than those without correction at the same stage.

heading: 6 Discussion
text: In this paper, we propose a sample-efficient model-based method EfficientZero. It achieves superhuman performance on the Atari games with as little as 2 hours of the gameplay experience and state-of-the-art performance on some DMControl tasks. Apart from the full results, we do detailed ablation studies to examine the effectiveness of the proposed components. This work is one step towards running RL in the physical world with complex sensory inputs. In the future, we plan to extend it to more directions, such as a better design for the continuous action space. And we also plan to study the acceleration of MCTS and how to combine this framework with life-long learning.
Acknowledgments and Disclosure of Funding
This work is supported by the Ministry of Science and Technology of the People’s Republic of China, the 2030 Innovation Megaprojects “Program on New Generation Artificial Intelligence” (Grant No. 2021AAA0150000).

Review summaries:
This paper presents an algorithm called EfficientZero, a sample efficient model-based RL algorithm built upon MuZero. The authors demonstrate the efficiency of the proposed method via atari games, and call for the research of MCTS based RL algorithms in a wider community. The core contribution lies in the adaption of MuZero to settings where data is limited. The paper is overall interesting, easy to follow, and technically sound. 

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.

This paper presents EfficientZero - a variant of MuZero that targets data efficiency in the low-data regime. The authors identify 3 factors that hinder the data efficiency of MuZero in the low-data regime, namely lack of supervision on learning the model, compounding errors in value estimation, and lack of off-policy correction with stale data. To provide richer supervision for learning the model, EfficientZero employs an auxiliary self-supervised objective similar to SimSiam. To combat compounding errors, EfficientZero directly learns the "value prefix" rather than per-step rewards. Finally, EfficientZero utilizes the model to address off-policy correction. EfficientZero achieves a new state of the art on the Atari 100k benchmark, outperforming existing methods by a large margin. Ablation studies demonstrate the importance of all three proposed components - removing any of them results in significant performance drop.

The authors propose EfficientZero, a sample efficient model-based visual RL algorithm built on MuZero. By making three changes, namely using self-supervised learning to learn the environment model, predicting the value f, and using the model to correct off-policy value targets, EfficientZero achieves state-of-the-art performance on the Atari 100k benchmark. The authors also provide an efficient implementation for future research.