[Paper Review] Reinforcement Learning with Augmented Data

Michael Laskin, Kimin Lee+4 more|arXiv (Cornell University)|Apr 30, 2020

Reinforcement Learning in Robotics54 references246 citations

TL;DR

RAD adds data augmentations to RL training, boosting data-efficiency and generalization across pixel-based and state-based inputs without changing the underlying RL algorithm.

ABSTRACT

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.

Motivation & Objective

• Motivate data-efficient and generalizable RL from visual observations.
• Investigate the effectiveness of diverse data augmentations in RL without extra losses.
• Show that augmentations improve performance on pixel-based and state-based benchmarks.
• Establish RAD as a simple plug-and-play module compatible with common RL methods.

Proposed method

• Apply stochastic data augmentations to input observations during RL training.
• Use augmentations consistently across frame stacks for pixel inputs and across time for state inputs.
• Plug RAD into base RL algorithms (SAC for off-policy; PPO for on-policy) without altering their core losses.
• Explore ten augmentations (crop, translate, window, grayscale, cutout, cutout-color, flip, rotate, random convolution, color jitter) for images; introduce random amplitude scaling for proprioceptive inputs.
• Evaluate on DMControl (pixel) and OpenAI ProcGen (generalization) as well as OpenAI Gym state-based tasks.
• Provide open-source RAD code implementation.

Experimental results

Research questions

• RQ1 Can data augmentations improve data-efficiency of RL with pixel inputs without changing the underlying algorithm?
• RQ2 Which augmentations most effectively improve RL performance and generalization across benchmarks?
• RQ3 Do augmentations extend benefits to state-based (proprioceptive) RL settings beyond pixel-based inputs?
• RQ4 How do augmentations affect representation learning and generalization to unseen environments?

Key findings

• RAD achieves state-of-the-art data-efficiency and final performance on all evaluated DMControl environments with pixel inputs.
• RAD improves the data-efficiency of pixel-based SAC by about a 4x gain in the tested regimes without auxiliary losses.
• RAD matches or surpasses many state-based baselines on DMControl environments, suggesting broad applicability to proprioceptive inputs.
• Random crop and random translate are among the most impactful augmentations for pixel inputs.
• RAD significantly improves test-time generalization on OpenAI ProcGen benchmarks.
• A novel random amplitude scaling augmentation improves state-based RL performance and robustness to input noise.