summaries: 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.

summaries: 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.

summaries: This paper explores the application of the MuZero agent for tasks which require generalization across environments, namely ProcGen and MetaWorld.  On 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.   Results 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.

summaries: 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.

abstractText: 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 [64], 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 [9]. However, we find that these factors do not always provide the same benefits for task generalization in Meta-World [81], indicating that transfer remains a challenge. 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.

authors:
Rishabh Agarwal
Max Schwarzer
Pablo Samuel Castro
Aaron Courville
Marc G Bellemare

title: Deep reinforcement learning at the edge of the statistical precipice

venue: NeurIPS,

authors:
Ankesh Anand
Evan Racah
Sherjil Ozair
Yoshua Bengio
Marc-Alexandre Côté
R Devon Hjelm

title: Unsupervised state representation learning in atari

venue: In NeurIPS,

authors:
OpenAI: Marcin Andrychowicz
Bowen Baker
Maciek Chociej
Rafal Jozefowicz
Bob Mc- Grew
Jakub Pachocki
Arthur Petron
Matthias Plappert
Glenn Powell
Alex Ray

title: Learning dexterous in-hand manipulation

venue: The International Journal of Robotics Research,

authors:
Andrea Banino
Adrià Puidomenech Badia
Jacob Walker
Tim Scholtes
Jovana Mitrovic
Charles Blundell

title: Coberl: Contrastive bert for reinforcement learning

venue: arXiv preprint arXiv:2107.05431,

authors:
Victor Bapst
Alvaro Sanchez-Gonzalez
Carl Doersch
Kimberly Stachenfeld
Pushmeet Kohli
Peter Battaglia
Jessica Hamrick

title: Structured agents for physical construction

venue: In ICML,

authors:
Ting Chen
Simon Kornblith
Mohammad Norouzi
Geoffrey Hinton

title: A simple framework for contrastive learning of visual representations

title: On the measure of intelligence

venue: arXiv preprint arXiv:1911.01547,

authors:
Karl Cobbe
Oleg Klimov
Chris Hesse
Taehoon Kim
John Schulman

title: Quantifying generalization in reinforcement learning

authors:
Karl Cobbe
Chris Hesse
Jacob Hilton
John Schulman

title: Leveraging procedural generation to benchmark reinforcement learning

authors:
Karl Cobbe
Jacob Hilton
Oleg Klimov
John Schulman

title: Phasic policy gradient

title: Efficient selectivity and backup operators in Monte-Carlo tree search

venue: In International conference on computers and games,

title: The nature of explanation, volume 445

venue: CUP Archive,

authors:
Will Dabney
Georg Ostrovski
David Silver
Rémi Munos

title: Implicit quantile networks for distributional reinforcement learning

authors:
Peter Dayan
Geoffrey E Hinton
Radford M Neal
Richard S Zemel

title: The helmholtz machine

venue: Neural computation,

authors:
Andrea Dittadi
Frederik Träuble
Manuel Wüthrich
Felix Widmaier
Peter Gehler
Ole Winther
Francesco Locatello
Olivier Bachem
Bernhard Schölkopf
Stefan Bauer

title: Representation learning for out-of-distribution generalization in reinforcement learning

venue: arXiv preprint arXiv:2107.05686,

authors:
Josip Djolonga
Jessica Yung
Michael Tschannen
Rob Romijnders
Lucas Beyer
Alexander Kolesnikov
Joan Puigcerver
Matthias Minderer
Alexander D’Amour
Dan Moldovan

title: On robustness and transferability of convolutional neural networks

venue: In CVPR,

authors:
Yan Duan
John Schulman
Xi Chen
Peter L Bartlett
Ilya Sutskever
Pieter Abbeel

title: Fast reinforcement learning via slow reinforcement learning

venue: arXiv preprint arXiv:1611.02779,

authors:
Chelsea Finn
Sergey Levine

title: Deep visual foresight for planning robot motion

venue: In ICRA,

authors:
Chelsea Finn
Pieter Abbeel
Sergey Levine

title: Model-agnostic meta-learning for fast adaptation of deep networks

authors:
Dibya Ghosh
Jad Rahme
Aviral Kumar
Amy Zhang
Ryan P Adams
Sergey Levine

title: Why generalization in rl is difficult: Epistemic pomdps and implicit partial observability

venue: Advances in Neural Information Processing Systems,

authors:
Karol Gregor
Danilo Jimenez Rezende
Frederic Besse
Yan Wu
Hamza Merzic
Aaron van den Oord

title: Shaping belief states with generative environment models for rl

authors:
Jean-Bastien Grill
Florent Altché
Yunhao Tang
Thomas Hubert
Michal Valko
Ioannis Antonoglou
Rémi Munos

title: Monte-carlo tree search as regularized policy optimization

authors:
Christopher Grimm
André Barreto
Satinder Singh
David Silver

title: The value equivalence principle for model-based reinforcement learning

authors:
Arthur Guez
Mehdi Mirza
Karol Gregor
Rishabh Kabra
Sébastien Racanière
Théophane Weber
David Raposo
Adam Santoro
Laurent Orseau
Tom Eccles

title: An investigation of model-free planning

