PaperlayerView exact revision on arXiv

Source-linked AI summary

Reinforcement Learning with Augmented Data

Michael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, Aravind Srinivas

arXiv:2004.14990v5cs.LGstat.ML

TL;DR

RL methods remain limited by data-efficiency and generalization, particularly for visual observations. RAD adds data augmentation directly to RL pipelines across pixel-based and state-based inputs, without changing the underlying algorithm. The approach improves benchmark performance, data-efficiency, and test-time generalization, while showing weaker gains on ProcGen's structurally varying Jumper task.

  • Problem

    RL still faces poor data-efficiency and generalization, challenges exacerbated by high-dimensional, partially observable pixel observations.

  • Method

    RAD is a plug-and-play module that applies data augmentations to RL observations, including novel random translation and random amplitude scaling.

  • Results

    RAD outperforms prior baselines across DeepMind Control and OpenAI Gym, improves ProcGen test-time generalization, and achieves significant pixel-based gains without changing the RL algorithm.

  • Takeaways & Limitations

    Data augmentation alone can improve RL data-efficiency and generalization across pixel-based and state-based benchmarks while remaining simple, efficient, and broadly applicable.

  • Takeaways & Limitations

    RAD's augmentation gains are less significant on ProcGen Jumper, whose structural map-layout generalization may require recurrent policies.

Abstract

from arXiv · show

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 https://www.github.com/MishaLaskin/rad.

1 Introduction

RL from visual observations has achieved strong results, but remains limited by poor data-efficiency and generalization. RAD addresses both challenges by applying data augmentation directly to RL inputs without changing the underlying algorithm.

  • RL from visual observations combines convolutional neural networks with modern credit-assignment methods to achieve strong results.
  • Poor data-efficiency and generalization remain major RL challenges, especially with high-dimensional, partially observable pixel inputs.Real-world RL successes can require months of data collection or training.
  • Data augmentation improves representation consistency across multiple views and has supported data-efficiency and generalization in computer vision.
  • RAD applies data augmentation to RL input observations so agents learn from multiple views of the same input.
  • RAD is presented as the first extensive RL data-augmentation study requiring no changes to the RL algorithm or additional domain assumptions beyond observation type.
  • RAD outperforms prior state-of-the-art baselines on pixel-based DeepMind Control and state-based OpenAI Gym benchmarks.The study reports improved test-time generalization on several OpenAI ProcGen environments and introduces random translation and random amplitude scaling.

2 Related work

Prior RL work used data augmentation for generalization and representation learning, but lacked an extensive benchmark study across widely used settings. RAD instead applies augmentations directly within RL without auxiliary losses or algorithm changes.

  • Computer-vision data augmentation introduced invariance priors that improved classification performance.
  • Prior RL studies used domain randomization, cutout, and random convolution to improve transfer or generalization.
  • CURL combined data augmentation with contrastive and reinforcement learning losses to learn contrastive representations.
  • RAD directly uses augmented observations for reinforcement learning without an auxiliary loss.
  • Unlike DrQ's SAC-specific approach, RAD can plug into on-policy PPO and off-policy SAC without changing the underlying algorithm.

3 Background

The paper frames RL as reward maximization in a Markov Decision Process and describes pixel observations as indirect state representations. It then outlines SAC and PPO as off-policy and on-policy baselines.

  • An RL agent operates in a Markov Decision Process defined by states, actions, transitions, and a discount factor.
  • The agent uses experience to maximize expected discounted rewards without prior knowledge of the transition distribution.
  • Pixel-based RL receives high-dimensional image observations that indirectly represent the underlying state.
  • Soft Actor-Critic: SAC learns a policy and critic using reward and policy-entropy objectives, with critic updates based on replay-buffer transitions.
  • Soft Actor-Critic: SAC estimates next-state values using sampled actions and a slowly updated target critic, while learning the policy from the soft-Q function.
  • Proximal policy optimization: PPO learns continuous or discrete policies from action advantages and a clipped-ratio loss over recent on-policy experience.

4 Reinforcement learning with augmented data

