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

Nicklas Hansen1 Hao Su1 Xiaolong Wang1
1University of California, San Diego
<EMAIL_ADDRESS><EMAIL_ADDRESS>
###### 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.222Website and code is available at:
https://nicklashansen.github.io/SVEA.

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

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
actor-critic 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]
$\mathcal{M}=\langle\mathcal{S},\mathcal{A},\mathcal{P},r,\gamma\rangle$,
where $\mathcal{S}$ is the state space, $\mathcal{A}$ is the action space,
$\mathcal{P}\colon\mathcal{S}\times\mathcal{A}\mapsto\mathcal{S}$ is the state
transition function that defines a conditional probability distribution
$\mathcal{P}\left(\cdot|\mathbf{s}_{t},\mathbf{a}_{t}\right)$ over all
possible next states given a state $\mathbf{s}_{t}\in\mathcal{S}$ and action
$\mathbf{a}_{t}\in\mathcal{A}$ taken at time $t$,
$r\colon\mathcal{S}\times\mathcal{A}\mapsto\mathbb{R}$ is a reward function,
and $\gamma\in[0,1)$ is the discount factor. Because image observations only
offer partial state observability [30], we define a state $\mathbf{s}_{t}$ as
a sequence of $k+1$ consecutive frames
$(\mathbf{o}_{t},\mathbf{o}_{t-1},\dots,\mathbf{o}_{t-k}),~{}\mathbf{o}\in\mathcal{O}$,
where $\mathcal{O}$ is the high-dimensional image space, as proposed in Mnih
et al. [43]. The goal is then to learn a policy
$\pi\colon\mathcal{S}\mapsto\mathcal{A}$ that maximizes discounted return
$R_{t}=\mathbb{E}_{\Gamma\sim\pi}[\sum_{t=1}^{T}\gamma^{t}r(\mathbf{s}_{t},\mathbf{a}_{t})]$
along a trajectory
$\Gamma=(\mathbf{s}_{0},\mathbf{s}_{1},\dots,\mathbf{s}_{T})$ obtained by
following policy $\pi$ from an initial state $\mathbf{s}_{0}\in\mathcal{S}$ to
a state $\mathbf{s}_{T}$ with state transitions sampled from $\mathcal{P}$,
and $\pi$ is parameterized by a collection of learnable parameters $\theta$.
For clarity, we hereon generically denote parameterization with subscript,
e.g. $\pi_{\theta}$. We further aim to learn parameters $\theta$ s.t.
$\pi_{\theta}$ 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
$\overline{\mathcal{M}}=\langle\overline{\mathcal{S}},\mathcal{A},\mathcal{P},r,\gamma\rangle$,
where states $\overline{\mathbf{s}}_{t}\in\overline{\mathcal{S}}$ are
constructed from observations
$\overline{\mathbf{o}}_{t}\in\overline{\mathcal{O}},~{}\mathcal{O}\subseteq\overline{\mathcal{O}}$
of a perturbed observation space $\overline{\mathcal{O}}$ (e.g. unseen
visuals), and $\overline{\mathcal{M}}\sim\mathbb{M}$ for a space of MDPs
$\mathbb{M}$.

Deep $Q$-Learning. Common model-free off-policy RL algorithms aim to estimate
an optimal state-action value function
$Q^{*}\colon\mathcal{S}\times\mathcal{A}\mapsto\mathbb{R}$ as
$Q_{\theta}(\mathbf{s},\mathbf{a})\approx
Q^{*}(\mathbf{s},\mathbf{a})=\max_{\pi_{\theta}}\mathbb{E}\left[R_{t}|\mathbf{s}_{t}=\mathbf{s},\mathbf{a}_{t}=\mathbf{a}\right]$
using function approximation. In practice, this is achieved by means of the
single-step Bellman residual
$\left(r(\mathbf{s}_{t},\mathbf{a}_{t})+\gamma\max_{\mathbf{a}^{\prime}_{t}}Q^{\textnormal{tgt}}_{\psi}(\mathbf{s}_{t+1},\mathbf{a}^{\prime}_{t})\right)-Q_{\theta}(\mathbf{s}_{t},\mathbf{a}_{t})$
[62], where $\psi$ parameterizes a target state-action value function
$Q^{\textnormal{tgt}}$. We can choose to minimize this residual (also known as
the temporal difference error) directly wrt $\theta$ using a mean squared
error loss, which gives us the objective

$\mathcal{L}_{Q}(\theta,\psi)=\mathbb{E}_{\mathbf{s}_{t},\mathbf{a}_{t},\mathbf{s}_{t+1}\sim\mathcal{B}}\left[\frac{1}{2}\left[\left(r(\mathbf{s}_{t},\mathbf{a}_{t})+\gamma\max_{\mathbf{a}^{\prime}_{t}}Q^{\textnormal{tgt}}_{\psi}(\mathbf{s}_{t+1},\mathbf{a}^{\prime}_{t})\right)-Q_{\theta}(\mathbf{s}_{t},\mathbf{a}_{t})\right]^{2}\right]\,,$
(1)

where $\mathcal{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
$\mathbf{a}_{t}=\arg\max_{\mathbf{a}_{t}}Q_{\theta}(\mathbf{s}_{t},\mathbf{a}_{t})$.
While $Q^{\textnormal{tgt}}=Q$ and periodically setting
$\psi\longleftarrow\theta$ 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]:

$\psi_{n+1}\longleftarrow(1-\zeta)\psi_{n}+\zeta\theta_{n}$ (2)

for an iteration step $n$ and a momentum coefficient $\zeta\in(0,1]$, and
others [25, 23, 53, 14, 28]. As computing
$\max_{\mathbf{a}^{\prime}_{t}}Q^{\textnormal{tgt}}_{\psi}(\mathbf{s}_{t+1},\mathbf{a}^{\prime}_{t})$
in Eq. 1 is intractable for large and continuous action spaces, a number of
prominent actor-critic algorithms that additionally learn a policy
$\pi_{\theta}(\mathbf{s}_{t})\approx\arg\max_{\mathbf{a}_{t}}Q_{\theta}(\mathbf{s}_{t},\mathbf{a}_{t})$
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_{\theta}$ and a stochastic policy
$\pi_{\theta}$ (and optionally a temperature parameter), where $Q_{\theta}$ is
optimized using a variant of the objective in Eq. 1 and $\pi_{\theta}$ is
optimized using a $\gamma$-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_{\theta}$, 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_{\theta}$ for an MDP $\mathcal{M}$ that generalizes to
novel MDPs $\overline{\mathcal{M}}\sim\mathbb{M}$ with unseen visuals, and we
leverage data augmentation as an optimality-invariant state transformation
$\tau$ to induce a bisimulation relation [34, 17] between a state $\mathbf{s}$
and its transformed (augmented) counterpart
$\mathbf{s}^{\textnormal{aug}}=\tau(\mathbf{s},\nu)$ with parameters
$\nu\sim\mathcal{V}$.

######  Definition 1 (Optimality-Invariant State Transformation [33]).

Given an MDP $\mathcal{M}$, a state transformation
$\tau\colon\mathcal{S}\times\mathcal{V}\mapsto\mathcal{S}$ is an optimality-
invariant state transformation if
$Q(\mathbf{s},\mathbf{a})=Q(\tau(\mathbf{s},\nu),\mathbf{a})~{}~{}\forall~{}\mathbf{s}\in\mathcal{S},~{}\mathbf{a}\in\mathcal{A},~{}\nu\in\mathcal{V}$,
where $\nu\in\mathcal{V}$ parameterizes the transformation $\tau$.