authors:
Zhaohan Daniel Guo
Mohammad Gheshlaghi Azar
Bilal Piot
Bernardo A Pires
Rémi Munos

title: Neural predictive belief representations

venue: arXiv preprint arXiv:1811.06407,

authors:
Zhaohan Daniel Guo
Bernardo Avila Pires
Bilal Piot
Jean-Bastien Grill
Florent Altché
Rémi Munos
Mohammad Gheshlaghi Azar

title: Bootstrap latent-predictive representations for multitask reinforcement learning

authors:
David Ha
Jürgen Schmidhuber

title: Recurrent world models facilitate policy evolution

authors:
Danijar Hafner
Timothy Lillicrap
Mohammad Norouzi
Jimmy Ba

title: Mastering atari with discrete world models

venue: In ICLR,

title: Analogues of mental simulation and imagination in deep learning

venue: Current Opinion in Behavioral Sciences,

authors:
Jessica B Hamrick
Abram L Friesen
Feryal Behbahani
Arthur Guez
Fabio Viola
Sims Witherspoon
Thomas Anthony
Lars Buesing
Petar Veličković
Théophane Weber

title: On the role of planning in model-based deep reinforcement learning

authors:
Nicolas Heess
Greg Wayne
David Silver
Timothy Lillicrap
Yuval Tassa
Tom Erez

title: Learning continuous control policies by stochastic value gradients

venue: arXiv preprint arXiv:1510.09142,

authors:
Matteo Hessel
Ivo Danihelka
Fabio Viola
Arthur Guez
Simon Schmitt
Laurent Sifre
Theophane Weber
David Silver
Hado van Hasselt

title: Muesli: Combining improvements in policy optimization

authors:
Thomas Hubert
Julian Schrittwieser
Ioannis Antonoglou
Mohammadamin Barekatain
Simon Schmitt
David Silver

title: Learning and planning in complex action spaces

authors:
Maximilian Igl
Kamil Ciosek
Yingzhen Li
Sebastian Tschiatschek
Cheng Zhang
Sam Devlin
Katja Hofmann

title: Generalization in reinforcement learning with selective noise injection and information bottleneck

authors:
Michael Janner
Justin Fu
Marvin Zhang
Sergey Levine

title: When to trust your model: Model-based policy optimization

authors:
Minqi Jiang
Edward Grefenstette
Tim Rocktäschel

title: Prioritized level replay

authors:
Niels Justesen
Ruben Rodriguez Torrado
Philip Bontrager
Ahmed Khalifa
Julian Togelius
Sebastian Risi

title: Illuminating generalization in deep reinforcement learning through procedural level generation

venue: arXiv preprint arXiv:1806.10729,

authors:
Lukasz Kaiser
Mohammad Babaeizadeh
Piotr Milos
Blazej Osinski
Roy H Campbell
Konrad Czechowski
Dumitru Erhan
Chelsea Finn
Piotr Kozakowski
Sergey Levine

title: Modelbased reinforcement learning for atari

authors:
Thomas Kipf
Elise van der Pol
Max Welling

title: Contrastive learning of structured world models

authors:
Louis Kirsch
Sjoerd van Steenkiste
Jürgen Schmidhuber

title: Improving generalization in meta reinforcement learning using learned objectives

authors:
Levente Kocsis
Csaba Szepesvári

title: Bandit based Monte-Carlo planning

venue: In European conference on machine learning,

title: muzero-pytorch. https://github.com/koulanurag/ muzero-pytorch, 2020

venue: Cited on page

authors:
Brenden M Lake
Tomer D Ullman
Joshua B Tenenbaum
Samuel J Gershman

title: Building machines that learn and think like people

venue: Behavioral and brain sciences,

authors:
Michael Laskin
Kimin Lee
Adam Stooke
Lerrel Pinto
Pieter Abbeel
Aravind Srinivas

title: Reinforcement learning with augmented data

authors:
Bogdan Mazoure
Remi Tachet des Combes
Thang Doan
Philip Bachman
R Devon Hjelm

title: Deep reinforcement and infomax learning

authors:
Bogdan Mazoure
Ahmed M Ahmed
Patrick MacAlpine
R Devon Hjelm
Andrey Kolobov

title: Cross-trajectory representation learning for zero-shot generalization in RL

venue: arXiv preprint arXiv:2106.02193,

authors:
Russell Mendonca
Oleh Rybkin
Kostas Daniilidis
Danijar Hafner
Deepak Pathak

title: Discovering and achieving goals via world models

venue: In Thirty-Fifth Conference on Neural Information Processing Systems,

authors:
Jovana Mitrovic
Brian McWilliams
Jacob Walker
Lars Buesing
Charles Blundell

title: Representation learning via invariant causal mechanisms

authors:
Thomas M Moerland
Joost Broekens
Catholijn M Jonker

title: Model-based reinforcement learning: A survey

venue: arXiv preprint arXiv:2006.16712,

authors:
Sharada Mohanty
Jyotish Poonganam
Adrien Gaidon
Andrey Kolobov
Blake Wulfe
Dipam Chakraborty
Gražvydas Šemetulskis
João Schapke
Jonas Kubilius
Jurgis Pašukonis

