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

Abstract:
While agents trained by Reinforcement Learning (RL) can solve increasingly challenging tasks directly from visual observations, generalizing learned skills to novel environments remains very challenging. Extensive use of data augmentation is a promising technique for improving generalization in RL, but it is often found to decrease sample efficiency and can even lead to divergence. In this paper, we investigate causes of instability when using data augmentation in common off-policy RL algorithms. We identify two problems, both rooted in high-variance Q-targets. Based on our findings, we propose a simple yet effective technique for stabilizing this class of algorithms under augmentation. We perform extensive empirical evaluation of image-based RL using both ConvNets and Vision Transformers (ViT) on a family of benchmarks based on DeepMind Control Suite, as well as in robotic manipulation tasks. Our method greatly improves stability and sample efficiency of ConvNets under augmentation, and achieves generalization results competitive with state-of-the-art methods for image-based RL in environments with unseen visuals. We further show that our method scales to RL with ViT-based architectures, and that data augmentation may be especially important in this setting.†

heading: 1 Introduction
text: 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.

heading: 2 Related Work
text: 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.

heading: 3 Preliminaries
text: 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 cannot 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.

heading: 5 Method
text: We propose SVEA: Stabilized Q-Value Estimation under Augmentation, a general framework for visual generalization in RL by use of data augmentation. SVEA applies data augmentation in a novel learning framework leveraging two data streams – with and without augmented data, respectively. Our method is compatible with any standard off-policy RL algorithm without changes to the underlying neural network that parameterizes the policy, and it requires no additional forward passes, auxiliary tasks, nor learnable parameters. While SVEA in principle does not make any assumptions about the structure of states st 2 S , we here describe our method in the context of image-based RL.

heading: 5.1 Architectural Overview
text: An overview of the SVEA architecture is provided in Figure 2. Our method leverages properties of common neural network architectures used in off-policy RL without introducing additional learnable parameters. We subdivide the neural network layers and corresponding learnable parameters of a state-action value function into sub-networks f✓ (denoted the state encoder) and Q✓ (denoted the Q-function) s.t qt , Q✓(f✓(st),at) is the predicted Q-value corresponding to a given stateaction pair (st,at). We similarly define the target state-action value function s.t. q tgt t
, r(st,at) + maxa0t Q tgt (f tgt (st+1),a0) is the target Q-value for (st,at), and we define parameters as an exponential moving average of ✓ as in Eq. 2. Depending on the choice of underlying algorithm, we may choose to additionally learn a parameterized policy ⇡✓ that shares encoder parameters with Q✓ and selects actions at ⇠ ⇡✓(·|f✓(st)).
To circumvent erroneous bootstrapping from augmented data (as discussed in Section 4), we strictly apply data augmentation in Q-value estimation of the current state st, without applying data augmentation to the successor state st+1 used in Eq. 1 for bootstrapping with Q tgt
(and ⇡✓ if applicable), which addresses Pitfall 1. If ⇡✓ is learned (i.e., SVEA is implemented with an actor-critic algorithm), we also optimize it strictly from unaugmented data. To mitigate over-regularization in optimization of f✓ and Q✓ (Pitfall 2), we further employ a modified Q-objective that leverages both augmented and unaugmented data, which we introduce in the following section.

