Automatic Data Augmentation for Generalization in Reinforcement Learning

1 Introduction . Generalization to new environments remains a major challenge in deep reinforcement learning ( RL ) . Current methods fail to generalize to unseen environments even when trained on similar settings [ 19 , 51 , 71 , 11 , 21 , 12 , 60 ] . This indicates that standard RL agents memorize specific trajectories rather than learning transferable skills . Several strategies have been proposed to alleviate this problem , such as the use of regularization [ 19 , 71 , 11 , 28 ] , data augmentation [ 11 , 44 , 69 , 38 , 41 ] , or representation learning [ 72 , 74 ] . In this work , we focus on the use of data augmentation in RL . We identify key differences between supervised learning and reinforcement learning which need to be taken into account when using data augmentation in RL . More specifically , we show that a naive application of data augmentation can lead to both theoretical and practical problems with standard RL algorithms , such as unprincipled objective estimates and poor performance . As a solution , we propose Data-regularized Actor-Critic or DrAC , a new algorithm that enables the use of data augmentation with actor-critic algorithms in a theoretically sound way . Specifically , we introduce two regularization terms which constrain the agent ’ s policy and value function to be invariant to various state transformations . Empirically , this approach allows the agent to learn useful behaviors ( outperforming strong RL baselines ) in settings in which a naive use of data augmentation completely fails or converges to a sub-optimal policy . While we use Proximal Policy Optimization ( PPO , Schulman et al . [ 56 ] ) to describe and validate our approach , the method 1in June 2020 , at the time of making this work publicly available on arXiv . it has since been surpassed . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . can be easily integrated with any actor-critic algorithm with a discrete stochastic policy such as A3C [ 49 ] , SAC [ 24 ] , or IMPALA [ 17 ] . The current use of data augmentation in RL either relies on expert knowledge to pick an appropriate augmentation [ 11 , 44 , 38 ] or separately evaluates a large number of transformations to find the best one [ 69 , 41 ] . In this paper , we propose three methods for automatically finding a useful augmentation for a given RL task . The first two learn to select the best augmentation from a fixed set , using either a variant of the upper confidence bound algorithm ( UCB , Auer [ 2 ] ) or meta-learning ( RL2 , Wang et al . [ 66 ] ) . We refer to these methods as UCB-DrAC and RL2-DrAC , respectively . The third method , Meta-DrAC , directly meta-learns the weights of a convolutional network , without access to predefined transformations ( MAML , Finn et al . [ 20 ] ) . Figure 1 gives an overview of UCB-DrAC . We evaluate these approaches on the Procgen generalization benchmark [ 12 ] which consists of 16 procedurally generated environments with visual observations . Our results show that UCB-DrAC is the most effective among these at finding a good augmentation , and is comparable or better than using DrAC with the best augmentation from a given set . UCB-DrAC also outperforms baselines specifically designed to improve generalization in RL [ 28 , 44 , 41 ] on both train and test . In addition , we show that our agent learns policies and representations that are more invariant to changes in the environment which do not alter the reward or transition function ( i.e . they are inconsequential for control ) , such as the background theme . To summarize , our work makes the following contributions : ( i ) we introduce a principled way of using data augmentation with actor-critic algorithms , ( ii ) we propose a practical approach for automatically selecting an effective augmentation in RL settings , ( iii ) we show that the use of data augmentation leads to policies and representations that better capture task invariances , and ( iv ) we demonstrate state-of-the-art results on the Procgen benchmark and outperform popular RL methods on four DeepMind Control tasks with natural and synthetic distractors . 2 Background . We consider a distribution q ( m ) of Partially Observable Markov Decision Processes ( POMDPs , [ 6 ] ) m ∈ M , with m defined by the tuple ( Sm , Om , A , Tm , Rm , γ ) , where Sm is the state space , Om is the observation space , A is the action space , Tm ( s′|s , a ) is the transition function , Rm ( s , a ) is the reward function , and γ is the discount factor . During training , we restrict access to a fixed set of POMDPs , Mtrain = { m1 , . . . , mn } , where mi ∼ q , ∀ i = 1 , n. The goal is to find a policy πθ which maximizes the expected discounted reward over the entire distribution of POMDPs , J ( πθ ) = Eq , π , Tm , pm [ ∑T t=0 γ tRm ( st , at ) ] . In practice , we use the Procgen benchmark which contains 16 procedurally generated games . Each game corresponds to a distribution of POMDPs q ( m ) , and each level of a game corresponds to a POMDP sampled from that game ’ s distribution m ∼ q . The POMDP m is determined by the seed ( i.e . integer ) used to generate the corresponding level . Following the setup from Cobbe et al . [ 12 ] , agents are trained on a fixed set of n = 200 levels ( generated using seeds from 1 to 200 ) and tested on the full distribution of levels ( generated by sampling seeds uniformly at random from all computer integers ) . Proximal Policy Optimization ( PPO , Schulman et al . [ 56 ] ) is an actor-critic algorithm that learns a policy πθ and a value function Vθ with the goal of finding an optimal policy for a given POMDP . PPO alternates between sampling data through interaction with the environment and maximizing a clipped surrogate objective function using stochastic gradient ascent . One component of the PPO objective is the policy gradient term JPG , which is estimated using importance sampling : JPG ( θ ) = ∑ a∈A πθ ( a|s ) Âθold ( s , a ) = Ea∼πθold [ πθ ( a|s ) πθold ( a|s ) Âθold ( s , a ) ] , ( 1 ) where Â ( · ) is an estimate of the advantage function , πθold is the behavior policy used to collect trajectories ( i.e . that generates the training distribution of states and actions ) , and πθ is the policy we want to optimize ( i.e . that generates the true distribution of states and actions ) . See Appendix A for a full description of PPO . 3 Automatic Data Augmentation for RL . 3.1 Data Augmentation in RL . Image augmentation has been successfully applied in computer vision for improving generalization on object classification tasks [ 58 , 9 , 8 , 39 ] . As noted by Kostrikov et al . [ 38 ] , those tasks are invariant to certain image transformations such as rotations or flips , which is not always the case in RL . For example , if your observation is flipped , the corresponding reward will be reversed for the left and right actions and will not provide an accurate signal to the agent . While data augmentation has been previously used in RL settings without other algorithmic changes [ 11 , 69 , 41 ] , we argue that this approach is not theoretically sound . If transformations are naively applied to observations in PPO ’ s buffer , as done in Laskin et al . [ 41 ] , the PPO objective changes and equation ( 1 ) is replaced by JPG ( θ ) = Ea∼πθold [ πθ ( a|f ( s ) ) πθold ( a|s ) Âθold ( s , a ) ] , ( 2 ) where f : S × H → S is the image transformation . However , the right hand side of the above equation is not a sound estimate of the left hand side because πθ ( a|f ( s ) ) 6= πθ ( a|s ) , since nothing constrains πθ ( a|f ( s ) ) to be close to πθ ( a|s ) . Note that in the on-policy case when θ = θold , the ratio used to estimate the advantage should be equal to one , which is not necessarily the case when using equation ( 2 ) . In fact , one can define certain transformations f ( · ) that result in an arbitrarily large ratio πθ ( a|f ( s ) ) /πθold ( a|s ) . Figure 3 shows examples where a naive use of data augmentation prevents PPO from learning a good policy in practice , suggesting that this is not just a theoretical concern . In the following section , we propose an algorithmic change that enables the use of data augmentation with actor-critic algorithms in a principled way . 3.2 Policy and Value Function Regularization . Inspired by the recent work of Kostrikov et al . [ 38 ] , we propose two novel regularization terms for the policy and value functions that enable the proper use of data augmentation for actor-critic algorithms . Our algorithmic contribution differs from that of Kostrikov et al . [ 38 ] in that it constrains both the actor and the critic , as opposed to only regularizing the Q-function . This allows our method to be used with a different ( and arguably larger ) class of RL algorithms , namely those that learn a policy and a value function . Following Kostrikov et al . [ 38 ] , we define an optimality-invariant state transformation f : S×H → S as a mapping that preserves both the agent ’ s policy π and its value function V such that V ( s ) = V ( f ( s , ν ) ) and π ( a|s ) = π ( a|f ( s , ν ) ) , ∀s ∈ S , ν ∈ H , where ν are the parameters of f ( · ) , drawn from the set of all possible parametersH . To ensure that the policy and value functions are invariant to such transformation of the input state , we propose an additional loss term for regularizing the policy , Gπ = KL [ πθ ( a|s ) | πθ ( a|f ( s , ν ) ) ] , ( 3 ) as well as an extra loss term for regularizing the value function , GV = ( Vφ ( s ) − Vφ ( f ( s , ν ) ) ) 2 . ( 4 ) Thus , our data-regularized actor-critic method , or DrAC , maximizes the following objective : JDrAC = JPPO − αr ( Gπ +GV ) , ( 5 ) where αr is the weight of the regularization term ( see Algorithm 1 ) . The use of Gπ and GV ensures that the agent ’ s policy and value function are invariant to the transformations induced by various augmentations . Particular transformations can be used to impose certain inductive biases relevant for the task ( e.g . invariance with respect to colors or translations ) . In addition , Gπ and GV can be added to the objective of any actor-critic algorithm with a discrete stochastic policy ( e.g . A3C , TRPO , ACER , SAC , or IMPALA ) without any other changes . Note that when using DrAC , as opposed to the method proposed by Laskin et al . [ 41 ] , we still use the correct importance sampling estimate of the left hand side objective in equation ( 1 ) ( instead of a wrong estimate as in equation ( 2 ) ) . This is because the transformed observations f ( s ) are only used to compute the regularization losses Gπ and GV , and thus are not used for the main PPO objective . Without these extra terms , the only way to use data augmentation is as explained in Section 3.1 , which leads to inaccurate estimates of the PPO objective . Hence , DrAC benefits from the regularizing effect of using data augmentation , while mitigating adverse consequences on the RL objective . See Appendix B for a more detailed explanation of why a naive application of data augmentation with certain policy gradient algorithms is theoretically unsound . Algorithm 1 DrAC : Data-regularized Actor-Critic applied to PPO Black : unmodified actor-critic algorithm . Cyan : image transformation . Red : policy regularization . Blue : value function regularization . 1 : Hyperparameters : image transformation f , regularization loss coefficient αr , minibatch size M , replay buffer size T , number of updates K. 2 : for k = 1 , . . . , K do 3 : Collect a new set of transitions D = { ( si , ai , ri , si+1 ) } Ti=1 using πθ . 4 : for j = 1 , . . . , ∗ TM do 5 : { ( si , ai , ri , si+1 ) } Mi=1 ∼ D . Sample a minibatch of transitions 6 : for i = 1 , . . . , M do 7 : νi ∼ H . Sample the augmentation parameters 8 : π̂i ← πφ ( ·|si ) . Compute the policy targets 9 : V̂i ← Vφ ( si ) . Compute the value function targets 10 : end for 11 : Gπ ( θ ) = 1 M ∑M i=1 KL [ π̂i | πθ ( ·|f ( si , νi ) ) ] . Regularize the policy 12 : GV ( φ ) = 1 M ∑M i=1 ( V̂i − Vφ ( f ( si , νi ) ) ) 2 . Regularize the value function 13 : JDrAC ( θ , φ ) = JPPO ( θ , φ ) − αr ( Gπ ( θ ) +GV ( φ ) ) . Compute the total loss function 14 : θ ← argmaxθ JDrAC . Update the policy 15 : φ← argmaxφ JDrAC . Update the value function 16 : end for 17 : end for

Summary: This paper presents a method for improving generalization in reinforcement learning using data augmentation. It identifies a key theoretical limitation of previous work in the naive application of data augmentation to actor critic methods and proposes a solution that avoids biased estimation of the policy gradient objective while still allowing for regularization of the policy and value function. The resulting loss is paired with three distinct approaches for identifying an optimal data augmentation strategy. The best of these approaches, a bandit algorithm, obtains SOTA performance (relative to published works) on the Procgen environment.