title: Measuring sample efficiency and generalization in reinforcement learning benchmarks: Neurips 2020 procgen benchmark

venue: arXiv preprint arXiv:2103.15332,

authors:
Anusha Nagabandi
Ignasi Clavera
Simin Liu
Ronald S Fearing
Pieter Abbeel
Sergey Levine
Chelsea Finn

title: Learning to adapt in dynamic, real-world environments through meta-reinforcement learning

authors:
Alex Nichol
Vicki Pfau
Christopher Hesse
Oleg Klimov
John Schulman

title: Gotta learn fast: A new benchmark for generalization in RL

venue: arXiv preprint arXiv:1804.03720,

authors:
Hyeonwoo Noh
Seunghoon Hong
Bohyung Han

title: Learning deconvolution network for semantic segmentation

venue: In ICCV,

authors:
Junhyuk Oh
Satinder Singh
Honglak Lee

title: Value prediction network

authors:
Aaron van den Oord
Yazhe Li
Oriol Vinyals

title: Representation learning with contrastive predictive coding

venue: arXiv preprint arXiv:1807.03748,

authors:
Sherjil Ozair
Yazhe Li
Ali Razavi
Ioannis Antonoglou
Aäron van den Oord
Oriol Vinyals

title: Vector quantized models for planning

venue: arXiv preprint arXiv:2106.04615,

authors:
Alec Radford
Jong Wook Kim
Chris Hallacy
Aditya Ramesh
Gabriel Goh
Sandhini Agarwal
Girish Sastry
Amanda Askell
Pamela Mishkin
Jack Clark

title: Learning transferable visual models from natural language supervision

authors:
Roberta Raileanu
Rob Fergus

title: Decoupling value and policy for generalization in reinforcement learning

authors:
Roberta Raileanu
Max Goldstein
Denis Yarats
Ilya Kostrikov
Rob Fergus

title: Automatic data augmentation for generalization in deep reinforcement learning

authors:
Kate Rakelly
Aurick Zhou
Chelsea Finn
Sergey Levine
Deirdre Quillen

title: Efficient offpolicy meta-reinforcement learning via probabilistic context variables

venue: In IMCL,

authors:
Danilo J Rezende
Ivo Danihelka
George Papamakarios
Nan Rosemary Ke
Ray Jiang
Theophane Weber
Karol Gregor
Hamza Merzic
Fabio Viola
Jane Wang

title: Causally correct partial models for reinforcement learning

venue: arXiv preprint arXiv:2002.02836,

authors:
Sebastian Risi
Julian Togelius

title: Increasing generality in machine learning through procedural content generation

venue: Nature Machine Intelligence,

title: Curious model-building control systems

venue: In Proc. international joint conference on neural networks, pp. 1458–1463,

authors:
Julian Schrittwieser
Thomas Hubert
Amol Mandhane
Mohammadamin Barekatain
Ioannis Antonoglou
David Silver

title: Online and offline reinforcement learning by planning with a learned model

venue: arXiv preprint arXiv:2104.06294,

authors:
John Schulman
Filip Wolski
Prafulla Dhariwal
Alec Radford
Oleg Klimov

title: Proximal policy optimization algorithms

authors:
Max Schwarzer
Ankesh Anand
Rishab Goel
R Devon Hjelm
Aaron Courville
Philip Bachman

title: Data-efficient reinforcement learning with self-predictive representations

authors:
Ramanan Sekar
Oleh Rybkin
Kostas Daniilidis
Pieter Abbeel
Danijar Hafner
Deepak Pathak

title: Planning to explore via self-supervised world models

authors:
David Silver
Thomas Hubert
Julian Schrittwieser
Ioannis Antonoglou
Matthew Lai
Arthur Guez
Marc Lanctot
Laurent Sifre
Dharshan Kumaran
Thore Graepel

title: A general reinforcement learning algorithm that masters chess, shogi, and go through self-play

venue: Science,

title: Dyna, an integrated architecture for learning, planning, and reacting

venue: ACM Sigart Bulletin,

authors:
Team
Adam Stooke
Anuj Mahajan
Catarina Barros
Charlie Deck
Jakob Bauer
Jakub Sygnowski
Maja Trebacz
Max Jaderberg
Michael Mathieu

title: Openended learning leads to generally capable agents

venue: arXiv preprint arXiv:2107.12808,

authors:
Josh Tobin
Rachel Fong
Alex Ray
Jonas Schneider
Wojciech Zaremba
Pieter Abbeel

title: Domain randomization for transferring deep neural networks from simulation to the real world

venue: In IROS,

authors:
Jane X Wang
Zeb Kurth-Nelson
Dhruva Tirumala
Hubert Soyer
Joel Z Leibo
Remi Munos
Charles Blundell
Dharshan Kumaran
Matt Botvinick

title: Learning to reinforcement learn

venue: arXiv preprint arXiv:1611.05763,

authors:
Tingwu Wang
Xuchan Bao
Ignasi Clavera
Jerrick Hoang
Yeming Wen
Eric Langlois
Shunshi Zhang
Guodong Zhang
Pieter Abbeel
Jimmy Ba

