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 that utilizes data augmentation for image-based reinforcement learning. The data augmentation is used to regularize the policy and function approximation in the proposed method. In addition, a method for automatically identifying effective ways of data augmentation is proposed. The experimental results show that the proposed method outperforms the baseline methods.