Procedural generalization by planning with self-supervised world models
Authors: Ankesh Anand
Jacob C Walker
Yazhe Li
Eszter Vértes
Julian Schrittwieser
Sherjil Ozair
Theophane Weber
Jessica B Hamrick
Venue: ICLR 2022 Poster (2022)
Keywords: Self-Supervised Learning
Model-Based RL
Generalization in RL

Abstract: One of the key promises of model-based reinforcement learning is the ability to generalize using an internal model of the world to make predictions in novel environments and tasks. However, the generalization ability of model-based agents is not well understood because existing work has focused on model-free agents when benchmarking generalization. Here, we explicitly measure the generalization ability of model-based agents in comparison to their model-free counterparts. We focus our analysis on MuZero (Schrittwieser et al., 2020), a powerful model-based agent, and evaluate its performance on both procedural and task generalization. We identify three factors of procedural generalization---planning, self-supervised representation learning, and procedural data diversity---and show that by combining these techniques, we achieve state-of-the art generalization performance and data efficiency on Procgen (Cobbe et al., 2019). However, we find that these factors do not always provide the same benefits for the task generalization benchmarks in Meta-World (Yu et al., 2019), indicating that transfer remains a challenge and may require different approaches than procedural generalization. Overall, we suggest that building generalizable agents requires moving beyond the single-task, model-free paradigm and towards self-supervised model-based agents that are trained in rich, procedural, multi-task environments.

Published as a conference paper at ICLR 2022
PROCEDURAL GENERALIZATION BY PLANNING WITH SELE-SUPERVISED WORLD MODELS
Ankesh Anand" Jacob Walkerl, Yazhe Lit, Eszter Vértest, Julian Schrittwieser, Sherjil Ozair, Théophane Weber, Jessica B. Hamrick DeepMind, London, UK
ABSTRACT
One of the key promises of model-based reinforcement learning is the ability to generalize using internal model of the world to make predictions in novel envi- ronments and tasks. However, the generalization ability of model-basee agents is not well understood because existing work has focused on model free agents when benchmarking generalization. Here, we explicitly measure the generalization abil- ity of model-based agents in comparison their model-free counterparts. We fo- cus our analysis on MuZero [60], a powerful model-based agent, and evaluate its performance both procedural and task generalization. We identify three factors of procedural generalization planaiiiaationplaminnnnn planning, self-supe representation learning. and procedural data diversity-and show that by combining these techniques, we achieve state-of-the art generalization performance and data efficiency on Proc- gen [9]. However, we find that these factors do not always provide the same ben- efits for the task generalization benchmarks in Meta-Won World [74], indicating that transfer remains a challenge and may require different approaches than procedu- ral generalization. Overall, we suggest that building generalizable agents requires moving beyond the single-task, model-free paradigm and towards self-supervised model-basee agents that are trained in rich, procedural, multi task- environments.
INTRODUCTION
The ability to generalize to previously unseen situations or tasks using an internal model of the world is hallmark capability of human general intelligence [12, 41] and is thought by many to be of central importance in machine intelligence as well [14, 26, 59, 66]. Although significant strides have been made in model-based systems in recent years [28], the most popular model-based
benchmarks consist of identical training and testing environments [e.g. 27, 69) and do not measure or optimize for for generalization at all. While plenty of other work model-based RL does measure generalization [c.g. 18, 48, 70, 76]. each approach is typically evaluated on a bespoke task, making difficult to ascertain the state of generalization in model- RL more broadly.
Model-free RL, like model-based RL, has also suffered from both the "train=test" paradigm and a lack of standardization around how to measure generalization. In response, recent papers have discussed what generalization in RL means and how to measure it [7, 8, 36, 49, 71], and others have proposed new environments such as Procgen [9] and Meta- World [74] as benchmarks focusing on measuring generalization. While popular in the model-free community [e.g. 47, 73, 78), these benchmarks have not yet been widely adopted in the model-based setting. It is therefore unclear whether model-based methods outperform model-free approaches when it comes to generalization, how well model-based methods perform on standardized benchmarks, and whether popular model- free algorithmic improvements such self-supervision [15, 44] or procedural data diversity [67, 77] yield the same benefits for generalization in model-based agents.
In this paper, investigate three factors of generalization in model-based RL: planning, self- represtnttin learning, and procedural data diversity. We analyze these methods through
WWrk done while visiting from Mila, University of Montreal.
I Joint first authors.
I Equal contribution.

---

Published as a conference paper at ICLR 2022
variety of modifications and ablations to MuZero Reanalyse [60, 61], ofate-of-the-art model- based algorithm. To assess generalization performance, we test our variations of MuZero on two types of generalization (See Figure 1): procedural and task. Procedural generalization involves evaluating agents on unseen configurations of an environment (e.g., changes in observation render- ing, map terrain changes, new goal locations) while keeping the reward function largely the same. Task generalization, in contrast, involves evaluating agents' adaptability to unseen reward functions ("tasks"" within the same environment. We focus on two benchmarks designed for both types of generalization, Procgen [9] and Meta- World [74].
Our results broadly indicate that self-supervised, model-based agents hold promise in making progress to- wards better generalization. We find that (1) MuZero achieves state-of- the-art performance on Procgen and the procedural and multi-task Meta- World benchmarks (ML-1 and ML- 45 train), outperforming a controlled model-free baseline; (2) MuZero's performance and data efficiency can be improved with the incorpora- tion of self-supervised representa- tion learning: and (3) that with self- supervision, less data diversity is re- quired to achieve good performance. However, (4) these ideas help less for task generalization on the ML.45 test set from Meta- World, suggesting that different forms of generalization may require different approaches.
Train
Test
Procedural Generalization
Generalization
Figure I: Two different kinds of generalization, using Proc- gen and Meta- as examples. Procedural generaliza- tion involves evaluating on unseen environment configura- tions, whereas task generalization evaluates adaptability to unseen tasks (reward functions).
2
MOTIVATION AND BACKGROUND
2.1 GENERALIZATION IN RL
We are interested in the setting where an agent is trained on set of Perain MDPs drawn IID from the same distribution, Marain {M: ] Pezsin where M; ~ p(M). The agent is then tested on another set of Thest MDPs drawn IID' from p, but disjoint from the training set: Mresst {M;) Jiteet where M2 ~ P(M) such that M, & Merain V MMy. Generalization, then, is how well the agent performs in expectation on the test MDPs after training (analogous to how well a supervised model performs on the test set). Qualitatively, generalization difficulty can be roughly seen as a function of the number (Atrain) of distinct MDPs seen during training (what we will refer to as data diversity), as well the breadth or amount variation in P(M) itself. Intuitively, the amount of diversity in the training set is small (i.e. low values of Merain) relative to the breadth of p( p(M,, then this poses more difficult generalization challenge. As Merain - Do, we will eventually enter a regime where p(M) is densely sampled, and generalization should become much easier.
Given the above definition of generalization, we distinguish between two qualitative types of distri- butions MDPs. In procedural generalization, all MDPs with non-zero probability under p(M) share the same underlying logic to their dynamics (e.g.. that walls are impassable) and rewards (e.g.. that coins are rewarding), but differ in how the environment is laid out (e.g., different mazes) and in how it is rendered (e.g., color or background). In task generalization, MDPs under p(M) share the same dynamics and rendering, but differ in the reward function (e.g.. picking an object up vs. pushing it), which may be parameterized (e.g. specifying goal location).
'eehhicalll, the different tasks contained in were designed by hand. However, they qualitatively involve the same types of objects and level of complexity, feel this fits approximately into this regime.
(see Section 3) are not quite drawn IID as they
2

---