title: Benchmarking model-based reinforcement learning

venue: arXiv preprint arXiv:1907.02057,

authors:
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

title: Imagination-augmented agents for deep reinforcement learning

title: Muzero general: Open reimplementation of muzero

authors:
Sam Witty
Jun Ki Lee
Emma Tosch
Akanksha Atrey
Michael Littman
David Jensen

title: Measuring and characterizing generalization in deep reinforcement learning

venue: arXiv preprint arXiv:1812.02868,

authors:
Denis Yarats
Rob Fergus
Alessandro Lazaric
Lerrel Pinto

title: Reinforcement learning with prototypical representations

authors:
Weirui Ye
Shaohuai Liu
Thanard Kurutach
Pieter Abbeel
Yang Gao

title: Mastering atari games with limited data

authors:
Tianhe Yu
Saurabh Kumar
Abhishek Gupta
Sergey Levine
Karol Hausman
Chelsea Finn

title: Gradient surgery for multi-task learning

authors:
Tianhe Yu
Deirdre Quillen
Zhanpeng He
Ryan Julian
Karol Hausman
Chelsea Finn
Sergey Levine

title: Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning

venue: In CoRL,

authors:
Vinicius Zambaldi
David Raposo
Adam Santoro
Victor Bapst
Yujia Li
Igor Babuschkin
Karl Tuyls
David Reichert
Timothy Lillicrap
Edward Lockhart

title: Relational deep reinforcement learning

venue: arXiv preprint arXiv:1806.01830,

authors:
Amy Zhang
Sainbayar Sukhbaatar
Adam Lerer
Arthur Szlam
Rob Fergus

title: Composable planning with attributes

authors:
Amy Zhang
Yuxin Wu
Joelle Pineau

title: Natural environment benchmarks for reinforcement learning

venue: arXiv preprint arXiv:1811.06032,

heading: 1 INTRODUCTION

text:
The ability to generalize to previously unseen situations or tasks using an internal model of the world is a hallmark capability of human general intelligence [12, 43] and is thought by many to be of central importance in machine intelligence as well [14, 27, 63, 70]. Although significant strides have been made in model-based systems in recent years [29], the most popular model-based benchmarks consist of identical training and testing environments [e.g. 28, 74] and do not measure or optimize for for generalization at all. While plenty of other work in model-based RL does measure generalization [e.g. 18, 51, 75, 83], each approach is typically evaluated on a bespoke task, making it difficult to ascertain the state of generalization in model-based 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, 37, 52, 77], and others have proposed new environments such as Procgen [9] and Meta-World [81] as benchmarks focusing on measuring generalization. While popular in the model-free community [e.g. 50, 80, 85], 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 modelfree algorithmic improvements such as self-supervision [15, 46] or procedural data diversity [72, 84] yield the same benefits for generalization in model-based agents.
In this paper, we investigate three factors of generalization in model-based RL: planning, selfsupervised representation learning, and procedural data diversity. We analyze these methods through a variety of modifications and ablations to MuZero Reanalyse [64, 65], a state-of-the-art modelbased 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 rendering, map or terrain changes, or 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 [81].
Our results broadly indicate that selfsupervised, model-based agents hold promise in making progress towards 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 incorporation of self-supervised representation learning; and (3) that with self-supervision, less data diversity is required to achieve good performance. However, (4) these ideas help less for task generalization on the ML45 test set from Meta-World, indicating that task transfer remains a challenge.

heading: 2 MOTIVATION AND BACKGROUND

heading: 2.1 GENERALIZATION IN RL

text:
We are interested in the setting where an agent is trained on a set of ntrain MDPs drawn IID from the same distribution, Mtrain = {Mi}ntraini=1 , whereMi ∼ p(M). The agent is then tested on another set of ntest MDPs drawn IID1 from p, but disjoint from the training set: Mtest = {Mj}ntestj=1 where Mj ∼ p(M) such thatMj 6∈Mtrain ∀Mj . 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 (ntrain) of distinct MDPs seen during training (what we will refer to as data diversity), as well as the breadth or amount of variation in p(M) itself. Intuitively, if the amount of diversity in the training set is small (i.e. low values of ntrain) relative to the breadth of p(M), then this poses a more difficult generalization challenge. As ntrain →∞, 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 distributions over 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 a goal location). While we have found this distinction to be conceptually useful, we do note that the boundaries between procedural and task generalization are not always exact: new tasks may imply new environment conditions (e.g. new objects), and new layouts may imply different rewards (e.g. goal locations).
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, 34, 44] and auxiliary self-supervised objectives [45, 46]. Other approaches including better hyperparameter tuning
1Technically, the different tasks contained in Meta-World (see Section 3) are not quite drawn IID as they were designed by hand. In particular, while the test tasks share the same robot arm and physics as the train tasks, they do introduce new objects. For example, on ML-10, 3/5 test tasks have new objects, and on ML-45 1/5 test tasks have a new object. However, the types of objects are still similar to those seen during training, as is the level of complexity, so we feel this fits approximately into this regime.
[50], curriculum learning strategies [36], and stronger architectural inductive biases [5, 24, 82] 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 it challenging to know which are the most beneficial for model-based RL. Here we focus explicitly on model-based agents, and evaluate three factors in combination: planning, self-supervision, and data diversity.
Task generalization Generalizing to new tasks or reward functions has been a major focus of meta-RL [19, 60], 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 51], 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 [78] or dynamics [68] in an exploratory pre-training phase, and then use them for transfer. Among these, Sekar et al. [68] focus on the model-based setting but require access to the reward function during evaluation. In Section 4.2, we take a similar approach of pre-training representations and using them for task transfer, with the goal of evaluating whether planning and self-supervised representation learning might assist in this.

