Procedural generalization by planning with self-supervised world models

1 INTRODUCTION . 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 . 2 MOTIVATION AND BACKGROUND . 2.1 GENERALIZATION IN RL . 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 . 2.2 FACTORS OF GENERALIZATION . 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 .

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