[Paper Review] Automatic Data Augmentation for Generalization in Deep Reinforcement Learning

Roberta Răileanu, Max A. Goldstein+3 more|arXiv (Cornell University)|Jun 23, 2020

Reinforcement Learning in Robotics67 references52 citations

TL;DR

The paper introduces DrAC, a data-regularized actor-critic framework, plus three automatic augmentation strategies (UCB-DrAC, RL2-DrAC, Meta-DrAC) to automatically select effective data augmentations for RL tasks, achieving state-of-the-art generalization on Procgen and strong results on DeepMind Control with distractors.

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.

Motivation & Objective

• Address generalization gaps in deep RL due to overfitting to training environments.
• Propose a theoretically sound data augmentation framework for actor-critic methods.
• Develop regularization terms to enforce policy and value-invariance to state transformations.
• Automatically select effective augmentations via UCB, RL2 meta-learning, or CNN-weight learning.
• Demonstrate state-of-the-art Procgen performance and robustness to irrelevant environment changes.

Proposed method

• Introduce Data-regularized Actor-Critic (DrAC) with two regularization terms: policy regularization and value-function regularization.
• Use an optimality-invariant state transformation f(s, ν) to enforce invariance: V(s)=V(f(s,ν)) and π(a|s)=π(a|f(s,ν)).
• Maintain a standard actor-critic objective (PPO) and subtract regularization losses G_π and G_V weighted by α_r.
• Provide three automatic augmentation strategies: UCB-DrAC (bandit-based selection), RL2-DrAC (meta-learning selection), Meta-DrAC (weights of a CNN augmentation).
• Approximate the augmentation selection as a non-stationary bandit or meta-learning problem while updating the agent concurrently.
• Demonstrate invariance and robustness through cycle-consistency and JSD analyses.

Experimental results

Research questions

• RQ1 Can data augmentation be used safely with actor-critic RL algorithms without breaking the objective estimates?
• RQ2 Can we automatically identify task-specific augmentations that improve generalization in RL?
• RQ3 Do policy and value function regularizations to transformations improve stability and performance with augmented observations?
• RQ4 How do automatic augmentation methods (UCB-DrAC, RL2-DrAC, Meta-DrAC) compare in Procgen and DM Control with distractors?
• RQ5 Does the learned representation become more invariant to irrelevant visual changes (e.g., background)?

Key findings

• UCB-DrAC achieves state-of-the-art performance on Procgen, outperforming several baselines and matching or surpassing the best-task augmentation.
• Regularizing both the policy and value function is crucial; DrAC outperforms variants that regularize only one component.
• Automatic augmentation with UCB-DrAC provides robust, stable performance across games, often exceeding fixed-augmentation baselines.
• On DeepMind Control with distractors, UCB-DrAC consistently outperforms PPO and RAD in challenging background settings.
• Across Procgen, UCB-DrAC yields lower background sensitivity (higher cycle-consistency) and better invariance in representations.