Stabilizing Deep Q-Learning with ConvNets and Vision Transformers under Data Augmentation

1 Introduction . Reinforcement Learning ( RL ) from visual observations has achieved tremendous success in various applications such as video-games [ 43 , 4 , 70 ] , robotic manipulation [ 37 ] , and autonomous navigation [ 42 , 83 ] . However , it is still very challenging for current methods to generalize the learned skills to novel environments , and policies trained by RL can easily overfit to the training environment [ 81 , 13 ] , especially for high-dimensional observation spaces such as images [ 8 , 58 ] . Increasing the variability in training data via domain randomization [ 66 , 50 ] and data augmentation [ 57 , 35 , 33 , 51 ] has demonstrated encouraging results for learning policies invariant to changes in environment observations . Specifically , recent works on data augmentation [ 35 , 33 ] both show improvements in sample efficiency from simple cropping and translation augmentations , but the studies also conclude that additional data augmentation in fact decrease sample efficiency and even cause divergence . While these augmentations have the potential to improve generalization , the increasingly varied data makes the optimization more challenging and risks instability . Unlike supervised learning , balancing the trade-off between stability and generalization in RL requires substantial trial and error . In this paper , we illuminate causes of instability when applying data augmentation to common off-policy RL algorithms [ 43 , 38 , 15 , 18 ] . Based on our findings , we provide an intuitive method for stabilizing this class of algorithms under use of strong data augmentation . Specifically , we find two main causes of instability in previous work ’ s application of data augmentation : ( i ) indiscriminate application of data augmentation resulting in high-variance Q-targets ; and ( ii ) that Q-value estimation strictly from augmented data results in over-regularization . †Website and code is available at : https : //nicklashansen.github.io/SVEA . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) To address these problems , we propose SVEA : Stabilized Q-Value Estimation under Augmentation , a simple yet effective framework for data augmentation in off-policy RL that greatly improves stability of Q-value estimation . Our method consists of the following three components : Firstly , by only applying augmentation in Q-value estimation of the current state , without augmenting Q-targets used for bootstrapping , SVEA circumvents erroneous bootstrapping caused by data augmentation ; Secondly , we formulate a modified Q-objective that optimizes Q-value estimation jointly over both augmented and unaugmented copies of the observations ; Lastly , for SVEA implemented with an actorcritic algorithm , we optimize the actor strictly on unaugmented data , and instead learn a generalizable policy indirectly through parameter-sharing . Our framework can be implemented efficiently without additional forward passes nor introducing additional learnable parameters . We perform extensive empirical evaluation on the DeepMind Control Suite [ 64 ] and extensions of it , including the DMControl Generalization Benchmark [ 21 ] and the Distracting Control Suite [ 60 ] , as well as a set of robotic manipulation tasks . Our method greatly improve Q-value estimation with ConvNets under a set of strong data augmentations , and achieves sample efficiency , asymptotic performance , and generalization that is competitive or better than previous state-of-the-art methods in all tasks considered , at a lower computational cost . Finally , we show that our method scales to RL with Vision Transformers ( ViT ) [ 10 ] . We find that ViT-based architectures are especially prone to overfitting , and data augmentation may therefore be a key component for large-scale RL . 2 Related Work . Representation Learning . Learning visual invariances using data augmentation and self-supervised objectives has proven highly successful in computer vision [ 46 , 45 , 82 , 74 , 68 , 65 , 75 , 27 , 7 ] . For example , Chen et al . [ 7 ] perform an extensive study on data augmentation ( e.g . random cropping and image distortions ) for contrastive learning , and show that representations pre-trained with such transformations transfer effectively to downstream tasks . While our work also uses data augmentation for learning visual invariances , we leverage the Q-objective of deep Q-learning algorithms instead of auxiliary representation learning tasks . Visual Learning for RL . Numerous methods have been proposed with the goal of improving sample efficiency [ 29 , 56 , 68 , 76 , 40 , 59 , 61 , 54 , 77 ] of image-based RL . Recently , using self-supervision to improve generalization in RL has also gained interest [ 80 , 47 , 55 , 1 , 22 , 21 , 72 ] . Notably , Zhang et al . [ 80 ] and Agarwal et al . [ 1 ] propose to learn behavioral similarity embeddings via auxiliary tasks ( bisimulation metrics and contrastive learning , respectively ) , and Hansen et al . [ 21 ] learn visual invariances through an auxiliary prediction task . While these results are encouraging , it has also been shown in [ 29 , 40 , 22 , 79 , 41 ] that the best choice of auxiliary tasks depends on the particular RL task , and that joint optimization with sub-optimally chosen tasks can lead to gradient interference . We achieve competitive sample-efficiency and generalization results without the need for carefully chosen auxiliary tasks , and our method is therefore applicable to a larger variety of RL tasks . Data Augmentation and Randomization for RL . Our work is directly inspired by previous work on generalization in RL by domain randomization [ 66 , 50 , 48 , 52 , 6 ] and data augmentation [ 36 , 9 , 71 , 35 , 33 , 51 , 61 , 21 ] . For example , Tobin et al . [ 66 ] show that a neural network trained for object localization in a simulation with randomized visual augmentations improves real world generalization . Similarly , Lee et al . [ 36 ] show that application of a random convolutional layer to observations during training improve generalization in 3D navigation tasks . More recently , extensive studies on data augmentation [ 35 , 33 ] have been conducted with RL , and conclude that , while small random crops and translations can improve sample efficiency , most data augmentations decrease sample efficiency and cause divergence . We illuminate main causes of instability , and propose a framework for data augmentation in deep Q-learning algorithms that drastically improves stability and generalization . Improving Deep Q-Learning . While deep Q-learning algorithms such as Deep Q-Networks ( DQN ) [ 43 ] have achieved impressive results in image-based RL , the temporal difference objective is known to have inherent instabilities when used in conjunction with function approximation and off-policy data [ 63 ] . Therefore , a variety of algorithmic improvements have been proposed to improve convergence [ 24 , 73 , 25 , 23 , 53 , 38 , 15 , 14 , 28 ] . For example , Hasselt et al . [ 24 ] reduce overestimation of Q-values by decomposing the target Q-value estimation into action selection and action evaluation using separate networks . Lillicrap et al . [ 38 ] reduce target variance by defining the target Q-network as a slow-moving average of the online Q-network . Our method also improves Q-value estimation , but we specifically address the instability of deep Q-learning algorithms on augmented data . 3 Preliminaries . Problem formulation . We formulate the interaction between environment and policy as a Markov Decision Process ( MDP ) [ 2 ] M = hS , A , P , r , i , where S is the state space , A is the action space , P : S ⇥ A 7 ! S is the state transition function that defines a conditional probability distribution P ( ·|st , at ) over all possible next states given a state st 2 S and action at 2 A taken at time t , r : S⇥A 7 ! R is a reward function , and 2 [ 0 , 1 ) is the discount factor . Because image observations only offer partial state observability [ 30 ] , we define a state st as a sequence of k + 1 consecutive frames ( ot , ot 1 , . . . , ot k ) , o 2 O , where O is the high-dimensional image space , as proposed in Mnih et al . [ 43 ] . The goal is then to learn a policy ⇡ : S 7 ! A that maximizes discounted return Rt = E ⇠⇡ [ P T t=1 t r ( st , at ) ] along a trajectory = ( s0 , s1 , . . . , sT ) obtained by following policy ⇡ from an initial state s0 2 S to a state sT with state transitions sampled from P , and ⇡ is parameterized by a collection of learnable parameters ✓ . For clarity , we hereon generically denote parameterization with subscript , e.g . ⇡✓ . We further aim to learn parameters ✓ s.t . ⇡✓ generalizes well ( i.e. , obtains high discounted return ) to unseen MDPs , which is generally unfeasible without further assumptions about the structure of the space of MDPs . In this work , we focus on generalization to MDPs M = hS , A , P , r , i , where states st 2 S are constructed from observations ot 2 O , O ✓ O of a perturbed observation space O ( e.g . unseen visuals ) , and M ⇠M for a space of MDPs M. Deep Q-Learning . Common model-free off-policy RL algorithms aim to estimate an optimal stateaction value function Q⇤ : S ⇥ A 7 ! R as Q✓ ( s , a ) ⇡ Q⇤ ( s , a ) = max⇡✓ E [ Rt|st = s , at = a ] using function approximation . In practice , this is achieved by means of the single-step Bellman residual ⇣ r ( st , at ) + maxa0t Q tgt ( st+1 , a0t ) ⌘ Q✓ ( st , at ) [ 62 ] , where parameterizes a target state-action value function Qtgt . We can choose to minimize this residual ( also known as the temporal difference error ) directly wrt ✓ using a mean squared error loss , which gives us the objective LQ ( ✓ , ) = Est , at , st+1⇠B '' 1 2 ✓ r ( st , at ) + max a0t Q tgt ( st+1 , a 0 t ) ◆ Q✓ ( st , at ) 2 # , ( 1 ) where B is a replay buffer with transitions collected by a behavioral policy [ 39 ] . From here , we can derive a greedy policy directly by selecting actions at = argmaxat Q✓ ( st , at ) . While Qtgt = Q and periodically setting ✓ exactly recovers the objective of DQN [ 43 ] , several improvements have been proposed to improve stability of Eq . 1 , such as Double Q-learning [ 24 ] , Dueling Q-networks [ 73 ] , updating target parameters using a slow-moving average of the online Q-network [ 38 ] : n+1 ( 1 ⇣ ) n + ⇣✓n ( 2 ) for an iteration step n and a momentum coefficient ⇣ 2 ( 0 , 1 ] , and others [ 25 , 23 , 53 , 14 , 28 ] . As computing maxa0t Q tgt ( st+1 , a0t ) in Eq . 1 is intractable for large and continuous action spaces , a number of prominent actor-critic algorithms that additionally learn a policy ⇡✓ ( st ) ⇡ argmaxat Q✓ ( st , at ) have therefore been proposed [ 38 , 15 , 18 ] . Soft Actor-Critic ( SAC ) [ 18 ] is an off-policy actor-critic algorithm that learns a state-action value function Q✓ and a stochastic policy ⇡✓ ( and optionally a temperature parameter ) , where Q✓ is optimized using a variant of the objective in Eq . 1 and ⇡✓ is optimized using a -discounted maximum-entropy objective [ 84 ] . To improve stability , SAC is also commonly implemented using Double Q-learning and the slow-moving target parameters from Eq . 2 . We will in the remainder of this work describe our method in the context of a generic off-policy RL algorithm that learns a parameterized state-action value function Q✓ , while we in our experiments discussed in Section 6 evaluate of our method using SAC as base algorithm . 4 Pitfalls of Data Augmentation in Deep Q-Learning In this section , we aim to illuminate the main causes of instability from naïve application of data augmentation in Q-value estimation . Our goal is to learn a Q-function Q✓ for an MDP M that generalizes to novel MDPs M ⇠M with unseen visuals , and we leverage data augmentation as an optimality-invariant state transformation ⌧ to induce a bisimulation relation [ 34 , 17 ] between a state s and its transformed ( augmented ) counterpart saug = ⌧ ( s , ⌫ ) with parameters ⌫ ⇠ V . Definition 1 ( Optimality-Invariant State Transformation [ 33 ] ) . Given an MDP M , a state transformation ⌧ : S ⇥ V 7 ! S is an optimality-invariant state transformation if Q ( s , a ) = Q ( ⌧ ( s , ⌫ ) , a ) 8 s 2 S , a 2 A , ⌫ 2 V , where ⌫ 2 V parameterizes the transformation ⌧ . Following our definitions of M , M from Section 3 , we can further extend the concept of optimalityinvariant transformations to MDPs , noting that a change of state space ( e.g . perturbed visuals ) itself can be described as a transformation ⌧ : S ⇥ V 7 ! S with unknown parameters ⌫ 2 V . If we choose the set of parameters V of a state transformation ⌧ to be sufficiently large such that it intersects with V with high probability , we can therefore expect to improve generalization to state and observation spaces not seen during training . However , while naïve application of data augmentation as in previous work [ 35 , 33 , 61 , 54 ] may potentially improve generalization , it can be harmful to Q-value estimation . We hypothesize that this is primarily because it dramatically increases the size of the observed state space , and consequently also increases variance Var [ Q ( ⌧ ( s , ⌫ ) ) ] Var [ Q ( s ) ] , ⌫ ⇠ V when V is large . Concretely , we identify the following two issues : Pitfall 1 : Non-deterministic Q-target . For deep Q-learning algorithms , previous work [ 35 , 33 , 61 , 54 ] applies augmentation to both state saug t , ⌧ ( st , ⌫ ) and successor state saugt+1 , ⌧ ( st+1 , ⌫0 ) where ⌫ , ⌫ 0 ⇠ V . Compared with DQN [ 43 ] that uses a deterministic ( more precisely , periodically updated ) Q-target , this practice introduces a non-deterministic Q-target r ( st , at ) + maxa0t Q tgt ( saug t+1 , a 0 t ) depending on the augmentation parameters ⌫0 . As observed in the original DQN paper , high-variance target values are detrimental to Q-learning algorithms , and may cause divergence due to the “ deadly triad ” of function approximation , bootstrapping , and off-policy learning [ 63 ] . This motivates the work to introduce a slowly changing target network , and several other works have refined the Q-target update rule [ 38 , 15 ] to further reduce volatility . However , because data augmentation is inherently non-deterministic , it can greatly increase variance in Q-target estimation and exacerbates the issue of volatility , as shown in Figure 1 ( top ) . This is particularly troubling in actor-critic algorithms such as DDPG [ 38 ] and SAC [ 18 ] , where the Q-target is estimated from ( st+1 , a0 ) , a0 ⇠ ⇡ ( ·|st+1 ) , which introduces an additional source of error from ⇡ that is non-negligible especially when st+1 is augmented . Pitfall 2 : Over-regularization . Data augmentation was originally introduced in the supervised learning regime as a regularizer to prevent overfitting of high-capacity models . However , for RL , even learning a policy in the training environment is hard . While data augmentation may improve generalization , it greatly increases the difficulty of policy learning , i.e. , optimizing ✓ for Q✓ and potentially a behavior network ⇡✓ . Particularly , when the temporal difference loss from Eq . 1 can not be well minimized , the large amount of augmented states dominate the gradient , which significantly impacts Q-value estimation of both augmented and unaugmented states . We refer to this issue as over-regularization by data augmentation . Figure 1 ( bottom ) shows the mean difference in Q-predictions made with augmented vs. unaugmented data in fully converged DrQ [ 33 ] agents trained with shift augmentation . Augmentations such as affine-jitter , random convolution , and random overlay incur large differences in estimated Q-values . While such difference can be reduced by regularizing the optimization with each individual augmentation , we emphasize that even the minimal shift augmentation used throughout training incurs non-zero difference . Since is commonly chosen to be a moving average of ✓ as in Eq . 2 , such differences caused by over-regularization affect Q✓ and Q tgt equally , and optimization may therefore still diverge depending on the choice of data augmentation . As such , there is an inherent trade-off between accurate Q-value estimation and generalization when using data augmentation . In the following section , we address these pitfalls .

Summary: This work proposes a simple technique to improve the robustness and sample-efficiency of reinforcement learning. The technique can be applied to any algorithm via a simple modification of the input pipeline to said algorithm's Q-learning subroutine. The claim is that this modification reduces the variance of Q-targets and therefore stabilizes learning.