Automatic Data Augmentation for Generalization in Deep Reinforcement Learning

TL;DR

This work tackles generalization in deep reinforcement learning by enabling principled data augmentation within actor-critic methods. It introduces Data-regularized Actor-Critic (DrAC), which adds policy and value regularizers to enforce invariance to state transformations, allowing augmentation to improve learning without biasing the objective. The authors propose three automatic augmentation strategies—UCB-DrAC, RL2-DrAC, and Meta-DrAC—evaluated on Procgen and DeepMind Control with distractors, achieving state-of-the-art generalization and robustness (e.g., ~40% test performance gains on Procgen) while producing more invariant representations. The approach yields strong practical impact by providing a principled, automatic augmentation framework applicable to a range of actor-critic algorithms, with code available for reproducibility ($J_{\mathrm{DrAC}} = J_{\mathrm{PPO}} - \alpha_{r}(G_{\pi} + G_{V})$).

Abstract

Deep reinforcement learning (RL) agents often fail to generalize to unseen scenarios, even when they are trained on many instances of semantically similar environments. Data augmentation has recently been shown to improve the sample efficiency and generalization of RL agents. However, different tasks tend to benefit from different kinds of data augmentation. In this paper, we compare three approaches for automatically finding an appropriate augmentation. These are combined with two novel regularization terms for the policy and value function, required to make the use of data augmentation theoretically sound for certain actor-critic algorithms. We evaluate our methods on the Procgen benchmark which consists of 16 procedurally-generated environments and show that it improves test performance by ~40% relative to standard RL algorithms. Our agent outperforms other baselines specifically designed to improve generalization in RL. In addition, we show that our agent learns policies and representations that are more robust to changes in the environment that do not affect the agent, such as the background. Our implementation is available at https://github.com/rraileanu/auto-drac.

Figures (13)

• Figure 1: Overview of UCB-DrAC. A UCB bandit selects an image transformation (e.g. random-conv) and applies it to the observations. The augmented and original observations are passed to a regularized actor-critic agent (i.e. DrAC) which uses them to learn a policy and value function which are invariant to this transformation.
• Figure 2: Average return on DMC tasks with natural video backgrounds with mean and standard deviation computed over 5 seeds. UCB-DrAC outperforms PPO and RAD with the best augmentations.
• Figure 3: Comparison between RAD and DrAC with the same augmentations, grayscale and random convolution, on the test environments of Chaser (left), Miner (center), and StarPilot (right). While DrAC's performance is comparable or better than PPO's, not using the regularization terms, i.e. using RAD, significantly hurts performance relative to PPO. This is because, in contrast to DrAC, RAD does not use a principled (importance sampling) estimate of PPO's objective.
• Figure 4: Cumulative number of times UCB selects each augmentation over the course of training for Ninja (a) and Dodgeball (c). Train and test performance for PPO, DrAC with the best augmentation for each game (color-jitter and crop, respectively), and UCB-DrAC for Ninja (b) and Dodgeball (d). UCB-DrAC finds the most effective augmentation from the given set and reaches the performance of DrAC. Our methods improve both train and test performance.
• Figure 5: Behavior of UCB for different values of its exploration coefficient c on Dodgeball. When c is too small, UCB might converge to a suboptimal augmentation. On the other hand, when c is too large, UCB might take too long to converge.