heading: 2.2 FACTORS OF GENERALIZATION

text:
Planning Model-based RL is an active area of research [29, 49] with the majority of work focusing on gains in data efficiency [27, 28, 35, 38, 65], though it is often motivated by a desire for better zero- or few-shot generalization as well [18, 30, 51, 68, 75]. In particular, model-based techniques may benefit data efficiency and generalization in three distinct ways. First, model learning can act as an auxiliary task and thus aid in learning representations that better capture the structure of the environment and enable faster learning [21]. Second, the learned model can be used to select actions on-the-fly via MPC [18], enabling faster adaptation in the face of novelty. Third, the model can also be used to train a policy or value function by simulating training data [70] or constructing more informative losses [22, 31], again enabling faster learning (possibly entirely in simulation). A recent state-of-the-art agent, MuZero [64], combines all of these techniques: model learning, MPC, simulated training data, and model-based losses.2 We focus our analysis on 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 modelbased RL: all that should be required is for models to learn the task dynamics, not necessarily environment dynamics [23]. 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 [46, 67, 78]. 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 selfpredictive. Reconstruction losses involve directly predicting future observations [e.g. 18, 28, 38, 75]. Contrastive objectives set up a classification task to determine whether a future frame could result from the current observation [4, 39, 55, 79]. Finally, self-predictive losses involve having agents predict their own future latent states [26, 67].
Data diversity Several works have shown that exposing a deep neural network to a diverse training distribution can help its representations better generalize to unseen situations [16, 57, 62]. A dense sampling over a diverse training distribution can also act as a way to sidestep the out-of-distribution generalization problem [57]. 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 learned in simulation to the real world [3, 72]. Guez et al. [24] also showed that in the context of zero-shot generalization, the amount of data diversity can have interesting interactions with other
2In 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 be sufficient, as in Hamrick et al. [30], Hessel et al. [32].
architectural choices such as model size. Here, we take a cue from these findings and explore how important procedural data diversity is in the context of self-supervised, model-based agents.

heading: 2.3 MUZERO

text:
We evaluate generalization with respect to a state-of-the-art model-based agent, MuZero [64]. 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 [65], planning [69], and model-based representation learning [23, 54]. However, while these components contribute to strong performance across a wide range of domains, other work has suggested that MuZero does not necessarily achieve perfect generalization on its own [30]. It therefore serves as a strong baseline but with clear room for improvement.
MuZero learns an implicit (or value-equivalent, see Grimm et al. [23]) world model by simply learning to predict future rewards, values and actions. It then plans with Monte-Carlo Tree Search (MCTS) [41, 11] over this learned model to select actions for the environment. Data collected from the environment, as well as the results of planning, are then further used to improve the learned reward function, value function, and policy. More specifically, MuZero optimizes the following loss at every time-step t, applied to a model that is unrolled 0 . . .K steps into the future: lt(θ) =∑K k=0(l k π + l k v + l k r ) = ∑K k=0(CE(π̂
k, πk) + CE(v̂k, vk) + CE(r̂k, rk)), where π̂k, v̂k and r̂k 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: πk is the improved policy generated by the search tree, vk is an n-step return bootstrapped by a target network, and rk 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. [64].
To enable better sample reuse and improved data efficiency, we use the Reanalyse version of MuZero [65]. 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 [33] that modifies MuZero to plan over sampled actions instead. See Section A.2 for more details on MuZero.

heading: 3 EXPERIMENTAL DESIGN

text: 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 it improve procedural generalization, and to what extent does it improve task generalization?

heading: 3.1 ENVIRONMENTS

text:
Procgen Procgen [9] is 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, 36, 58]. 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, we train an agent for 30M environment frames. We use the same network architecture and hyper-parameters that Schrittwieser et al. [65] used for Atari and perform no environment specific tuning. We report mean normalized scores across all games as in Cobbe et al. [9]; see Section A.1.
Meta-World Meta-World [81] is a suite of 50 different tasks on a robotic SAWYER arm, making it more suitable to test task generalization. Meta-World has three benchmarks focused on generalization: ML-1, ML-10, and ML-45. ML-1 consists of three goal-conditioned tasks where the objective is to generalize to unseen goals during test time on the same task. ML-10 and ML-45 require generalization to five completely new tasks (reward functions) at test time after training on either 10
or 45 tasks, respectively. For Meta-World, we trained Sampled MuZero [33] for 50M environment frames. We measure performance in terms of average episodic success rate across task(s).

heading: 3.2 FACTORS OF GENERALIZATION