Published as a conference paper at ICLR 2022
Procedural generalization Many recent works attempt to improve procedural generalization in different ways. For example, techniques that have been successful in supervised learning have also been shown to help procedural generalization in RL, including regularization [9, 33, 42] and aux- iliary self-supervised objectives [43, 44]. Other approaches including better hyperparameter tuning [47], curriculum learning strategies [35], and stronger architectural inductive biases [5, 23, 75] may also be effective. However, these various ideas are often only explored in the model-free setting, and are not usually evaluated in combination, making challenging know which are the most beneficial for model- based RL. Here we focus explicitly on model-based based agents, and evaluate three factors in combination: planning, self supervsurrviiinn and data diversity.
Task generalization Generalizing to new tasks or reward functions has been a major focus of meta-RL. [19, 57], in which an agent is trained on a dense distribution of tasks during training and a meta-objective encourages few-shot generalization to new tasks. However, meta-RL works have primarily focused on model-free algorithms [though see 48), and the distribution from which train/test MDPs are drawn is typically quite narrow and low dimensional. Another approach to task generalization has been to first train task-agnostic representations [72] or dynamics [64] in an exploratory pre-training phase, and then use them for transfer. Among these, Sekar et al. [64] focus on the model-based based setting but require access to the reward function during evaluation. In Section 4.2, we take a similar approach of pre-training models and using them for task transfer, with the goal of evaluating whether planning and self-supervised learning might assist in this.
2.2 FACTORS OF GENERALIZATION
Planning Model-based RL is an active area of research [28, 46] with the majority of work focusing on gains in data efficiency [26, 27, 34, 37, 61], though is often motivated by a desire for better zero- or few-shot generalization as well [18, 29, 48, 64, 70]. In particular, model-based techniques may benefit data efficiency and generalization in three distinct ways. First, model learning can act auxiliary task and thus aid in learning representations that better capture the structure of the environment and enable faster learning [20]. Second, the learned model can be used to select actions on-the-fly fly via MPC [18], enabling faster adaptation in the face of novelty. Third, the model can also be used to train policy or value function by simulating training data [66] or constructing more informative losses [21, 30), again enabling faster learning (possibly entirely in simulation). A recent state-of-the-an agent, MuZero [60], combines all of these techniques: model learning, MPC, simulated training data, and model-based based losses. 2 We focus our analysis MuZero for this reason, and compare it to baselines that do not incorporate model-based components.
Self-supervision Model learning is itself a form of self-supervision, leveraging an assumption about the structure of MDPs in order to extract further learning signal from collected data than is possible via rewards alone. However, recent work has argued that this is unnecessary for model- based RL: all that should be required is for models to learn the task dynamics, not necessarily environment dynamics [22]. Yet even in the context of model- free RL, exploiting the structure of the dynamics has been shown to manifest in better learning efficiency, generalization, and representation learning (44, 63, 72]. Here, we aim to test the impact of self-supervision in model-based agents by focusing on three popular classes of self-supervised losses: reconstruction, contrastive, and self- predictive. Reconstruction losses involve directly predicting future observations [c.g. 18, 27, 70]. Contrastive objectives set up classification task to determine whether future frame could result from the current observation [4, 38, 52]. Finally, self- self-predictive losses involve having agents predict their own future latent states [25, 63].
Data diversity Several works have shown that exposing deep neural network to diverse training
distribution can help its representations better generalize to unseen situations [16, 54). A dense sampling over diverse training distribution can also act as a way to sidestep the dus-of-distiiiutiin generalization problem [54]. In robotics, collecting and training over a diverse range of data has been found to be critical particularly in the sim2real literature where the goal is to transfer a policy
"In this paper, we refer to planning as some combination of MPC, simulated training data, and model-based losses (excluding model learning iteself). We are agnostic to the question of how deep or far into the future the model must be used; indeed, it may be the case that in the environments we test, very shallow or even single-step planning may he sufficient, as in Hamrick et al. [29], Hessel et al. [31].
3

---

