Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels

By Ilya Kostrikov, Denis Yarats, Rob Fergus. 2020

946 citations recorded by the source corpus.

We propose a simple data augmentation technique that can be applied to standard model-free reinforcement learning algorithms, enabling robust learning directly from pixels without the need for auxiliary losses or pre-training. The approach leverages input perturbations commonly used in computer vision tasks to regularize the value function. Existing model-free approaches, such as Soft Actor-Critic (SAC), are not able to train deep networks effectively from image pixels. However, the addition of our augmentation method dramatically improves SAC's performance, enabling it to reach state-of-the-art performance on the DeepMind control suite, surpassing model-based (Dreamer, PlaNet, and SLAC) methods and recently proposed contrastive learning (CURL). Our approach can be combined with any model-free reinforcement learning algorithm, requiring only minor modifications. An implementation can be found 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

Habitat 2.0: Training Home Assistants to Rearrange their Habitat (2021, background): “[101] Ilya Kostrikov, Denis Yarats, and Rob Fergus.”
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].”
Anatomy of Catastrophic Forgetting: Hidden Representations and Task Semantics (2020, background): “While the past few years have seen the development of increasingly versatile machine learning systems capable of learning complex tasks [1, 2, 3, 4, 5], catastrophic forgetting remains a core challenge to designing machine learning systems.”
Provable Representation Learning for Imitation with Contrastive Fourier Features (2021, background): “We emphasize that our own contrastive objectives are distinct from typical approaches in image domains [15] and popular in image-based RL [26, 37, 36], which use prior knowledge to generate pairs of similar images (e.”
SoftGym: Benchmarking Deep Reinforcement Learning for Deformable Object Manipulation (2020, methodology): “Recent works [51, 52, 53] show evidence that the gap between image-based RL and state-based RL can be closed on a range of tasks with the data augmentation in reinforcement learning.”
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.”
Reinforcement Learning with Augmented Data (2020, methodology): “Concurrent and independent to our work, DrQ [23] uses data augmentations and weighted Q-functions in conjunction with the off-policy RL algorithm SAC [35] to achieve state-of-the-art data-efﬁciency results on the DeepMind Control Suite.”
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.”
Mastering Atari Games with Limited Data (2021, background): “(4) DrQ [21], which adds data augmentations to the input images while learning the original RL objective.”
Data-Efficient Reinforcement Learning with Self-Predictive Representations (2020, methodology): “We follow DrQ (Yarats et al., 2021) by: using the 3-layer convolutional encoder from Mnih et al.”
TD-MPC2: Scalable, Robust World Models for Continuous Control (2023, methodology): “Additionally, we also compare against current state-of-the-art visual RL methods (4) CURL (Srinivas et al., 2020), an extension of SAC that uses a contrastive auxiliary objective, and (5) DrQ-v2 , a model-free RL algorithm that uses data augmentation.”
Meta-Learning Requires Meta-Augmentation (2020, methodology): “Recent work in reinforcement learning [18, 21] also demonstrate large increases in reward from applying image augmentations.”
R3M: A Universal Visual Representation for Robot Manipulation (2022, methodology): “Such methods include contrastive learning with data augmentation [25, 26, 27], dynamics prediction [28, 29], bi-simulation [30], temporal or goal distance [31, 32], or domain specific information [33].”
Temporal Difference Learning for Model Predictive Control (2022, methodology): “We adopt action repeat hyperparameters for DMControl from previous work (Hafner et al., 2019; Kostrikov et al., 2020) for state-based experiments as well as the DMControl 100k benchmark; we list all values below.”
Transformers are Sample Efficient World Models (2022, methodology): “With minimal tuning, IRIS outperforms a line of recent methods (Kaiser et al., 2020; Hessel et al., 2018; Laskin et al., 2020; Yarats et al., 2021; Schwarzer et al., 2021) for sample-efﬁcient RL in the Atari 100k benchmark (Kaiser et al., 2020).”
Domain Generalization with MixStyle (2021, background): “Very recent studies (Laskin et al., 2020; Kostrikov et al., 2020) have shown that it is useful to combine a diverse set of label-preserving transformations, such as rotation, shifting and Cutout.”