text:
Planning To evaluate the contribution of planning (see Section 2.2), we compared the performance of a vanilla MuZero Reanalyse agent with a Q-Learning agent. We designed the Q-Learning agent to be as similar to MuZero as possible: for example, it shares the same codebase, network architecture, and replay strategy. The primary difference is that the Q-Learning agent uses a Q-learning loss instead of the MuZero loss in Section 2.3 to train the agent, and uses -greedy instead of MCTS to act in the environment. See Section A.3 for further details.
Self-supervision We looked at three self-supervised methods as auxiliary losses on top of MuZero: image reconstruction, contrastive learning [25], and self-predictive representations [67]. We do not leverage any domain-specific data-augmentations for these self-supervised methods.
Reconstruction. Our approach for image reconstruction differs slightly from the typical use of mean reconstruction error over all pixels. In this typical setting, the use of averaging implies that the decoder can focus on the easy to 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 max reconstruction error over all pixels. See Section A.4 for details.
Contrastive. We also experiment with a temporal contrastive objective which treats pairs of observations close in time as positive examples and un-correlated timestamps as negative examples [2, 55, 24]. Our implementation of the contrastive objective is action-conditioned and uses the MuZero dynamics model to predict future embeddings. Similar to ReLIC [48], we also use a KL 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 reconstruction loss. See Section A.5 for details.
Self-predictive. Finally, we experiment with self-predictive representations (SPR) [67] which are trained by predicting future representations of the agent itself from a 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.6 for details.
Data diversity To evaluate the contribution of data diversity (see Section 2.1), we ran experiments on Procgen in which we varied the number of levels seen during training (either 10, 100, 500, or∞). In all cases, we evaluated on the infinite test split. Meta-World does not expose a way to modify the amount of data diversity, therefore we did not analyze this factor in our Meta-World experiments.

heading: 4 RESULTS

text: 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 as medians across seeds.

heading: 4.1 PROCEDURAL GENERALIZATION

text:
Overall results on Procgen are shown in Figure 2 and Table A.2. MuZero achieves a meannormalized test score of 0.50, which is slightly higher than the earlier state-of-the-art UCBDrAC+PLR, which was specifically designed for procedural generalization tasks [36]. MuZero is also much more data efficient (30M vs. 200M frames). MuZero also outperforms our Q-Learning baseline which gets a score of 0.36. Performance is further improved with self-supervision, indicating that both planning and self-supervised model learning are important for generalization.
Effect of planning While not reaching the same level of performance of MuZero, the Q-Learning baseline performs quite well and achieves a mean normalized score of 0.36, matching performance of other specialized model-free methods such as PLR [36]. By modifying the Q-Learning baseline to learn a 5-step value equivalent model (similar to MuZero), we find its performance further improves to 0.45, though does not quite catch up to MuZero itself (see Figure 3, left). This suggests that while simply learning a value-equivalent 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.8 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 a 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-predictive representations both substantially improve over MuZero’s normalized score of 0.50 to 0.63 (contrastive) or 0.64 (SPR), which are new state-of-the-art scores on Procgen. The reconstruction loss also provides benefits but of a lesser magnitude, improving MuZero’s performance from 0.50 to 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-supervision, we performed a 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 representations. We trained each agent on 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-grained details such as the position of the characters and enemies. In Chaser, for example, the model augmented with a reconstruction loss is able to to predict the position of the character across multiple time steps, while MuZero’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 value-equivalent models [23], where the idea is that world models do not need to capture full environment dynamics—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 [56] and causality [61].
Effect of data diversity Overall, we find that increased data diversity improves procedural generalization in Procgen, as shown in Figure 5. Specifically, at 10 levels, self-supervision barely improves the test performance on Procgen over MuZero. But as we increase the number of levels to 100, 500, or∞, we observe a substantial improvement when using self-supervision. This is an interesting finding which suggests that methods that improve generalization might show promise only when we evaluate them on environments with a lot of inherent diversity, such as procedurally generated environments. Consequently, evaluating methods solely on single task settings (as is common in RL) might lead us to overlook innovations which might have a big impact on generalization.
Not only do we find that MuZero with self-supervision achieves better final generalization performance than the baselines, we also observe that self-supervision improves generalization performance
even when controlling for training performance. To see this, we visualized generalization performance as a function of training performance (see Figure A.4). 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-supervision generally achieving stronger generalization performance. This suggests that in addition to improving data efficiency, self-supervision leads to more robust representations—a feature that again might be overlooked if not measuring generalization.
4.2 TASK GENERALIZATION MAML RL2 QL MZ MZ+Recon ML-10 train 44.4% 86.9% 85.2% 97.6% 97.8%
zero-shot - - 6.8% 26.5% 25.0% few-shot 31.6% 35.8% - - - finetune - - 80.1% 94.1% 97.5%
ML-45 train 40.7% 70% 55.9% 77.2% 74.9% zero-shot - - 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, fewshot, finetuning) on the ML-10 and ML-45 task generalization benchmarks of Meta-World. Shown are baseline results on MAML [19] and RL2 [17] from the Meta-World paper [80] as well as our results with Q-Learning and MuZero. Note that MAML and RL2 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 we have shown, planning, selfsupervision, and data diversity all play an important role in procedural generalization. We next investigated whether this trend holds for task generalization, too. To test this, we pre-trained the Q-Learning and MuZero agents with and without selfsupervision 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 as in Yu et al. [81].
Effect of planning 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 [81]. These existing approaches are meta-learning methods, and are therefore not directly comparable in terms of 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 baselines). MuZero also outperforms Q-Learning in terms of zero-shot test performance, indicating a positive benefit of planning for task generalization.
We also looked at whether the pre-trained representations or dynamics would yield more dataefficient task transfer by fine-tuning on the test tasks for 10M environment steps. Figure 6 shows the results, measured as cumulative regret (see also Figure A.9 for learning curves). Using pre-trained representations or dynamics does enable both Q-Learning and MuZero to outperform corresponding
agents trained from scratch—evidence for weak positive transfer to unseen tasks for these agents. MuZero also exhibits better data efficiency than Q-Learning, again showing a benefit for planning.
Effect of self-supervision and data diversity While MuZero performs better than Q-Learning both in the zero-shot and fine-tuning regime, reconstruction and other forms of self-supervision do not improve performance and may even decrease it (see Table 1 and Figure A.9). In terms of data diversity, ML-45 involves a wider set of training tasks than ML-10; however, the test tasks for these two benchmarks are different. The level of data diversity is also relatively small compared to Procgen, where we can train on 100s of levels. These issues prevent us from directly measuring an effect of data diversity, leaving open the question as to whether data diversity has similar benefits for task generalization as for procedural generalization.
Summary Our results indicate that while MuZero excels at representing variation across tasks seen during training, there is room for improvement to better transfer this knowledge to unseen reward functions. We hypothesize this is due to a combination of factors. First, the data that its model is trained on is biased towards the training tasks, and may not be sufficient for learning a globallyaccurate world model suitable for planning in different tasks. Incorporating a more exploratory pretraining phase [e.g. 47, 68] might help to alleviate this problem. Second, because the model relies on task-specific gradients, it may over-represent training task-specific features (perhaps suggesting that value-equivalent models [23] may be poorly suited to task generalization). It is possible that increasing data diversity (e.g., letting an agent self-generate a wide number of tasks [71]) would prevent such overfitting. Third, during finetuning, the agent must still discover what the reward function is, even if it already knows the dynamics. We suspect that the exploration required to do this is the primary bottleneck for task transfer, rather than the model representation itself. Metalearning methods like those benchmarked by Yu et al. [81] may implicitly learn better exploration policies [73], suggesting a possible reason for their stronger performance. Finally, MuZero only gets 15 frames of history, and better memory mechanisms [20] might also facilitate better task inference.

