Local Feature Swapping for Generalization in Reinforcement Learning

1 INTRODUCTION . Advances made in deep learning have opened the way to many applications in computer vision such as classification , object recognition , or image segmentation . The powerful representation capabilities of deep neural networks paved the way for many successes in deep reinforcement learning with the design of agents able to take decisions directly from pixels ( Mnih et al. , 2013 ; 2015 ) . However , the sensitivity of neural networks to the distribution of training data strongly affects their generalization abilities . Neural networks are intrinsically designed to memorize the data they are trained upon , since their fitting implies empirical risk minimization ( Vapnik , 1992 ) ( they minimize the empirical average prediction error over a large training dataset ) . Therefore , they are prone to prediction errors on unseen samples . RL agents also suffer from this handicap and tend to memorize training trajectories , rather than general skills and features leading to transferable policies . This phenomenon , usually known as overfitting , takes a double sense in RL . Generalization in RL implies the ability to generalize across states ( as in supervised learning ) , but also across environments . It is only recently that several environments with different configurations for training and testing have emerged and received a lot of attention ( Nichol et al. , 2018 ; Justesen et al. , 2018 ; Zhang et al. , 2018a ; Cobbe et al. , 2019 ; 2020 ) , shedding light on the generalization issue which remained mostly overlooked , and confirming the poor generalization ability of current algorithms . Strategies to achieve good generalization and avoid overfitting in deep learning fall into three categories of regularization : explicit regularization ( e.g. , loss penalization , weight decay ) , implicit regularization via the architecture and optimization ( e.g. , dropout , batch-normalization , batch size selection , momentum , early stopping ) , or implicit regularization by enhancement of the input data ( data augmentation ) . Direct application of these strategies to deep RL agents has demonstrated some improvements in agent generalization abilities in some environments , but much progress remains to be made to integrate RL systems in real-world applications . In this work , we address the observational overfitting issue introduced in Song et al . ( 2019 ) which considers a zero-shot generalization RL setting . An agent is trained on a specific distribution of environments ( for example some levels of a platform video game ) and tested on similar environments sharing the same high-level goal and dynamics but with different layouts and visual attributes ( background or assets ) . We argue that a structured stochastic permutation of features of an RL agent during training leads to state-of-theart generalization performance for vision-based policies . We introduce an efficient regularization technique based on Channel-consistent LOcal Permutations ( CLOP ) of the feature maps that mitigates overfitting . We implement it as an intermediate layer in feed-forward neural networks , and demonstrate its effectiveness on several reinforcement and supervised learning problems . This paper is organized as follows . Section 2 presents the necessary background on generalization in supervised learning and reinforcement learning . Section 3 reviews recent work in the literature that allows for a critical look at our contribution and put it in perspective . Section 4 introduces the CLOP technique and the corresponding layer . Section 5 empirically evaluates agents using CLOP against state-of-the-art generalization methods and discusses their strengths , weaknesses , and variants . Section 6 summarizes and concludes this paper . 2 WHAT IS GENERALIZATION ? . Generalization in supervised learning ( SL ) . LetX be an input space of descriptors andY an output space of labels . A SL problem is defined by a distribution p ( x , y ) of elements of X × Y , and a loss L ( ŷ , y ) which measures how different ŷ ∈ Y and y ∈ Y are . Then , for a given function f intended to capture the mapping from X to Y underlying the p distribution , one defines the ( expected ) risk as R ( f ) = Ep [ L ( f ( x ) , y ) ] . Since the true p is generally unknown , one can not directly minimize the risk in search for the optimal f . Given a training set S = { ( xi , yi ) } ni=1 of n items in X × Y drawn i.i.d . according to p ( x , y ) , empirical risk minimization ( Vapnik , 1992 ) seeks to minimize RS ( f ) = 1/n ∑n i=1 [ L ( f ( xi ) , yi ) ] . The ability of f to generalize to unseen samples is then defined by the generalization gap R ( f ) − RS ( f ) , which is often evaluated by approaching R ( f ) as the empirical risk RT ( f ) over a test set T = { ( xi , yi ) } n ′ i=1 also drawn i.i.d . from p ( x , y ) . Closing the generalization gap can be attempted through structural risk minimization , which modifies RS so as to include a regularization penalty into the optimization , or more generally by the introduction of inductive biases ( Mitchell , 1980 ) . Reinforcement learning ( RL ) . RL ( Sutton & Barto , 2018 ) considers the problem of learning a decision making policy for an agent interacting over multiple time steps with a dynamic environment . At each time step , the agent and environment are described through a state s ∈ S , and an action a ∈ A is performed ; then the system transitions to a new state s′ according to probability T ( s′|s , a ) , while receiving reward R ( s , a ) . The tuple M = ( S , A , T , R ) forms a Markov Decision Process ( Puterman , 2014 , MDP ) , which is often complemented with the knowledge of an initial state distribution p0 ( s ) . A decision making policy parameterized by θ is a function πθ ( a|s ) mapping states to distributions over actions . Training a reinforcement learning agent consists in finding the policy that maximizes the discounted expected return : J ( πθ ) = E [ ∑∞ t=0 γ tR ( st , at ) ] . Generalization in RL . Departing from the rather intuitive definition of generalization in SL , the idea of generalization in RL may lead to misconceptions . One could expect a policy that generalizes well to perform well across environments . It seems important to disambiguate this notion and stress out the difference between generalization in the SL sense ( which is a single MDP problem ) , and domain generalization or robustness . For instance , one could expect a policy that learned to play a certain level of a platform game to be able to play on another level , provided the key game features ( e.g . platforms , enemies , treasures ) remain similar enough and the game dynamics are the same . This type of generalization benchmark is typically captured by procedurally generated environments , such as Procgen ( Cobbe et al. , 2020 ) . Finding a policy that yields a guaranteed minimal performance among a set of MDPs ( sharing common state and action spaces ) is the problem of solving Robust MDPs ( Iyengar , 2005 ) . Similarly , that of policy optimization over a distribution over MDPs ( sharing common states and actions too ) is that of domain generalization ( Tobin et al. , 2017 ) . We argue that these problems are unrelated and much harder than vanilla generalization in RL . Robust MDPs induce a maxπminT , R problem , domain generalization is a maxπ ET , R one , while vanilla generalization as understood in SL remains a maxπ problem that we try to solve for all states within a single MDP ( and not only the ones in the training set ) . Consequently , the ability to generalize in RL is a broader notion than structural risk minimization , which , unfortunately , uses the same name of generalization . In a given game , finding a policy that plays well on all levels is still solving the very same MDP . The underlying transition and reward models are the same and only the parts of the state space explored in each level are different . Note that these explored subsets of the state space may have non-empty intersections ( e.g . the final part of two different levels might be common ) and the optimal distribution over actions for a given state is unique . In this work we focus on generalization as the problem of preventing overfitting ( and thus reducing the generalization gap ) within a single MDP , which we call observational overfitting ( Song et al. , 2019 ) . Observational overfitting ( OO ) in RL . The inability to generalize might be caused by overfitting to the partially explored environment dynamics ( Rajeswaran et al. , 2017 ) , or to some misleading signal that is correlated with progress but does not generalize to new levels ( Machado et al. , 2018 ; Song et al. , 2019 ) . Therefore , preventing OO in RL remains the ability to generalize across states , within the same MDP . Beyond the ability of RL agents to memorize good actions in explored states , it boils down to their ability to capture rules that extend to unencountered states , just as in structural risk minimization . As in SL , many policies might fit the observed data , but few might generalize to the true mechanisms of the underlying MDP . Song et al . ( 2019 ) propose to capture OO through a framework where a unique latent state space is transformed into a variety of observation spaces . Observation functions are built by combining useful features with purely decorative and unimportant ones which vary from one observation function to the next . They suppose the observation functions are drawn from a certain distribution ( over procedural generation parameters for instance ) and define the corresponding distribution over MDPs . In turn , they define the risk and the generalization gap with respect to this distribution over MDPs . We slightly depart from their derivation and argue this distinction is unnecessary : what is being solved is the unique MDP defined by the underlying dynamics and the projection of latent states into observations . The risk is thus defined with respect to the distribution over observations induced by the distributions over initial states and observation functions . Overall , this allows capturing the problem of generalization gap minimization in RL and underpins the developments proposed in the remainder of this paper . 3 RELATED WORK . In the following paragraphs , we cover essential works which aim to improve the generalization abilities of neural networks in both supervised learning and reinforcement learning . In supervised learning , the process of modifying a learning algorithm with the objective to reduce its test error while preserving its train error is known as regularization ( Goodfellow et al. , 2016 ) . Direct , or explicit , regularization can be achieved by adding a regularizer term into the loss function , such as an L2 penalty on networks parameters ( Plaut et al. , 1986 ; Krogh & Hertz , 1992 ) . A second , implicit , regularization strategy consists in feature-level manipulations like dropout ( Srivastava et al. , 2014 ) , drop-connect ( Wan et al. , 2013 ) or batch-normalization ( Ioffe & Szegedy , 2015 ) . Regularization can also be achieved implicitly by directly augmenting the training data with perturbations such as adding Gaussian noise ( Bishop , 1995 ) , or , in the case of visual inputs , random cropping and flipping ( Krizhevsky et al. , 2012 ; Szegedy et al. , 2017 ) or removing structured parts of the image ( DeVries & Taylor , 2017 ) . Another efficient augmentation strategy , called label smoothing ( Szegedy et al. , 2016 ) , consists in penalizing overconfident predictions of neural networks by perturbing the labels . Combining perturbation of both inputs and outputs , Zhang et al . ( 2018c ) ; Yun et al . ( 2019 ) , produce synthetic data and labels using two different samples and their corresponding labels . Many studies have highlighted the limited ability of RL agents to generalize to new scenarios ( Farebrother et al. , 2018 ; Packer et al. , 2018 ; Zhang et al. , 2018b ; Song et al. , 2019 ; Cobbe et al. , 2019 ) . Using saliency maps , Song et al . ( 2019 ) exhibit that RL agents trained from pixels in environments with rich and textured observations , such as platform video games ( e.g . Sonic ( Nichol et al. , 2018 ) ) , focus on elements of the scenery correlated with in-game progress but which lead to poor generalization in later situations . One of the identified counter-measures to overfitting in RL consists in applying standard methods in supervised learning . Cobbe et al . ( 2019 ; 2020 ) demonstrated the contribution of classical supervised learning regularization techniques like weight decay , dropout , or batch-normalization to generalization in procedurally generated environments . Similar to dataaugmentation in supervised learning , Raileanu et al . ( 2020 ) ; Laskin et al . ( 2020 ) ; Yarats et al . ( 2020 ) apply visual-data-augmentation on the observations provided by the environment to train robust agents . Igl et al . ( 2019 ) use a selective noise injection and information bottleneck to regularize their agent . Wang et al . ( 2020 ) propose mixreg , a direct application of mixup ( Zhang et al. , 2018c ) in RL , combining two randomly sampled observations , and training the RL agent using their interpo- lated supervision signal . Lee et al . ( 2020 ) use a random convolution layer ahead of the network architecture , to modify the color and texture of the visual observations during training . Tobin et al . ( 2017 ) tackle the sim-to-real problem , using domain randomization on visual inputs to bridge the gap between simulation and reality in robotics . Raileanu & Fergus ( 2021 ) dissociate the optimization process of the policy and value function represented by separate networks and introduce an auxiliary loss that encourages the representation to be invariant to task-irrelevant properties of the environment . Another recent strategy consists in learning representations that are invariant to visual changes . Higgins et al . ( 2017 ) use a two-stage learning process : first they extract disentangled representations from random observation and then they exploit these representations to train an RL agent . Zhang et al . ( 2020 ) use bisimulation metrics to quantify behavioral similarity between states and learn robust task-relevant representations . Wang et al . ( 2021b ) extract , without supervision , the visual foreground to provide background invariant inputs to the policy learner . Our work focuses on the OO situations , where an agent overfits visual observations features that are irrelevant to the latent dynamics of the MDP . Our method approaches the augmentation of data by noise injection at the feature level , thus avoiding the computationally costly operations of highdimensional image transformation in regular data augmentation . By directly modifying the encountered features ’ spatial localization , the CLOP layer aims to remove the correlations between spurious features and rewards .

Summary: This paper proposes a new regularization scheme, which involves randomly permuting nearby (in terms of spatial locations in the downsized feature tensor) depth blocks. The authors show its effectiveness in supervised learning tasks and establish a SOTA-like result over Procgen generalization, and also provide several empirical ablations and visualizations over the method to understand its inner workings. The paper shows that regularization which involves the locality of objects is particularly effective.