Following our definitions of $\mathcal{M},\overline{\mathcal{M}}$ from Section
3, we can further extend the concept of optimality-invariant transformations
to MDPs, noting that a change of state space (e.g. perturbed visuals) itself
can be described as a transformation
$\overline{\tau}\colon\mathcal{S}\times\overline{\mathcal{V}}\mapsto\overline{\mathcal{S}}$
with unknown parameters $\overline{\nu}\in\overline{\mathcal{V}}$. If we
choose the set of parameters $\mathcal{V}$ of a state transformation $\tau$ to
be sufficiently large such that it intersects with $\overline{\mathcal{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
$\operatorname{Var}\left[Q(\tau(\mathbf{s},\nu))\right]\geq\operatorname{Var}\left[Q(\mathbf{s})\right],~{}\nu\sim\mathcal{V}$
when $\mathcal{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
$\mathbf{s}^{\textnormal{aug}}_{t}\triangleq\tau(\mathbf{s}_{t},\nu)$ and
successor state
$\mathbf{s}^{\textnormal{aug}}_{t+1}\triangleq\tau(\mathbf{s}_{t+1},\nu^{\prime})$
where $\nu,\nu^{\prime}\sim\mathcal{V}$. Compared with DQN [43] that uses a
deterministic (more precisely, periodically updated) $Q$-target, this practice
introduces a non-deterministic $Q$-target
$r(\mathbf{s}_{t},\mathbf{a}_{t})+\gamma\max_{\mathbf{a}^{\prime}_{t}}Q^{\textnormal{tgt}}_{\psi}(\mathbf{s}^{\textnormal{aug}}_{t+1},\mathbf{a}^{\prime}_{t})$
depending on the augmentation parameters $\nu^{\prime}$. 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
$(\mathbf{s}_{t+1},\mathbf{a}^{\prime}),~{}\mathbf{a^{\prime}}\sim\pi(\cdot|\mathbf{s}_{t+1})$,
which introduces an additional source of error from $\pi$ that is non-
negligible especially when $\mathbf{s}_{t+1}$ is augmented.

Figure 1: (Top) Mean $Q$-target variance of DrQ [33] and SVEA (ours), both
trained with conv augmentation [36]. (Bottom) Mean difference in $Q$-value
estimation on augmented vs. non-augmented data. We measure mean absolute error
in $Q$-value estimation from converged DrQ agents (trained with shift
augmentation) on the same observations before and after augmentation. Both
figures are averages across 5 seeds for each of the 5 tasks from DMControl-GB.

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 $\theta$
for $Q_{\theta}$ and potentially a behavior network $\pi_{\theta}$.
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 $\psi$ is commonly
chosen to be a moving average of $\theta$ as in Eq. 2, such differences caused
by over-regularization affect $Q_{\theta}$ and $Q^{\textnormal{tgt}}_{\psi}$
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.

Figure 2: Overview. An observation $\mathbf{s}_{t}$ is transformed by data
augmentation $\tau(\cdot,\nu),~{}\nu\sim\mathcal{V}$ to produce a view
$\mathbf{s}^{\textnormal{aug}}_{t}$. The $Q$-function $Q_{\theta}$ is then
jointly optimized on both augmented and unaugmented data wrt the objective in
Eq. 7, with the $Q$-target of the Bellman equation computed from an
unaugmented observation $\mathbf{s}_{t+1}$. We illustrate our data-mixing
strategy by the $\otimes$ operator.

##  5 Method

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
$\mathbf{s}_{t}\in\mathcal{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 sub-networks $f_{\theta}$ (denoted the state encoder) and
$Q_{\theta}$ (denoted the $Q$-function) s.t $q_{t}\triangleq
Q_{\theta}(f_{\theta}(\mathbf{s}_{t}),\mathbf{a}_{t})$ is the predicted
$Q$-value corresponding to a given state-action pair
$(\mathbf{s}_{t},\mathbf{a}_{t})$. We similarly define the target state-action
value function s.t. $q^{\textnormal{tgt}}_{t}\triangleq
r(\mathbf{s}_{t},\mathbf{a}_{t})+\gamma\max_{\mathbf{a}^{\prime}_{t}}Q^{\textnormal{tgt}}_{\psi}(f^{{}^{\textnormal{tgt}}}_{\psi}(\mathbf{s}_{t+1}),\mathbf{a}^{\prime})$
is the target $Q$-value for $(\mathbf{s}_{t},\mathbf{a}_{t})$, and we define
parameters $\psi$ as an exponential moving average of $\theta$ as in Eq. 2.
Depending on the choice of underlying algorithm, we may choose to additionally
learn a parameterized policy $\pi_{\theta}$ that shares encoder parameters
with $Q_{\theta}$ and selects actions
$\mathbf{a}_{t}\sim\pi_{\theta}(\cdot|f_{\theta}(\mathbf{s}_{t}))$.

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 $\mathbf{s}_{t}$, without applying data augmentation to the
successor state $\mathbf{s}_{t+1}$ used in Eq. 1 for bootstrapping with
$Q^{\textnormal{tgt}}_{\psi}$ (and $\pi_{\theta}$ if applicable), which
addresses Pitfall 1. If $\pi_{\theta}$ 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_{\theta}$ and
$Q_{\theta}$ (Pitfall 2), we further employ a modified $Q$-objective that
leverages both augmented and unaugmented data, which we introduce in the
following section.

###  5.2 Learning Objective

Our method redefines the temporal difference objective from Eq. 1 to better
leverage data augmentation. First, recall that
$q^{\textnormal{tgt}}_{t}=r(\mathbf{s}_{t},\mathbf{a}_{t})+\gamma\max_{\mathbf{a}^{\prime}_{t}}Q^{\textnormal{tgt}}_{\psi}(f^{{}^{\textnormal{tgt}}}_{\psi}(\mathbf{s}_{t+1}),\mathbf{a}^{\prime})$.
Instead of learning to predict $q^{\textnormal{tgt}}_{t}$ only from state
$\mathbf{s}_{t}$, we propose to minimize a nonnegative linear combination of
$\mathcal{L}_{Q}$ over two individual data streams,
${\color[rgb]{0,0.44140625,0.73828125}\mathbf{s}_{t}}$ and
${\color[rgb]{0.45,0.01,0.92}\mathbf{s}^{\textnormal{aug}}_{t}=\tau(\mathbf{s}_{t},\nu),~{}\nu\sim\mathcal{V}}$,
which we define as the objective

$\displaystyle\mathcal{L}^{\textbf{SVEA}}_{Q}(\theta,\psi)$
$\displaystyle\triangleq{\color[rgb]{0,0.44140625,0.73828125}\alpha\mathcal{L}_{Q}\left(\mathbf{s}_{t},q^{\textnormal{tgt}}_{t};\theta,\psi\right)}+{\color[rgb]{0.45,0.01,0.92}\beta\mathcal{L}_{Q}\left(\mathbf{s}^{\textnormal{aug}}_{t},q^{\textnormal{tgt}}_{t};\theta,\psi\right)}$
(3)
$\displaystyle=\mathbb{E}_{\mathbf{s}_{t},\mathbf{a}_{t},\mathbf{s}_{t+1}\sim\mathcal{B}}\left[{\color[rgb]{0,0.44140625,0.73828125}\alpha\left\|Q_{\theta}(f_{\theta}(\mathbf{s}_{t}),\mathbf{a}_{t})-q^{\textnormal{tgt}}_{t}\right\|^{2}_{2}}+{\color[rgb]{0.45,0.01,0.92}\beta\left\|Q_{\theta}(f_{\theta}(\mathbf{s}^{\textnormal{aug}}_{t}),\mathbf{a}_{t})-q^{\textnormal{tgt}}_{t}\right\|^{2}_{2}}\right]\,,$
(4)

where
${\color[rgb]{0,0.44140625,0.73828125}\alpha},{\color[rgb]{0.45,0.01,0.92}\beta}$
are constant coefficients that balance the ratio of the unaugmented and
augmented data streams, respectively, and $q^{\textnormal{tgt}}_{t}$ is
computed strictly from unaugmented data.
$\mathcal{L}^{\textbf{SVEA}}_{Q}(\theta,\psi)$ 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 $\alpha=\beta$, the objective
in Eq. 4 can be evaluated in a single, batched forward-pass by rewriting it
as:

$\displaystyle\mathbf{g}_{t}$
$\displaystyle=\left[{\color[rgb]{0,0.44140625,0.73828125}\mathbf{s}_{t}},{\color[rgb]{0.45,0.01,0.92}\tau(\mathbf{s}_{t},\nu)}\right]_{\textnormal{N}}$
(5) $\displaystyle h_{t}$
$\displaystyle=\left[q^{\textnormal{tgt}}_{t},q^{\textnormal{tgt}}_{t}\right]_{\textnormal{N}}$
(6) $\displaystyle\mathcal{L}^{\textbf{SVEA}}_{Q}(\theta,\psi)$
$\displaystyle=\mathbb{E}_{\mathbf{s}_{t},\mathbf{a}_{t},\mathbf{s}_{t+1}\sim\mathcal{B},~{}\nu\sim\mathcal{V}}\left[(\alpha+\beta)\left\|Q_{\theta}(f_{\theta}(\mathbf{g}_{t}),\mathbf{a}_{t})-h_{t}\right\|^{2}_{2}\right]\,,$
(7)

where $\left[\cdot\right]_{\textnormal{N}}$ is a concatenation operator along
the batch dimension $N$ for
$\mathbf{s}_{t},\mathbf{s}^{\textnormal{aug}}_{t}\in\mathbb{R}^{N\times
C\times H\times W}$ and $q^{\textnormal{tgt}}_{t}\in\mathbb{R}^{N\times 1}$,
which is illustrated as $\otimes$ in Figure 2. Empirically, we find
$\alpha=0.5,\beta=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 $\alpha,\beta$ and/or varying them as training
progresses could be interesting directions for future research. If the base
algorithm of choice learns a policy $\pi_{\theta}$, its objective
$\mathcal{L}_{\pi}(\theta)$ is optimized solely on unaugmented states
$\mathbf{s}_{t}$ without changes to the objective, and a stop-grad operation
is applied after $f_{\theta}$ to prevent non-stationary gradients of
$\mathcal{L}_{\pi}(\theta)$ from interfering with $Q$-value estimation, i.e.,
only the objective from Eq. 4 or optionally Eq. 7 updates $f_{\theta}$ using
stochastic gradient descent. As described in Section 5.1, parameters $\psi$
are updated using an exponential moving average of $\theta$ and a stop-grad
operation is therefore similarly applied after $Q^{\textnormal{tgt}}_{\psi}$.
We summarize our method for $\alpha=\beta$ applied to a generic off-policy
algorithm in Algorithm 1.

Algorithm 1 Generic SVEA off-policy algorithm ($\blacktriangleright$ naïve
augmentation, $\blacktriangleright$ our modifications)

1:$\theta,\theta_{\pi},\psi$: randomly initialized network parameters,
$\psi\longleftarrow\theta$ $\vartriangleright$ Initialize $\psi$ to be equal
to $\theta$

2:$\eta,\zeta$: learning rate and momentum coefficient

3:$\alpha,\beta$: loss coefficients, default: $(\alpha=0.5,\beta=0.5)$

4:for timestep $t=1...T$ do

5:act:

6:
$\mathbf{a}_{t}\sim\pi_{\theta}\left(\cdot|f_{\theta}(\mathbf{s}_{t})\right)$
$\vartriangleright$ Sample action from policy

7:
$\mathbf{s}^{\prime}_{t}\sim\mathcal{P}(\cdot|\mathbf{s}_{t},\mathbf{a}_{t})$
$\vartriangleright$ Sample transition from environment

8:
$\mathcal{B}\leftarrow\mathcal{B}\cup(\mathbf{s}_{t},\mathbf{a}_{t},r(\mathbf{s}_{t},\mathbf{a}_{t}),\mathbf{s}^{\prime}_{t})$
$\vartriangleright$ Add transition to replay buffer

9:update:

10:
$\\{\mathbf{s}_{i},\mathbf{a}_{i},r(\mathbf{s}_{i},\mathbf{a}_{i}),\mathbf{s}^{\prime}_{i}~{}|~{}i=1...N\\}\sim\mathcal{B}$
$\vartriangleright$ Sample batch of transitions

11:
$\mathbf{s}_{i}=\tau(\mathbf{s}_{i},\nu_{i}),~{}\mathbf{s}^{\prime}_{i}=\tau(\mathbf{s}^{\prime}_{i},\nu^{\prime}_{i}),~{}\nu_{i},\nu^{\prime}_{i}\sim\mathcal{V}$
$\blacktriangleright$ Naïve application of data augmentation

12: for transition $i=1..N$ do

13:
$\theta_{\pi}\longleftarrow\theta_{\pi}-\eta\nabla_{\theta_{\pi}}\mathcal{L}_{\pi}\left(\mathbf{s}_{i};\theta_{\pi}\right)$
(if applicable) $\vartriangleright$ Optimize $\pi_{\theta}$ with SGD

14:
$q^{\textnormal{tgt}}_{i}=r(\mathbf{s}_{i},\mathbf{a}_{i})+\gamma\max_{\mathbf{a}^{\prime}_{i}}Q^{\textnormal{tgt}}_{\psi}(f^{\textnormal{tgt}}_{\psi}(\mathbf{s}^{\prime}_{i}),\mathbf{a}^{\prime}_{i})$
$\vartriangleright$ Compute $Q$-target

15:
$\mathbf{s}^{\textnormal{aug}}_{i}=\tau(\mathbf{s}_{i},\nu_{i}),~{}\nu_{i}\sim\mathcal{V}$
$\blacktriangleright$ Apply stochastic data augmentation

16:
$\mathbf{g}_{i}=\left[\mathbf{s}_{i},\mathbf{s}^{\textnormal{aug}}_{i}\right]_{\textnormal{N}},~{}h_{i}=\left[q^{\textnormal{tgt}}_{i},q^{\textnormal{tgt}}_{i}\right]_{\textnormal{N}}$
$\blacktriangleright$ Pack data streams

17:
$\theta\longleftarrow\theta-\eta\nabla_{\theta}\mathcal{L}^{\textbf{SVEA}}_{Q}\left(\mathbf{g}_{i},h_{i};\theta,\psi\right)$
$\blacktriangleright$ Optimize $f_{\theta}$ and $Q_{\theta}$ with SGD

18: $\psi\longleftarrow(1-\zeta)\psi+\zeta\theta$ $\vartriangleright$ Update
$\psi$ using EMA of $\theta$

Figure 3: Experimental setup. Agents are trained in a fixed environment and
are expected to generalize to novel environments with e.g. random colors,
backgrounds, and camera poses.

##  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\times 84\times 3$ (and
$96\times 96\times 3$ in ViT experiments). In the DMControl-GB 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.

Figure 4: Data augmentations. Training performance of SVEA (top) and DrQ
(bottom) under 6 common data augmentations. Mean of 5 seeds. Red line at $800$
return is for visual guidance only. We omit visualization of std. deviations
for clarity, but provide per-augmentation comparisons to DrQ (including std.
deviations) across all tasks in Appendix B, and test performances in Appendix
C. Figure 5: Training and test performance. We compare SVEA to DrQ with and
without random convolution augmentation, as well as a set of ablations. Data-
mixing only indiscriminately applies our data-mixing strategy to all data
streams, and $(\alpha=0,\beta=1)$ only augments $Q$-predictions but without
data-mixing. We find both components to contribute to SVEA’s success. Top:
episode return on the training environment during training. Bottom:
generalization measured by episode return on the color_hard benchmark of
DMControl-GB. Mean of 5 seeds, shaded area is $\pm 1$ std. deviation.

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.

###  6.1 Stability and Generalization on DMControl

Table 1: Comparison to state-of-the-art. Test performance (episode return) of
methods trained in a single, fixed environment and evaluated on (i) randomized
colors, and (ii) natural video backgrounds from DMControl-GB. Results for
CURL, RAD, PAD, and SODA are obtained from [21] and we report mean and std.
deviation over 5 seeds. DrQ corresponds to our SAC base algorithm using random
shift augmentation. SVEA matches or outperforms prior methods in all tasks
considered.

DMControl-GB | CURL | RAD | DrQ | PAD | SODA | SODA | SVEA | SVEA
---|---|---|---|---|---|---|---|---
(random colors) |  |  |  |  | (conv) | (overlay) | (conv) | (overlay)
walker, | $445$ | $400$ | $520$ | $468$ | $697$ | $692$ | $\mathbf{760}$ | 749
walk | $\scriptstyle{\pm 99}$ | $\scriptstyle{\pm 61}$ | $\scriptstyle{\pm 91}$ | $\scriptstyle{\pm 47}$ | $\scriptstyle{\pm 66}$ | $\scriptstyle{\pm 68}$ | $\mathbf{\scriptstyle{\pm 145}}$ | $\scriptstyle{\pm 61}$
walker, | $662$ | $644$ | $770$ | $797$ | $930$ | $893$ | $\mathbf{942}$ | $933$
stand | $\scriptstyle{\pm 54}$ | $\scriptstyle{\pm 88}$ | $\scriptstyle{\pm 71}$ | $\scriptstyle{\pm 46}$ | $\scriptstyle{\pm 12}$ | $\scriptstyle{\pm 12}$ | $\mathbf{\scriptstyle{\pm 26}}$ | $\scriptstyle{\pm 24}$
cartpole, | $454$ | $590$ | $586$ | $630$ | $831$ | $805$ | $\mathbf{837}$ | 832
swingup | $\scriptstyle{\pm 110}$ | $\scriptstyle{\pm 53}$ | $\scriptstyle{\pm 52}$ | $\scriptstyle{\pm 63}$ | $\scriptstyle{\pm 21}$ | $\scriptstyle{\pm 28}$ | $\mathbf{\scriptstyle{\pm 23}}$ | $\scriptstyle{\pm 23}$
ball_in_cup, | $231$ | $541$ | $365$ | $563$ | $892$ | $949$ | $\mathbf{961}$ | $959$
catch | $\scriptstyle{\pm 92}$ | $\scriptstyle{\pm 29}$ | $\scriptstyle{\pm 210}$ | $\scriptstyle{\pm 50}$ | $\scriptstyle{\pm 37}$ | $\scriptstyle{\pm 19}$ | $\mathbf{\scriptstyle{\pm 7}}$ | $\scriptstyle{\pm 5}$
finger, | $691$ | $667$ | $776$ | $803$ | $901$ | $793$ | $\mathbf{977}$ | $972$
spin | $\scriptstyle{\pm 12}$ | $\scriptstyle{\pm 154}$ | $\scriptstyle{\pm 134}$ | $\scriptstyle{\pm 72}$ | $\scriptstyle{\pm 51}$ | $\scriptstyle{\pm 128}$ | $\mathbf{\scriptstyle{\pm 5}}$ | $\scriptstyle{\pm 6}$
DMControl-GB | CURL | RAD | DrQ | PAD | SODA | SODA | SVEA | SVEA
(natural videos) |  |  |  |  | (conv) | (overlay) | (conv) | (overlay)
walker, | $556$ | $606$ | $682$ | $717$ | $635$ | $768$ | $612$ | $\mathbf{819}$
walk | $\scriptstyle{\pm 133}$ | $\scriptstyle{\pm 63}$ | $\scriptstyle{\pm 89}$ | $\scriptstyle{\pm 79}$ | $\scriptstyle{\pm 48}$ | $\scriptstyle{\pm 38}$ | $\scriptstyle{\pm 144}$ | $\mathbf{\scriptstyle{\pm 71}}$
walker, | $852$ | $745$ | $873$ | $935$ | $903$ | $955$ | $795$ | $\mathbf{961}$
stand | $\scriptstyle{\pm 75}$ | $\scriptstyle{\pm 146}$ | $\scriptstyle{\pm 83}$ | $\scriptstyle{\pm 20}$ | $\scriptstyle{\pm 56}$ | $\scriptstyle{\pm 13}$ | $\scriptstyle{\pm 70}$ | $\mathbf{\scriptstyle{\pm 8}}$
cartpole, | $404$ | $373$ | $485$ | $521$ | $474$ | $758$ | $606$ | $\mathbf{782}$
swingup | $\scriptstyle{\pm 67}$ | $\scriptstyle{\pm 72}$ | $\scriptstyle{\pm 105}$ | $\scriptstyle{\pm 76}$ | $\scriptstyle{\pm 143}$ | $\scriptstyle{\pm 62}$ | $\scriptstyle{\pm 85}$ | $\mathbf{\scriptstyle{\pm 27}}$
ball_in_cup, | $316$ | $481$ | $318$ | $436$ | $539$ | $\mathbf{875}$ | $659$ | $871$
catch | $\scriptstyle{\pm 119}$ | $\scriptstyle{\pm 26}$ | $\scriptstyle{\pm 157}$ | $\scriptstyle{\pm 55}$ | $\scriptstyle{\pm 111}$ | $\mathbf{\scriptstyle{\pm 56}}$ | $\scriptstyle{\pm 110}$ | $\scriptstyle{\pm 106}$
finger, | $502$ | $400$ | $533$ | $691$ | $363$ | $695$ | $764$ | $\mathbf{808}$
spin | $\scriptstyle{\pm 19}$ | $\scriptstyle{\pm 64}$ | $\scriptstyle{\pm 119}$ | $\scriptstyle{\pm 80}$ | $\scriptstyle{\pm 185}$ | $\scriptstyle{\pm 97}$ | $\scriptstyle{\pm 86}$ | $\mathbf{\scriptstyle{\pm 33}}$

Figure 6: (Left) Comparison with additional DrQ baselines. We compare SVEA
implemented with DrQ [K=1,M=1] as base algorithm to DrQ with varying values of
its $K,M$ hyperparameters. All methods use the conv augmentation (in addition
to shift augmentation used by DrQ). Results are averaged over 5 seeds for each
of the 5 tasks from DMControl-GB [21] and shaded area is $\pm 1$ std.
deviation across seeds. Increasing values of $K,M$ improve sample efficiency
of DrQ, but at a high computational cost; DrQ uses approx. 6x wall-time to
match the sample efficiency of SVEA. (Right) DistractingCS. Episode return as
a function of randomization intensity at test-time, aggregated across 5 seeds
for each of the 5 tasks from DMControl-GB. See Appendix E for per-task
comparison.

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 $\mathbf{27}$ out of $\mathbf{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 $\mathbf{12}$ out of $\mathbf{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 $\mathbf{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.

Figure 7: (Top) ViT architecture. Observations are divided into $144$ non-
overlapping space-time patches and linearly projected into tokens. Each token
uses a learned positional encoding and we also use a learnable class token as
in [10]. The ViT encoder consists of 4 stacked Transformer encoders [69].
(Bottom) RL with a ViT encoder. Training and test performance of SVEA and DrQ
using ViT encoders. We report results for three tasks and test performance is
evaluated on the color_hard benchmark of DMControl-GB. Mean of 5 seeds, shaded
area is $\pm 1$ std. deviation.

###  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\times 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 ($\mathbf{94}$ test return on
color_hard for DrQ with ViT vs. $\mathbf{569}$ with a ConvNet on the Walker,
walk task). SVEA achieves comparable sample efficiency and improves
generalization by $\mathbf{706\%}$ and $\mathbf{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 $\mathbf{877}$ on the color_hard benchmark of the
challenging Walker, walk task (vs. $\mathbf{760}$ using CNNs). SVEA might
therefore be a promising technique for future 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.

Table 2: Generalization in robotic manipulation. Task success rates of SVEA
and DrQ with CNN and ViT encoders in the training environment, as well as
aggregated success rates across 25 different test environments with randomized
camera pose, colors, lighting, and background. Mean of 5 seeds.

Robotic | Arch. | reach | reach | mv.tgt. | mv.tgt. | push | push
---|---|---|---|---|---|---|---
manipulation | (encoder) | (train) | (test) | (train) | (test) | (train) | (test)
DrQ | CNN | $\mathbf{1.00}$ | $0.60$ | $\mathbf{1.00}$ | $0.69$ | $\mathbf{0.76}$ | $0.26$
DrQ + conv | CNN | $0.59$ | $0.77$ | $0.60$ | $0.89$ | $0.13$ | $0.12$
SVEA w/ conv | CNN | $\mathbf{1.00}$ | $\mathbf{0.89}$ | $\mathbf{1.00}$ | $\mathbf{0.96}$ | $0.72$ | $\mathbf{0.48}$
DrQ | ViT | $0.93$ | $0.14$ | $\mathbf{1.00}$ | $0.16$ | $0.73$ | $0.05$
DrQ + conv | ViT | $0.26$ | $0.67$ | $0.48$ | $\mathbf{0.82}$ | $0.08$ | $0.07$
SVEA w/ conv | ViT | $\mathbf{0.98}$ | $\mathbf{0.71}$ | $\mathbf{1.00}$ | $0.81$ | $\mathbf{0.82}$ | $\mathbf{0.17}$

Figure 8: Robotic manipulation. Agents are trained in a fixed environment and
evaluated on challenging environments with randomized colors, lighting,
background, and camera pose.

Figure 9: Stability with a CNN encoder. Training performance (episode return)
of SVEA and DrQ in 3 robotic manipulation tasks. Mean and std. deviation of 5
seeds. Success rates are shown in Table 2.

Figure 10: Stability with a ViT encoder. Training performance (episode return)
of SVEA and DrQ in 3 robotic manipulation tasks. Mean and std. deviation of 5
seeds. Success rates are shown in Table 2. DrQ is especially unstable under
augmentation when using a ViT encoder.

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

## References

  * Agarwal et al. [2021] Rishabh Agarwal, Marlos C. Machado, P. S. Castro, and Marc G. Bellemare.  Contrastive behavioral similarity embeddings for generalization in reinforcement learning.  _ArXiv_ , abs/2101.05265, 2021.
  * Bellman [1957] Richard Bellman.  A markovian decision process.  _Journal of Mathematics and Mechanics_ , 6(5):679–684, 1957.
  * Bender et al. [2021] Emily M. Bender, Timnit Gebru, Angelina McMillan-Major, and Shmargaret Shmitchell.  On the dangers of stochastic parrots: Can language models be too big?  In _Proceedings of the 2021 ACM Conference on Fairness, Accountability, and Transparency_. Association for Computing Machinery, 2021.
  * Berner et al. [2019] Christopher Berner, Greg Brockman, Brooke Chan, Vicki Cheung, Przemyslaw Debiak, Christy Dennison, David Farhi, Quirin Fischer, et al.  Dota 2 with large scale deep reinforcement learning.  _ArXiv_ , abs/1912.06680, 2019.
  * Brown et al. [2020] T. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, J. Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, et al.  Language models are few-shot learners.  _arXiv_ , abs/2005.14165, 2020.
  * Chebotar et al. [2019] Yevgen Chebotar, A. Handa, Viktor Makoviychuk, M. Macklin, J. Issac, Nathan D. Ratliff, and D. Fox.  Closing the sim-to-real loop: Adapting simulation randomization with real world experience.  _2019 International Conference on Robotics and Automation (ICRA)_ , pages 8973–8979, 2019.
  * Chen et al. [2020] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton.  A simple framework for contrastive learning of visual representations, 2020.
  * Cobbe et al. [2019] K. Cobbe, Oleg Klimov, Christopher Hesse, Taehoon Kim, and John Schulman.  Quantifying generalization in reinforcement learning.  In _Icml_ , 2019.
  * Cobbe et al. [2018] Karl Cobbe, Oleg Klimov, Chris Hesse, Taehoon Kim, and John Schulman.  Quantifying generalization in reinforcement learning, 2018.
  * Dosovitskiy et al. [2020] A. Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, M. Dehghani, Matthias Minderer, et al.  An image is worth 16x16 words: Transformers for image recognition at scale.  _ArXiv_ , abs/2010.11929, 2020.
  * Ebert et al. [2018] Frederik Ebert, Chelsea Finn, Sudeep Dasari, Annie Xie, Alex Lee, and Sergey Levine.  Visual foresight: Model-based deep reinforcement learning for vision-based robotic control, 2018.
  * Espeholt et al. [2018] Lasse Espeholt, Hubert Soyer, R. Munos, K. Simonyan, V. Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, et al.  Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures.  _ArXiv_ , abs/1802.01561, 2018.
  * Farebrother et al. [2018] Jesse Farebrother, Marlos C. Machado, and Michael H. Bowling.  Generalization and regularization in dqn.  _ArXiv_ , abs/1810.00123, 2018.
  * Fortunato et al. [2018] Meire Fortunato, M. G. Azar, Bilal Piot, Jacob Menick, Ian Osband, A. Graves, Vlad Mnih, R. Munos, et al.  Noisy networks for exploration.  _ArXiv_ , abs/1706.10295, 2018.
  * Fujimoto et al. [2018] Scott Fujimoto, H. V. Hoof, and D. Meger.  Addressing function approximation error in actor-critic methods.  _ArXiv_ , abs/1802.09477, 2018.
  * Gidaris et al. [2018] Spyros Gidaris, Praveer Singh, and Nikos Komodakis.  Unsupervised representation learning by predicting image rotations, 2018\.
  * Givan et al. [2003] R. Givan, T. Dean, and M. Greig.  Equivalence notions and model minimization in markov decision processes.  _Artificial Intelligence_ , 147:163–223, 2003.
  * Haarnoja et al. [2018] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine.  Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor.  _arXiv preprint arXiv:1801.01290_ , 2018.
  * Hafner et al. [2019] Danijar Hafner, T. Lillicrap, Ian S. Fischer, R. Villegas, David R Ha, H. Lee, and James Davidson.  Learning latent dynamics for planning from pixels.  _ArXiv_ , abs/1811.04551, 2019.
  * Hafner et al. [2020] Danijar Hafner, T. Lillicrap, Jimmy Ba, and Mohammad Norouzi.  Dream to control: Learning behaviors by latent imagination.  _ArXiv_ , abs/1912.01603, 2020.
  * Hansen and Wang [2021] Nicklas Hansen and Xiaolong Wang.  Generalization in reinforcement learning by soft data augmentation.  In _International Conference on Robotics and Automation_ , 2021.
  * Hansen et al. [2021] Nicklas Hansen, Rishabh Jangir, Yu Sun, Guillem Alenyà, Pieter Abbeel, Alexei A. Efros, Lerrel Pinto, and Xiaolong Wang.  Self-supervised policy adaptation during deployment.  In _International Conference on Learning Representations_ , 2021.
  * Hasselt et al. [2016a] H. V. Hasselt, A. Guez, Matteo Hessel, V. Mnih, and D. Silver.  Learning values across many orders of magnitude.  In _Nips_ , 2016a.
  * Hasselt et al. [2016b] H. V. Hasselt, A. Guez, and D. Silver.  Deep reinforcement learning with double q-learning.  In _Aaai_ , 2016b.
  * Hausknecht and Stone [2015] M. Hausknecht and P. Stone.  Deep recurrent q-learning for partially observable mdps.  In _AAAI Fall Symposia_ , 2015.
  * He et al. [2016] Kaiming He, X. Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  _2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 770–778, 2016.
  * He et al. [2020] Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross B. Girshick.  Momentum contrast for unsupervised visual representation learning.  _2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 9726–9735, 2020.
  * Hessel et al. [2018] Matteo Hessel, Joseph Modayil, H. V. Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, et al.  Rainbow: Combining improvements in deep reinforcement learning.  In _Aaai_ , 2018.
  * Jaderberg et al. [2016] Max Jaderberg, Volodymyr Mnih, Wojciech Marian Czarnecki, Tom Schaul, Joel Z Leibo, David Silver, and Koray Kavukcuoglu.  Reinforcement learning with unsupervised auxiliary tasks, 2016.
  * Kaelbling et al. [1998] Leslie Pack Kaelbling, Michael L. Littman, and Anthony R. Cassandra.  Planning and acting in partially observable stochastic domains.  _Artificial Intelligence_ , 1998.
  * Kalashnikov et al. [2018] D. Kalashnikov, A. Irpan, P. Pastor, J. Ibarz, A. Herzog, Eric Jang, Deirdre Quillen, Ethan Holly, et al.  Qt-opt: Scalable deep reinforcement learning for vision-based robotic manipulation.  _ArXiv_ , abs/1806.10293, 2018.
  * Kingma and Ba [2015] Diederik P. Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _CoRR_ , abs/1412.6980, 2015.
  * Kostrikov et al. [2020] Ilya Kostrikov, Denis Yarats, and Rob Fergus.  Image augmentation is all you need: Regularizing deep reinforcement learning from pixels.  _International Conference on Learning Representations_ , 2020.
  * Larsen and Skou [1989] K. G. Larsen and A. Skou.  Bisimulation through probabilistic testing (preliminary report).  In _Proceedings of the 16th ACM SIGPLAN-SIGACT Symposium on Principles of Programming Languages_ , 1989.
  * Laskin et al. [2020] Michael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas.  Reinforcement learning with augmented data.  _arXiv preprint arXiv:2004.14990_ , 2020.
  * Lee et al. [2019] Kimin Lee, Kibok Lee, Jinwoo Shin, and Honglak Lee.  A simple randomization technique for generalization in deep reinforcement learning.  _ArXiv_ , abs/1910.05396, 2019.
  * Levine et al. [2016] Sergey Levine, Chelsea Finn, Trevor Darrell, and Pieter Abbeel.  End-to-end training of deep visuomotor policies.  _The Journal of Machine Learning Research_ , 17(1):1334–1373, 2016.
  * Lillicrap et al. [2016] T. Lillicrap, J. Hunt, A. Pritzel, N. Heess, T. Erez, Y. Tassa, D. Silver, and Daan Wierstra.  Continuous control with deep reinforcement learning.  _CoRR_ , abs/1509.02971, 2016.
  * Lin [2004] L. J. Lin.  Self-improving reactive agents based on reinforcement learning, planning and teaching.  _Machine Learning_ , 8:293–321, 2004.
  * Lin et al. [2019] Xingyu Lin, Harjatin Singh Baweja, George Kantor, and David Held.  Adaptive auxiliary task weighting for reinforcement learning.  _Advances in neural information processing systems_ , 32, 2019.
  * Lyle et al. [2021] Clare Lyle, Mark Rowland, Georg Ostrovski, and Will Dabney.  On the effect of auxiliary tasks on representation dynamics.  In _Aistats_ , 2021.
  * Mirowski et al. [2017] P. Mirowski, Razvan Pascanu, Fabio Viola, Hubert Soyer, Andy Ballard, Andrea Banino, Misha Denil, R. Goroshin, et al.  Learning to navigate in complex environments.  _ArXiv_ , abs/1611.03673, 2017.
  * Mnih et al. [2013] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller.  Playing atari with deep reinforcement learning.  _arXiv preprint arXiv:1312.5602_ , 2013.
  * Nair et al. [2018] Ashvin V Nair, Vitchyr Pong, Murtaza Dalal, Shikhar Bahl, Steven Lin, and Sergey Levine.  Visual reinforcement learning with imagined goals.  In _Advances in Neural Information Processing Systems_ , pages 9191–9200, 2018.
  * Noroozi and Favaro [2016] Mehdi Noroozi and Paolo Favaro.  Unsupervised learning of visual representations by solving jigsaw puzzles.  In _European Conference on Computer Vision_ , pages 69–84. Springer, 2016.
  * Pathak et al. [2016] Deepak Pathak, Philipp Krahenbuhl, Jeff Donahue, Trevor Darrell, and Alexei A Efros.  Context encoders: Feature learning by inpainting.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 2536–2544, 2016.
  * Pathak et al. [2019] Deepak Pathak, Dhiraj Gandhi, and A. Gupta.  Self-supervised exploration via disagreement.  _ArXiv_ , abs/1906.04161, 2019.
  * Peng et al. [2018] Xue Bin Peng, Marcin Andrychowicz, Wojciech Zaremba, and Pieter Abbeel.  Sim-to-real transfer of robotic control with dynamics randomization.  _2018 IEEE International Conference on Robotics and Automation (ICRA)_ , May 2018.
  * Pinto and Gupta [2016] Lerrel Pinto and Abhinav Gupta.  Supersizing self-supervision: Learning to grasp from 50k tries and 700 robot hours.  In _2016 IEEE international conference on robotics and automation (ICRA)_ , pages 3406–3413. Ieee, 2016.
  * Pinto et al. [2017] Lerrel Pinto, Marcin Andrychowicz, Peter Welinder, Wojciech Zaremba, and Pieter Abbeel.  Asymmetric actor critic for image-based robot learning.  _arXiv preprint arXiv:1710.06542_ , 2017.
  * Raileanu et al. [2020] Roberta Raileanu, M. Goldstein, Denis Yarats, Ilya Kostrikov, and R. Fergus.  Automatic data augmentation for generalization in deep reinforcement learning.  _ArXiv_ , abs/2006.12862, 2020.
  * Ramos et al. [2019] Fabio Ramos, Rafael Possas, and Dieter Fox.  Bayessim: Adaptive domain randomization via probabilistic inference for robotics simulators.  _Robotics: Science and Systems XV_ , Jun 2019.
  * Schaul et al. [2016] Tom Schaul, John Quan, Ioannis Antonoglou, and D. Silver.  Prioritized experience replay.  _CoRR_ , abs/1511.05952, 2016.
  * Schwarzer et al. [2020] Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip Bachman.  Data-efficient reinforcement learning with self-predictive representations.  _arXiv preprint arXiv:2007.05929_ , 2020.
  * Sekar et al. [2020] Ramanan Sekar, Oleh Rybkin, Kostas Daniilidis, Pieter Abbeel, Danijar Hafner, and Deepak Pathak.  Planning to explore via self-supervised world models, 2020.
  * Shelhamer et al. [2017] Evan Shelhamer, Parsa Mahmoudieh, Max Argus, and Trevor Darrell.  Loss is its own reward: Self-supervision for reinforcement learning.  _ArXiv_ , abs/1612.07307, 2017.
  * Shorten and Khoshgoftaar [2019] Connor Shorten and T. Khoshgoftaar.  A survey on image data augmentation for deep learning.  _Journal of Big Data_ , 6:1–48, 2019.
  * Song et al. [2020] Xingyou Song, Yiding Jiang, Stephen Tu, Yilun Du, and Behnam Neyshabur.  Observational overfitting in reinforcement learning.  _ArXiv_ , abs/1912.02975, 2020.
  * Srinivas et al. [2020] Aravind Srinivas, Michael Laskin, and Pieter Abbeel.  Curl: Contrastive unsupervised representations for reinforcement learning.  _arXiv preprint arXiv:2004.04136_ , 2020.
  * Stone et al. [2021] Austin Stone, Oscar Ramirez, K. Konolige, and Rico Jonschkowski.  The distracting control suite - a challenging benchmark for reinforcement learning from pixels.  _ArXiv_ , abs/2101.02722, 2021.
  * Stooke et al. [2020] Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin.  Decoupling representation learning from reinforcement learning.  _ArXiv_ , abs/2004.1499, 2020.
  * Sutton [2005] R. Sutton.  Learning to predict by the methods of temporal differences.  _Machine Learning_ , 3:9–44, 2005.
  * Sutton and Barto [2018] Richard S. Sutton and Andrew G. Barto.  _Reinforcement Learning: An Introduction_.  The MIT Press, second edition, 2018.
  * Tassa et al. [2018] Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, et al.  Deepmind control suite.  Technical report, DeepMind, 2018.
  * Tian et al. [2019] Yonglong Tian, Dilip Krishnan, and Phillip Isola.  Contrastive multiview coding.  _arXiv preprint arXiv:1906.05849_ , 2019.
  * Tobin et al. [2017] Josh Tobin, Rachel Fong, Alex Ray, Jonas Schneider, Wojciech Zaremba, and Pieter Abbeel.  Domain randomization for transferring deep neural networks from simulation to the real world.  _2017 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)_ , Sep 2017.
  * Todorov et al. [2012] Emanuel Todorov, Tom Erez, and Yuval Tassa.  Mujoco: A physics engine for model-based control.  In _2012 IEEE/RSJ International Conference on Intelligent Robots and Systems_ , pages 5026–5033, 2012.  doi: 10.1109/iros.2012.6386109.
  * van den Oord et al. [2018] Aaron van den Oord, Yazhe Li, and Oriol Vinyals.  Representation learning with contrastive predictive coding, 2018.
  * Vaswani et al. [2017] Ashish Vaswani, Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  _arXiv_ , abs/1706.03762, 2017.
  * Vinyals et al. [2019] Oriol Vinyals, I. Babuschkin, Wojciech Czarnecki, Micha
"el Mathieu, Andrew Dudzik, J. Chung, D. Choi, Richard Powell, et al.
Grandmaster level in starcraft ii using multi-agent reinforcement learning.
_Nature_ , pages 1–5, 2019.

  * Wang et al. [2020] K. Wang, Bingyi Kang, Jie Shao, and Jiashi Feng.  Improving generalization in reinforcement learning with mixture regularization.  _ArXiv_ , abs/2010.10814, 2020.
  * Wang et al. [2021] Xudong Wang, Long Lian, and Stella X. Yu.  Unsupervised visual attention and invariance for reinforcement learning.  _ArXiv_ , abs/2104.02921, 2021.
  * Wang et al. [2016] Ziyu Wang, Tom Schaul, Matteo Hessel, H. V. Hasselt, Marc Lanctot, and N. D. Freitas.  Dueling network architectures for deep reinforcement learning.  _ArXiv_ , abs/1511.06581, 2016.
  * Wu et al. [2018] Zhirong Wu, Yuanjun Xiong, Stella X Yu, and Dahua Lin.  Unsupervised feature learning via non-parametric instance discrimination.  In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_ , pages 3733–3742, 2018.
  * Xu et al. [2020] Zhenlin Xu, Deyi Liu, Junlin Yang, and M. Niethammer.  Robust and generalizable visual representation learning via random convolutions.  _ArXiv_ , abs/2007.13003, 2020.
  * Yarats et al. [2019] Denis Yarats, Amy Zhang, Ilya Kostrikov, Brandon Amos, Joelle Pineau, and Rob Fergus.  Improving sample efficiency in model-free reinforcement learning from images, 2019.
  * Yarats et al. [2021] Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto.  Reinforcement learning with prototypical representations.  _arXiv preprint arXiv:2102.11271_ , 2021.
  * Young et al. [2020] Sarah Young, Dhiraj Gandhi, Shubham Tulsiani, Abhinav Gupta, Pieter Abbeel, and Lerrel Pinto.  Visual imitation made easy.  _CoRL_ , 2020.
  * Yu et al. [2020] Tianhe Yu, Saurabh Kumar, A. Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn.  Gradient surgery for multi-task learning.  _ArXiv_ , abs/2001.06782, 2020.
  * Zhang et al. [2020] A. Zhang, Rowan McAllister, R. Calandra, Y. Gal, and Sergey Levine.  Learning invariant representations for reinforcement learning without reconstruction.  _ArXiv_ , abs/2006.10742, 2020.
  * Zhang et al. [2018] C. Zhang, Oriol Vinyals, R. Munos, and S. Bengio.  A study on overfitting in deep reinforcement learning.  _ArXiv_ , abs/1804.06893, 2018.
  * Zhang et al. [2016] Richard Zhang, Phillip Isola, and Alexei A Efros.  Colorful image colorization.  In _European conference on computer vision_ , pages 649–666. Springer, 2016.
  * Zhu et al. [2017] Yuke Zhu, Roozbeh Mottaghi, Eric Kolve, Joseph J Lim, Abhinav Gupta, Li Fei-Fei, and Ali Farhadi.  Target-driven visual navigation in indoor scenes using deep reinforcement learning.  In _Icra_ , pages 3357–3364. Ieee, 2017.
  * Ziebart et al. [2008] Brian D Ziebart, Andrew Maas, J Andrew Bagnell, and Anind K Dey.  Maximum entropy inverse reinforcement learning.  In _Proceedings of the 23rd National Conference on Artificial Intelligence_ , volume 3, 2008.

##  Appendix A Ablations

We ablate the design choices of SVEA and compare both training and test
performance to DrQ and RAD. Results are shown in Table 3. We find that our
proposed formulation of SVEA outperforms the test performance of all other
variants, and by a large margin (method 2). Using a ViT encoder (method 1)
instead of a CNN further improves both the training and test performance of
SVEA, whereas the test performance of DrQ decreases by a factor of 5 when
using a ViT encoder (method 7). This indicates that ViT encoders overfit
heavily to the training environment without the strong augmentation of SVEA.
We observe that both DrQ and RAD are unstable under strong augmentation
(method 10 and 12). While the test performance of DrQ does not benefit from
using a ViT encoder, we observe a slight improvement in training performance
(method 7), similar to that of SVEA.

Table 3: Ablations. We vary the following choices: (i) architecture of the
encoder; (ii) our proposed objective $\mathcal{L}^{\textbf{SVEA}}_{Q}$ as
opposed to $\mathcal{L}_{Q}$ or a mix-all objective that uses two data-streams
for both $Q$-predictions, $Q$-targets, and $\pi$; (iii) using strong
augmentation (conv) in addition to the random shift augmentation used by
default (abbreviated as Str. aug.); and (iv) whether the target is augmented
or not (abbreviated as Aug. tgt.). We report mean episode return in the
training and test environments (color_hard) of the Walker, walk task. Method 1
and 2 are the default formulations of SVEA using ViT and CNN encoders,
respectively, method 7 and 8 are the default formulations of DrQ using Vit and
CNN encoders, respectively, and method 11 is the default formulation of RAD
that uses a random crop augmentation and is implemented using a CNN encoder.
Mean and std. deviation. of 5 seeds.

| Method | Arch. | Objective | Str. aug. | Aug. tgt. | Train return | Test return
---|---|---|---|---|---|---|---
1 | SVEA | ViT | SVEA | ✓ | ✗ | $918\scriptstyle{\pm 57}$ | $\mathbf{877\scriptstyle{\pm 54}}$
2 | $-$ | CNN | SVEA | ✓ | ✗ | $833\scriptstyle{\pm 91}$ | $760\scriptstyle{\pm 145}$
3 | $-$ | CNN | SVEA | ✓ | ✓ | $872\scriptstyle{\pm 53}$ | $605\scriptstyle{\pm 148}$
4 | $-$ | CNN | mix-all | ✓ | ✓ | $\mathbf{927\scriptstyle{\pm 24}}$ | $599\scriptstyle{\pm 214}$
5 | $-$ | CNN | $Q$ | ✓ | ✗ | $596\scriptstyle{\pm 55}$ | $569\scriptstyle{\pm 139}$
6 | $-$ | CNN | $Q$ | ✗ | ✗ | $771\scriptstyle{\pm 317}$ | $498\scriptstyle{\pm 196}$
7 | DrQ | ViT | $Q$ | ✗ | ✓ | $\mathbf{920\scriptstyle{\pm 36}}$ | $94\scriptstyle{\pm 18}$
8 | $-$ | CNN | $Q$ | ✗ | ✓ | $892\scriptstyle{\pm 65}$ | $520\scriptstyle{\pm 91}$
9 | $-$ | ViT | $Q$ | ✓ | ✓ | $286\scriptstyle{\pm 225}$ | $470\scriptstyle{\pm 67}$
10 | $-$ | CNN | $Q$ | ✓ | ✓ | $560\scriptstyle{\pm 158}$ | $\mathbf{569\scriptstyle{\pm 139}}$
11 | RAD | CNN | $Q$ | ✗ | ✓ | $\mathbf{883\scriptstyle{\pm 23}}$ | $\mathbf{400\scriptstyle{\pm 61}}$
12 | $-$ | CNN | $Q$ | ✓ | ✓ | $260\scriptstyle{\pm 201}$ | $246\scriptstyle{\pm 184}$

##  Appendix B Stability under Data Augmentation

Figure 11 compares the sample efficiency and stability of SVEA and DrQ under
each of the 6 considered data augmentations for 5 tasks from DMControl. We
observe that SVEA improves stability in all 27 instances where DrQ is impaired
by data augmentation. Stability of DrQ under data augmentation is found to be
highly sensitive to both the choice of augmentation and the particular task.
For example, the DrQ + aug baseline is relatively unaffected by a majority of
data augmentations in the Walker, stand task, while we observe significant
instability across all data augmentations in the Cartpole, swingup task. Our
results therefore indicate that SVEA can be a highly effective method for
eliminating the need for costly trial-and-error associated with application of
data augmentation.

Figure 11: Stability under data augmentation. Training performance measured by
episode return of SVEA and DrQ under 6 common data augmentations (using
ConvNets). We additionally provide reference curves for DrQ without additional
augmentation. Mean of 5 seeds, shaded area is $\pm 1$ std. deviation. SVEA
obtains similar sample efficiency to DrQ without augmentation, while the
sample efficiency of DrQ + aug is highly dependent on the task and choice of
augmentation.

##  Appendix C Data Augmentation in RL

Application of data augmentation in image-based RL has proven highly
successful [9, 36, 35, 33, 61, 21, 51] in improving generalization by
regularizing the network parameterizing the $Q$-function and policy $\pi$.
However, not all augmentations are equally effective. Laskin et al. [35] and
Kostrikov et al. [33] find small random crops and random shifts (image
translations) to greatly improve sample efficiency of image-based RL, but they
empirically offer no significant improvement in generalization to other
environments [21, 60]. On the other hand, augmentations such as random
convolution [36] have shown great potential in improving generalization, but
is simultaneously found to cause instability and poor sample efficiency [35,
21]. In this context, it is useful to distinguish between weak augmentations
such as small random translations that improve sample efficiency due to their
regularization, and strong augmentations such as random convolution that
improve generalization at the expense of sample efficiency. In this work, we
focus on stabilizing deep $Q$-learning under strong data augmentation with the
goal of improving generalization.

Figure 12: Generalization depends on the choice of data augmentation. A
comparison of SVEA implemented using each of the 6 data augmentations
considered in this work (using ConvNets). SVEA exhibits comparable stability
and sample efficiency for all augmentations, but generalization ability is
highly dependent on the choice of augmentation. Top: episode return on the
training environment during training. Bottom: generalization measured by
episode return on the color_hard benchmark of DMControl-GB. Mean of 5 seeds,
shaded area is $\pm 1$ std. deviation.

Figure 12 shows training and test performance of SVEA implemented using each
of the 6 data augmentations considered in this work. SVEA exhibits comparable
stability and sample efficiency for all augmentations, but we find that
generalization ability on the color_hard benchmark of DMControl-GB is highly
dependent on the choice of augmentation. Generally, we observe that
augmentations such as conv, overlay, and affine-jitter achieve the best
generalization, but they empirically also cause the most instability in our
DrQ + aug baseline as shown in Figure 11.

Figure 13 provides a comprehensive set of samples for each of the data
augmentations considered in this study: random shift [33], random convolution
(denoted conv) [36], random overlay [21], random cutout [9], Gaussian blur,
random affine-jitter, and random rotation [35]. We emphasize that the random
convolution augmentation is not a convolution operation, but rather
application of a randomly initialized convolutional layer as in the original
proposal [36]. As in previous work [35, 33, 21] that applies data augmentation
to image-based RL, we either clip values or apply a logistic function,
whichever is more appropriate, to ensure that output values remain within the
$[0,1)$ interval that unaugmented observations are normalized to. Each of the
considered data augmentations are applied to the walker and cartpole
environments and are representative of the Walker, walk, Walker, stand,
Cartpole, swingup, and Cartpole, balance tasks. To illustrate the diversity of
augmentation parameters associated with a given transformation, we provide a
total of 6 samples for each data augmentation in each of the two environments.
Note that, while random shift has been shown to improve sample efficiency in
previous work, it provides very subtle randomization. Stronger and more varied
augmentations such as random convolution, random overlay, and affine-jitter
can be expected to improve generalization to a larger set of MDPs, but naïve
application of these data augmentations empirically results in optimization
difficulties and poor sample efficiency.

(a) No augmentation (walker).

(b) No augmentation (cartpole).

(c) Random shift (walker).

(d) Random shift (cartpole).

(e) Random convolution (walker).

(f) Random convolution (cartpole).

(g) Random overlay (walker).

(h) Random overlay (cartpole).

(i) Random cutout (walker).

(j) Random cutout (cartpole).

(k) Random blur (walker).

(l) Random blur (cartpole).

(m) Random affine-jitter (walker).

(n) Random affine-jitter (cartpole).

(o) Random rotation (walker).

(p) Random rotation (cartpole).

Figure 13: Data augmentation. Visualizations of all data augmentations
considered in this study. Left column contains samples from the Walker, walk
and Walker, stand tasks, and right column contains samples from the Cartpole,
swingup and Cartpole, balance tasks. Figure 14: DistractingCS. Episode return
as a function of randomization intensity, for each of the 5 tasks from
DMControl-GB (using ConvNets). Mean of 5 seeds. We find that the difficulty of
DistractingCS varies greatly between tasks, but SVEA consistently outperforms
DrQ in terms of generalization across all intensities and tasks, except for
Ball in cup, catch at the highest intensity.

##  Appendix D Choice of Base Algorithm

In the majority of our results, we implement SVEA using SAC as base algorithm,
implemented with a random shift augmentation as proposed by DrQ [33]. We now
consider an additional set of experiments where we instead implement SVEA
using RAD [35] as base algorithm, which proposes to add a random cropping to
SAC (in place of random shift). Training and test performances on color_hard
are shown in Figure 15. We find SVEA to provide similar performance and
benefits in terms of sample efficiency and generalization as when using DrQ as
base algorithm. RAD likewise has similar performance to DrQ without use of
strong augmentation, however, we observe that RAD generally is more unstable
than DrQ when additionally using conv augmentation, and the relative
improvement of SVEA is therefore comparably larger in our RAD experiments.

Figure 15: Choice of base algorithm. We compare SVEA implemented with RAD as
base algorithm to instances of RAD with and without random convolution
augmentation (using ConvNets). Top: episode return on the training environment
during training. Bottom: generalization measured by episode return on the
color_hard benchmark of DMControl-GB. Mean of 5 seeds, shaded area is $\pm 1$
std. deviation. SVEA improves generalization in all instances.

##  Appendix E Test Environments

Figure 16 provides visualizations for each of the two generalization
benchmarks, DMControl Generalization Benchmark [21] and Distracting Control
Suite [60], used in our experiments. Agents are trained in a fixed training
environment with no visual variation, and are expected to generalize to novel
environments of varying difficulty and factors of variation. The color_hard,
video_easy, and video_hard benchmarks are from DMControl Generalization
Benchmark, and we further provide samples from the Distracting Control Suite
(DistractingCS) benchmark for intensities $I=\\{0.1,0.2,0.5\\}$. While methods
are evaluated on a larger set of intensities, we here provide samples deemed
representative of the intensity scale. We note that the DistractingCS
benchmark has been modified to account for action repeat (frame-skip).
Dynamically changing the environment at each simulation step makes the
benchmark disproportionally harder for tasks that use a large action repeat,
e.g. Cartpole tasks. Therefore, we choose to modify the DistractingCS
benchmark and instead update the distractors every second simulation step,
corresponding to the lowest action repeat used (2, in Finger, spin). This
change affects both SVEA and baselines equally. Figure 14 shows generalization
results on DistractingCS for each task individually. We find that the
difficulty of DistractingCS varies greatly between tasks, but SVEA
consistently outperforms DrQ in terms of generalization across all intensities
and tasks.

(a) Training environment (walker).

(b) Training environment (cartpole).

(c) color_hard environment (walker).

(d) color_hard environment (cartpole).

(e) video_easy environment (walker).

(f) video_easy environment (cartpole).

(g) video_hard environment (walker).

(h) video_hard environment (cartpole).

(i) DistractingCS for intensity $0.1$ (walker).

(j) DistractingCS for intensity $0.1$ (cartpole).

(k) DistractingCS for intensity $0.2$ (walker).

(l) DistractingCS for intensity $0.2$ (cartpole).

(m) DistractingCS for intensity $0.5$ (walker).

(n) DistractingCS for intensity $0.5$ (cartpole).

Figure 16: Test environments. Samples from each of the two generalization
benchmarks, DMControl Generalization Benchmark [21] and Distracting Control
Suite [60], considered in this study. In our experiments, agents are trained
in a fixed training environment with no visual variation, and are expected to
generalize to novel environments of varying difficulty and factors of
variation. Table 4: Comparison to state-of-the-art. Test performance (episode
return) of methods trained in a fixed environment and evaluated on the
video_hard benchmark from DMControl-GB. In this setting, the entirety of the
floor and background is replaced by natural videos; see Figure 16 for samples.
Results for CURL, RAD, PAD, and SODA are obtained from Hansen and Wang [21]
and we report mean and std. deviation of 5 runs. We compare SVEA and SODA
using the overlay augmentation, since this is the augmentation for which the
strongest results are reported for SODA in Hansen and Wang [21]. SVEA achieves
competitive results in all tasks considered, though there is still some room
for improvement before the benchmark is saturated.

DMControl-GB | CURL | RAD | DrQ | PAD | SODA | SVEA
---|---|---|---|---|---|---
(video_hard) |  |  |  |  | (overlay) | (overlay)
walker, | $58$ | $56$ | $104$ | $93$ | $\mathbf{381}$ | $377$
walk | $\scriptstyle{\pm 18}$ | $\scriptstyle{\pm 9}$ | $\scriptstyle{\pm 22}$ | $\scriptstyle{\pm 29}$ | $\mathbf{\scriptstyle{\pm 72}}$ | $\scriptstyle{\pm 93}$
walker, | $45$ | $231$ | $289$ | $278$ | $771$ | $\mathbf{834}$
stand | $\scriptstyle{\pm 5}$ | $\scriptstyle{\pm 39}$ | $\scriptstyle{\pm 49}$ | $\scriptstyle{\pm 72}$ | $\scriptstyle{\pm 83}$ | $\mathbf{\scriptstyle{\pm 46}}$
cartpole, | $114$ | $110$ | $138$ | $123$ | $\mathbf{429}$ | $393$
swingup | $\scriptstyle{\pm 15}$ | $\scriptstyle{\pm 16}$ | $\scriptstyle{\pm 9}$ | $\scriptstyle{\pm 24}$ | $\mathbf{\scriptstyle{\pm 64}}$ | $\scriptstyle{\pm 45}$
ball_in_cup, | $115$ | $97$ | $92$ | $66$ | $327$ | $\mathbf{403}$
catch | $\scriptstyle{\pm 33}$ | $\scriptstyle{\pm 29}$ | $\scriptstyle{\pm 23}$ | $\scriptstyle{\pm 61}$ | $\scriptstyle{\pm 100}$ | $\mathbf{\scriptstyle{\pm 174}}$
finger, | $27$ | $34$ | $71$ | $56$ | $302$ | $\mathbf{335}$
spin | $\scriptstyle{\pm 21}$ | $\scriptstyle{\pm 11}$ | $\scriptstyle{\pm 45}$ | $\scriptstyle{\pm 18}$ | $\scriptstyle{\pm 41}$ | $\mathbf{\scriptstyle{\pm 58}}$

##  Appendix F Additional Results on DMControl-GB

Table 1 contains results for the color_hard and video_easy generalization
benchmarks from DMControl-GB. We here provide additional results for the
video_hard benchmark, and note that we leave out the color_easy benchmark
because it is already considered solved by previous work [21, 72]. Results are
shown in Table 4. SVEA also achieves competitive performance across all 5
tasks in the video_hard benchmark.

##  Appendix G Robotic Manipulation Tasks

We conduct experiments with a set of three goal-conditioned robotic
manipulation tasks: (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. We implement the tasks
using MuJoCo [67] for simulation, and we use a simulated Kinova Gen3 robotic
arm. The initial configuration of gripper, cube, and goal is randomized, the
agent uses 2D positional control, and policies are trained using dense
rewards. For reach and reach moving target, at each time step there is a
reward penalty proportional to the euclidean distance between the gripper and
the goal, and there is a reward bonus of $+1$ when the distance is within a
fixed threshold corresponding to the radius of the red mark. We use the same
reward structure for the push task, but use the euclidean distance between the
cube and the goal.

Each episode consists of $50$ time steps, which makes $50$ an upper bound on
episode return, while there is no strict lower bound. All observations are
stacks of three RGB frames of size $84\times 84\times 3$, and the agent has no
access to state information. During training, the camera position is fixed,
and the camera orientation follows the gripper. During testing, the camera
orientation still follows the gripper, but we additionally randomize the
camera position, as well as colors, lighting, and background of the
environment. Samples for the robotic manipulation tasks – both in the training
environment and the randomized test environments – are shown in Figure 17. We
use a binary threshold for measuring task success: the episode is considered a
success if the environment is in success state (i.e. either the gripper or the
cube is within a fixed distance to the center of the red mark) for at least
$50\%$ of all time steps in the two reaching tasks, and $25\%$ for the push
task. This is to ensure that the success rate of a random policy does not get
inflated by trajectories that coincidentally visit a success state for a small
number of time steps, e.g. passing through the goal with the gripper.

(a) Training environment (reach).

(b) Training environment (push).

(c) Test environments (reach).

(d) Test environments (push).

Figure 17: Environments for robotic manipulation. Training and test
environments for our robotic manipulation experiments. Agents are trained in a
fixed environment, and are expected to generalize to the unseen test
environments with randomized camera pose, colors, lighting, and backgrounds.

##  Appendix H Implementation Details

In this section, we provide extensive implementation details for our
experimental setup, including network architecture, hyperparameters, as well
as design choices specific to our ViT experiments.

Network architecture. For experiments in DMControl [64], we adopt our network
architecture from Hansen and Wang [21], without any changes to the
architecture nor hyperparameters to ensure a fair comparison. The shared
encoder $f_{\theta}$ is implemented as an 11-layer CNN encoder that takes a
stack of RGB frames rendered at $84\times 84\times 3$ and outputs features of
size $32\times 21\times 21$, where $32$ is the number of channels and
$21\times 21$ are the dimensions of the spatial feature maps. All
convolutional layers use $32$ filters and $3\times 3$ kernels. The first
convolutional layer uses a stride of 2, while the remaining convolutional
layers use a stride of 1. Following previous work on image-based RL for
DMControl tasks [76, 59, 35, 33, 21], the shared encoder is followed by
independent linear projections for the actor and critic of the Soft Actor-
Critic [18] base algorithm used in our experiments, and the actor and critic
modules each consist of three fully connected layers with hidden dimension
$1024$. Training takes approximately 24 hours on a single NVIDIA V100 GPU. For
simplicity, we choose to apply the same experimental setup for robotic
manipulation.

Hyperparameters. Whenever applicable, we adopt hyperparameters from Hansen and
Wang [21]. We detail hyperparameters relevant to our experiments in Table 5;
ViT hyperparameters are discussed in the following. We use the default SVEA
loss coefficients $\alpha=0.5,\beta=0.5$ in all experiments using a CNN
encoder.

Table 5: Hyperparameters used in experiments on DMControl and robotic
manipulation.

Hyperparameter | Value
---|---
Frame rendering | $84\times 84\times 3$
Stacked frames | 3
Random shift | Up to $\pm 4$ pixels
| Action repeat
---
| 2 (finger)
---
8 (cartpole)
4 (otherwise)
Discount factor $\gamma$ | 0.99
Episode length | 1,000
Learning algorithm | Soft Actor-Critic (SAC)
Number of frames | 500,000
Replay buffer size | 500,000
Optimizer ($\theta$) | Adam ($\beta_{1}=0.9,\beta_{2}=0.999$)
Optimizer ($\alpha$ of SAC) | Adam ($\beta_{1}=0.5,\beta_{2}=0.999$)
Learning rate ($\theta$) | 1e-3
Learning rate ($\alpha$ of SAC) | 1e-4
Batch size | 128
SVEA coefficients | $\alpha=0.5,\beta=0.5$
$\psi$ update frequency | 2
|  $\psi$ momentum coefficient
---
| 0.05 (encoder)
---
0.01 (critic)

RL with Vision Transformers. We adopt a similar experimental setup for our
experiments with Vision Transformers (ViT) [10] as replacement for the CNN
encoder used in the rest of our experiments, but with minimal changes to
accommodate the new encoder. Our ViT encoder takes as input a stack of RGB
frames rendered at $96\times 96\times 3$ (versus $84\times 84\times 3$ for
CNN) and uses a total of $144$ non-overlapping $8\times 8\times 3k$ (where $k$
is the number of frames in a frame stacked observation) image patches from a
spatial grid evenly placed across the image observation. All patches are
projected into $128$-dimensional embeddings, and all embeddings are then
forwarded as tokens for the ViT encoder. Following the original ViT
implementation, we use learned positional encodings as well as a learnable
class token. Our encoder consists of $4$ stacked Transformer [69] encoders,
each using Multi-Head Attention with $8$ heads. We use the ViT encoder as a
drop-in replacement to CNN and optimize it jointly together with the
$Q$-function using the Adam [32] optimizer and with no changes to the learning
rate. We do not pre-train the parameters of the ViT encoder, and we do not use
weight decay. Similar to the CNN encoder, we find ViT to benefit from random
shift augmentation and we therefore apply it by default in all methods.
Training takes approximately 6 days on a single NVIDIA V100 GPU. See Table 6
for an overview of hyperparameters specific to the ViT encoder experiments.

Table 6: Hyperparameters used in our ViT experiments.

Hyperparameter | Value
---|---
Frame rendering | $96\times 96\times 3$
Random shift | Up to $\pm 6$ pixels
Patch size | $8\times 8\times 3k$
Number of patches | $144$
Embedding dimensionality | $128$
Number of layers | $4$
Number of attention heads | $8$
Score function | Scaled dot-product
Number of frames | 300,000
Replay buffer size | 300,000
Batch size | $512$
| SVEA coefficients
---
|  $\alpha=1,\beta=1$ (walker)
---
$\alpha=1,\beta=0.5$ (push)
$\alpha=0.5,\beta=0.5$ (otherwise)
Number of ViT parameters | $489,600$

Implementation of data augmentation in SVEA vs. previous work. Previous work
[35, 33, 61, 54] applies augmentation to both state
$\mathbf{s}^{\textnormal{aug}}_{t}=\tau(\mathbf{s}_{t},\nu)$ and successor
state
$\mathbf{s}^{\textnormal{aug}}_{t+1}=\tau(\mathbf{s}_{t+1},\nu^{\prime})$
where $\nu,\nu^{\prime}\sim\mathcal{V}$. As discussed in Section C, this is
empirically not found to be an issue in application of weak augmentation such
as random shift [33]. However, previous work finds that strong augmentation,
e.g. random convolution [36, 35], can cause instability and poor sample
efficiency. We apply random shifts in SVEA and all baselines by default, and
aim to stabilize learning under strong augmentation. As such, we generically
refer to observations both with and without the random shift operation as
unaugmented, and instead refer to observations as augmented after application
of one of the 6 augmentations considered in our study. While we provide the
full SVEA algorithm in Algorithm 1, we here provide supplementary Python-like
pseudo-code for the update rules of SVEA as well as generic off-policy actor-
critic algorithms both with and without naïve application of data
augmentation.

Generic off-policy actor-critic algorithm. We here provide a reference
implementation for a generic base algorithm from which we will implement our
proposed SVEA framework for data augmentation.

def update(state, action, reward, next_state): """Generic off-policy actor-
critic RL algorithm""" next_action = actor(next_state) q_target = reward +
critic_target(next_state, next_action) q_prediction = critic(state, action)
update_critic(q_prediction, q_target) update_actor(state)
update_critic_target()

Naïve application of data augmentation. A natural way to apply data
augmentation in off-policy RL algorithms is to augment both $\mathbf{s}_{t}$
and $\mathbf{s}_{t+1}$. Previous work on data augmentation in off-policy RL
[35, 33, 61, 54] follow this approach. However, this is empirically found to
be detrimental to sample efficiency and stability under strong data
augmentation. In the following, we generically refer to application of data
augmentation as an aug operation.

def update_aug(state, action, reward, next_state): """Generic off-policy
actor-critic RL algorithm that uses data augmentation""" state = aug(state)
next_state = aug(next_state) next_action = actor(next_state) q_target = reward
+ critic_target(next_state, next_action) q_prediction = critic(state, action)
update_critic(q_prediction, q_target) update_actor(state)
update_critic_target()

SVEA. Our proposed method, SVEA, does not apply strong augmentation to
$\mathbf{s}_{t+1}$ nor $\mathbf{s}_{t}$ when used for policy learning. SVEA
jointly optimizes the $Q$-function over two data streams with augmented and
unaugmented data, respectively, which can be implemented efficiently for
$\alpha=\beta$ as in Algorithm 1. The following pseudo-code assumes
$\alpha=0.5,\beta=0.5$.

def update_svea(state, action, reward, next_state): """SVEA update for a
generic off-policy actor-critic RL algorithm""" update_actor(state)
next_action = actor(next_state) # [B, A] q_target = reward +
critic_target(next_state, next_action) # [B, 1] svea_state =
concatenate(state, aug(state), dim=0) # [2*B, C, H, W] svea_action =
concatenate(action, action, dim=0) # [2*B, A] svea_q_target =
concatenate(q_target, q_target, dim=0) # [2*B, 1] svea_q_prediction =
critic(svea_state, svea_action) # [2*B, 1] update_critic(svea_q_prediction,
svea_q_target) update_critic_target() where $B$ is the batch size, $C$ is the
number of input channels, $H$ and $W$ are the dimensions of observations, and
$A$ is the dimensionality of the action space. For clarity, we omit
hyperparameters such as the discount factor $\gamma$, learning rates, and
update frequencies.

##  Appendix I Task Descriptions

We experiment on tasks from DMControl [64] as well as a set of robotic
manipulation tasks that we implement using MuJoCo [67]. DMControl tasks are
selected based on previous work on both sample efficiency [19, 76, 59, 20] and
generalization [21, 60, 72], and represent a diverse and challenging skill set
in the context of image-based RL. Our set of robotic manipulation tasks are
designed to represent fundamental visuomotor skills that are widely applied in
related work on robot learning [37, 50, 11, 44, 22, 78]. We here provide a
unified overview of the tasks considered in our study and their properties;
see Section G for a detailed discussion of the robotic manipulation
environment. All tasks emit observations $\mathbf{o}\in\mathbb{R}^{84\times
84\times 3}$ that are stacked as states $\mathbf{s}\in\mathbb{R}^{84\times
84\times 9}$.

  * •

Walker, walk ($\mathbf{a}\in\mathbb{R}^{6}$). A planar walker that is rewarded
for walking forward at a target velocity. Dense rewards.

  * •

Walker, stand ($\mathbf{a}\in\mathbb{R}^{6}$). A planar walker that is
rewarded for standing with an upright torso at a constant minimum height.
Dense rewards.

  * •

Cartpole, swingup ($\mathbf{a}\in\mathbb{R}$). Swing up and balance an
unactuated pole by applying forces to a cart at its base. The agent is
rewarded for balancing the pole within a fixed threshold angle. Dense rewards.

  * •

Cartpole, balance ($\mathbf{a}\in\mathbb{R}$). Balance an unactuated pole by
applying forces to a cart at its base. The agent is rewarded for balancing the
pole within a fixed threshold angle. Dense rewards.

  * •

Ball in cup, catch ($\mathbf{a}\in\mathbb{R}^{2}$). An actuated planar
receptacle is to swing and catch a ball attached by a string to its bottom.
Sparse rewards.

  * •

Finger, spin ($\mathbf{a}\in\mathbb{R}^{2}$). A manipulation problem with a
planar 3 DoF finger. The task is to continually spin a free body. Sparse
rewards.

  * •

Robot, reach ($\mathbf{a}\in\mathbb{R}^{2}$). A manipulation problem with a
simulated Kinova Gen3 robotic arm. The task is to move the gripper to a
randomly initialized goal position. Dense rewards.

  * •

Robot, reach moving target ($\mathbf{a}\in\mathbb{R}^{2}$). A manipulation
problem with a simulated Kinova Gen3 robotic arm. The task is to continuously
track a randomly initialized goal with the gripper. The goal moves in a zig-
zag pattern at a random constant speed. Dense rewards.

  * •

Robot, push ($\mathbf{a}\in\mathbb{R}^{2}$). A manipulation problem with a
simulated Kinova Gen3 robotic arm. The task is to push a cube to a goal
position. All positions are randomly initialized. Dense rewards.

##  Appendix J Broader Impact

The discipline of deep learning – and reinforcement learning in particular –
is rapidly evolving, which can in part be attributed to better algorithms [43,
38, 18], neural network architectures [26, 69, 12], and availability of data
[49, 31], but advances are also highly driven by increased computational
resources and larger models such as GPT-3 [5] in natural language processing
and ViT [10] in computer vision. As a result, both computational and economic
requirements for training and deploying start-of-the-art models are increasing
at an unprecedented rate [3]. While we are concerned by this trend, we remain
excited about the possibility of reusing and re-purposing large learned models
(in the context of RL: policies and value functions) that learn and generalize
far beyond the scope of their training environment. A greater reuse of learned
policies can ultimately decrease overall computational costs, since new models
may need to be trained less frequently. As researchers, we are committed to
pursue research that is to the benefit of society. We strive to enable reuse
of RL policies through extensive use of data augmentation, and we firmly
believe that our contribution is an important step towards that goal. Our
method is empirically found to reduce the computational cost (in terms of both
stability, sample efficiency, and total number of gradient steps) of training
RL policies under augmentation, which is an encouraging step towards learning
policies that generalize to unseen environments. By extension, this promotes
policy reuse, and may therefore be a promising component both for reducing
costs and for improving generalization of large-scale RL that the field
appears to be trending towards.