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