heading: 5.2 Learning Objective
text: Our method redefines the temporal difference objective from Eq. 1 to better leverage data augmentation. First, recall that qtgt
t = r(st,at) + maxa0t Q
tgt (f tgt (st+1),a0). Instead of learning to predict
q tgt t
only from state st, we propose to minimize a nonnegative linear combination of LQ over two individual data streams, st and s aug t = ⌧(st, ⌫), ⌫ ⇠ V , which we define as the objective
L SVEA Q (✓, ) , ↵LQ st, q tgt t ; ✓,
+ LQ saug t , q tgt t ; ✓,
(3)
= Est,at,st+1⇠B h ↵ Q✓(f✓(st),at) qtgtt 2 2 +
Q✓(f✓(saugt ),at) q tgt t 2 2
i , (4)
where ↵, are constant coefficients that balance the ratio of the unaugmented and augmented data streams, respectively, and qtgt
t is computed strictly from unaugmented data. LSVEA Q (✓, ) serves as a
data-mixing strategy that oversamples unaugmented data as an implicit variance reduction technique.
As we will verify empirically in Section 6, data-mixing is a simple and effective technique for variance reduction that works well in tandem with our proposed modifications to bootstrapping. For ↵ = , the objective in Eq. 4 can be evaluated in a single, batched forward-pass by rewriting it as:
gt = [st, ⌧(st, ⌫)]N (5) ht = ⇥ q tgt t , q tgt t ⇤ N (6)
L SVEA Q (✓, ) = Est,at,st+1⇠B, ⌫⇠V h (↵+ ) kQ✓(f✓(gt),at) htk 2 2 i , (7)
where [·]N is a concatenation operator along the batch dimension N for st, s aug t 2 RN⇥C⇥H⇥W and qtgt t 2 RN⇥1, which is illustrated as ⌦ in Figure 2. Empirically, we find ↵ = 0.5, = 0.5 to be both effective and practical to implement, which we adopt in the majority of our experiments. However, more sophisticated schemes for selecting ↵, and/or varying them as training progresses could be interesting directions for future research. If the base algorithm of choice learns a policy ⇡✓, its objective L⇡(✓) is optimized solely on unaugmented states st without changes to the objective, and a stop-grad operation is applied after f✓ to prevent non-stationary gradients of L⇡(✓) from interfering with Q-value estimation, i.e., only the objective from Eq. 4 or optionally Eq. 7 updates f✓ using stochastic gradient descent. As described in Section 5.1, parameters are updated using an exponential moving average of ✓ and a stop-grad operation is therefore similarly applied after Qtgt .
We summarize our method for ↵ = applied to a generic off-policy algorithm in Algorithm 1.
Algorithm 1 Generic SVEA off-policy algorithm (I naïve augmentation, I our modifications) ✓, ✓⇡, : randomly initialized network parameters, ✓ B Initialize to be equal to ✓ ⌘, ⇣: learning rate and momentum coefficient ↵, : loss coefficients, default: (↵ = 0.5, = 0.5)
1: for timestep t = 1...T do act: 2: at ⇠ ⇡✓ (·|f✓(st)) B Sample action from policy 3: s0
t ⇠ P(·|st,at) B Sample transition from environment
4: B B [ (st,at, r(st,at), s0t) B Add transition to replay buffer update: 5: {si,ai, r(si,ai), s0i | i = 1...N} ⇠ B B Sample batch of transitions 6: si = ⌧(si, ⌫i), s0i = ⌧(s 0 i , ⌫ 0 i ), ⌫i, ⌫0i ⇠ V I Naïve application of data augmentation 7: for transition i = 1..N do 8: ✓⇡ ✓⇡ ⌘r✓⇡L⇡ (si; ✓⇡) (if applicable) B Optimize ⇡✓ with SGD 9: qtgt
i = r(si,ai) + maxa0i Q
tgt (f tgt (s0 i ),a0 i ) B Compute Q-target
10: saug i = ⌧(si, ⌫i), ⌫i ⇠ V I Apply stochastic data augmentation 11: gi = ⇥ si, s aug i ⇤ N , hi = ⇥ q tgt i , q tgt i ⇤ N I Pack data streams 12: ✓ ✓ ⌘r✓LSVEAQ (gi, hi; ✓, ) I Optimize f✓ and Q✓ with SGD 13: (1 ⇣) + ⇣✓ B Update using EMA of ✓
6 Experiments
We evaluate both sample efficiency, asymptotic performance, and generalization of our method and a set of strong baselines using both ConvNets and Vision Transformers (ViT) [10] in tasks from DeepMind Control Suite (DMControl) [64] as well as a set of robotic manipulation tasks. DMControl offers challenging and diverse continuous control tasks and is widely used as a benchmark for image-based
RL [19, 20, 76, 59, 35, 33]. To evaluate generalization of our method and baselines, we test methods under challenging distribution shifts (as illustrated in Figure 3) from the DMControl Generalization Benchmark (DMControl-GB) [21], the Distracting Control Suite (DistractingCS) [60], as well as distribution shifts unique to the robotic manipulation environment. Code is available at https://github.com/nicklashansen/dmcontrol-generalization-benchmark.
Setup. We implement our method and baselines using SAC [18] as base algorithm, and we apply random shift augmentation to all methods by default. This makes our base algorithm equivalent to DrQ [33] when K=1,M=1; we refer to the base algorithm as unaugmented and consider stability under additional data augmentation. We use the same network architecture and hyperparameters for all methods (whenever applicable), and adopt the setup from Hansen and Wang [21]. Observations are stacks of 3 RGB frames of size 84⇥ 84⇥ 3 (and 96⇥ 96⇥ 3 in ViT experiments). In the DMControlGB and DistractingCS benchmarks, all methods are trained for 500k frames and evaluated on all 5 tasks from DMControl-GB used in prior work, and we adopt the same experimental setup for robotic manipulation. See Appendix H for hyperparameters and further details on our experimental setup.
Baselines and data augmentations. We benchmark our method against the following strong baselines: (1) CURL [59], a contrastive learning method for RL; (2) RAD that applies a random crop; (3) DrQ that applies a random shift; (4) PAD [22] that adapts to test environments using self-supervision; (5) SODA [21] that applies data augmentation in auxiliary learning; as well as a number of ablations. We compare to the K=1,M=1 setting of DrQ by default, but also provide comparison to varying K,M . We experiment with a diverse set of data augmentations proposed in previous work on RL and computer vision, namely random shift [33], random convolution (denoted conv) [36], random overlay [21], random cutout [9], Gaussian blur, random affine-jitter, and random rotation [35, 16]. We provide samples for all data augmentations in Appendix C and test environments in Appendix E.

heading: 6.1 Stability and Generalization on DMControl
text: Stability. We evaluate the stability of SVEA and DrQ under 6 common data augmentations; results are shown in Figure 4. While the sample efficiency of DrQ degrades substantially for most augmentations, SVEA is relatively unaffected by the choice of data augmentation and improves sample efficiency in 27 out of 30 instances. While the sample efficiency of DrQ can be improved by increasing its K,M parameters, we find that DrQ requires approx. 6x wall-time to match the sample efficiency of SVEA; see Figure 6 (left). We further ablate each component of SVEA and report both training and test curves in Figure 5; we find that both components are key to SVEA’s success. Because we empirically find the conv augmentation to be particularly difficult to optimize, we provide additional stability experiments in Section 6.2 and 6.3 using this augmentation. See Appendix A for additional ablations.
Generalization. We compare the test performance of SVEA to 5 recent state-of-the-art methods for image-based RL on the color_hard and video_easy benchmarks from DMControl-GB (results in Table 1), as well as the extremely challenging DistractingCS benchmark, where camera pose, background, and colors are continually changing throughout an episode (results in Figure 6 (right)). We here use conv and overlay augmentations for fair comparison to SODA, and we report additional results on the video_hard benchmark in Appendix F. SVEA outperforms all methods considered in 12 out of 15 instances on DMControl-GB, and at a lower computational cost than CURL, PAD, and SODA that all learn auxiliary tasks. On DistractingCS, we observe that SVEA improves generalization by 42% at low intensity, and its generalization degrades significantly slower than DrQ for high intensities. While generalization depends on the particular choice of data augmentation and test environments, this is an encouraging result considering that SVEA enables efficient policy learning with stronger augmentations than previous methods.
6.2 RL with Vision Transformers
Vision Transformers (ViT) [10] have recently achieved impressive results on downstream tasks in computer vision. We replace all convolutional layers from the previous experiments with a 4-layer ViT encoder that operates on raw pixels in 8⇥ 8 space-time patches, and evaluate our method using data augmentation in conjunction with ViT encoders. Importantly, we design the ViT architecture such that it roughly matches our CNN encoder in terms of learnable parameters. The ViT encoder is trained from scratch using RL, and we use the same experimental setup as in our ConvNet experiments. In particular, it is worth emphasizing that both our ViT and CNN encoders are trained using Adam [32] as optimizer and without weight decay. See Figure 7 (top) for an architectural overview, and refer to Appendix H for additional implementation details.
Our training and test results are shown in Figure 7 (bottom). We are, to the best of our knowledge, the first to successfully solve image-based RL tasks without CNNs. We observe that DrQ overfits significantly to the training environment compared to its CNN counterpart (94 test return on color_hard for DrQ with ViT vs. 569 with a ConvNet on the Walker, walk task). SVEA achieves comparable sample efficiency and improves generalization by 706% and 233% on Walker, walk and Cartpole, swingup, respectively, over DrQ, while DrQ + conv remains unstable. Interestingly, we observe that our ViT-based implementation of SVEA achieves a mean episode return of 877 on the color_hard benchmark of the challenging Walker, walk task (vs. 760 using CNNs). SVEA might therefore be a promising technique for fu-
ture research on RL with CNN-free architectures, where data augmentation appears to be especially important for generalization. We provide additional experiments with ViT encoders in Section 6.3 and make further comparison to ConvNet encoders in Appendix A.
6.3 Robotic Manipulation
We additionally consider a set of goal-conditioned robotic manipulation tasks using a simulated Kinova Gen3 arm: (i) reach, a task in which the robot needs to position its gripper above a goal indicated by a red mark; (ii) reach moving target, a task similar to (i) but where the robot needs to follow a red mark moving continuously in a zig-zag pattern at a random velocity; and (iii) push, a task in which the robot needs to push a cube to a red mark. The initial configuration of gripper, object, and goal is randomized, the agent uses 2D positional control, and policies are trained using dense rewards. Observations are stacks of RGB frames with no access to state information. Training and test environments are shown in Figure 8. See Appendix G for further details and environment samples.
Results are shown in Figure 9 and Figure 10. For both CNN and ViT encoders, SVEA trained with conv augmentation has similar sample efficiency and training performance as DrQ trained without augmentation, while DrQ + conv exhibits poor sample efficiency and fails to solve the push task. Generalization results are shown in Table 2. We find that naïve application of data augmentation has a higher success rate in test environments than DrQ, despite being less
successful in the training environment, which we conjecture is because it is optimized only from augmented data. Conversely, SVEA achieves high success rates during both training and testing.
Conclusion. SVEA is found to greatly improve both stability and sample efficiency under augmentation, while achieving competitive generalization results. Our experiments indicate that our method scales to ViT-based architectures, and it may therefore be a promising technique for large-scale RL experiments where data augmentation is expected to play an increasingly important role.
Broader Impact. While our contribution aims to reduce computational cost of image-based RL, we remain concerned about the growing ecological and economical footprint of deep learning – and RL in particular – with increasingly large models such as ViT; see Appendix J for further discussion.
Acknowledgments and Funding Transparency. This work was supported by grants from DARPA LwLL, NSF CCF-2112665 (TILOS), NSF 1730158 CI-New: Cognitive Hardware and Software Ecosystem Community Infrastructure (CHASE-CI), NSF ACI-1541349 CC*DNI Pacific Research Platform, NSF grant IIS-1763278, NSF CCF-2112665 (TILOS), as well as gifts from Qualcomm, TuSimple and Picsart.

Review summaries:
(Auxiliary Review only -- the authors need not respond to this).  This work proposes issues with data augmentation in RL and proposes a couple of techniques to address those issues. It applies augmentation only on the current state and not future state to address erroneous bootstrapping. It also proposes a Q objective that uses both the augmented and un-augmented data. Also, for actor critic experiments, the actor is optimised only with un-augmented data and learn a  policy that is general via sharing parameters between the networks. These are simple techniques that help the paper demonstrate good results in a range of tasks.  This is a clear paper with good set of experiments. The problem has been described clearly and the work proposes simple refinements to the data augmentation strategies used in RL. While I am convinced of the quality of the method itself and the experiments, this work could have benefitted from more rigour with the baselines to make a stronger case for the method. Hence, I am offering a score of 6 noting that the paper is marginally above acceptance threshold.

This paper studies the problem of high-variance Q-targets under data augmentation, and proposes a method to reduce the variance of the targets called SVEA. The key idea contains three parts: (1) Apply augmentation to the Q-value estimation of the current state and not the next state Q-targets for bootstrapping, and (2) modify the objective to include both the targets using augmented and original observations and (3) optimize the actor using only the unaugmented data for actor-critic algorithms, with parameter sharing.   The paper performed experiments in DeepMind control suite and its variants (with raw pixel inputs). Compared to other data augmentation methods (e.g. DrQ) their proposed SVEA method is robust to the type of data augmentation and is more sample efficient with a better final performance. Some ablations were also presented among the above components to indicate the usefulness of all components. Generalization performance was demonstrated in the DistractingCS and robotic manipulation environments. The method also appears to help when using a Vision Transformers architecture.  

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.

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].  

The paper presents a way of improving training stability of DrQ (an image-based RL algorithm that uses data augmentation) by applying image augmentation more carefully. The authors propose to input a mix of augmented and un-augmented observations into the critic, and stop augmenting the actor's inputs. On the considered benchmarks the proposed method demonstrates better generalization results.