RAD processes sampled observations with stochastic augmentations before RL training, applying the same randomization across stacked frames or time. The method covers image and state-based inputs, including novel translation and amplitude-scaling augmentations.

  • RAD evaluates stochastic image augmentations for both off-policy SAC and on-policy PPO.
  • Training samples observations from a replay buffer or recent trajectory and augments them within each minibatch.
  • Augmentations vary randomly across a minibatch but remain consistent across stacked frames or time to preserve temporal information.
  • Augmentations of image-based input: The image study investigates crop, translate, window, grayscale, cutout, cutout-color, flip, rotate, random convolution, and color jitter.
  • Augmentations of image-based input: Random translation places the full image at a random location within a larger frame, while cropping extracts a random patch.
  • Extension to state-based inputs: For state-based inputs, random amplitude scaling multiplies states by a uniform random variable, while Gaussian noise adds a Gaussian vector.Both augmentations use the same randomization for current and next states; amplitude scaling preserves intrinsic information such as input sign.

5 Experimental results

RAD improves data-efficiency and generalization across pixel-based and state-based RL benchmarks, with especially strong results from random crop, translation, and amplitude scaling. Its benefits extend to unseen ProcGen environments, although gains are smaller for structural navigation generalization in Jumper.

  • Data-efficiency on pixel-based inputs: RAD achieves state-of-the-art performance on all 6 DMControl environments at both 100k and 500k environment steps.The benchmark reports averages across 10 seeds for the six main environments.
  • Data-efficiency on pixel-based inputs: 4x improvement over pixel SAC is achieved on both DMControl100k and DMControl500k solely through data augmentation.RAD does not use a forward model or auxiliary task for this comparison.
  • Augmentation effects: Random crop is the most effective augmentation by a large margin on Walker Walk, while random translation or crop has the highest impact on final performance.The translation component accounts for the benefit of cropping, whereas random information masking has little effect; two additional pixels can sometimes suffice.
  • Generalization on ProcGen: 55.8% relative gain over pixel-based PPO is obtained by RAD with random crop on ProcGen BigFish.Random crop and cutout also significantly improve generalization on BigFish and StarPilot.
  • Generalization on ProcGen: RAD trained with 100 levels outperforms pixel-based PPO trained with 200 levels on both BigFish and StarPilot.This comparison indicates that augmentation can be more effective than simply increasing the number of training environments.
  • Generalization limitations: Augmentation gains are less significant on Jumper because its structural generalization across map layouts likely requires recurrent policies.Jumper is a navigation task, unlike the visual generalization settings where augmentation improves performance more clearly.
  • Data-efficiency on state-based inputs: RAD is state-of-the-art on 4 of 6 OpenAI Gym environments and improves the average return over POPLIN-P by 1.7x in Walker.Random amplitude scaling is hypothesized to promote robustness to input noise while preserving sign and relative differences between state values.
  • Data-efficiency on state-based inputs: RAD improves RL with data augmentation for both visual and state inputs, indicating that the approach is not specific to pixel-based observations.The paper frames this as generality from incorporating inductive biases such as amplitude invariance through augmentation.

6 Conclusion

RAD is a simple plug-and-play module that uses data augmentations to improve RL data-efficiency and generalization without changing the underlying algorithm. Its implementation is efficient and open-sourced for future research and real-world applications.

  • RAD enhances reinforcement learning methods with data augmentations without modifying the underlying RL algorithm.
  • RAD significantly improves data-efficiency and generalization for pixel-based RL on DeepMind Control Suite and OpenAI ProcGen benchmarks.
  • The implementation is simple, efficient, and open-sourced to support future research and real-world RL applications.

7 Broader Impact

The paper situates RAD within broader concerns about computational complexity, accessibility, and real-world RL safety. It connects data augmentation to prior methods for diversifying training data and improving representations, while noting important practical and ethical boundaries.

  • Increasing model complexity and compute requirements can raise energy demands and make AI research less accessible.
  • The paper states that Deep RL remains far from solving real-world-complexity robotics problems and emphasizes complementary safety research.
  • Data augmentation and related techniques diversify RL training data through domain randomization, hindsight relabeling, and synthetic rollouts.
  • CURL uses augmented views with an auxiliary contrastive loss, whereas RAD applies augmentations directly within RL without an auxiliary loss.
  • RAD extends prior RL generalization work by investigating additional augmentations, including random crop and color jitter, across more tasks.
  • Attention maps show random crop focusing encoders on robot bodies while reducing attention to irrelevant scene details such as background stars.

