A Recipe for Unbounded Data Augmentation in
Visual Reinforcement Learning

Visual Reinforcement Learning (RL) has a myriad of real-world applications (Mnih et al., 2013; Levine et al., 2016; Pinto & Gupta, 2016; Kalashnikov et al., 2018; Berner et al., 2019; Vinyals et al., 2019), and visual $Q$-learning algorithms are especially enticing because of their potential for high data-efficiency. However, they are very prone to overfitting on their training distribution due to the combination of flexible representation, high-dimensional data, and limited visual diversity in training environments (Peng et al., 2018; Cobbe et al., 2019; Julian et al., 2020).

Data augmentation is a widely used technique for learning visual invariances in supervised learning (Noroozi & Favaro, 2016; Tian et al., 2019; Chen et al., 2020), but has been found to cause training instabilities when applied to visual RL (Lee et al., 2019; Laskin et al., 2020; Hansen & Wang, 2021). Prior work, SVEA (Hansen et al., 2021b), found that a more selective application of data augmentation in the critic update of actor-critic algorithms (Lillicrap et al., 2015; Haarnoja et al., 2018) improved training stability significantly. The actor (policy) – which shares its visual backbone with the critic ($Q$-function) – is then optimized solely from unaugmented observations. By sharing their visual backbone, the actor indirectly benefits from the learned invariances.

In this work, we revisit the data augmentation recipe proposed in SVEA, and discover an assumption that limits its practicality to augmentations of a photometric (color or light altering) nature. SVEA assumes that an encoder’s output embedding can become fully invariant to input augmentations. If an encoder’s output is fully invariant to input augmentations, then an actor, only trained on unaugmented observations, can become robust to input augmentations indirectly through a shared actor-critic encoder. However, this leads to two key failure cases: (i) the output of a convolutional neural network (CNN) encoder can not become invariant to input geometric augmentations e.g., rotation or translation; (see Figure 1) and (ii) the encoder and critic are trained end-to-end, thus, part of the invariance may be off-loaded to the critic regardless of the type of augmentation.

To address these limitations, we propose SADA: Stabilized Actor-Critic under Data Augmentation, a generalized data augmentation recipe that supports a wide variety of augmentations. Instead of only augmenting critic inputs, SADA augments both actor and critic inputs, but does so carefully to avoid training instabilities: (1) in actor updates, only the policy input is augmented while the $Q$-function input is left unaugmented, (2) in critic updates, only the online $Q$-function input is augmented while the target $Q$-function input is unaugmented, and (3) we jointly optimize components on both augmented and unaugmented data. Importantly, SADA requires no additional forward passes, losses, or parameters.

To stress-test our method, we propose DMC-GB2, an extension of the DeepMind Control Suite Generalization Benchmark (Hansen & Wang, 2021) that encompasses a wider and more challenging collection of test environments than existing benchmarks. We benchmark methods across DMC-GB2, tasks from Meta-World (Yu et al., 2020), and the Distracting Control Suite (Stone et al., 2021), and find that SADA greatly improves training stability and generalization of RL agents under a diverse set of augmentations.

The practice of learning visual invariances by augmenting data is ubiquitous in machine learning literature, and has been studied extensively in the context of supervised and self-supervised learning algorithms for computer vision problems (Noroozi & Favaro, 2016; Wu et al., 2018; Oord et al., 2018; Tian et al., 2019; Chen et al., 2020; He et al., 2022). More recently, use of augmentation has also been popularized in the context of visual RL. However, there is mounting evidence that much of the wisdom and practices developed in other areas (e.g. computer vision) do not translate to RL problems, presumably due to differences in learning objectives, datasets, and network architectures used. For example, while machine learning literature commonly considers a fixed dataset, RL algorithms are often trained on a non-stationary data distribution (replay buffer) that changes throughout training, and incoming data is typically a function of the current (behavioral) policy. As a result, RL datasets are often small and have limited diversity. This section provides an overview of prior work that leverages data augmentation to improve data-efficiency and generalization.

Improving data-efficiency with data augmentation. Much of the existing literature on visual RL leverages weak data augmentation (e.g. random crop or image shift) as a regularizer when data is limited, i.e., when data-efficiency is critical (Srinivas et al., 2020; Laskin et al., 2020; Kostrikov et al., 2020; Stooke et al., 2020; Yarats et al., 2021; Hansen et al., 2023), without particular emphasis on generalization or robustness to changes in the environment. For example, seminal works RAD (Laskin et al., 2020) and DrQ (Kostrikov et al., 2020) demonstrate that randomly cropping or shifting images, respectively, by a few pixels greatly improves data-efficiency and training stability of $Q$-learning algorithms – even when agents are trained and tested in the same environment. However, Laskin et al. (2020) simultaneously find that other types of augmentation (rotation, random convolution, masking) lead to training instabilities and a substantial decrease in data-efficiency.