Published as a conference paper at ICLR 2022
learned in simulation to the real world [3, 67]. Guez al. [23] also showed that in the context of zero-shot generalization, the amount of data diversity can have interesting interactions with other architectural choices such as model size. Here, we take a cue from these findings and explore how important procedural data diversity in the context of self-supervised, model-based agents.
2.3 MUZEEO
We evaluate generalization with respect to a state-of-the-art model-based agent, MuZero [60]. MuZero is an appealing candidate for investigating generalization because it already incorporates many ideas that are thought to be useful for generalization and transfer, including replay [61] J. plan- ning [65], and model-based representation learning [22. 51J. However, while these components contribute to strong performance across wide range of domains, other work has suggested that MuZero does not necessarily achieve perfect generalization on its own [29]. therefore serves as a strong baseline but with clear room for improvement.
MuZero learns an implicit (or value-cquivalent, see Grimm et al. [22]) world model by simply learning to predict future rewards, values and actions. It then plans with Monte-Carlo Tree Search (MCTS) [40, 11] over this learned model to select actions for the environment. Data collected from the environment, as well the results of planning, are then further used improve the learned reward function, value function, and policy. More specifically, MuZero optimizes the following loss every time-step L, applied
model that is unrolled O...K K steps into the future: (0) Ex + 1 + 15) = Ex=of (CE(#*, œ²) + CE((*, v²) + CE(f"; ,r")), where a 6* and f are respectively the policy, value and reward prediction produced by the k-step unrolled model. The targets for these predictions are drawn from the corresponding time-step t + k of the real trajectory: T is the improved policy generated by the search tree, v² is n-step return bootstrapped by a target network, and r* is the true reward. As MuZero uses a distributional approach for training the value and reward functions [13], their losses involve computing the cross-entropy (CE); this also typically results in better representation learning. For all our experiments, we specifically parameterize v and r as categorical distributions similar to the Atari experiments in Schrittwieser et al. [60].
To enable better sample reuse and improved data efficiency, we use the Reanalyse version of MuZero [61]. Reanalyse works by continuously re-running MCTS on existing data points, thus computing new improved training targets for the policy and value function. It does not change the loss function described in Section 2.3. In domains that use continuous actions (such as Meta-World), we use Sampled MuZero [32] that modifies MuZero to plan over sampled actions instead. See Section A. for more details on MuZero.
3 EXPERIMENTAL DESIGN
We analyze three potential drivers of generalization (planning, self-supervision, and data diversity) across two different environments. For each algorithmic choice, we ask: to what extent does improve procedural generalization, and to what extent does it improve task generalization?
3.1 ENVIRONMENTS
Procgen Procgen [9] a suite of 16 different Atari-like games with procedural environments (e.g. game maps, terrain, and backgrounds), and was explicitly designed as a setting in which to test procedural generalization. It has also been extensively benchmarked [9, 10, 35, 55). For each game, Procgen allows choosing between two difficulty settings (easy or hard) as well as the number of levels seen during training. In our experiments, we use the "hard" difficulty setting and vary the numbers of training levels to test data diversity (see below). For each Procgen game, agent for 30M environment frames. We use the same network architecture and hyper-parameters that Schrittwieser et al. [61] used for Atari and perform no environment specific tuning. We report mean normalized scores across all games as in Cobbe et al. [9]. Following the recommendation of Agarwal al. [1], report the min-max normalized scores across all games instead of PPO- normalized scores. Thus, the normalized score for each game is computed as (score- min) "maz manj" where min and max are the minimum and maximum scores possible per game as reported in [9]. We then average the normalized scores across all games report the mean normalized score.
train an
4

---

Published as a conference paper at ICLR 2022
ninerr
nnn
plundor
Starploc
PLR (200M) (200M) (200M)
Moni GL
helz
Lumper
peoper
Mean Epicode Return
Mcan Normalized Score
dimber
colnrun
dodgebia
TruRsos
0.4
0.2--
bipftsh
bossfgttt
cavelyeer
chaser
0.0-
10M 2000
10M 200
15M 20M
10M
10M
20M
Environment fraies
Environment frames
(a) Individual scores across all 16 Procgen games
(b) Mean normalized Score
Figure 2: The impact of planning and self-supervision on procedural generalization in Procgen (hard difficulty, 500 train levels). We plot the zero-shot evaluation performance on unseen levels for each agent throughout training. The Q-Learning agent (QL) is replica of the MuZero (MZ) with its model-based components removed. MZ+Contr MuZero agent augmented with a tem- poral contrastive self-supervised loss that is action-conditioned (we study other losses in Figure 3). We observe that both planning and supervuprvis improve procedural generalization on Proegen. Comparing with existing state-of-the-a art methods which were trained for 200M frames on the right (PPO [62], PLR [35], and UCB-DrAC+PLR [56, 35], data from [35]), we note that MuZero itself exceeds state-of-the-art performance after being trained on only 30M frames. For all plots, dark lines indicate median performance across 3 seeds and the shaded regions denote the min and max performance across seeds. For training curves see Figure A.4, for additional metrics see Figure A.1.
Meta-World Meta-World [74] a suite of different tasks on robotic SAWYER arm, making more suitable to test task generalization. Meta-World has three benchmarks focused on general- MI.-45. MI.-1 consists of goal-conditioned tasks where the objective to generalize to unseen goals during test time. ML-10 and ML-45 require generalization to com- pletely new tasks (reward functions) at test time after training on either or 45 tasks, respectively. Meta-World exposes both state and pixel observations, as well as dense and sparse versions of the reward functions. In our experiments, we use the v2 version of Meta- World, dense rewards, and the corner3 camera angle for pixel observations. For Meta- Meta-World, we trained Sampled MuZero [32] for 50M environment frames. We measure performance in terms of average episodic success rate across task(s). Note that this measure different from task rewards, which are dense.
ization: MI.-1, MI-10,
3.2 FACTORS GENERALIZATION
Planning To evaluate the contribution of planning (see Section 2.2), we compared the performance of vanilla MuZero Reanalyse agent with a Q-Learning agent. We designed the Q-Learning agent to be similar MuZero as possible: for example, shares the same codebase, network architecture, and replay strategy. The primary difference is that the Q-Learning agent uses Q-learning loss instead of the MuZero loss Section 2.3 to train the agent, and uses e-greedy instead of MCTS to act in the environment. See Section A.2 for further details.
Self-supervision We looked at three self-supervised methods as auxiliary losses on top of MuZero: image reconstruction, contrastive learning [24], and self prediciciitive representations [63]. We do not leverage any domain-specifit data-augmentations for these self-supervised methods.
Reconstruction. Our approach for image reconstruction differs slightly from the typical use of mean reconstruction eтror over all pixels. In this typical setting, the use averaging implies that the decoder focus on the easy model parts of the observation and still perform well. To encourage the decoder to model all including the hardest to reconstruct pixels, we add an additional loss term which corresponds to the reconstruction error over all pixels. See Section A.3 for details.
5

---

Published as a conference paper at ICLR 2022
Moan Normalizos lrod Score
0.6
NZ 0.4
Normallzes
Modil Hoddl-Rocan
OL
QL- Model
0.2
MZ Contr
0.0
10M
20M
30M
38M
201
J0N
Environment frames
Environment frames
Figure 3: Evaluation performance on Procgen (hard, 500 train levels). On the left, we ablate the effectiveness of planning. The Q-Learning agent (QL) is a replica of MuZero (MZ) without model- based components. We then add model to this agent (QL+Model) (see Section A.2) to disentangle the effects of the model-based representation learning from planning in the full MuZero model (MZ). On the right, we ablate the effect of self- supervuuerviioon with three different losses: Contrastive (Contr), SeffPPredictive (SPR), and Image Reconstruction (Recon). We also include Q-Learning agent with reconstruction (QL+Model+ performance across 3 seeds (5 seeds for MZ and MZ+Recon) and the shaded regions denote the min and max performance across seeds. For corresponding training curves Figure A.5.
as a baseline. For all plots, dark lines indicate median
Contrastive. We also experiment with a temporal contrastive objective which treats pairs of ob- servations close in time as positive examples and un-correlated timestamps as negative examples [2, 52, 23]. Our implementation of the contrastive objective is action-conditionel and uses the MuZero dynamics model to predict future embeddings. Similar to ReLIC [45], we also use KI. regularization term in addition to the contrastive loss. The contrastive loss operates entirely in the latent space and does not need decoding to the pixel space at each prediction step, unlike the recon- struction loss. See Section A.4 for details.
Self-predictive. Finally, we experiment with self-predictive representations (SPR) [63] which are trained by predicting future representations of the agent itself from target network using the MuZero dynamics model. Similar to the contrastive loss, this objective operates in the latent space but does not need negative examples. See Section A.5 for details.
Data diversity To evaluate the contribution of data diversity (see Section 2.1), we ran experiments on Proegen in which we varied number of levels seen during training (either 10, 100, 500, or oc). In all cases, evaluated the infinite test split. Meta- World does not expose way modify the amount of data diversity, therefore we did not analyze this factor our Meta- World experiments.
4 RESULTS
We ran all experiments according to the design described in Section 3. Unless otherwise specified, all reported results are on the test environments of Procgen and Meta-World (results on the training environments are also reported in the Appendix) and are computed medians across seeds.
4.1 PROCEDURAL GENERALIZATION
Overall results Procgen shown in Figure 2 and Table A.2. MuZero achieves mean- normalized test score of 0.50, which is slightly higher than earlier saute-o--the-e methods like UCB-DrAd UCB-DrAC+PLR which was specifically designed for procedural generalization tasks [35], while also being much more data efficient (30M frames vs. 200M frames). MuZero also outperforms our Q-Learning baseline which gets score of 0.36. Performance is further improved by the addition of self-supervision, indicating that both planning and self-supervised model learning are important.
Effect of planning While not reaching the same level of performance MuZero, the Q-Learning baseline performs quite well and achieves mean normalized score of 0.36, matching performance

---

Published as a conference paper at ICLR 2022
Chaser
Climber
Chaser
C Iimbe
k-5
k-5
kw0
Decoding Oboeevvation
MuZero
MuZero Reconstruction
Figure 4: Qualitative comparison of the information encoded in the embeddings learned by MuZero with and without the auxiliary pixel reconstruction loss. For MuZero, embeddings are visualized by learning a standalone pixel decoder trained with MSE. Visualized are environment frames (top row) and decoded frames (bottom row) for two games (Chaser and Climber), for embeddings at the current time step (k 0) and 5 steps into the future (k = 5). Colored circles highlight important entities that are or are not well captured (bluc=captured, yellow=so-so red=missing).
of other specialized model-free methods such as PLR [35]. By modifying the Q-Learning baseline to learn a 5-step value equivalent model (similar to MuZero), we find its performance further improves 0.45, though does not quite catch up to MuZero itself (see Figure left). This suggests that while simply learning a value-equivalen model can bring representational benefits, the best results come from also using this model for action selection and/or policy optimization.
We also tested the effect of planning in the Meta- World ML-1 benchmark from states. Table A.5 and Figure A.7 show the results. We find that both the Q-Learning and MuZero agents achieve perfect or near-perfect generalization performance on this task, although MuZero is somewhat more stable and data-efficient. The improvement of MuZero in stability or data efficiency again suggests that model-learning and planning can play role in improving performance.
Effect of self-supervision Next, we looked at how well various self-supervised auxiliary losses improve over MuZero's value equivalent model on Procgen (Figure 3. right). We find that contrastive learning and self- ceff-rredictive representations both substantially improve over MuZero's normalized score of 0.50 0.63 (contrastive) or 0.64 (SPR), which are new state-of-the-art scores Proc- gen. The reconstruction loss also provides benefits but of a lesser magnitude, improving MuZero's 'ss performance from 0.50 0.57. All three self-supervised losses also improve the data efficiency of generalization, Of note is the fact that a MuZero agent with the contrastive loss can match the final performance of the baseline MuZero agent using only a third of the data (10M environment frames).
To further tease apart the difference between MuZero with and without self- seff-supervision, we per- formed qualitative comparison between reconstructed observations of MuZero with and without image reconstruction. The vanilla MuZero agent was modified to include image reconstruction, but with a stop gradient on the embedding so that the reconstructions could not influence the learned rep- resentations. We trained each agent two games, Chaser and Climber. As can be seen in Figure 4, it appears that the primary benefit brought by self-supervision is in learning a more accurate world model and capturing more fine-gra tne-ggained details such as the position of the characters and enemies. In Chaser, for example, the model augmented with a reconstruction loss able to to predict the position of the character across multiple time steps, while MuZero's 's embedding only retains precise information about the character at the current time step. This is somewhat surprising in light of the arguments made for eeuiveeeui models [22], where the idea is that world models do need to capture full environment dynamics-one only what is needed for the task. Our results suggest that in more complex, procedurally generated environments, it may be challenging to learn even the task dynamics from reward alone without leveraging environment dynamics, too. Further gains in model quality might be gained by also properly handling stochasticity (53) and causality [58].
Effect of data diversity Overall, we find that increased data diversity improves procedural gen- eralization in Procgen, as shown in Figure 5. Specifically, at 10 levels, self-supervision barely improves the test performance procgen over MuZero. But as we increase the number of levels to
7

---

Published as a conference paper at ICLR 2022
Co+Co
NZ+Rccon MZ+
10 training locels
training lecels
500 training levels
All trakting lwvels
Mean Normalized Score
10M
20M
30M
10M
20M
10M
20M 30M 10M
20M 30M
Enviranment frames
Figure 5: Interaction of self-supervision and data diversity on procedural generalization. Each plot shows generalization performance as a function of environment frames for different numbers of training levels. With only 10 levels, self-supervision does not bring much benefit over vanilla MuZero. Once the training set includes at least 100 levels, there is large improvement with self- supervised learning both in terms of data efficiency and final performance. For all plots, dark lines indicate median performance across seeds and shading indicates min/max seeds. See also Figure A.6 for corresponding training curves.
100, 500, or oo, we observe substantial improvement when using self-supervision. This is an in- teresting finding which suggests that methods that improve generalization might show promise only when we evaluate them on environments with lot of inherent diversity, such as procedurally gener- ated environments. Consequently, evaluating methods solely on single task settings (as is common in RI.) might lead us to overlook innovations which might have big impact on generalization.
Not only do we find that MuZero with self- supervisoon achieves better final generalization perfor- mance than the baselines, we also observe that self-supervision improves generalization performance even when controlling for training performance. To see this, we visualized generalization perfor- mance as a function of training performance (see Figure A.3). The figure shows that even when two agents are equally strong (according to the rewards achieved during training), they differ at test time, with those trained with self- suffssuperviion generally achieving stronger generalization perfor- mance. This suggests that in addition to improving data efficiency, supe-supeoviiion leads more robust representations a feature that again might be overlooked if not measuring generalization.
4.2 TASK GENERALIZATION
MAML
RL 2
QL
MZ
MZ+Recon
ML-10
44.4%
86.9%
85.2%
97.6%
97.8%
As we have shown, planning, self- supervision, and data diversity all play an important role in procedu- ral generalization. We next inves- tigated whether this trend holds for task generalization, too. To test this, pre-trained the Q-Learning and MuZero agents with and without self- supervision on the ML-10 and ML-45 training sets of Meta-World and then evaluated zero-shot test performance. In all experiments, we trained agents from pixels using the dense setting of the reward functions. We report the agent's success rate computed the same way in Yu et al. [74].
zero-shot
6.8%
26.5 %
25.0%
few-shot
31.6%
35.8%
finetune
80.1%
94.1%
97.5%
ML-45
train zero-shot
40.7%
70%
55.9%
77.2% 2
74.9%
10.8%
17.7%
18.5%
few-shot
39.9%
33.3%
finetune
78.1%
76.7%
81.7%
Table 1: Train and various test success rates (zero-shot, few- shot, finetuning) on the ML-10 and ML-45 task generaliza- tion benchmarks of Meta- World, Shown are baseline results on MAMI. [19] and RI.2 2 [17] from the Meta- paper [73] well as our results with Q-Learning and MuZero. Note that MAML and RL² were trained for around 400M environment steps from states, whereas MuZero was trained from pixels for 50M steps on train tasks and 10M steps on test tasks for fine-tuning.
As shown in Table 1, we can observe that MuZero reaches better training performance on ML- 10 (97.6%) and ML-45 (77.2%) compared to existing state-of-the-art agents [74]. These existing approaches are meta-learning methods, and are therefore not directly comparable in terms test performance due to different data budgets allowed at test time; although MuZero does not reach the same level of performance as these agents at test, we find it compelling that it succeeds as often as it does, especially after being trained for only 50M environment steps (compared to 400M for the
8

---

Published as a conference paper at ICLR 2022
ML10-Test
ML45-Test
Total rearet
QL
QL (finetune)
MZ (fnetune)
MZ+Recon (fintune)
10M
Environment frames
Environment frames
Figure 6: Finetuning performance ML-10 and MI,45, shown cumulative regret the suc- cess rate (lower is better). Both the pre-trained Q-Leerniig and MuZero agents have lower regret than corresponding agents trained from scratch. MuZero also achieves lower regret than Q-Learning, indicating a positive benefit of planning (though the difference is small on ML-45), Self-supervision (pixel reconstruction) does not provide any additional benefits. Solid lines indicate median perfor- mance across seeds, and shading indicates min/max seeds. See also Figure A.8.
baselines). MuZero also outperforms Q-Learning in terms zero-shot performance, indicating positive benefit of planning for task generalization. However, reconstruction and other forms of self-supervision do not improve performance and may even decrease (see also Figure A.8).
We also looked at whether the pre-trained representations or dynamics would assist in more data- efficient task transfer by fine-tuning the agents on the test tasks for 10M environment steps. Figure shows the results, measured as cumulative regret (see also Figure A.8 for learning curves). Using pre-trained representations or dynamics does enable both Q-Learning and MuZero to outperform corresponding agents trained from scratch-evidesd these agents. Additionally, MuZero exhibits better data efficiency than Q-Learning again showing benefit for planning. However, self- superviseii again does not yield improvements in fine-tuning, and as before may hurt in some cases (Figure A.8). This indicates that while MuZero (with or without self-supervision) excels representing variation across tasks seen during training, there is room for improvement to better transfer this knowledge to unseen reward functions.
for weak positive transfer to unseen tasks for
We hypothesize that MuZero only exhibits weak positive transfer unseen tasks due to combina- tion of factors. First, the data that its model is trained on is biased towards the training tasks, and thus may not be sufficient for learning globally-accurate world model that is suitable for planning in different tasks. Incorporating more exploratory pre-training phase [e.g. 64] might help to alleviate this problem. Second, because the model relies on task-specific gradients, the model may over- represent features that are important to the training tasks (perhaps suggesting that value-equivalent models [22] may be poorly suited task generalization). Third, during finetuning, the agent must still discover what the reward function is, even if it already knows the dynamics. It is possible that the exploration required to do this the primary bottleneck for task transfer, rather than the model representation itself.
CONCLUSION
In this paper, we systematically investigate how well modern model-based methods perform at hard generalization problems, and whether self-supervision of such methods. We find that the case of procedural generalization in Procgen, both model-based learning and self-supervision have additive benefits and result in state-of-the-a performance on test levels with remarkable data efficiency. In the case of task generalization in Meta- Hera-World, we find that while a model-based agent does exhibit weak positive transfer to unseen tasks, auxiliary self- supervision does not provide any additional benefit, suggesting that having access to a good world model not always sufficient for good generalization [29]. Indeed, we suspect that to succeed at task generalization, model-based methods must be supplemented with more sophisticated online exploration strategies, such as those learned via meta-learning [17, 39, 68] or by leveraging world models in other ways [64]. Overall, conclude that self-supervised model-based methods are promising starting point for developing agents that generalize better, particularly when trained in rich, procedural, and multi-task environments.
improve the generalization performance

---

Published as a conference paper at ICLR 2022
REPRODUCIBILITY STATEMENT
In this paper. we have taken multiple steps to ensure that the algorithms, experiments, and results are as reproducible as possible. Both environments used in our experiments, Proegen and Meta- World, are publicly available. In order to produce reliable results, we ran multiple seeds in our key experiments. Our choice of agent is based on MuZero, which is published work [32, 60. 61]. In addition, we have included an extensive appendix (Appendix A) which describes the architectural details of MuZero, the exact hyperparameters used for experiments, and a detailed description of the Q-Learning agents and the self-supervised losses.
ACKNOWLEDGEMENTS
We would like thank Jovana Mitrovié, Ivo Danihelka, Feryal Behbahani, Abe Friesen, Peter Battaglia and many others for helpful suggestions and feedback on this work.
REFERENCES
[1] Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron Courville, and Mare Belle- mare. Deep reinforcement learning the edge of the statistical precipice. In NeurlPS, 2021. Cited on pages 4 and 19.
[2] Ankesh Anand, Evan Racah, Sherjil Ozair, Yoshua Bengio, Marc-Alexandre Côté, and R De- von Hjelm. Unsupervised state representation learning in atari. In NeurIPS, 2019. Cited on page 6.
[3] OpenAl: Marcin Andrychowicz, Bowen Baker, Maciek Chociej, Rafal Jozefowicz, Bob Mc- Grew, Jakub Pachocki, Arthur Petron, Matthias Plappert, Glenn Powell, Alex Ray, et al. Learn- ing dexterous in-hand manipulation. The International Journal of Robotics Research, 2020. Cited on page 4.
[4] Andrea Banino, Adrià Puidomenech Badia, Jacob Walker, Tim Scholtes, Jovana Mitrovic, and Charles Blundell. Coberl: Contrastive bert for reinforcement learning. arXiv preprint arXiv:2107.05431, 2021. Cited pages and 16.
[5] Victor Bapst, Alvaro Sanchez-Gonzalez, Carl Doersch, Kimberly Stachenfeld, Pushmeet Kohli, Peter Battaglia, and Jessica Hamrick. Structured agents for physical construction. In ICML, 2019. Cited on page 3.
[6] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In ICML, 2020. Cited on page 16.
[7] François Chollet. On the measure intelligence. arXiv preprint arXiv:1911.015477 2019.
Cited on page 1.
[8] Karl Cobbe, Oleg Klimov, Chris Hesse, Taehoon Kim, and John Schulman. Quantifying gen- eralization in reinforcement learning. In ICML, 2019. Cited on page
[9] Karl Cobbe, Chris Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation benchmark reinforcement learning. In ICML, 2020. Cited pages 2, 3, and
[10] Karl Cobbe, Jacob Hilton, Oleg Klimov, and John Schulman. Phasic policy gradient. In ICML, 2021. Cited on page 4.
[11] Rémi Coulom. Efficient selectivity and backup operators in Monte-Carke tree search. In Inter- national conference on computers and games, 2006. Cited on page 4.
[12] Kenneth James Williams Craik. The nature of explanation, volume 445. CUP Archive, 1952. Cited on page
[13] Will Dabney, Georg Ostrovski, David Silver, and Rémi Munos. Implicit quantile networks for distributional reinforcement learning. In ICML, 2018. Cited on page 4.
10

---

Published as a conference paper at ICLR 2022
[14] Peter Dayan, Geoffrey E Hinton, Radford M Neal, and Richard S Zemel. The helmholtz machine. Neural computation, (15):89-9044) 1995. Cited on page 1.
[15] Andrea Dittadi, Frederik Träuble, Manuel Wüthrich, Felix Widmaier, Peter Gehler, Ole Winther, Francesco Locatello, Olivier Bachem, Bernhard Schölkopf, and Stefan Bauer. Rep- resentation learning for out-of-distribution generalization in reinforcement learning. arXiv preprint arXivv2107.055666, 2021. Cited
page 1.
[16] Josip Djolonga, Jessica Yung. Michael Tschannen, Rob Romijnders, Lucas Beyer, Alexander Kolesnikov, Joan Puigcerver, Matthias Minderer, Alexander D' Amour, Dan Moldovan, et al. On robustness and transferability of convolutional neural networks. In CVPR, 2021. Cited on page 3.
[17] Yan Duan, John Schulman, Xi Chen, Peter L. Bartlett, Ilya Sutskever, and Pieter Abbeel. RI 2: Fast reinforcement learning via slow reinforcement learning. arXiv preprint arXiv: 2016. Cited on pages 8, 9, and 19.
[18] Chelsea Finn and Sergey Levine. Deep visual foresight for planning robot motion. In ICRA, 2017. Cited on pages and 3.
[19] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adap- tation of deep networks. In ICML, 2017. Cited on pages and 8.
[20] Karol Gregor, Danilo Jimenez Rezende, Frederic Besse, Yan Wu, Hamza Merzic, and Aaron van den Oord. Shaping belief states with generative environment models for rl. In NeurlPS, 2019. Cited on page 3.
[21] Jean-Bastien Grill, Florent Altché, Yunhao Tang, Thomas Hubert, Michal Valko, Ioannis Antonoglou, and Rémi Munos. Monte-carlo tree search as regularized policy optimization. In ICML, 2020. Cited on 3.ge
[22] Christopher Grimm, André Barreto, Satinder Singh, and David Silver. The value equivalence principle for model-based reinforcement learning. In NeurlPS, 2020. Cited on pages 3, 4, 7, and 9.
[23] Arthur Guez, Mehdi Mirza, Karol Gregor, Rishabh Kabra, Sébastien Racanière, Théophane Weber, David Raposo, Adam Santoro, Laurent Orscau, Tom Eccles, et al. An investigation of model-free planning. In JCML, 2019. Cited on pages 3, 4, and 6.
[24] Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, Bilal Piot, Bernardo A Pires, and Rémi Munos. Neural predictive belief representations. arXiv preprint arXiv:1811.06407, 2018. Cited on pages 5 and 16.
[25) Zhaohan Daniel Guo, Bernardo Avila Pires, Bilal Piot, Rémi Munos, and Mohammad Gheshlaghi Azar. Bootstrap preeiteeee representations for multitask reinforcement learning. ICMI,, 2020. Cited on page 3.
Grill, Florent Altché,
[26] David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In NeurlPS, 2018. Cited on pages and
[27] Danijar Hafner, Timothy Lillicrap. Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In /CLR, 2021. Cited on pages I and 3.
[28) Jessica B Hamrick. Analogues of mental simulation and imagination in deep learning. Current Opinion in Behavioral Sciences, 29:8-16 16, 2019. Cited on pages and
[29] Jessica B Hamrick, Abram I. Friesen, Feryal Behbahani, Arthur Guez, Fabio Viola, Sims With- erspoon, Thomas Anthony, Lars Buesing, Petar Veličković, and Théophane Weber. On the role of planning in model-based deep reinforcement learning. In ICLR, 2021. Cited on pages 3, 4, and 9.
[30] Nicolas Heess, Greg Wayne, David Silver, Timothy Lillicrap. Yuval Tassa, and Tom Erez. Learning continuous control policies by stochastic value gradients. arXiv preprint 100.09 142, 2015. Cited page

---

Published as a conference paper at ICLR 2022
[31] Matteo Hessel, Ivo Danihelka, Fabio Viola, Arthur Guez, Simon Schmitt, Laurent Sifre, Theo- phane Weber, David Silver, and Hado van Hasselt. Muesli: Combining improvements in policy optimization. In ICML, 2021. Cited on page 3.
[32] Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Mohammadamin Barekatain, Si- mon Schmitt, and David Silver. Learning and planning in complex action spaces. In ICML, 2021. Cited on pages 4, 5, 10, and 16.
[33] Maximilian Igl. Kamil Ciosek, Yingzhen Li, Sebastian Tschiatschek, Cheng Zhang, Sam De- vlin, and Katja Hofmann. Generalization in reinforcement learning with selective noise injec- tion and information bottleneck. In NeurIPS, 2019. Cited on page
[34] Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model-based policy optimization. In NeurIPS, 2019. Cited on page 3.
[35] Minqi Jiang, Edward Grefenstette, and Tim Rocktäschel. Prioritized level replay. In ICML, 2021. Cited on pages 3, 4, 5, 6, 7, and 19.
[36] Niels Justesen, Ruben Rodriguez Torrado, Philip Bontrager, Ahmed Khalifa, Julian Togelius, and Sebastian Risi. Illuminating generalization in deep reinforcement learning through proce- dural level generation. arXiv preprint arXiv: 10729, 2018. Cited on page 1.
[37] Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Kon- rad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, el. Model- based reinforcement learning for atari. In ICLR, 2010. Cited on page 3.
[38] Thomas Kipf, Elise van der Pol, and Max Welling. Contrastive learning of structured world models. In ICLR, 2020. Cited on page 3.
[39] Louis Kirsch, Sjoerd van Steenkiste, and Jürgen Schmidhuber. Improving generalization in meta reinforcement learning using learned objectives. In ICLR, 2020. Cited page 9.
[40] Levente Koesis and Csaba Szepesvári. Bandit based Monte-Carke planning. In European conference on machine learning, 2006. Cited on page 4.
[41] Brenden M Lake, Tomer D Ullman, Joshua B Tenenbaum, and Samuel J Gershman. Building machines that learn and think like people. Behavioral and brain sciences, 40, 2017. Cited on page 1.
[42] Michael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Reinforcement learning with augmented data. In NeurIPS, 2020. Cited on page 3.
[43] Bogdan Mazoure, Remi Tachet des Combes, Thang Doan, Philip Bachman, and R Devon Hjelm. Deep reinforcement and infomax learning. In NeurIPS, 2020. Cited page 3.
[44] Bogdan Mazoure, Ahmed M Ahmed, Patrick MacAlpine, R Devon Hjelm, and Andrey Kolobov. Cross-trajectory representation learning for zero-shot generalization in RL. arXiv preprint arXiv:2106.021 02193, 2021. Cited on pages I and 3.
[45] Jovana Mitrovic, Brian McWilliams, Jacob Walker, Lars Buesing, and Charles Blundell. Rep- resentation learning via invariant causal mechanisms. In ICLR, 2021. Cited on pages and 16.
[46] Thomas M Moerland, Joost Broekens, and Catholijn M Jonker. Model- based reinforcement learning: A survey, arXiv preprint arXiv:2006. 167 2020. Cited on page 3.
[47] Sharada Mohanty, Jyotish Poonganam, Adrien Gaidon, Andrey Kolobov, Blake Wulfe, Di- pam Chakraborty, Gražvydas Šemetulskis, João Schapke, Jonas Kubilius, Jurgis Pašukonis, et al. Measuring sample efficiency and generalization in reinforcement learning benchmarks: Neurips 2020 procgen benchmark. arXiv preprint arXiv:2 15332, 2021. Cited on pages and 3.
[48] Anusha Nagabandi, Ignasi Clavera, Simin Liu, Ronald S Fearing, Pieter Abbeel, Sergey Levine, and Chelsea Finn. Learning to adapt in dynamic, real-world environments through meta-reinforceement learning. In ICLR, 2019. Cited on pages and 3.
12

---

Published as a conference paper at ICLR 2022
[49] Alex Nichol, Vicki Pfau, Christopher Hesse, Oleg Klimov, and John Schulman. Gotta learn fast: A new benchmark for generalization in RL. arXiv preprint arXiv:1804.03720, 2018. Cited on page 1.
[50]
Hyeonwoo Noh, Seunghoon Hong, and Bohyung Han. Learning deconvolution network for semantic segmentation. In ICCV, 2015. Cited on page 16.
[51] Junhyuk Oh, Satinder Singh, and Honglak Lee. Value prediction network. In NeurlPS, 2017. Cited on page 4.
[52] Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv: 2018. Cited on pages 3, 6, and 16.
[53] Sherjil Ozair, Yazhe Li, Ali Razavi, Joannis Antonoglou, Aäron van den Oord, and Oriol Vinyals. Vector quantized models for planning. arXiv preprint arXiv:2106.04615. 2021. Cited on page 7.
[54] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agar- wal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et Learning transferable visual models from natural language supervision. In ICML, 2021. Cited on page 3.
[55] Roberta Raileanu and Rob Fergus. Decoupling value and policy for generalization in rein- forcement learning. In ICML, 2021. Cited page 4.
[56] Roberta Raileanu, Max Goldstein, Denis Yarats, Ilya Kostrikov, and Rob Fergus. Automatic data augmentation for generalization in deep reinforcement learning. In NeurlPS, 2021. Cited on pages 5 and 19.
[57] Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient off- policy meta-reinforcement learning via probabilistic context variables. In IMCL, 2019. Cited on page 3.
[58] Danilo Rezende, Ivo Danihelka, George Papamakarios, Nan Rosemary Ke, Ray Jiang, Theo- phane Weber, Karol Gregor, Hamza Merzic, Fabio Viola, Jane Wang, et al. Causally correct partial models for reinforcement learning. arXiv preprint urXiv:2002.0288 02836, 2020. Cited on page
[59] Jürgen Schmidhuber. conference on neural networks, pp. 1458-1463, 1991. Cited on page 1.
Curious model-building control systems. In Proc. international joint
[60] Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mas- tering atari, go, chess and shogi by planning with a learned model. Nature, 2020. Cited on pages 2, 2, 4, and 10.
[61] Julian Schrittwieser, Thomas Hubert, Amol Mandhane, Mohammadamin Barekatain, Ioannis Antonoglou, and David Silver. Online and offline reinforcement learning by planning with a learned model. arXiv preprint arXiv:2104.06294, 2021. Cited on pages 2, 3, 4, 10, and 15.
[62] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.063477 2017. Cited on pages and 19.
[63] Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip Bachman. Data-efficient reinforcement learning with self-predictive representations. In ICLR, 2021. Cited on pages 3, 5, and 6.
[64] Ramanan Sekar, Oleh Rybkin, Kostas Daniilidis, Pieter Abbeel, Danijar Hafner, and Deepak Pathak. Planning explore via self-supervised world models. In ICMI., 2020. Cited on pages and 9.
[65] David Silver, Thomas Hubert, Julian Schrittwiesrr, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Mare Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, et al. A general rein- forcement learning algorithm that masters chess, shogi, and go through self-play. Science, 362 140-9):140-1144 2018. Cited page 4.
13

---

Published as a conference paper at ICLR 2022
[66] Richard S Sutton. Dyna, integrated architecture for learning, planning, and reacting. ACM Sigart Bulletin, 2(4): 163, 1991. Cited on pages and
[67] Josh Tobin, Rachel Fong, Alex Ray, Jonas Schneider, Wojeiech Zaremba, and Pieter Abbeel. Domain randomization for transferring deep neural networks from simulation to the real world. In IROS, 2017. Cited on pages and 4.
[68] Jane X Wang, Zeb Kurth-Nelson, Dhruva Tirumala, Hubert Soyer, Joel Z. Leibo, Remi Munos, Charles Blundell, Dharshan Kumaran, and Matt Botvinick. Learning to reinforcement learn. arXiv preprint arXi:: 2016. Cited page 9.
[69] Tingwu Wang. Xuchan Bao, Ignasi Clavera, Jerrick Hoang, Yeming Wen, Eric Langlois, Shun- shi Zhang, Guodong Zhang, Pieter Abbeel, and Jimmy Ba. Benchmarking model-based rein- forcement learning. arXiv preprint arXiv:
2019. Cited page
[70] Théophane Weber, Sébastien Racanière, David P Reichert, Lars Buesing, Arthur Guez, Danilo Jimenez Rezende, Adria Puigdomenech Badia, Oriol Vinyals, Nicolas Heess, Yujia Li, et al. imaiination-augmente agents for deep reinforcement learning. In NeurlPS, 2017.
Cited on pages and 3.
[71] Sam Witty, Jun Ki Lee, Emma Tosch, Akanksha Atrey, Michael Littman, and David Jensen. Measuring and characterizing generalization in deep reinforcement learning. arXiv preprint arXiv:///2.00886, 2018. Cited page 1.
[72] Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Reinforcement learning with prototypical representations. In ICML, 2021. Cited on page 3.
[73] Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn. Gradient surgery for multi-task learning. NeurIPS, 2020. Cited on pages I, 8, and
[74] Tianhe Yu, Deirdre Quillen, Zhampeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforce- ment learning. CoRL, 2020. Cited on pages 1, 2, 5, and
[75] Vinicius Zambaldi, David Raposo, Adam Santoro, Victor Bapst, Yujia Li, Igor Babuschkin, Karl Tuyls, David Reichert, Timothy Lillicrap, Edward Lockhart, al. Relational deep rein- forcement learning. arXiv preprint 1rXiv:1806.0/888 2018. Cited on page 3.
[76] Amy Zhang, Sainbayar Sukhbaatar, Adam Lerer, Arthur Szlam, and Rob Fergus. Composable planning with attributes. In ICML, 2018. Cited on page 1.
[77] Amy Zhang, Yuxin Wu, and Joelle Pincau. Natural environment benchmarks for reinforcement learning. arXiv preprint arXiv: /81 1.00032, 2018. Cited on page 1.
[78] Jin Zhang, Jianhao Wang. Hao Hu, Tong Chen, Yingfeng Chen, Changjie Fan, and Chongjie Zhang. Metacure: Meta reinforcement learning with empowerment-drine exploration. In ICML, 2021. Cited on page I.
14

---

Published as a conference paper at ICLR 2022
A
AGENT DETAILS
A.1 MUZERO
Network Architecture: We follow the same network architectures follow the ones used in MuZero ReAnalyse [61]. For pixel based inputs, the images are first sent through a convolutional stack that downsamples the
96 image down to an tensor (for Meta- Meta-World World). This tensor then serves as input to the encoder. Both the encoder and the dynamics model were implemented by ResNet with 10 blocks, each block containing 2 layers. For pixel-based inputs (Procgen and ML-45) each layer of the residual stack was convolutional with kernel size of 3x3 and 256 planes. For state based inputs (ML-1) each layer was fully-connected with hidden size of 512.
tensor (for Procgen)
16 x 16
Hyperparameters
We list major parametrsssssess used this work Table A.1.
Table A.1: Styper-paaameten.
HYPER-PARAMETER
VALUE (PROCGEN)
VALUE (META-WORLM))
TRAINING
Model Unroll Length
TD-Steps
5
ReAnalyse Fraction
0.945
0.95
Replay Size (in sequences)
50000
2000
MCTS
Number of Simulations
50
50
UCB-constant
1.25
1.25
Number of Samples
n/a
20
SELF-SUPERVISION
Reconstruction Loss Weight
1.0
1.0
Contrastive Loss Weight
1.0
0.1
SPR Loss Weight
10.0
OPTIMIZATION
Optimizer Initial Learning Rate
Adam
Adam 10-4
10-4
Batch Size
1024
1024
Additional Implementation Details: For Meta- World experiments, we provide the agent with past reward history well. We found this to be particularly helpful when training Meta-Worke since implicit task inference becomes easier. In both Procgen and Meta- Meta-World, the agent is given history of the last observations. The images are concatenated by channel and then input as one tensor to the encoder. For each game, we trained our model using 2 TPUV3-8 machines. A separate actor gathered environment trajectories with TPUv3-8 machine.
A.2 CONTROLLED MODEL-FREE BASELINE
Our controlled Q-Learning baseline begins with the same setup the MuZero agent (Section A.1) and modifies it few key ways to make it model-free rather than model-based.
The Q-Learning baseline uses n-step targets for action value function. Given trajectory (st, 02, rt r= the target action value is computed as follow
R-1
Qtarget a.)
Σ
r're4s + + max (e(+)t): a)
aEA
Where Q3 is the target network whose parameter & is updated every 100 training steps.
15