C Random translate ablations

The ablations examine how random translation works and show that its benefits arise mainly from translation rather than windowing or cropping. Even minimal translation improves performance over SAC without augmentation.

  • Random translation places a 100x100 image within larger frames of sizes 102, 104, 108, and 116 pixels.
  • Augmentation significantly improves performance over SAC without augmentation on Cartpole and Cheetah, even with the smallest 102-pixel frame.
  • Ablations of translation, cropping, and windowing find that the gains come primarily from translation.
  • RAD with translation produces more stable and efficient learning on most environments.

E Implementation details for DMControl

The DMControl implementation uses a shared encoder setup with mostly fixed hyperparameters, while ProcGen environments vary in objectives, layouts, themes, and entity dynamics. These settings define the benchmark context for the experiments.

  • DMControl experiments use the same encoder architecture as CURL and provide hyperparameters in Table 4.
  • Most DMControl hyperparameters remain unchanged across environments, except action repeat, learning rate, and batch size.
  • BigFish, StarPilot, Jumper, and Modified CoinRun vary in objectives, map structures, backgrounds, spawning patterns, and obstacle or agent appearances.

F.2 Implementation details

The implementation evaluates PPO and RAD on ProcGen settings and provides batched augmentation code, with RAD designed for efficient, frame-stack-consistent processing.

  • ProcGen experiments use the IMPALA CNN policy architecture and PPO with the hyperparameters proposed in prior work.
  • The augmentation code samples implement random cropping, cutout, and horizontal flipping on image batches.Cropping samples random spatial offsets, while cutout samples rectangular regions and fill values.
  • RAD’s primary implementation gain is efficient augmentation of stacked frame inputs consistently across each frame stack within a minibatch.The authors report that general-purpose single-frame augmentation APIs are unsuitable for this requirement.
  • RAD augmentations require nearly 2x less additional time than equivalent native PyTorch augmentations for 100k training steps.The comparison is reported in Table 6 as an average wall-clock result.
  • RAD and CURL both improve data-efficiency by enforcing consistency across augmented input observations, with RAD matching CURL and surpassing it in some environments.

I.2 Is data augmentation sufficient for RL from pixels?

RAD directly optimizes the RL objective on augmented views, whereas CURL adds an explicit contrastive objective; the paper therefore treats augmentation alone as promising but not universally sufficient.

  • RAD trains directly on multiple augmented observation views without an auxiliary loss, while CURL jointly optimizes reinforcement and contrastive objectives.
  • RAD’s stronger task-score performance over CURL is attributed to focusing optimization on task reward while implicitly enforcing consistency across augmented views.
  • CURL is described as a more general framework because its contrastive objective can operate without task or environment rewards.It can learn representations from high-dimensional observations gathered from random rollouts.
  • Representation-learning methods may be more important for sparse-reward, image-goal applications than for benchmarks centered on single- or multi-task reward optimization.
  • The paper concludes that RAD and representation-learning methods such as CURL can both be useful for data-efficient and generalizable RL research.

J.2 Experimental results on OpenAI Gym

On OpenAI Gym, random amplitude scaling is effective across almost all tested environments, with the single-variable version generally outperforming multivariate scaling and Gaussian noise.

  • Random amplitude scaling is effective in almost all OpenAI Gym environments.The result is reported for RAD variants evaluated across the benchmark.
  • RAS-S generally performs better than RAS-M because multivariate scaling can change the relative differences between state variables.
  • Batch normalization is reported to be less effective than RAD, indicating that RAD’s gains are not reproduced by normalization alone.
  • Figure 12 compares learning curves for RAS-S, RAS-M, and Gaussian noise on OpenAI Gym across four runs.Solid lines show means and shaded regions show standard deviations.
Loading 2004.14990v5…