heading: 5 CONCLUSION

text: In this paper, we systematically investigate how well modern model-based methods perform at hard generalization problems, and whether self-supervision can improve the generalization performance of such methods. We find that in the case of procedural generalization in Procgen, both model-based learning and self-supervision have additive benefits and result in state-of-the-art performance on test levels with remarkable data efficiency. In the case of task generalization in Meta-World, we find that while a model-based agent does exhibit weak positive transfer to unseen tasks, auxiliary selfsupervision does not provide any additional benefit, suggesting that having access to a good world model is not always sufficient for good generalization [30]. 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, 40, 73] or by leveraging world models in other ways [68]. Overall, we conclude that self-supervised model-based methods are a promising starting point for developing agents that generalize better, particularly when trained in rich, procedural, and multi-task environments.

heading: REPRODUCIBILITY STATEMENT

text: 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, Procgen and MetaWorld, 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 [33, 64, 65] with several open source implementations [42, 76, 79]. We have included an extensive appendix (Appendix A) which describes the architectural details of MuZero, the exact hyperparameters used for all experiments, and a detailed description of the Q-Learning agents and the self-supervised losses. In addition, we plan to release the code for our self-supervised losses along with a minimal MuZero pseudocode in Python.

heading: A AGENT DETAILS

heading: A.1 ENVIRONMENT DETAILS

text:
Procgen Following the recommendation of Agarwal et al. [1], we 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)(max−min) , 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 to report the mean normalized score.
Meta-World Meta-World exposes both state and pixel observations, as well as dense and sparse versions of the reward functions. We use the v2 version of Meta-World, dense rewards, and the corner3 camera angle for pixel observations.

heading: A.2 MUZERO

text:
Network Architecture: We follow the same network architectures follow the ones used in MuZero ReAnalyse [65]. For pixel based inputs, the images are first sent through a convolutional stack that downsamples the 96 × 96 image down to an 8 × 8 tensor (for Procgen) or a 16 × 16 tensor (for Meta-World). This tensor then serves as input to the encoder. Both the encoder and the dynamics model were implemented by a 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 a kernel size of 3x3 and 256 planes. For state based inputs (ML-1) each layer was fully-connected with a hidden size of 512.
Hyperparameters We list major hyper-parameters used in this work in Table A.1.
Additional Implementation Details: For Meta-World experiments, we provide the agent with past reward history as well. We found this to be particularly helpful when training on Meta-World since implicit task inference becomes easier. In both Procgen and Meta-World, the agent is given a history of the 15 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 1 TPUv3-8 machine.

