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: The paper addresses the question of zero-shot or fast generalization of value-based deep reinforcement learning agents using visual inputs to image transformation which do not change the (conceptual) state of the environment. The long term goal is therefore training agents which learn efficient, high return control policies both which are robust to irrelevant variations in high dimensional input observations. The ultimate goal seems to be generalization of learned skills to novel environments, but as far as I can tell this particular evaluation philosophy is implemented in experiments.  The paper reviews previous attempts to learn robust state representations to visual variations and formulates clear hypotheses for explaining the causes of difficulties which prior methods faced and were hindered by. The said hypotheses suggest a simplified approach (SVEA) which is empirically developed and evaluated in the remainder of the paper. Control tasks from widely used benchmarks are used to evaluate the method, somewhat less usual is the focus on vision-based high-dimensional versions of such tasks, instead of standard domains such as Atari games and benchmarks for generalization such as Natural Atari [1], CoinRun/ProcGen [2].