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: To enable robust policy learning with image observations, the paper proposes a simple data augmentation technique that can be used with existing model-free reinforcement learning algorithms. It defines a notion of optimality invariant state transformation which preserves the Q function. An example of such transformations can be random image translations. It uses these transformations to (i) transform the input images, (ii) average the target Q values, and (iii) average the Q function themselves. Using this simple technique, they are able to get SOTA on DM control tasks and Atari 100k benchmark. On DM control tasks, it’s able to outperform SAC trained on state representations. Additionally, the paper provides ablation studies on different image transformations and robustness analysis with respect to hyperparameter settings.