heading: A.3 CONTROLLED MODEL-FREE BASELINE

text:
Our controlled Q-Learning baseline begins with the same setup as the MuZero agent (Section A.2) and modifies it in a 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 a trajectory {st, at, rt}Tt=0, the target action value is computed as follow
Qtarget(st, at) = n−1∑ i=0 γirt+i + γ nmax a∈A Qξ(st+n, a) (1)
Where Qξ is the target network whose parameter ξ is updated every 100 training steps.
In order to make the model architecture most similar to what is used in the MuZero agent, we decompose the action value function into two parts: a reward prediction r̂ and a value prediction V , and model these two parts separately. The total loss function is, therefore, Ltotal = Lreward + Lvalue. The reward loss is exactly the same as that of MuZero. For the value loss, we can decompose Equation 1 in the same way:
Qtarget(st, at) = r̂t + γVtarget(s)
= n−1∑ i=0 γirt+i + γ nmax a∈A ( r̂t+n + γVξ(s ′ ) )
=⇒ Vtarget(s) = n−1∑ i=1 γi−1rt+i + γ n−1 max a∈A ( r̂t+n + γVξ(s ′ ) )
(2)
Since the reward prediction should be taken care of by Lreward and it usually converges fast, we assume r̂t = rt and the target is simplified to Equation 2. We can then use this value target to compute the value loss Lvalue = CE(Vtarget(s), V (s)). In Meta-World, since it has a continuous action space, maximizing over the entire action space is infeasible. We follow the Sampled Muzero approach [33] and maximize only over the sampled actions.
The Q-Learning baseline uses 5-step targets for computing action values. This is the same as in MuZero, but unlike MuZero, the Q-Learning baseline only trains the dynamics for a single step. However, we also provide results for a 5-step dynamics function which we call QL+Model and QL+Model+Recon which further adds an auxiliary for reconstruction.

heading: A.4 MUZERO + RECONSTRUCTION

text:
The decoder architecture mirrors the ResNet encoder with 10 blocks, but with each convolutional layer replaced by a de-convolutional layer [53]. In addition to 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 to reconstruct pixels, and in turn makes the reconstructions sharper. The reconstruction loss can be used as follows:
L(X,Y ) =  1 |X| ∑ i,j,f `(Xijf , Yijf ) +max i,j,f `(Xijf , Yijf ),
where `(x, y) = (x − y)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 f is the feature index, and where |X| is the size of the array.

heading: A.5 MUZERO + CONTRASTIVE

text:
Our implementation of the contrastive loss does not rely on any image augmentations but focuses on the task of predicting future embeddings of the image encoder [4, 55]. Howerver, in this case, the
dynamics function produces a predicted vector xn at each step n in the unroll, making the contrastive task action-conditional similar to CPC|Action [25]. We utilize the same target network used in train temporal difference in MuZero to compute our target vectors. We then attempt to match xn 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 a classification problem, we also need negative examples. For each pair xn and yn, the sets {xi}Ni=0,i6=n and {yi}Ni=0,i6=n serve as the negative examples with N the set of randomly sampled examples in the minibatch. We randomly sample 10% of the examples to constituteN . As in previous work [4, 6, 48] we use a cosine distance between xi and yi and feed this distance metric through a softmax function. We use the loss function in ReLIC [48] and CoBERL [4] and use a temperature of 0.1.
The 256-dimensional vector xn is derived from a two layer critic function which is first a convolutional layer with stride 1 and kernel size 3 on top of the state of the dynamics function. This is then fed into a fully connected layer to produce xn. 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 n. The encoding is then fed into the critic network to 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...n − 1) up to the corresponding step n, 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.

heading: A.6 MUZERO + SPR

text: The implementation of SPR is similar to the contrastive objective except that it does not use negative examples. Thus it only uses a mean squared error instead of InfoNCE. The architecture for SPR is nearly identical to what is described in Section A.5. The same layers mentioned compute predicted and target vectors xn and yn respectively. However, SPR relies on an additional projection layer for learning. We thus add a projection layer pn of 256 dimensions on top of the critic function mentioned in Section A.5. Instead of computing a loss between xn and yn, we follow the original formulation of SPR and compute mean squared error between L2-normalized pn and L2-normalized yn.

heading: B ADDITIONAL RESULTS

heading: B.1 COMPARISON OF Q-LEARNING AND MUZERO ON PROCGEN

text:
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 QLearning agent, is reported in Table A.2 and visualized in Figure A.3. Both incorporating the 5-step model and adding reconstruction loss helps the Q-Learning agent’s performance.
When using model-based planning, it’s important to ensure that the model is “correct”. Otherwise, planning on an “incorrect” model could hurts the agent performance, as we can see for the last few games in Figure A.3. 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-supervision, MuZero’s model is poor, resulting in compounding model errors during planning. Figure 4 provides a qualitative analysis of MuZero’s model in two of these games, confirming this hypothesis.

title: PROCEDURAL GENERALIZATION BY PLANNING WITH SELF-SUPERVISED WORLD MODELS