Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels

1 INTRODUCTION . Sample-efficient deep reinforcement learning ( RL ) algorithms capable of directly training from image pixels would open up many real-world applications in control and robotics . However , simultaneously training a convolutional encoder alongside a policy network is challenging when given limited environment interaction , strong correlation between samples and a typically sparse reward signal . Limited supervision is a common problem across AI and two approaches are commonly taken : ( i ) training with an additional auxiliary losses , such as those based on self-supervised learning ( SSL ) and ( ii ) training with data augmentation . A wide range of auxiliary loss functions have been proposed to augment supervised objectives , e.g . weight regularization , noise injection ( Hinton et al. , 2012 ) , or various forms of auto-encoder ( Kingma et al. , 2014 ) . In RL , reconstruction losses ( Jaderberg et al. , 2017 ; Yarats et al. , 2019 ) or SSL objectives ( Dwibedi et al. , 2018 ; Srinivas et al. , 2020 ) are used . However , these objectives are unrelated to the task at hand , thus have no guarantee of inducing an appropriate representation for the policy network . SSL losses are highly effective in the large data regime , e.g . in domains such as vision ( Chen et al. , 2020 ; He et al. , 2019 ) and NLP ( Collobert et al. , 2011 ; Devlin et al. , 2018 ) where large ( unlabeled ) datasets are readily available . However , in sample-efficient RL , training data is more limited due to restricted interaction between the agent and the environment , limiting their effectiveness . Data augmentation methods are widely used in vision and speech domains , where output-invariant perturbations can easily be applied to the labeled input examples . Surprisingly , data augmentation ∗Equal contribution . Author ordering determined by coin flip . Both authors are corresponding . has received little attention in the RL community . In this paper we propose augmentation approaches appropriate for sample-efficient RL and comprehensively evaluate them . The key idea of our approach is to use standard image transformations to perturb input observations , as well as regularizing the Q-function learned by the critic so that different transformations of the same input image have similar Q-function values . No further modifications to standard actor-critic algorithms are required . Our study is , to the best of our knowledge , the first careful examination of image augmentation in sample-efficient RL . The main contributions of the paper are as follows : ( i ) the first to demonstrate that data augmentation greatly improves performance when training model-free RL algorithms from images ; ( ii ) introducing a natural way to exploit MDP structure through two mechanisms for regularizing the value function , in a manner that is generally applicable to model-free RL and ( iii ) setting a new state-of-the-art performance on the standard DeepMind control suite ( Tassa et al. , 2018 ) , closing the gap between learning from states , and Atari 100k ( Kaiser et al. , 2019 ) benchmarks . 2 RELATED WORK . Data Augmentation in Computer Vision Data augmentation via image transformations has been used to improve generalization since the inception of convolutional networks ( Becker & Hinton , 1992 ; Simard et al. , 2003 ; LeCun et al. , 1989 ; Ciresan et al. , 2011 ; Ciregan et al. , 2012 ) . Following AlexNet ( Krizhevsky et al. , 2012 ) , they have become a standard part of training pipelines . For object classification tasks , the transformations are selected to avoid changing the semantic category , i.e . translations , scales , color shifts , etc . While a similar set of transformations are potentially applicable to control tasks , the RL context does require modifications to be made to the underlying algorithm . Data augmentation methods have also been used in the context of self-supervised learning . Dosovitskiy et al . ( 2016 ) use per-exemplar perturbations in a unsupervised classification framework . More recently , several approaches ( Chen et al. , 2020 ; He et al. , 2019 ; Misra & van der Maaten , 2019 ) have used invariance to imposed image transformations in contrastive learning schemes , producing state-of-the-art results on downstream recognition tasks . By contrast , our scheme addresses control tasks , utilizing different types of invariance . Data Augmentation in RL In contrast to computer vision , data augmentation is rarely used in RL . Certain approaches implicitly adopt it , for example Levine et al . ( 2018 ) ; Kalashnikov et al . ( 2018 ) use image augmentation as part of the AlexNet training pipeline without analysing the benefits occurring from it , thus being overlooked in subsequent work . HER ( Andrychowicz et al. , 2017 ) exploits information about the observation space by goal and reward relabeling , which can be viewed as a way to perform data augmentation . Other work uses data augmentation to improve generalization in domain transfer ( Cobbe et al. , 2018 ) . However , the classical image transformations used in vision have not previously been shown to definitively help on standard RL benchmarks . Concurrent with our work , RAD ( Laskin et al. , 2020 ) performs an exploration of different data augmentation approaches , but is limited to transformations of the image alone , without the additional augmentation of the Q-function used in our approach . Moreover , RAD can be regarded as a special case of our algorithm . Multiple follow ups to our initial preprint appeared on ArXiv ( Raileanu et al. , 2020 ; Okada & Taniguchi , 2020 ) , using similar techniques on other tasks , thus supporting the effectiveness and generality of data augmentation in RL . Continuous Control from Pixels There are a variety of methods addressing the sample-efficiency of RL algorithms that directly learn from pixels . The most prominent approaches for this can be classified into two groups , model-based and model-free methods . The model-based methods attempt to learn the system dynamics in order to acquire a compact latent representation of high-dimensional observations to later perform policy search ( Hafner et al. , 2018 ; Lee et al. , 2019 ; Hafner et al. , 2019 ) . In contrast , the model-free methods either learn the latent representation indirectly by optimizing the RL objective ( Barth-Maron et al. , 2018 ; Abdolmaleki et al. , 2018 ) or by employing auxiliary losses that provide additional supervision ( Yarats et al. , 2019 ; Srinivas et al. , 2020 ; Sermanet et al. , 2018 ; Dwibedi et al. , 2018 ) . Our approach is complementary to these methods and can be combined with them to improve performance . 3 BACKGROUND . Reinforcement Learning from Images We formulate image-based control as an infinite-horizon partially observable Markov decision process ( POMDP ) ( Bellman , 1957 ; Kaelbling et al. , 1998 ) . An POMDP can be described as the tuple ( O , A , p , r , γ ) , where O is the high-dimensional observation space ( image pixels ) , A is the action space , the transition dynamics p = Pr ( o′t|o≤t , at ) capture the probability distribution over the next observation o′t given the history of previous observations o≤t and current action at , r : O ×A → R is the reward function that maps the current observation and action to a reward rt = r ( o≤t , at ) , and γ ∈ [ 0 , 1 ) is a discount factor . Per common practice ( Mnih et al. , 2013 ) , throughout the paper the POMDP is converted into an MDP ( Bellman , 1957 ) by stacking several consecutive image observations into a state st = { ot , ot−1 , ot−2 , . . . } . For simplicity we redefine the transition dynamics p = Pr ( s′t|st , at ) and the reward function rt = r ( st , at ) . We then aim to find a policy π ( at|st ) that maximizes the cumulative discounted return Eπ [ ∑∞ t=1 γ trt|at ∼ π ( ·|st ) , s′t ∼ p ( ·|st , at ) , s1 ∼ p ( · ) ] . Soft Actor-Critic The Soft Actor-Critic ( SAC ) ( Haarnoja et al. , 2018 ) learns a state-action value function Qθ , a stochastic policy πθ and a temperature α to find an optimal policy for an MDP ( S , A , p , r , γ ) by optimizing a γ-discounted maximum-entropy objective ( Ziebart et al. , 2008 ) . θ is used generically to denote the parameters updated through training in each part of the model . Deep Q-learning DQN ( Mnih et al. , 2013 ) also learns a convolutional neural net to approximate Q-function over states and actions . The main difference is that DQN operates on discrete actions spaces , thus the policy can be directly inferred from Q-values . In practice , the standard version of DQN is frequently combined with a set of refinements that improve performance and training stability , commonly known as Rainbow ( van Hasselt et al. , 2015 ) . For simplicity , the rest of the paper describes a generic actor-critic algorithm rather than DQN or SAC in particular . Further background on DQN and SAC can be found in Appendix A . 4 SAMPLE EFFICIENT REINFORCEMENT LEARNING FROM PIXELS . 4.1 OPTIMALITY INVARIANT IMAGE TRANSFORMATIONS FOR Q FUNCTION . We first introduce a general framework for regularizing the value function through transformations of the input state . For a given task , we define an optimality invariant state transformation f : S×T → S as a mapping that preserves the Q-values Q ( s , a ) = Q ( f ( s , ν ) , a ) for all s ∈ S , a ∈ A and ν ∈ T . where ν are the parameters of f ( · ) , drawn from the set of all possible parameters T . One example of such transformations are the random image translations successfully applied in the previous section . For every state , the transformations allow the generation of several surrogate states with the same Q-values , thus providing a mechanism to reduce the variance of Q-function estimation . In particular , for an arbitrary distribution of states µ ( · ) and policy π , instead of using a single sample s∗ ∼ µ ( · ) , a∗ ∼ π ( ·|s∗ ) estimation of the following expectation E s∼µ ( · ) a∼π ( ·|s ) [ Q ( s , a ) ] ≈ Q ( s∗ , a∗ ) we generate K samples via random transformations and obtain an estimate with lower variance E s∼µ ( · ) a∼π ( ·|s ) [ Q ( s , a ) ] ≈ 1 K K∑ k=1 Q ( f ( s∗ , νk ) , ak ) where νk ∈ T and ak ∼ π ( ·|f ( s∗ , νk ) ) . This suggests two distinct ways to regularize Q-function . First , we use the data augmentation to compute the target values for every transition tuple ( si , ai , ri , s′i ) as yi = ri + γ 1 K K∑ k=1 Qθ ( f ( s ′ i , ν ′ i , k ) , a ′ i , k ) where a ′ i , k ∼ π ( ·|f ( s′i , ν′i , k ) ) ( 1 ) where ν′i , k ∈ T corresponds to a transformation parameter of s′i . Then the Q-function is updated using these targets through an SGD update using learning rate λθ θ ← θ − λθ∇θ 1 N N∑ i=1 ( Qθ ( f ( si , νi ) , ai ) − yi ) 2 . ( 2 ) In tandem , we note that the same target from Equation ( 1 ) can be used for different augmentations of si , resulting in the second regularization approach θ ← θ − λθ∇θ 1 NM N , M∑ i=1 , m=1 ( Qθ ( f ( si , νi , m ) , ai ) − yi ) 2 . ( 3 ) When both regularization methods are used , νi , m and ν′i , k are drawn independently .

Summary: This paper investigates data augmentation in the context of RL and proposes a novel augmentation algorithm to enabling robust learning directly from pixels without the need for auxiliary losses or pre-training. The authors propose to average both the Q function and its target over multiple image transformations. The experiments on DeepMind control suite and Atari 100k benchmark show that their method outperforms previous model-free, model-based and contrastive learning approaches.