Stabilizing Deep Q-Learning with ConvNets and Vision Transformers under Data Augmentation
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 efﬁciency and can even lead to divergence.
In this paper, we investigate causes of instability when using data augmentation in common offpolicy RL algorithms.
We identify two problems, both rooted in high-variance Q -targets.
Based on our ﬁndings, 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 efﬁciency of ConvNets under augmentation, and achieves generalization results competitive with state-of-the-art methods for image-based RL.
We further show that our method scales to RL with
ViT-based architectures, and that data augmentation may be especially important in this setting.
1. Introduction
Reinforcement Learning (RL) from visual observations has achieved tremendous success in various applications such as video-games (Mnih et al., 2013; Berner et al., 2019;
Vinyals et al., 2019), robotic manipulation (Levine et al., 2016) and autonomous navigation (Mirowski et al., 2017;
Zhu et al., 2017).
However, it is still very challenging for current methods to generalize the learned skills to novel environments, and policies trained by RL can easily overﬁt to the training environment (Zhang et al., 2018;
Farebrother et al., 2018), especially for high-dimensional observation spaces, e.g. images (Cobbe et al., 2019; Song et al., 2020).
Increasing variability in training data via domain randomization (Tobin et al., 2017; Pinto et al., 2017) and data augmentation (Shorten & Khoshgoftaar, 2019; Laskin et al., 2020;
Kostrikov et al., 2020; Raileanu et al., 2020) has demonstrated encouraging results for learning policies invariant to environment changes.
Speciﬁcally, recent works on data augmentation (Laskin et al., 2020; Kostrikov et al., 2020;
Hansen & Wang, 2021) show improvements in sample efﬁciency from simple cropping and translation augmentations, but also conclude that stronger data augmentation in fact decreases sample efﬁciency and even cause divergence.
While these augmentations have the potential to improve generalization, 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 theoretically grounded causes of instability when applying data augmentation to common off-policy RL algorithms (Mnih et al., 2013;
Lillicrap et al., 2016; Fujimoto et al., 2018;
Haarnoja et al., 2018).
Based on our ﬁndings, we provide an intuitive method for stabilizing this class of algorithms under use of strong data augmentation.
Speciﬁcally, we ﬁnd 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.
To address these problems, we propose SVEA :
S tabilized
Q - V alue E stimation under A ugmentation, 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 modiﬁed 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 indi- rectly through parameter-sharing.
Our framework can be implemented efﬁciently without additional forward passes nor introducing additional learnable parameters.
We perform extensive empirical evaluation on the DeepMind Control Suite (Tassa et al., 2018) and extensions of it, including the DMControl Generalization Benchmark (Hansen & Wang, 2021) and the Distracting Control Suite (Stone et al., 2021), as well as a set of robotic manipulation tasks.
Our method successfully stabilizes Q -value estimation with ConvNets under a set of strong data augmentations, and achieves sample efﬁciency, asymptotic performance, and generalization that is competitive or better than previous state-of-the-art methods.
Finally, we show that our method scales to RL with Vision Transformers (ViT) (Dosovitskiy et al., 2020).
We ﬁnd that ViT-based architectures are especially prone to overﬁtting, 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 (Pathak et al., 2016; Noroozi & Favaro, 2016; Zhang et al., 2016; Wu et al., 2018; van den Oord et al., 2018; Tian et al., 2019;
Xu et al., 2020;
He et al., 2020; Chen et al., 2020).
For example, Chen et al.
(Chen et al., 2020) 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 efﬁciency (Jaderberg et al., 2016; Shelhamer et al., 2017; van den
Oord et al., 2018; Yarats et al., 2019; Lin et al., 2019;
Srinivas et al., 2020; Stooke et al.;
Schwarzer et al., 2020;
Yarats et al., 2021) of image-based RL.
Recently, using self-supervision to improve generalization in RL has also gained interest (Zhang et al., 2020; Pathak et al., 2019;
Sekar et al., 2020; Agarwal et al., 2021; Hansen et al., 2020;
Hansen & Wang, 2021; Wang et al., 2021).
Notably, Zhang et al.
(Zhang et al., 2020) and Agarwal et al.
(Agarwal et al., 2021) propose to learn behavioral similarity embeddings via auxiliary tasks (bisimulation metrics and contrastive learning, respectively), and Hansen et al.
(Hansen & Wang, 2021) learn visual invariances through an auxiliary prediction task.
While these results are encouraging, it has also been shown in (Jaderberg et al., 2016; Lin et al., 2019;
Hansen et al., 2020; Yu et al., 2020; Lyle et al., 2021) 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 sampleefﬁciency 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 (Tobin et al., 2017; Pinto et al., 2017; Peng et al., 2018; Ramos et al., 2019; Chebotar et al., 2019) and data augmentation (Lee et al., 2019;
Cobbe et al., 2018; Wang et al., 2020; Laskin et al., 2020;
Kostrikov et al., 2020; Raileanu et al., 2020; Stooke et al.
Hansen & Wang, 2021).
For example, Tobin et al.
(Tobin et al., 2017) show that a neural network trained for object localization in a simulation with randomized visual augmentations improves real world generalization.
Similarly, Lee et al.(Lee et al., 2019) 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 (Laskin et al., 2020;
Kostrikov et al., 2020) have been conducted with RL, and conclude that, while small random crops and translations can improve sample efﬁciency, most data augmentations decrease sample efﬁciency 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) (Mnih et al., 2013) 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 (Sutton & Barto, 2018).
Therefore, a variety of algorithmic improvements have been proposed to improve convergence (Hasselt et al., 2016b;
Wang et al., 2016; Hausknecht & Stone, 2015;
Hasselt et al., 2016a; Schaul et al., 2016; Lillicrap et al., 2016; Fujimoto et al., 2018;
Fortunato et al., 2018; Hessel et al., 2018).
For example, Hasselt et al.
(Hasselt et al., 2016b) reduce overestimation of Q -values by decomposing the target Q -value estimation into action selection and action evaluation using separate networks.
Lillicrap et al.
(Lillicrap et al., 2016) reduce target variance by deﬁning the target Q -network as a slow-moving average of the online Q -network.
Our method also improves Q -value estimation, but we speciﬁcally 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) (Bellman, 1957)
M = (cid:104)S , A , P , r, γ (cid:105) , where
S is the state space, A is the action space, P : S × A (cid:55)→ S is the state transition function that deﬁnes a conditional probability distribution P ( ·| s t , a t ) over all possible next states given a state s t ∈ S and action a t ∈ A taken at time t , r :
S × A
(cid:55)→ R is a reward function, and γ ∈ [0 , 1) is the discount factor.
Because image observations only offer partial state observability (Kaelbling et al., 1998), we deﬁne a state s t as a sequence of k +consecutive frames ( o t , o t − 1 , . . .
, o t − k ) , o ∈ O , where
O is the high-dimensional image space, as is common practice (Mnih et al., 2013).
The goal is then to learn a policy π : S (cid:55)→
A that maximizes discounted return
R t =
E Γ ∼ π [ (cid:80) Tt =1 γ t r ( s t , a t )] along a trajectory Γ =
( s 0 , s 1 , . .
, s T ) obtained by following policy π from an initial state s 0 ∈ S to a state s T 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 novel MDPs, which is generally unfeasible without further assumptions about the structure of the space of MDPs.
In this work, we therefore consider generalization to MDPs M =
(cid:104)S , A , P , r, γ (cid:105) , where states s t ∈ S are constructed from observations o t ∈
O , O ⊆ O of a novel observation space O , and M ∼ M for a space of
MDPs M .
Deep Q -Learning.
Common model-free off-policy RL algorithms aim to estimate an optimal state-action value function Q ∗ :
S × A (cid:55)→ R as Q θ ( s , a ) ≈ Q ∗ ( s , a ) =
max π θ E
[ R t | s t = s , a t = a ] using function approximation.
In practice, this is achieved by means of the single-step
Bellman residual (cid:16) r ( s t , a
[Equation]
+ γ
[Equation]
[Equation]
Q θ ( s t , a t ) (Sutton, 2005), where ψ parameterizes a target state-action value function Q tgt .
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
[Equation]
[Equation]
where B is a replay buffer with transitions collected by a behavioral policy (Lin, 2004).
From here, we can derive a greedy policy directly by selecting actions a t = arg max a t Q θ ( s t , a t ) .
While Q tgt = Q and periodically setting ψ ←− θ exactly recovers the objective of DQN (Mnih et al., 2013), several improvements have been proposed to improve stability of Eq. 1, such as Double Q-learning (Hasselt et al., 2016b),
Dueling Q -networks (Wang et al., 2016), updating target parameters using a slow-moving average of the online Q -network (Lillicrap et al., 2016):
[Equation]
for an iteration step n and a momentum coefﬁcient
[Equation]
0 , 1] , and others (Hausknecht & Stone, 2015; Hasselt et al., 2016a; Schaul et al., 2016;
Fortunato et al., 2018; Hessel et al., 2018).
As computing max a (cid:48)
t Q tgt ψ ( s t +1 , a (cid:48) t ) in Eq.is intractable for large and continuous action spaces, a number of prominent actor-critic algorithms that additionally learn a policy π θ ( s t )
≈ arg max a t Q θ ( s t , a t ) have therefore been proposed (Lillicrap et al., 2016; Fujimoto et al., 2018;
Haarnoja et al., 2018).
Soft Actor-Critic (SAC) (Haarnoja et al., 2018) is an offpolicy 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 (Ziebart et al., 2008).
To improve stability, SAC is also commonly implemented using Double Q-learning and the slow-moving target parameters from Eq.
We will in the remainder of this work describe our method in the context of a generic offpolicy 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.
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 Qvalue estimation.
Our goal is to learn a Q -function Q θ for an
MDP M that generalizes to novel MDPs M ∼ M , and we leverage data augmentation as an optimality-invariant state transformation τ to induce a bisimulation relation (Larsen & Skou, 1989;
Givan et al., 2003) between a state s and its transformed (augmented) counterpart s aug = τ ( s , ν ) with parameters ν ∼ V .
Deﬁnition 1 (Optimality-Invariant State Transformation (Kostrikov et al., 2020)) .
Given an MDP M , a state transformation τ : S × V (cid:55)→ S is an optimality-invariant state transformation if Q ( s , a ) =
Q ( τ ( s , ν ) , a ) ∀ s ∈ S , a ∈
A , ν ∈ V , where ν ∈ V parameterizes τ .
Following our deﬁnitions of M , M from Section 3, we can further extend the concept of optimality-invariant transformations to MDPs, noting that a change of state space itself can be described as a transformation τ : S × V (cid:55)→ S with unknown parameters ν ∈ V .
If we choose the set of parameters V of a state transformation τ to be sufﬁciently large, we can therefore expect to improve generalization to state spaces not seen during training.
However, while na¨ıve application of data augmentation as in (Laskin et al., 2020; Kostrikov et al., 2020; Stooke et al.;
Schwarzer et al., 2020) may potentially improve generalization, it can be harmful to Q -value estimation.
We hypothesize that this
[Figure] a n [Figure] o r 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 Qlearning algorithms, previous work (Laskin et al., 2020;
Kostrikov et al., 2020; Stooke et al.; Schwarzer et al., 2020) applies augmentation to both state s aug t (cid:44)
τ
( s t , ν ) and successor state s aug t +1 (cid:44) τ ( s t +1 , ν (cid:48) ) where ν, ν (cid:48) ∼
V .
Compared with DQN (Mnih et al., 2013) that uses a deterministic (more precisely, periodically updated) Qtarget, this practice introduces a non-deterministic Q -target r ( s t , a t ) + γ max a (cid:48)
t Q tgt ψ ( s aug t +1 , a (cid:48) t ) depending on the augmentation parameters ν (cid:48) .
As observed in the original DQN paper, high-variance target values are detrimental to Qlearning algorithms, and may cause divergence due to the “deadly triad” of function approximation, bootstrapping, and off-policy learning (Sutton & Barto, 2018).
This motivates the work to introduce a slowly changing target network, and several other works have reﬁned the Q -target update rule (Lillicrap et al., 2016; Fujimoto et al., 2018) to further reduce volatility.
However, because data augmentation is inherently non-deterministic, it greatly increases variance in Q -target estimation and exacerbates the issue of volatility.
This is particularly troubling in actor-critic algorithms such as DDPG (Lillicrap et al., 2016) and SAC (Haarnoja et al., 2018), where the Q -target is estimated from ( s t +1 , a (cid:48) ) , a (cid:48) ∼ π ( ·| s t +1 ) , which introduces an additional source of error from π that is non-negligible especially when s t +1 is augmented.
Pitfall 2: Over-regularization.
Data augmentation was originally introduced in the supervised learning regime as a regularizer to prevent overﬁtting 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 difﬁculty 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 signiﬁcantly impacts Q -value estimation of both augmented and unaugmented states.
We refer to this issue as over-regularization by data augmentation.
Figure 1 shows the mean difference in Q -predictions made with augmented vs. unaugmented data in fully converged DrQ (Kostrikov et al., 2020)
agents trained with shift augmentation.
Augmentations such as afﬁne-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-negligible difference.
Since ψ is commonly chosen to be a moving average of θ as in Eq.
3, 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.
5. Method
We propose SVEA : S tabilized Q - V alue E stimation under
A ugmentation, a general framework for generalization by data augmentation in RL.
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 does not make any assumptions about the structure of states s t ∈ S , we here describe our method in the context of image-based RL.
5.1. Architectural Overview
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 subnetworks f θ (denoted the state encoder ) and Q θ (denoted the Q - function )
s.t q t (cid:44) Q
θ ( f θ ( s t ) , a t ) is the predicted Qvalue corresponding to a given state-action pair ( s t , a t ) .
We similarly deﬁne the target state-action value function s.t.
q tgt t (cid:44) r
( s t , a t )
+ γ max a (cid:48)
t Q tgt ψ ( f tgt ψ ( s t +1 ) , a (cid:48) ) is the target
Q -value for ( s t , a t ) , and we deﬁne parameters ψ as an exponential moving average of θ as in Eq. 3.
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 a t ∼ π θ ( ·| f θ ( s t )) .
To circumvent erroneous bootstrapping from augmented [Figure] data (as discussed in Section 4), we strictly apply data augmentation in Q -value estimation of the current state s t , without applying data augmentation to the successor state s t +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 novel Q -objective that leverages both augmented and unaugmented data, which we introduce in the following section.
5.2. Learning Objective
Our method redeﬁnes the temporal difference objective from
Eq.
1 to better leverage data augmentation.
First, recall that q tgt t = r
( s t , a t )
+ γ max a (cid:48)
t Q tgt ψ ( f tgt ψ ( s t +1 ) , a (cid:48) ) .
Instead of learning to predict q tgt t only from state s t , we propose to minimize a linear combination of L Q over two individual data streams, s t and s aug t =
τ ( s t , ν ) , ν ∼ V , which we deﬁne as the objective
[Equation]
where α, β are constant coefﬁcients that balance the ratio of the unaugmented and augmented data streams, respectively, and q tgt t is computed strictly from unaugmented data.
( θ, ψ ) serves as a data-mixing strategy that oversamples unaugmented data as an implicit variance reduction technique.
As we will verify empirically in Section 6, datamixing is a simple and effective technique for variance reduction that works well in tandem with our proposed modiﬁcations to bootstrapping.
For α = β , the objective in
Eq.
5 can be evaluated in a single, batched forward-pass by rewriting it as:
[Equation]
where [ · ] N is a concatenation operator along the batch dimension N for s t , s aug t ∈ R N × C × H × W and q tgt t ∈ R N × 1 , which is illustrated as ⊗ in Figure 2.
Empirically, we ﬁnd α = 0 .
5 , β = 0 .
5 to be both effective and practical to implement.
If the base algorithm of choice learns a policy π θ , its objective L π ( θ ) is optimized solely on unaugmented states s t 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.
5 or optionally Eq. 8 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 Q tgt ψ .
We summarize our method for α = β applied to a generic off-policy algorithm in Algorithm 1.
6. Experiments
We evaluate both sample efﬁciency, asymptotic performance, and generalization of our method and a set of strong baselines in tasks from DeepMind Control Suite (DMControl)
(Tassa et al., 2018) 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 (Hafner et al., 2019; 2020; Yarats et al., 2019; Srinivas et al., 2020;
Laskin et al., 2020; Kostrikov et al., 2020).
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) (Hansen &
Wang, 2021), the Distracting Control Suite (DistractingCS) (Stone et al., 2021), as well as distribution shifts unique to the robotic manipulation environment.
Code is avail- (cid:66) ζ α, β : loss coefﬁcients, default: ( α = 0 .
5 , β = 0 . 5) 1: for timestep t = 1 ...
T do act: 2: a t ∼ π θ ( ·| f θ ( s t )) (cid:66) Sample action from policy 3: s (cid:48)
t ∼ P ( ·| s t , a t ) (cid:66)
Sample transition from environment 4:
B ← B ∪ ( s t , a t , r ( s t , a t ) , s (cid:48) t ) (cid:66) Add transition to replay buffer update: 5: { s i , a i , r ( s i , a i ) , s (cid:48) i | i
= 1 ...N } ∼ B (cid:66)
Sample batch of transitions 6: s i =
τ ( s i , ν i ) , s (cid:48)
i =
τ ( s (cid:48)
i , ν (cid:48) i ) , ν i , ν (cid:48)
i ∼ V (cid:73)
Na¨ıve application of data augmentation 7: for transition i = 1 ..
N do 8: θ π ←− θ π − η ∇ θ π L π ( s i ; θ π ) (if applicable) (cid:66)
Optimize π θ with SGD 9: q tgt i =
r ( s i , a i ) + γ max a (cid:48) i Q tgt ψ ( f tgt ψ ( s (cid:48) i ) , a (cid:48) i ) (cid:66) Compute Q -target 10: s aug i = τ ( s i , ν i ) , ν i ∼ V (cid:73) Apply stochastic data augmentation 11: g i = (cid:2)
s i ,
s aug
i (cid:3) N , h i = (cid:2) q tgt i , q tgt i (cid:3) N (cid:73) Pack data streams 12: θ ←− θ − η ∇ θ L SVEA Q ( g i ,
h
i ; θ, ψ ) (cid:73)
Optimize f θ and Q θ with SGD 13: ψ ←−
[Equation]
(cid:66) Update ψ using EMA of θ test training colors natural videos camera poses able at https://github.com/nicklashansen/ dmcontrol-generalization-benchmark .
Setup.
We implement our method and baselines using SAC (Haarnoja et al., 2018) as base algorithm, and we apply random shift augmentation to all methods by default, which makes our base algorithm equivalent to DrQ (Kostrikov et al., 2020).
Network architecture and hyperparameters are adopted from (Hansen & Wang, 2021), and observations are stacks of 3 RGB frames of size 84 × 84 × 3 .
In our
DMControl-GB and DistractingCS experiments, all methods are trained for 500k frames 2 and evaluated on the full set of tasks proposed in DMControl-GB.
For simplicity, we adopt the same experimental setup for robotic manipulation.
Baselines and data augmentations.
We benchmark our method against the following strong baselines: (1) CURL (Srinivas et al., 2020), a contrastive learning method for
RL; (2) RAD that applies a random crop; (3) DrQ that applies a random shift; (4) PAD
(Hansen et al., 2020)
[Table]
[Table] [Table] that adapts to test environments using self-supervision; (5)
SODA (Hansen & Wang, 2021) that applies data augmentation in auxiliary learning; as well as a number of ablations.
We experiment with a diverse set of data augmentations proposed in previous work on RL and computer vision, namely random shift (Kostrikov et al., 2020), random convolution (denoted conv ) (Lee et al., 2019), random overlay (Hansen & Wang, 2021), random cutout (Cobbe et al., 2018), Gaussian blur , random afﬁne-jitter , and random rotation (Laskin et al., 2020; Gidaris et al., 2018).
6.1. Stability and Generalization on DMControl
We evaluate sample efﬁciency, asymptotic performance, and generalization of SVEA, DrQ, and a set of ablations across all 5 tasks from DMControl-GB, and report training and test curves in Figure 4.
SVEA outperforms all baselines on the test environment – often by a large margin – and maintains a sample efﬁciency comparable to DrQ trained without strong augmentation, while DrQ degrades substantially when using the additional conv augmentation.
We examine the reason for SVEA’s success with the following two ablations: a data-mixing only variant that applies our data-mixing strategy in both Q θ , Q tgt ψ , and π θ , as well as a variant of SVEA that only applies augmentation in Q θ but does not apply data-mixing.
We ﬁnd that both components of SVEA are necessary to achieve both stability and the generalization beneﬁts of strong data augmentation.
We further evaluate the stability of SVEA and DrQ undercommon data augmentations; results are shown in Figure 5.
While the sample efﬁciency of DrQ degrades substantially for most augmentations, SVEA is relatively unaffected by the choice of data augmentation and improves sample efﬁ-
[Figure] [Figure]
ciency in all 18 instances .
See supplementary material for a per-augmentation comparison.
Because we empirically ﬁnd the conv augmentation to be particularly difﬁcult to optimize, we provide additional stability experiments in Section 6.2 and 6.3 using this augmentation.
To benchmark the generalization ability of SVEA, we compare its test performance to 5 recent state-of-the-art methods for image-based RL on the challenging color hard and video easy benchmarks from DMControl-GB, and report results in Table 1.
All methods use the same architecture and hyperparameters whenever applicable, and we here use conv and overlay augmentations for fair comparison to
SODA.
SVEA outperforms all methods considered in 9 out of 10 instances, and at a signiﬁcantly lower computational cost than CURL, PAD, and SODA that learn auxiliary tasks.
We further evaluate generalization on DistractingCS, an extremely challenging benchmark for generalization where camera pose, background, and colors are continually changing throughout an episode. Figure 6 (top) shows a compari-
Test
) natural video backgrounds as visual distraction. Results for CURL, RAD, PAD, and SODA are obtained from (Hansen & Wang, 2021) and we report mean and std. deviation of 5 seeds.
SVEA achieves competitive results in all tasks. [Table] son of SVEA with conv and overlay augmentations to DrQ over a range of gradually increasing intensity of distractions, averaged across all 5 tasks from DMControl-GB.
SVEA improves generalization by 42 % at low intensity, and degrades signiﬁcantly slower than DrQ for high intensities.
6.2. RL with Vision Transformers
Vision Transformers (ViT) (Dosovitskiy et al., 2020) 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
[Figure]
with ViT encoders.
Results are shown in Figure 6 (bottom) [Figure]
We are, to the best of our knowledge, the ﬁrst to successfully solve image-based RL tasks without CNNs.
We observe that DrQ overﬁts signiﬁcantly to the training environment compared to its CNN counterpart.
SVEA achieves comparable sample efﬁciency 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 a CNN).
SVEA might therefore be a promising technique for future RL studies with CNN-free architectures, where data augmentation appears to be especially important.
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 conﬁguration
[Figure]
[Figure]
[Figure] [Table] [Table] [Figure]
[Table]
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 7.
Results are shown in Figure 8 and Table 2.
SVEA trained with conv augmentation exhibits similar stability and sample efﬁciency as DrQ trained without, while
DrQ + conv is found to have poor sample efﬁciency and fails to solve the challenging push task.
SVEA outperforms both baselines in terms of generalization.
Interestingly, we ﬁnd that na¨ıve application of data augmentation has a higher success rate in test environments than the DrQ baseline, 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 efﬁciency 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 largescale RL experiments where data augmentation is expected to play an increasingly important role.