Improving generalization with data augmentation. Visual generalization is a challenging but increasingly important problem in RL due to its limited data diversity. Multiple prior works aim to improve the training stability and generalization of RL algorithms by, e.g., proposing new types of augmentation (Lee et al., 2019; Wang et al., 2020; Hansen & Wang, 2021; Hansen et al., 2021b; Zhang & Guo, 2021; Wang et al., 2023), or introducing new (auxiliary) objectives (Raileanu et al., 2020; Hansen et al., 2021a; Wang et al., 2021; Fan et al., 2021; Yuan et al., 2022; Yang et al., 2024). For example, Lee et al. (2019) augment high-frequency content in observations using random convolution, Hansen & Wang (2021) randomly overlay observations with out-of-domain images, and Yang et al. (2024) adapt to camera changes at test-time using an auxiliary self-supervised objective and augmented data. Perhaps most importantly, SVEA (Hansen et al., 2021b) investigate why strong augmentations (such as those used in the aforementioned works) often destabilize training in an RL context, and propose an alternative method of applying augmentations that mitigate these instabilities. Our work builds upon SVEA and demonstrates that – while SVEA is robust to photometric augmentations – it largely fails when applied to (equally important) geometric augmentations.

We recommend the survey by Kirk et al. (2023) for a more comprehensive overview of prior work.

Visual Reinforcement Learning (RL) formulates interaction between an agent and its environment as a Partially Observable Markov Decision Process (POMDP) (Kaelbling et al., 1998). A POMDP can be formalized as a tuple $(\mathcal{S},\mathcal{O},\mathcal{A},\mathcal{T},R,\gamma)$, where $\mathcal{S}$ is an unobservable state space, $\mathbf{o}\in\mathcal{O}$ are observations from the environment (e.g., RGB images), $\mathbf{a}\in\mathcal{A}$ are actions, $\mathcal{T\colon\mathcal{S}\times\mathcal{A}\mapsto\mathcal{S}}$ is a transition function, $r$ is a task reward from a reward function $R\colon\mathcal{S}\times\mathcal{A}\mapsto\mathbb{R}$, and $\gamma$ is a discount factor. Throughout this work, we approximate the unobservable states $\mathbf{s}\in\mathcal{S}$ by defining observations as a stack of the three most recent RGB frames $\mathbf{o}_{t}\doteq\{\mathbf{x}_{t},\mathbf{x}_{t-1},\mathbf{x}_{t-2}\}$ for frames $\mathbf{x}_{t:t-2}$ at time $t$ (Mnih et al., 2013). The goal is then to learn a policy $\pi\colon\mathcal{O}\mapsto\mathcal{A}$ such that the discounted sum of rewards $\mathbb{E}_{\pi}\left[\sum_{t=0}^{\infty}\gamma^{t}r_{t}\right]$ is maximized (in expectation) when following the policy $\pi$.

$Q$-Learning algorithms developed for visual RL generally estimate the optimal state-action value function $Q^{*}\colon\mathcal{O}\times\mathcal{A}\mapsto\mathbb{R}$ with a neural network (denoted the critic). This is achieved by dynamic programming using the single-step Bellman error $Q(\mathbf{o}_{t},\mathbf{a}_{t})-y_{t}$ where $y_{t}$ is the temporal difference (TD) target $y_{t}\doteq r_{t}+\gamma Q(\mathbf{o}_{t+1},\mathbf{a}_{t+1}),~\mathbf{a}_{t+1}\sim\pi(\cdot|\mathbf{o}_{t+1})$. In practice, the $Q$-network used to compute $y_{t}$ is usually chosen to be an exponential moving average of the $Q$-function being learned (Lillicrap et al., 2015; Haarnoja et al., 2018). The policy $\pi$ is obtained by taking the action $\mathbf{a}_{t}\approx\arg\max_{\mathbf{a}_{t}}Q(\mathbf{o}_{t},\mathbf{a}_{t})~\forall\mathbf{o}_{t}$ in the current dataset (replay buffer), which is typically estimated by training a separate actor network when $\mathcal{A}$ is continuous. These two components – the actor and the critic – are iteratively updated by collecting data in the environment, appending it to a replay buffer $\mathcal{D}$, and optimizing $Q,\pi$ with the following objectives using stochastic gradient descent:

where gradients of the first objective are computed wrt. $Q$ only, and gradients of the second objective are computed wrt. $\pi$ only. When learning from images, observations are commonly encoded using a shared convolutional encoder $f$ such that $Q,\pi$ are redefined as $Q(f(\mathbf{o}_{t}),\mathbf{a})$ and $\pi(f(\mathbf{o}_{t}))$, with $f$ only being updated by the critic objective $\mathcal{L}_{Q}$. Due to the recurrent and self-referential nature of Equations 1-2, $Q$-learning algorithms are often more data-efficient than other algorithm classes, but are very prone to training instabilities – especially when data augmentation is applied to observations.