---

Published as a conference paper at ICLR 2022
In order to make the model architecture most similar to what is used in the MuZero agent. decompose the action value function into two parts: a reward prediction and a value prediction V. and model these two parts separately. The total loss function is, therefore, Liotal Lreward + Leatue- The reward loss is exactly the same as that of MuZero. For the value loss, we can decompose Equation in the same way:
Qtarget (4) = Pr + (sa
n-1
Y Tott + 7 max (fetn + (fetn aEA
77Ve(s))
Viarget (5)
rots + + n max (eem + +Ve(ss))
(2)
Since the reward prediction should be taken care of by Lreward and it usually converges fast, assume ir = Tr and the target is simplified to Equation 2. We can then use this value target to compute the value loss Lvaiue CE(Viargee (s), Vs)).
In Meta-World, since it has a continuous action space, maximizing over the entire action space is infeasible. We follow the Sampled Muzero approach [32] and maximize only over the sampled actions.
in The Q-Larniing baseline uses 5-step targets for computing action values. This is the MuZero, but unlike MuZero, the Q-Learning baseline only trains the dynamics for a single step. However, we also provide results for 5-step dynamics function which we call ML+Model and QL+Model+Recon which further adds an auxiliary for reconstruction.
A.3 MUZERO RECONSTRUCTION
The decoder architecture mirrors the ResNet encoder with 10 blocks, but with each convolutional layer replaced by de-convolutional layer [50]. In addition the regular mean reconstruction loss, we add an additional max-loss term which computes the maximum reconstruction loss across all pixels. This incentives the reconstruction objective to not neglect the hardest reconstruct pixels, and in turn makes the reconstructions sharper. The reconstruction loss can be used as follows:
1 X
((Xis,, (osss.Yss) Yar) 1.5.J
((X,Y)
+ max ((xir, Yij)), 1.3
(z 2² is the element-wise squared distance between each feature (i.e. RGB) value in the image array, where i and j are the pixel indices, where is the feature index, and where [X] is the size of the array.
where f(x, 3)
A.4 MUZERO + CONTRASTIVE
Our implementation the contrastive loss does not rely on any image augmentations but focuses on the task of predicting future embeddings of the image encoder [4, 52]. Howerver, in this case, the dynamics function produces a predicted vector H. at each step n in the unroll, making the contrastive task action-conditional similar to APtiontion [24]. We utilize the same target network used in train temporal difference in MuZero to compute our target vectors. We then attempt to match 2,, with the corresponding target vector Yn directly computed from the observation at that timestep. We compute the target vector yn with the same target network weights used to compute the temporal difference loss in MuZero. Since the contrastive loss is ultimately classification problem, also need negative examples. For each pair I,, and yn. the sets {a;] {m:} and {yi J serve as the negative examples with the set of randomly sampled examples in the minibatch. We randomly sample 10% of the examples to constitute N. As in previous work [4, 6, 45] we use cosine distance between x; and y; and feed this distance metric through softmax function. We use the loss function in ReLIC [45] and CoBERL [4] and use temperature of O.I.
16

---

Published as a conference paper at ICLR 2022
The 256-dimensional vector 2n is derived from two layer critic function which is first convolu- tional layer with stride and kernel size on top of the state of the dynamics function. This is then fed into fully connected layer to produce Zn. For each step n. the target vector yn. which serves as a positive example, is derived from the final convolutional layer (8) that is fed into the residual stack of the encoder. The vector is computed by inputting the encoder, using the target network weights, with the corresponding image at that future step The encoding is then fed into the critic network compute yn.
The encoder of MuZero uses 15 past images as history. However, when we compute the target vectors, our treatment of the encoding history is different from that of the agent; instead of stacking all of the historical images (n 15... 1) up to the corresponding step we simply replace the history stack with 15 copies of the image at the current step n. We found that this technique enhanced the performance of contrastive learning.
A.5 MUZERO SPR
The implementation SPR is similar to the contrastive objective except that it does not use negative examples. Thus it only uses nearly identical to what is described in Section A.4. The same layers mentioned compute predicted and target vectors In and Yn respectively. However, SPR relies on an additional projection layer for learning. We thus add a projection layer of 256 dimensions on top of the critic function mentioned in Section A.4. Instead of computing a loss between 2n and yn, we follow the original formulation of SPR and compute mean squared error between 12-normalized and L2- normalizzed Ya.
squared error instead of InfoNCE. The architecture for SPR is
17

---

Published as a conference paper at ICLR 2022
B ADDITIONAL RESULTS
B.1 COMPARISON OF Q-LEARNING AND MUZERO ON PROCGEN
As discussed in the main text, we trained a 5-step value-equivalent model on top of the Q-Learning agent. The performance of this agent with and without supervision, as well as the baseline Q- Learning agent, reported in Table A.2 and visualized in Figure A.2. Both incorporating the 5-step model and adding reconstruction loss helps the Q-Learning agent's 's performance.
When using model-based planning, it's important to ensure that the model "correct". Otherwise, planning on an "incorrect" model could hurts the agent performance, as we can see for the last few games in Figure A.2. Specifically, we can see that on these games, MuZero performs far worse than Q-Learning; yet, its performance is dramatically improved by the addition of the reconstruction loss. This suggests that without self- self-supervision, MuZero's 's model is poor, resulting in compounding model errors during planning. Figure 4 provides a qualitative analysis of MuZero's 's model in two of these games, confirming this hypothesis.
MuZero
Q-Learning
Game
Mode
MZ
MZ4Contr
MZ+ Recon
MZ+SPR
QL
QLeModel
MLe Recon
biglish
train
0.77
0.73
0.80
0.79
0.60
0.53
0.57
test
0.58
bosslight
train
lest
0.91
0.92
0.91
0.93
0.92
0.94
0.95
0.91
0.51
0.57
0.65
0.61
0.71
0.79
caveflyer
tmin
0.91
0,90
0.87
0.77
0.67
0.65
0.82
0:83
0.64
0.88
chaser
test
tmin
0.24
0.89
0.89
0.24
0.89
0.16
0.17
0.68
0.54
0.65
climber
train
0.42
0.72
0,96
0.87
0.76
0.92
0.83
0.91
0.27
0.42
0.58
0.10
0.6
0.81
trin
0.85
coiorun
test
0.64
0.68
0.58
066
dodgeball
test
train
0.77
0.78
0.29
0.30
0.74
0.75
0.41
0.42
0.58
0.54
0.60
0.64
0.59
0.62
fruitbot
train
0.33
0.41
0.55
0,52
0.69
0.56
0.61
0.65
heist
test
triin
-0:17
0.13
40.00
0.17
-0.18
0.18
-0166
0.13
-0.17
0.14
-017
0.14
0.5
0,55
tmin
0.75
0.89
0.87
0.90
0.54
0.61
0.82
jumper
0.66
laaer
test
rain
0.98
0.86
0.98
0.90
0.88
0.97
0.97
0.90
0.70
0.80
0.75
0.85
0.7
0.8
maze
train
-0.25
0.58
0.36
0..9
00.19
0.54
0331
0.44
0.26
00.11
0.39
0.30
-0.39
0.39
miner
lest
train
0.59
0.96
0.78
0.85
0.9
0.7
0.23
0.15
0.56
0.11
ninja
rest
trinn
01.84
0.98
0.88
0.7
0.9%
01883
0.98
0.85
0.70
0.52
0.77
0.61
0.78
083
plunder
train
0.38
0.97
0.31
0.32
0.41
0.50
test
0.94
0.04
0.40
starpilot
train
1.19
1.16
1,23
0.93
0.65
0.76
0.77
Lest
0.98
0.95
0.91
0.84
0.66
0.77
0.50
Average
train
0.72 (0.70/0.75)
0.50 (0.48/0.55)
0.78 (0. 70.77/00.77 79)
0.63 (0.61.0.55
0.75 (0.74H0.88)
0.57 (0.54/0.75)
0.79
0.64 (0.57/0.64)
(0.71.0.79)
01551 (0.49/0.51)
0.36 (0.35/0.37)
0.60
0.45 (0.4440.47)
(0.5940.60) 60)
0.66 (0.58X0.60 (0)6666
0.54 (0.45X0.54)
Table A.2: Final normalised training and test scores on Procgen when training on 500 levels for 30M environment frames, reporting the median across seeds (for MZ, MZ+Recon across 5 seeds). Min/max over seeds are shown in parentheses. Final scores for each seed were computed as means over data points from the last 1M frames. A subset of the experiments on Maze have terminated before 30M frames therefore we computed final scores at 27M for all agents on this game.
18

---

Published as a conference paper at ICLR 2022
Training levels
Mode
M7.
M7+Conir
MZ+Recon
M7+SPR
train
0.73
0.76
0.76
0.68
10
test
0.07
0.13
0.12
0.09
train
0.72
0.76
0.82
0.79
100
test
0.29
0.46
0.47
0.40
0.72
0.75
0.79
500
0.50
0.63
0.57
0.64
train
0.71
0.82
0.93
0.85
All
test
0.68
0.80
0,93
0.83
Table A.3: MuZero with self-supervised losses on Proegen. Final mean normalised training and test scores for different number of training levels. Reporting median values over seed on 10, 100 levels; on 500 levels seeds for MZ+Contr, MZ+SPR and 5 seeds for MZ, MZ+Recon; on infinite levels seed for MZ+Contr, MZ+SPR and seeds for MZ, MZ+Recon. Final scores for each seed were computed as means over data points from the last 1M frames.
Mode
PPO (200M)
PLR (200M)
UCB-DrAC+PLR (200M)
train
0.41
0.61
test
0.220
0.485
Table A.4: Mean-normalized scores on ProcGen for PPO [62], PLR [35] and UCB-DraC+PLR ([56, 35) used in Figure and Figure A.4. For all scores in this table, we the experimental data from Jiang et al. [35].
reach
push
pick place
Agent
Test
Test
Test
RL
100%
96.5%
98.5%
MuZero
100%
100%
100%
Q-Learning 97.5%
99.8%
99.9%
Table A.5: Zero-shot test performance on the ML-1 procedural generalization benchmark Meta- World. Shown baseline results RL3 [17] from the Meta-World paper [73] as well results on MuZero. Note that RL is a meta-learning method and was trained for 300M environment steps, while MuZero and Q-Learning do require meta-training were only trained for 50M steps and fine-tune for 10M. We average the success rate of the last 1M environment steps for each seed and report medians across three seeds for MuZero and Q-Learning (RL2 results from [17]).
Median
IQM
Mean
Optimality Gap
QL
MZ
MZ Recon
MZ+SPR
PPO
UCB-DrAC+PLR
0.4
0.6
0.8
0.30
0.45
0.60
0.75
0.30
0.45
0.60
0.45
0.60
0.75
Min-Max Normalized Score
Figure A.I: Additional metrics (proposed in Agarwal et al. [1]) indicating zero-shot test performance of different methods on ProcGen. IQM corresponds the Inter- Optimality Gap refers to the amount by which the algorithm fails to meet a minimum score of 1.0.
Inter-Quratie Quratile Mean among all runs, and
19

---

Published as a conference paper at ICLR 2022
Mo+Me
QL+Model-Recon
train eval
Normalized return
0.8
0.4
0.2
plunder
starpilot
dodgeball
ninja
bossfight
chaser
bigfish
caveflyer
leaper
maze
heist
coinrun
miner
fruitbot
climber
jumper
Figure A.2: Per-game breakdown of scores for MuZero and Q-Learning agents along with their
reconstruction variants. Shaded bars indicate training performance while solid bars indicate zero- shot test performance. We average the normalized return of the last 1M environment steps for each seed and report medians across 5 seeds for MZ and MZ+Recon and 3 seeds for all other agents.
10 training leves
training levels
5:00 training levels
training levels
0.8
Test Test normalised score
0.6
0.4
0
0.2
0.4
0.8
0.88
0.2
0.4
0.8
0.8
0.2
0.4 0.6
0.8
0.2
0.4
0.8
0.8
Training normalised
Figure A.3: Isolating the interaction between self-supervision and data diversity. Each plot shows the generalization performance as function of training performance for different numbers train- ing levels. Also shown is the unity line, where training and testing performance are equivalent (note that with infinite training levels, all lines collapse to the unity line, as in this case there is no differ- ence between train and test). Generally speaking, self-supervision results in stronger generalization than MuZero. Reporting median values over seed on 10, 100 levels; on 500 levels seeds for MZ+Contr, MZ+SPR and 5 seeds for MZ, MZ+Recon; infinite levels 1 seed for MZ+Contr, MZ+SPR and 2 seeds for MZ, MZ+Recon. For visibility, min/max seeds are not shown.
MuZero
Q-Learning
Game
Mode
MZ
MZ+Contr
MZ+Recon
MZ+SPR
QL
train 97.6%
78. 1%
97.8%
82.5%
85 55.2%
ML10
zero-shot test
26.5%
17.0%
25.0%
17.8%
6.8%
fine-tune test
94.1 1%
60.0%
97.5%
51.3%
80.1%
train
77.2%
57.5%
74.9%
54.4%
55.9%
ML45
zero-shot test
17.7%
12.5%
18.5%
16.0%
10.8%
fine-tune test
76.7%
77.3%
81.7%
77.2%
78.1%
Table A.6: Success rates MuZero (with and without self-supervision) and Q-Learning on ML-10 and ML-45. Training results are shown at 50M frames and fine-tuning results are shown at 10M frames. Reporting average of the last 1M frames and then medians across seeds (2 for MZ and Recn on MI.
20

---

Published as a conference paper at ICLR 2022
20-
PLR (200M) PPO (200M) UCS-DDAC+ (200M)
MZ4Contr
10
Aimper
Moan Episade Return
clmbbr
cottiun
dodgebol
Trubbo
Mightan
bomilight
Cavaltyer
30-
10-
2.5
20M
10M
2IM
35M
3016
10M Environment frames
Environment fromes
(a) Individual scores across all 16 Procgen games
(b) Mean normalized Score
Figure A.4: Training performance on Procgen. See Figure 2 in the main text for results on the test Reporting median values over seeds (5 seeds for MZ), with shaded regions indicating min/max seeds.
Mean Noraaizzd Score
0.75-
0.50
NZ
Mean Nor Norrmalize
Rocon
MZ+5PR
0.55
QLi Model
0.5-
N2nCo
MZ+Becon
10M Environment frames
3014
Environment frames
Figure A.5: Training performance Procgen. See Figure 3 in the main text for corresponding results on the test set. Left: the effect of planning. Right: the effect of self-supervision. Reporting median performance over 3 seeds (5 seeds for MZ and MZ+Recon), with shaded regions indicating min/max seeds.
MZ 5PR
10 training lnvels
100 training levels
500 traning levels
Al training levels
Mean Normalizedd Score
10M
20M
30M
10M
Environment frames
Figure A.6: Interaction of self-supervision and data diversity on training performance. See Figure in the main text for corresponding results on the test set. Reporting median values over 1 seed on 10, 100 levels: on 500 levels seeds for MZ+Contr, MZ+SPR and 5 seeds for MZ, MZ+Recon; on infinite levels seed for MZ+Contr, MZ+SPR and 2 seeds for MZ, MZ+Recon. Shaded regions indicate min/max seeds.
21

---

Published as a conference paper at ICLR 2022
push-v2
pick-place
reach-v2
1.0
1.0
0.5
0.a
0.6
0.4
MuZeer
0.0
0.0
167
167
167
goals the MI-1 procedural generalization benchmark of Meta-World. While both MuZero and Q-L_earnin achieve optimal performance, MuZero is stable and bit data-efficient than Q-Learning. Shown are medians across three seeds (for visibility, min/max seeds are not shown).
Figure A.7: Zero-shot success rates MuZero and Q-Learning
ML10-Train Tasks
ML10-Zero. shot Test Tasks
ML1 ML10-Finetune Test Tasks
QL (from scratch)
1.5
QL
Success rate
scratch)
MZ
MZ Recon
MZ+Contr
MZ+SPR
5.2
5.0
16M
20M
32M
1/M
I/M
201
30M
siM
T/M
invironment frames
Environment frames
Environment frames
MLAS-Finetune
QL (from scratch)
Succrss rate
MZ (from scratch)
MZ
4.4
MZ+Contr
MZ+SPA
209
Environment frames
Environ frames
Environment
Figure A.8: Train, zero-shot test, and fine-tuning performance ML-10 and ML-45. The top row shows ML-10, while the bottom row shows ML-45. The left column shows training performance, the middle column shows zero-shot test performance, and the right column shows fine-tuning perfor- mance. Reporting medians across seeds seeds for MZ and MZ+Recon on ML10), with shaded regions indicating min/max seeds. See also Figure 6 for plots showing cumulative regret during fine-tuning.
22


Reviews:
[{"review_id": "Rb-_MQYKlHe", "reviewer": "Reviewer_Nzu6", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper systematically investigates how MuZero, a SOTA model-based agent, performs on procedural generalization (in Procgen) and task generalization (in Meta-World). They identify three factors that improve procedural generalization through careful experiments. Next, they test their effect for task generalization and find only a weak positive transfer to unseen tasks. However, the experiments and results testing the effect of planning and data diversity for task generalization were hard to follow.", "summary_of_the_paper": "This paper measures the generalization ability of model-based agents, i.e. MuZero, in comparison to model-free agents. The authors identify three key factors for procedural generalization: planning, self-supervised representation learning, and procedural data diversity. However, they find that these factors do not necessarily provide the same benefit for task generalization. They argue for a move towards self-supervised model-based agents trained in rich, procedural, multi-task environments. ", "summary_of_the_review": "The authors provide a thorough evaluation of MuZero, a model-based agent, on procedural generalization in Procgen and task generalization in Meta-World. They identify three factors that improve procedural generation and demonstrate challenges in task generalization. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "PlSOYj_VnFW", "reviewer": "Reviewer_mZzY", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros:\n- very clear and well-written paper\n- very strong results on ProcGen\n- informative ablations and insights\n- novel application of model-based methods to procedurally generated environments\n\nCons:\n- limited novelty of methods \n- the authors do not mention code release in their reproducibility statement\n\nIn section 4.2, is there a reason why MuZero is only trained for 50M while the baselines are trained for 400M steps? It would be helpful to see how MuZero performs when it is run for 400M steps. Or is this too slow? If so, this should be noted somewhere. If it is too slow, then it would be helpful to compare to the baseline performance at 50M steps.", "summary_of_the_paper": "This paper explores the application of the MuZero agent for tasks which require generalization across environments, namely ProcGen and MetaWorld. \nOn ProcGen, they find that MuZero in its standard form performs on par or better than strong model-free methods. Furthermore, when combined with auxiliary self-supervised learning (SSL) losses, there is a significant jump in performance which achieves a new state of the art. The paper includes interesting control experiments disentangling the effects of different components. For example, it shows that both MuZero’s modified targets for the value functions, as well as the tree search for action selection, each separately contribute to performance. Another interesting finding is that adding auxiliary SSL objectives can help generalization performance on unseen environments, even when they do not improve performance on the training environments, which I found surprising but useful. \n\nResults are also reported on task generalization benchmarks from MetaWorld. Here the results are less strong, and self-supervision does not appear to help. There appears to be some transfer between tasks, but it is limited. ", "summary_of_the_review": "Overall this is a strong paper and I recommend acceptance. It investigates the application of a model-based RL (MBRL) to procedurally generated environments, which contrasts with most existing MBRL works which on run on singleton environments. In addition to strong results on ProcGen, the ablations are quite informative in understanding the effect of the different algorithmic components. The MetaWorld results are a bit disappointing, but these are nevertheless helpful to include.\n\nMy main issue with the paper is that the authors do not mention code release in their reproducibility statement. While the detailed appendix is appreciated, this is not a substitute for releasing code (I am also not aware of thoroughly tested open-source implementations of MuZero). I strongly urge the authors to make their code open source, otherwise it will be difficult for the research community to build on this work. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "NvfF-hd1BXJ", "reviewer": "Reviewer_hXT2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths: \n\n- The paper is very well written and the motivation and findings are clearly presented and easy to follow.\n- They provide a good ablation study of the different losses and explored a few interesting cases.\n\nWeaknesses:\n\n- The main weakness of the paper lies in that the scope is somewhat limited to an empirical study with only a small number of take-away learnings for the community, which are not necessarily very surprising. \n- Minor: “Although significant strides have been made in model-based systems in recent years [27], the most popular model-based benchmarks consist of identical training and testing environments [e.g. 26, 68] and do not measure or optimize for for generalization at all.” remove double “for”.\n", "summary_of_the_paper": "The authors present a systematic empirical study of the effect of planning and model learning on generalization performance, using the MuZero agent. They use 2 environments, Procgen and Meta-World, to respectively explore procedural generalization to new variants of the environment with the same reward structure and task generalisation to new structures of the reward function in the same environment. Their main contributions are their empirical results, specifically that additional reconstruction or self-supervised losses enable the MueZero agent to achieve state-of-the-art performance in procedural generalization, but are not enough to promote a similar increase in performance in task generalization tasks in Meta-World. Finally, they also explore the data diversity dimension, showing that having more diverse data during training can help procedural generalization even more.", "summary_of_the_review": "The work is very clearly presented, easy to understand and presents a number of ablation cases on some important considerations for RL agents, e.g. model-free, planning and model learning. The learnings from the work are clear, though not necessarily very surprising. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "2a177pxXfyZ", "reviewer": "Reviewer_KYiU", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths\n\n- While prior works have shown some sort of self-supervised representation learning to be important in learning world models [1], this work seems to be the first to evaluate its effect on generalization. Since multi-task RL and generalization seems to be an active research area, I think this is a valuable contribution.\n- The paper shows that some conclusions as to generalization performance cannot be drawn when using a small training set, which suggests researchers should evaluate their methods on diverse sets of environments to see whether the method affects generalization.\n\nWeaknesses\n\n- The paper talks about procedural generalization and task generalization, where the former corresponds to generalizing to unseen configurations of an environment with the same reward function and the latter corresponds to unseen reward functions in the same environment. Examples of each include ProcGen and ML-1 for procedural generalization and ML-10 and ML-45 for task generalization. To me, it is unclear that ML-10 and ML-45 exactly fit into the category of \"task generalization\", since there are different types of objects (different environments) in the different task instances.\n- ML-10 and ML-45 have significantly fewer training environments than the proc-gen tasks, where the paper claims that the improvement of self-supervised representation learning is less apparent with lower data diversity. Perhaps the results on task-generalization are simply due to this? I am curious as to whether there is truly a difference in generalization performance when measuring generalization to new reward functions rather than new environment configurations. Perhaps there is a better environment to test this in where the data diversity can be controlled for?\n\n[1] Babaeizadeh, Mohammad, et al. Models, Pixels, and Rewards: Evaluating Design Trade-offs in Visual Model-Based Reinforcement Learning. [https://arxiv.org/abs/2012.04603](https://arxiv.org/abs/2012.04603)", "summary_of_the_paper": "This paper evaluates how well model-based RL, specifically MuZero, generalizes in comparison to model-free RL. It empirically compares how planning, representation learning, and data diversity affect the generalization of agents. To evaluate the effect of planning, a Q-learning agent is constructed to be as similar as possible to MuZero without the MCTS. In experiments, it is found that planning, self-supervised representation learning (reconstruction, contrastive, self-predictive) and data diversity all improve generalization performance. However, results are not similar in the Meta World benchmarks, where self-supervision did not seem to improve results much. The paper concludes that self-supervision is a promising approach to improving the generalization of MBRL agents in procedural environments, but perhaps it does not improve task generalization.", "summary_of_the_review": "This paper shows that self-supervised representation learning not only improves the training performance of MBRL but also the generalization performance and also points out that researchers may want to evaluate their algorithms on tasks with high data diversity. However, I feel as though the argument that task generalization differs from procedural generalization is not well supported by the experiments. I do not recommend accepting this paper in its current form.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper evaluates the generalization capabilities of model-based agents, in particular, MuZero, compared with model-free agents. Reviewers agree that the paper is well-written and the topic is interesting. The ablation study is especially interesting, as it disentangles the effect of different algorithmic components. Some concerns are raised about the significance of this work, as the scope is limited to an empirical study and the results are not necessarily very surprising. 

Since the paper presents clear results on an important and relevant topic, I recommend acceptance.
Decision: Accept (Poster)