Reinforcement Learning with Augmented Data

By Michael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, A. Srinivas. 2020

798 citations recorded by the source corpus.

Learning from visual observations is a fundamental yet challenging problem in Reinforcement Learning (RL). Although algorithmic advances combined with convolutional neural networks have proved to be a recipe for success, current methods are still lacking on two fronts: (a) data-efficiency of learning and (b) generalization to new environments. To this end, we present Reinforcement Learning with Augmented Data (RAD), a simple plug-and-play module that can enhance most RL algorithms. We perform the first extensive study of general data augmentations for RL on both pixel-based and state-based inputs, and introduce two new data augmentations - random translate and random amplitude scale. We show that augmentations such as random translate, crop, color jitter, patch cutout, random convolutions, and amplitude scale can enable simple RL algorithms to outperform complex state-of-the-art methods across common benchmarks. RAD sets a new state-of-the-art in terms of data-efficiency and final performance on the DeepMind Control Suite benchmark for pixel-based control as well as OpenAI Gym benchmark for state-based control. We further demonstrate that RAD significantly improves test-time generalization over existing methods on several OpenAI ProcGen benchmarks. Our RAD module and training code are available at this https URL.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Reinforcement Learning, Representation & Self-Supervised Learning

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Reinforcement Learning

In Representation & Self-Supervised Learning

Citations in the corpus

What Matters in Learning from Offline Human Demonstrations for Robot Manipulation (2021, methodology): “To implement random pixel shifts [84–87] for input image observations, we take large random crops from the source images when feeding image observations to any network.”
Widening the Pipeline in Human-Guided Reinforcement Learning with Explanation and Context-Aware Data Augmentation (2020, methodology): “The context-agnostic augmentations we compare to include Gaussian blurring and random cropping, which result in state-of-the-art sample efﬁciency in recent works [20, 18].”
Visual Imitation Made Easy (2020, methodology): “To improve the performance of our networks with limited data, we experiment with using the following data augmentations in training [39, 40, 48]: • Color Jitter: Randomly adds up to ± 20% random noise to the brightness, contrast and saturation of each observation.”
CURL: Contrastive Unsupervised Representations for Reinforcement Learning (2020, background): “If one has access to a rich stream of rewards from the underlying environment and is interested in optimzing the performance in terms of average reward, RAD and DrQ are likely to work better than CURL.”
Meta-Learning Requires Meta-Augmentation (2020, background): “Recent work in reinforcement learning [18, 21] also demonstrate large increases in reward from applying image augmentations.”
Data-Efficient Reinforcement Learning with Self-Predictive Representations (2020, background): “Data augmentation has also been found to improve generalization of reinforcement learning methods (Combes et al., 2018; Laskin et al., 2020) in multi-task and transfer settings.”
A Survey on Masked Autoencoder for Self-supervised Learning in Vision and Beyond (2022, methodology): “The results indicate that CNN-based RAD [113] performs better on most image-based deep RL tasks, but reconstruction-based ViT models [1], [24] outperform RAD on some tasks.”
Automatic Data Augmentation for Generalization in Reinforcement Learning (2021, methodology): “More similar to our work, Cobbe et al. [11], Ye et al. [69] and Laskin et al. [41] add augmented observations to the training buffer of an RL agent.”
Automatic Data Augmentation for Generalization in Deep Reinforcement Learning (2020, methodology): “Recent results have shown that the use of transformations common in the supervised learning paradigm (such as cutout, color-jitter, or random shifts) can also be beneﬁcial in RL [11, 37, 32, 34].”
Learning Domain Invariant Representations in Goal-conditioned Block MDPs (2021, methodology): “Data augmentation via RAD provides marginal improvement over the vanilla Skew-Fit.”
Masked Visual Pre-training for Motor Control (2022, methodology): “We compare MVP to two state-of-the-art methods that train the vision encoder with environment data: CURL (Srinivas et al., 2020) and RAD (Laskin et al., 2020).”
Dreaming: Model-based Reinforcement Learning by Latent Imagination without Reconstruction (2020, methodology): “MFRL methods without representation learning : Recently proposed MFRL methods, which include reinforcement learning with augmented data (RAD) [28], data-regularized Q (DrQ) [29], and simple uniﬁed framework for reinforcement learning using ensembles (SUNRISE) [30], have achieved state-of-the-art res”
MimicGen: A Data Generation System for Scalable Robot Learning using Human Demonstrations (2023, methodology): “In fact, this can be complementary to MimicGen— we leverage pixel shift randomization [7,36–39] when training image-based agents on MimicGen data.”
Mastering Visual Continuous Control: Improved Data-Augmented Reinforcement Learning (2021, methodology): “Recently, it has been shown that data augmentation techniques, commonplace in Computer Vision, are also important for achieving the state-of-the-art performance in image-based RL [Yarats et al., 2021b, Laskin et al., 2020].”
Predictive Information Accelerates Learning in RL (2020, background): “[27, 28] showed that input perturbation facilitates actor-critic learning, and we show that it also improves CEB learning.”
Bisimulation Makes Analogies in Goal-Conditioned Reinforcement Learning (2022, methodology): “We use a similar encoder architecture to (Laskin et al., 2020b) except we use a six layer convolutional network.”
Deep Reinforcement Learning at the Edge of the Statistical Precipice (2021, methodology): “We compare SAC+AE [114], SLAC [58], Dreamer [37], CURL [98], RAD [57], DrQ [53], PISAC [60], SUNRISE [59], and CURL-D2RL [97].”
Contrastive Learning as Goal-Conditioned Reinforcement Learning (2022, methodology): “While such representations might emerge from end-to-end training [7, 79, 118, 125], prior work has found it necessary to equip RL algorithms with perception-specific loss functions [31, 43, 70, 74, 89, 91, 101, 140] or data augmentations [72, 74, 117, 136], effectively decoupling the representation ”