Understanding the Generalization Gap in Visual Reinforcement Learning

1 INTRODUCTION . Deep Reinforcement Learning ( RL ) has achieved tremendous success on several video game suites ( Mnih et al. , 2013 ; Vinyals et al. , 2019 ; Berner et al. , 2019 ) . However , state-of-the-art ( SOTA ) RL agents are usually trained and tested on the same game . It turns out that the learned policy fails to transfer to different games or even unseen levels of the same game ( Cobbe et al. , 2018 ; 2019 ) . In contrast , humans are remarkable at transferring to new tasks and environments . Behavior studies ( Dubey et al. , 2018 ) show that humans make extensive use of visual priors such as object permanence , sub-goals , and intuitive physics when learning new games . Without these priors , data efficiency of humans drops an order of magnitude ( Dubey et al. , 2018 ) . ( Sax et al. , 2018 ) learns visual priors with supervised learning on datasets with manual annotations of several semantically meaningful features . In practice however , acquiring such datasets is challenging . In the absence of labelled datasets and human supervision , past works have explored three main ways to incorporate visual priors to improve generalization . The first paradigm is data augmentation ( Michael et al. , 2020 ; Kostrikov et al. , 2020 ) where invariance in the learned representations is induced by training the agent on a large set of task-irrelevant variations of the observed data . Since the distinction between what is relevant vs irrelevant is task-dependent , the optimal choice of data augmentation is also task dependent . To avoid manual selection of data augmentations , UCB-DrAC ( Raileanu et al. , 2020 ) used upper confidence bound ( UCB ) algorithm ( Auer , 2002 ) for choosing optimal augmentations . However , it remains unclear whether choosing one out of many task-agnostic data augmentations , such as cropping , helps in closing the generalization gap . In this paper , we show that set of transformations used in state-of-the-art methods are insufficient at closing the generalization gap . On the other hand , we also show that task-informed data augmentations close the generalization gap , but require task variation information during training . How to construct task sensitive data augmentations remains an open question and an avenue of future research . We discuss this finding in Section 4 . The second paradigm for transfer learning is to leverage domain confusion ( Tzeng et al. , 2014 ; 2015 ; Hoffman et al. , 2013 ) , technique commonly used in computer vision to make the source data distribution indistinguishable from the target distribution . The main idea is to discourage learning of spurious domain/level-specific features by enforcing the constraint that it should not be possible to predict the identity of the domain/level from the policy features . iDAAC ( Raileanu & Fergus , 2021 ) used this technique to improve policy generalization . However , our investigation reveals that even policies that generalize well contain domain/level specific information . In hindsight , this is not surprising . On image classification tasks , the identity of the object doesnot change based on factors such as background clutter , size , etc . However , small changes in the size of obstacles or platforms can drastically change the optimal policy . Therefore , the policy must encode information about level-specific game layout , making application of domain invariance to policy learning challenging . Details of this investigation are provided in Section 5 . The third paradigm of improving generalization involves using auxiliary prediction tasks ( such as depth prediction , image reconstruction , contrastive learning ) to regularize the policy features to prevent learning spurious features ( Jaderberg et al. , 2016 ; Srinivas et al. , 2020 ) . In this work , we choose the auxiliary task of inverse model prediction ( Agrawal et al. , 2016 ) due to its simplicity . This improves policy generalization . We detail this finding in Section 6 . It is well-known that choice of hyper parameters has a substantial impact on the performance of RL agents . When it comes to policy generalization , discussions around hyperparameter selection have been limited . We show that state-of-the-art ( SOTA ) methods that improve policy generalization require careful hyperparameter selection for each task . When these methods are constrained to choose one set of hyperparameters across tasks , they donot outperform the base PPO algorithm . We detail this finding in Section 7 . Finally , we probe if deep RL algorithms learn spurious features that deters performance of the same task in new environments . Our experiments reveal that visual features learned by training on a limited set of levels of one game donot require adaption to achieve good performance on new levels . Infact , only the layers of the policy network that convert visual features into actions require finetuning . Details of this investigation are provided in Section 8 . In summary , our experiments on Procgen ( Cobbe et al. , 2019 ) , a suite of procedurally generated games reveal that ( i ) task-informed data augmentation closes the generalization gap ; ( ii ) learning levelinvariant ( or domain-invariant ) features is not necessary for good generalization ; ( iii ) regularizing the policy features using simple auxiliary tasks can improve generalization ; ( iv ) SOTA methods that improve policy generalization rely on careful hyperparameter selection ; ( v ) training on a limited number of levels of the same game doesnot result in learning of level-specific spurious features . When operating on new levels , good performance can be achieved by only fine-tuning two layers that transform these features into actions . 2 RELATED WORK . Prior works ( Cobbe et al. , 2018 ; 2019 ; Justesen et al. , 2018 ; Zhang et al. , 2018a ; b ; Juliani et al. , 2019 ; Rajeswaran et al. , 2017 ; Raileanu & Rocktäschel , 2020 ; Grigsby & Qi , 2020 ; Kuttler et al. , 2020 ; Farebrother et al. , 2018 ) have established the problem of overfitting in reinforcement learning . In response , multiple methods have been proposed to mitigate overfitting . ( Cobbe et al. , 2018 ; 2019 ) showed that classical techniques such as dropout ( Srivastava et al. , 2014 ) , L2 regularization , and batch normalization ( Ioffe & Szegedy , 2015 ) – originally developed for supervised learning – reduce the generalization gap in RL . Other works ( Michael et al. , 2020 ; Cobbe et al. , 2018 ; Ye et al. , 2019 ; Raileanu et al. , 2020 ; Wang et al. , 2020 ) have made use of data augmentation to learn policies that generalize well . Another popular approach is to use representation learning techniques such as variational information bottleneck ( Igl et al. , 2019 ) and bisimulation metrics ( Agarwal et al. , 2021 ; Zhang et al. , 2020 ) to help with the problem of overfitting in RL . ( Cobbe et al. , 2020 ; Raileanu & Fergus , 2021 ) learned a separate network for policy and value function and showed that it helps reduce the generalization gap in RL . ( Sax et al. , 2018 ) showed that using mid-level visual features optimized for segmentation , depth , keypoints , surface normal prediction and others improve transfer . In our work , rather than proposing a novel approach to solve overfitting , we analyze the limitations of the current techniques and propose avanues for improvements . 3 PRELIMINARIES . We consider a distribution of Partially Observable Markov Decision Process ( POMDPs ) p ( M ) such thatMi = ( Si , Oi , Ai , Ti , Ωi , Ri , γ ) ∼ p ( M ) . We can think of differentMi as instances of same task in different environments . Si is the state space , Oi is the observation space , Ai is the action space , Ti ( s ′|s , a ) is the transition function , Ωi ( o′|s , a ) is the observation function , Ri ( s , a ) is the reward function and γ is the discount factor . The goal is to find a policy πθ that maximizes the expected sum of discounted rewards over the distribution of POMDPs , J ( πθ ) = Eπ , p ( M ) [ ∑T−1 t=0 γ tRi ( st , at ) ] . During training , we only have access to a limited number of POMDPs M̂ = { Mi } Ni=1 . Our goal is to find a policy πθ that generalizes to new POMDPs sampled from p ( M ) . 3.1 ENVIRONMENT SETUP . We conduct experiments using the ProcGen suite ( Cobbe et al. , 2019 ) , a collection of 16 procedurally generated video games . Each game contains multiple levels wherein the agent needs to perform the same task , but the environment varies due to changes in position of objects , their textures and the background texture . We refer to variation in levels resulting solely from changes in texture of object , agent and the background as theme variation . We collectively refer to variation in levels resulting from changes in positions of the objects as layout variation . Figure 1 illustrates theme and layout variation for two games , Jumper and Climber . For each game , the first and the second level differ in theme whereas the first and the third level differ in both the theme and the layout . 3.2 METHODS . We use standard neural network architecture to represent the policy , a ∼ πθ ( a|s ) , consisting of : a visual encoder z = πencθ1 ( s ) and a policy head a ∼ π h θ2 ( a|z ) with parameters θ1 , θ2 respectively . Collectively , they are referred as θ = [ θ1 , θ2 ] . The architecture of the visual encoder is borrowed from IMPALA ( Espeholt et al. , 2018 ) and is a ResNet with 15 convolutional layers . It outputs a flattened feature map ( z ) that we refer to as the visual feature . These features are passed into a two 2-layer fully-connected neural networks with 256 hidden units for predicting the actions ( i.e. , the policy head ) and the value function ( i.e. , the value head , V hψ ( z ) ) . We analyze the generalization performance of PPO ( Schulman et al. , 2017 ) and three state-of-the-art ( SOTA ) methods : • UCB-DrAC ( Raileanu et al. , 2020 ) uses Upper Confidence Bound ( UCB ) to automatically select game specific data augmentations out of eight possible ones : crop , grayscale , cutout , cutout-color , flip , rotate , random convolution and color jitter . For brevity , we sometimes refer to UCB-DrAC as simply DrAC . • DAAC ( Raileanu & Fergus , 2021 ) trains two separate networks for representing the policy and the value function . The intuition is that in POMDPs , the value network will overfit , but the policy can still generalize . The seperation of networks prevents interference in visual features learned for representing the value function and the policy . PPG ( Cobbe et al. , 2020 ) also trains separate networks for actor and critic but mainly focuses on improved training of policies and not their generalization performance . • iDAAC ( Raileanu & Fergus , 2021 ) adds an auxiliary loss to the training objective of DAAC to encourage domain invariance . We use PPO ’ s pytorch implementation ( Kostrikov , 2018 ) , the official implementation for UCBDrAC ( Raileanu et al. , 2020 ) , and the official implementation for DAAC/iDAAC ( Raileanu & Fergus , 2021 ) . See Appendix B for the hyperparameters used .

Summary: This paper looks into why some of the recent methods for improving generalization in deep RL still do not perfectly generalize to new levels of a game. For data augmentation methods, they show that better augmentations which make use of the underlying process that generates the data / ground truth task variation. For approaches that learn domain invariant features, the authors claim that you actually need domain-dependent features to generalize well and that these methods may work well for different reasons (i.e. the additional signal and regularization from the auxiliary losses). They investigate approaches that use data augmentation and domain invariant features and provide some new insights. The authors also show that simple auxiliary tasks can improve generalization policies and that other approaches.