Image transformations. Throughout this work, we classify image transformations into two types: photometric and geometric transformations. Photometric transformations alter image color and lighting properties while preserving the spatial arrangement of pixels (e.g. random convolution, image overlay). Geometric transformations alter the spatial arrangement of pixels while keeping image color and lighting properties intact (e.g. rotation, shift). We visualize examples of photometric and geometric transformations in Figure 1.

We revisit common wisdom and practices when applying data augmentation in $Q$-learning algorithms, and discover that prior work makes an assumption that only holds for augmentations that are photometric in nature. We propose SADA: Stabilized Actor-Critic Learning under Data Augmentation, a generalized recipe for data augmentation that significantly improves the performance of a wider variety of augmentations. We start by outlining the assumptions of prior work, its implications, and then present our proposed solution.

We evaluate our method and baselines across 11 visual RL tasks from the DMControl (Tassa et al., 2018) and Meta-World-v2 (Yu et al., 2020) benchmarks and 12 test distributions from our proposed DMControl - Generalization Benchmark 2 (DMC-GB2). We additionally evaluate on the Distracting Control Suite (Stone et al., 2021) and provide the results in Appendix B.2. All methods are trained under strong augmentations in the training environments and evaluated in a zero-shot manner on their respective test distributions. See Figure 3 for a visualization of DMC-GB2 test environments. The full DMControl and Meta-World task lists are provided in Appendix A.2. Concretely, we aim to answer the following questions through experimentation:

• -

  Robustness. How does SADA compare to baselines in terms of overall augmentation robustness? In terms of geometric vs photometric robustness?

• -

  Analysis. Why do baselines fail to display geometric robustness, and how does SADA solve the problem? How do each of the SADA design choices affect results?

• -

  Generality. Can SADA be readily applied to other RL backbones and benchmarks?

Setup. We build on DrQ (Laskin et al., 2020) as our backbone algorithm, and use a fixed set of hyperparameters across all tasks and environments. All agents are trained for one million environments steps and use stacks of the three most recent RGB frames ($3\times\mathbb{R}^{\mathbf{(3\times 84\times 84)}})$ as observations. A full list of hyperparameters and training details is provided in Appendix A.

Test environments. As our experiments will reveal, previous methods largely fail to generalize to geometric changes, which has gone unnoticed due to existing benchmarks predominantly evaluating photometric robustness. Therefore, we propose the Deepmind Control Suite Generalization Benchmark 2 (DMC-GB2), an extension of DMC-GB (Hansen & Wang, 2021) to encompass a wider collection of photometric and geometric test distributions. In DMC-GB2, we provide geometric and photometric test sets. The geometric test set considers two types of geometric distributions – rotations and shifts – both individually and jointly, and at varying intensities categorized as easy and hard environments. The photometric test set considers a complementary setup for two types of photometric distributions – colors and videos. Detailed visualizations of all 12 DMC-GB2 test distributions is provided in Appendix C.2.

Data augmentation. We apply a weak augmentation of random shifting to all our inputs as conducted in our DrQ baseline, and consider it unaugmented. We further employ a set of strong augmentations, taking into account both geometric and photometric transformations. For geometric augmentations, we use random shift (Laskin et al., 2020), random rotation, and a combination consisting of random rotation followed by random shift. For photometric augmentations we use random convolution (Lee et al., 2019), random overlay (Hansen & Wang, 2021), and a combination consisting of random convolution followed by random overlay. We sample an augmentation from this set of six strong augmentations for each input sample in all our experiments unless stated otherwise. Detailed visualizations of all augmentations is provided in Appendix C.1.

Baselines. We benchmark our method against the following well-established baselines. 1) DrQ Laskin et al. (2020), a visual based Soft Actor Critic baseline that uses random shifts as the default augmentation for all inputs. 2) DrQ + Aug, a variant of DrQ implemented with a naive application of strong augmentations. 3) SVEA Hansen et al. (2021b), which builds on DrQ with a selective application of augmentation in the $Q$-function to increase robustness under strong augmentations.

Throughout this work, we give an overview of data augmentation within visual RL, highlighting the shortcomings of previous work, its implications, and presenting SADA, a generic data augmentation recipe for modern visual based reinforcement learning. We empirically prove SADA’s superiority to previous methods and provide a deep analysis of its effectiveness. Concurrently, we curated a comprehensive visual generalization benchmark, DMC-GB2, which we make publicly available at https://aalmuzairee.github.io/SADA, with the aim of furthering research efforts within visual RL.