PaperlayerView exact revision on arXiv

Source-linked AI summary

Automatic Data Augmentation for Generalization in Deep Reinforcement Learning

Roberta Raileanu, Max Goldstein, Denis Yarats, Ilya Kostrikov, Rob Fergus

arXiv:2006.12862v2cs.LGcs.AI

TL;DR

Deep RL agents can fail to generalize to unseen environments, while suitable augmentations vary by task. The paper combines actor-critic regularization with automatic augmentation selection or learning; UCB-DrAC reports a 40% test-performance improvement on Procgen and better invariance to control-irrelevant changes.

  • Problem

    RL agents can fail to generalize to unseen environments despite training on similar settings, and choosing an appropriate augmentation can require expert knowledge or evaluating many transformations.

    Source: §1 · p. 1

  • Method

    DrAC regularizes policy and value functions for transformation invariance, alongside UCB or RL2 selection from a fixed set or Meta-DrAC learning convolutional transformation weights.

    Source: §1 · p. 1

  • Results

    UCB-DrAC improves test performance by 40% on Procgen and is reported to set a new state-of-the-art there.

    Source: §6 · p. 8

  • Takeaways & Limitations

    The learned policies and representations are more invariant to spurious correlations between observations and rewards, including changes that do not affect control.

    Source: §1 · p. 2, §6 · p. 8

Abstract

from arXiv · show

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.

1. Introduction

Deep RL agents can fail to generalize beyond familiar environments, and data augmentation in RL requires care because task-relevant transformations differ from those in supervised learning. The paper proposes regularized actor-critic learning and three approaches for selecting or learning augmentations, reporting strong Procgen results and greater invariance to control-irrelevant changes.

1. Introduction

  • Deep RL agents may fail on unseen environments despite training on similar settings, suggesting they memorize trajectories rather than learn transferable skills.The paper focuses on data augmentation while noting that differences between supervised learning and RL must be considered.Source: §1 · p. 1
  • DrAC adds policy and value regularizers that make actor-critic use of data augmentation theoretically sound and constrain both functions toward transformation invariance.Naive augmentation can produce unprincipled objective estimates and poor performance; the approach is described with PPO and can be integrated with other actor-critic algorithms with discrete stochastic policies.Source: §1 · p. 1
  • Three automatic approaches are proposed: UCB and RL2 select from a fixed transformation set, while Meta-DrAC meta-learns convolutional-network weights.The approaches address prior reliance on expert choice or separate evaluation of many transformations.Source: §1 · p. 1
  • The approaches are evaluated on Procgen, a benchmark of 16 procedurally generated environments with visual observations..Source: §1 · p. 1
  • In the UCB-DrAC overview, a UCB bandit selects an image transformation, and DrAC receives both original and augmented observations to learn invariant policy and value functions.Source: Figure 1 · p. 2
  • UCB-DrAC outperforms generalization-focused baselines on Procgen train and test, and the paper reports state-of-the-art Procgen results.The paper also reports outperforming popular RL methods on four DeepMind Control tasks with natural and synthetic distractors.Source: §1 · p. 2
  • The learned policies and representations are more invariant to changes that do not affect reward or transitions, including background changes.Source: §1 · p. 2

2. Background

The background frames the objective as finding a policy that performs across a distribution of partially observable environments, despite training access being restricted to a fixed set. It describes Procgen's seeded levels and PPO's actor-critic optimization setup.

2. Background

  • The task is to maximize expected discounted reward over a distribution of POMDPs, while training uses a fixed set sampled from that distribution.Each POMDP is defined by state and observation spaces, actions, transitions, rewards, and a discount factor.Source: §2 · p. 2
  • Procgen levels correspond to POMDPs generated from game-specific distributions; agents train on 200 fixed levels and test on randomly sampled seeds.Source: §2 · p. 2
  • PPO learns a policy and value function by alternating environment sampling with stochastic-gradient maximization of a clipped surrogate objective.Its policy-gradient term uses importance sampling; the behavior policy generated the training state-action distribution, while the optimized policy generates the target distribution.Source: §2 · p. 2

3. Automatic Data Augmentation for RL

Naively augmenting PPO observations can make its objective estimate unsound and impair learning. DrAC adds policy and value regularization, then combines it with automatic methods for choosing or learning augmentations.

3.1. Data Augmentation in RL

  • Naively applying transformations to PPO observations can make its objective estimate unsound and prevent learning a good policy.The transformed-policy ratio need not equal one in the on-policy case and can be arbitrarily large.Source: §3.1 · p. 3

3.2. Policy and Value Function Regularization

  • DrAC regularizes both policy and value functions toward invariance, including for optimality-invariant transformations that preserve both.The regularizers can be added to actor-critic algorithms with discrete stochastic policies.Source: §3.2 · p. 3
  • DrAC uses transformed observations only for regularization, retaining the correct importance-sampling estimate for PPO’s main objective.This is intended to retain augmentation’s regularizing effect while mitigating adverse consequences for the RL objective.Source: §3.2 · pp. 3–4

3.3. Automatic Data Augmentation

  • UCB-DrAC selects from a fixed transformation set using UCB and scores each choice with a sliding window of recent returns.The recent-return window addresses the nonstationarity caused by the agent improving during training.Source: §3.3 · p. 4
  • A meta-learner can select an augmentation before each DrAC update and update from the resulting trajectories’ mean return.Source: §3.3 · p. 4
  • Meta-DrAC learns a convolutional transformation’s weights using meta-train and meta-test splits of DrAC’s buffer.Source: §3.3 · p. 4

4. Experiments

The experiments evaluate augmentation and regularization methods on Procgen and DeepMind Control tasks. They compare automatic augmentation selection, examine regularization effects, and assess robustness to irrelevant background changes.

4. Experiments

  • The evaluation covers 16 Procgen games and four pixel-based DeepMind Control tasks; Procgen agents train on 200 levels per game and are tested across the full level distribution.Methods use PPO as a base; DrAC adds actor and critic regularization to the transformation set used by RAD.Source: §4 · p. 4

4.1. Generalization Performance on Procgen

  • UCB-DrAC performs best on Procgen train and test and matches or exceeds DrAC with the best augmentation; adding PLR yields a reported 76% gain over PPO.The benchmark aggregates 16 tasks; regularized DrAC also outperforms RAD, which uses the best augmentation without the regularization losses.Source: Table 1 · p. 5, §4.1 · pp. 5–6

4.2. DeepMind Control with Distractors

  • On DeepMind Control tasks with natural video distractors, UCB-DrAC outperforms PPO and RAD; it also outperforms baselines with default and simple distractors.The evaluation uses Cartpole Balance, Finger Spin, Walker Walk, and Cheetah Run.Source: §4.2 · p. 6

4.3. Regularization Effect

  • Without policy and value regularization, RAD can perform drastically worse than vanilla RL, while DrAC is comparable to or better than PPO.The comparison uses grayscale and random-convolution augmentations on Chaser, Miner, and StarPilot.Source: §4.3 · p. 7

4.4. Automatic Augmentation

  • No single augmentation works best across Procgen games, and UCB-DrAC finds effective task-specific augmentations, reaching performance comparable to or better than using each game's best augmentation throughout training.Ninja and Jumper favor different transformations, while Miner, Climber, and Ninja share color-jitter as their best-performing augmentation.Source: §4.4 · p. 7, §4.5 · p. 7

4.5. Robustness Analysis

  • UCB-DrAC has lower policy JSD and higher representation cycle-consistency than PPO across changes to background themes, indicating greater robustness to task-irrelevant visual changes.The compared observations vary in background while retaining the same semantics; cycle-consistency uses unseen environments with the same layout.Source: §4.5 · pp. 7–8

5. Related Work

Prior work addresses deep RL generalization through regularization, representation methods, and data augmentation. This paper distinguishes its approach by addressing theoretical and practical issues in RL augmentation and by automatically selecting transformations.

Generalization in Deep RL

Data Augmentation

  • Naive RL data augmentation can create theoretical and practical issues, which the paper's algorithmic contributions address while retaining augmentation's regularization effect.The cited related work also notes that automatic augmentation methods designed for supervised learning cannot be applied to RL without algorithmic changes.Source: §5 · p. 8
  • Recent RL augmentation methods use fixed augmentations rather than automatically finding the most effective one (Cobbe et al., 2018; Lee et al., 2020; Laskin et al., 2020); this work automatically selects one, regularizes both actor and critic, and targets generalization over sample efficiency.Source: §5 · p. 8

6. Discussion

The discussion reports that UCB-DrAC improves Procgen training and test performance and achieves state-of-the-art results. The PPO appendix explains how importance sampling and clipping constrain policy updates.

6. Discussion

  • UCB-DrAC improves Procgen training performance by 19% and test performance by 40%, and sets a new state-of-the-art on the benchmark.The discussion also reports greater invariance of learned policies and representations to spurious correlations between observations and rewards.Source: §6 · p. 8

A. PPO

  • PPO alternates environment interaction for sampling data with stochastic-gradient optimization of an objective for a policy and value function.Its objective also includes an entropy bonus to aid exploration and a value-function loss.Source: A. PPO · p. 12
  • PPO estimates its policy-gradient objective using importance sampling when data come from a behavior policy different from the policy being optimized.The advantage estimate uses the behavior policy’s sampling distribution; the ratio between new and behavior policy probabilities supplies the importance weight.Source: A. PPO · p. 12
  • PPO clips the policy probability ratio to [1 − ϵ, 1 + ϵ] and takes the minimum of clipped and original objectives to discourage extreme policy updates.This simpler clipped surrogate implements a constraint inspired by TRPO, which limits policy changes by minimizing KL divergence between old and new policies.Source: A. PPO · p. 12
  • Adam rescales loss-component gradients, allowing DrAC to use one extra regularization coefficient rather than separate policy and value coefficients.The passage says this alleviates the burden of hyperparameter search.Source: A. PPO · p. 12

B. Naive Application of Data Augmentation in RL

Naively applying augmentation to importance-sampling policy-gradient objectives can be theoretically incorrect and can prevent PPO from learning well. The appendix also describes cycle-consistency as a measure of trajectory alignment in learned representations.

B. Naive Application of Data Augmentation in RL

  • The argument about naive augmentation applies to importance-sampling policy-gradient algorithms including TRPO, PPO, IMPALA, and ACER.PPO is used to explain the issue because these algorithms use importance sampling, typically when making multiple policy updates from the same data.Source: B. Naive Application of Data… · p. 13
  • Using augmented observations in PPO’s policy-gradient objective produces incorrect importance-sampling estimates because rewards and transitions are defined for true, not augmented, observations.DrAC leaves the policy-gradient objective on true observations and uses augmented observations in additional regularization losses; the passage says both proposed RAD ratios are incorrect.Source: B. Naive Application of Data… · p. 13

C. Cycle-Consistency

  • Cycle-consistency measures whether a trajectory point returns within one index of its starting point after a nearest-neighbor round trip in representation space.The passage states that this implies accurate alignment of two trajectories in hidden space; its three-way variant checks both paths through a third trajectory.Source: C. Cycle-Consistency · p. 13

D. Automatic Data Augmentation Algorithms

The paper presents automatic augmentation approaches that either select from a fixed transformation set or learn transformation weights. UCB and RL2 select augmentations, while Meta-DrAC meta-learns a convolutional transformation.

Algorithm 2 UCB-DrAC

  • UCB-DrAC selects an augmentation from a fixed set, updates the agent with one DrAC step, and tracks returns for each augmentation to update its Q-value.Its pseudocode initializes selection counts, Q-values, and FIFO return windows, and uses an exploration coefficient for UCB selection.Source: Algorithm 2 UCB-DrAC · p. 13
  • RL2-DrAC uses an LSTM selection policy that conditions on the previously chosen augmentation and its resulting average return, then learns from returns using REINFORCE.The selected augmentation is used for a one-step DrAC update before the resulting mean return rewards the selection policy.Source: Algorithm 2 UCB-DrAC · pp. 13–14

Algorithm 4 Meta-DrAC

  • Meta-DrAC learns convolutional augmentation weights by splitting experience into meta-training and meta-testing data and using MAML to maximize post-update return.For each DrAC update, trajectories are split 9:1; the learned CNN transforms observations, and the objective is the average return after a DrAC update.Source: Algorithm 2 UCB-DrAC · p. 14, Algorithm 4 Meta-DrAC · p. 14

E. Hyperparameters

The experiments use PPO-based methods and report selected hyperparameters for DrAC and its automatic-augmentation variants. The appendix also notes a trade-off in UCB’s exploration coefficient and a benchmark-wide limitation of a comparison baseline.

E. Hyperparameters

  • The best DrAC regularization coefficient found in the grid search was α_r = 0.1, and this value was used for all reported results.The search considered α_r values from 0.0001 to 1.0.Source: E. Hyperparameters · p. 15
  • UCB-DrAC used exploration coefficient c = 0.1 and sliding-window size K = 10 after grid searches over both settings.The tested values were c ∈ [0.0, 0.1, 0.5, 1.0, 5.0] and K ∈ [10, 50, 100].Source: E. Hyperparameters · p. 15
  • RL2-DrAC used hidden-state dimension 32, learning rate 5e −4, and entropy coefficient 1e −3; Meta-DrAC used learning rate 7e −4 and entropy coefficient 1e −2.These are the reported best values from the respective hyperparameter searches.Source: E. Hyperparameters · p. 15
  • Rand-FM used the authors’ recommended implementation settings, while IBAC-SNI used the authors’ selected best version evaluated on CoinRun.The paper reports that IBAC-SNI outperforms other methods on heist, maze, and miner but remains significantly worse than UCB-DrAC across Procgen.Source: E. Hyperparameters · p. 15
  • UCB’s exploration coefficient balances choosing high-Q augmentations with exploring alternatives: too little can settle on a suboptimal choice, while too much can delay convergence.The figure caption describes this behavior on Dodgeball.Source: Figure 5 · p. 16

F. Analysis of UCB’s Behavior

UCB’s selection speed depends on how distinguishable the augmentations’ effects are. Its exploration coefficient must balance finding a better augmentation against converging efficiently.

F. Analysis of UCB’s Behavior

  • UCB converges to the best augmentation earlier on Dodgeball, where augmentation effects differ significantly, than on Ninja, where it takes 15M training steps.Early discovery is associated with significant improvements over PPO on both train and test environments.Source: F. Analysis of UCB’s Behavior · p. 16
  • An exploration coefficient of 0.1 worked well across Procgen games; zero can lock UCB into a suboptimal augmentation, while an excessively large value can delay convergence.The coefficient balances exploration and exploitation of augmentations.Source: F. Analysis of UCB’s Behavior · p. 16

G. Procgen Benchmark

The Procgen materials present example levels, per-game best-augmentation tables, score breakdowns, and learning curves. The supplied captions describe these comparisons but do not provide score values or curve outcomes.

G. Procgen Benchmark

  • The Procgen benchmark illustrations show procedurally generated levels from 15 named environments.The caption lists StarPilot, CaveFlyer, Dodgeball, FruitBot, Chaser, Miner, Jumper, Leaper, Maze, BigFish, Heist, Climber, Plunder, Ninja, and BossFight.Source: Figure 6 · p. 17

H. Best Augmentations

  • Tables 4 and 5 identify the best augmentation type for each game, evaluated on test environments.The supplied table captions do not state the augmentation selected for any individual game.Source: Table 4 · p. 17, Table 5 · p. 17

I. Breakdown of Procgen Scores

  • Tables 6 and 7 report Procgen scores on train and test levels after 25M environment steps.The captions specify means and standard deviations over 10 runs, but the supplied text omits the score values.Source: Table 6 · p. 18, Table 7 · p. 18

J. Procgen Learning Curves

  • The learning-curve figures cover test and train performance for automatically selecting augmentations and compare PPO, DrAC, and actor-only or critic-only regularization ablations.The supplied captions name UCB-DrAC and RL2-DrAC among the automatic-selection approaches; Figures 9 and 10 specify test and train comparisons, respectively.Source: Figure 7 · p. 19, Figure 8 · p. 20, Figure 9 · p. 21

K. DeepMind Control Suite Experiments

The DeepMind Control experiments use selected PPO settings and compare performance across backgrounds with and without distractors. On simple distractor backgrounds, UCB-DrAC outperforms PPO and RAD with their best augmentations.

K. DeepMind Control Suite Experiments

  • The DMC experiments use a grid search for PPO hyperparameters, three stacked observation frames, and crop as the best augmentation among eight transformations.Action repeats are two for Walker Walk and Finger Spin, and four for the other tasks.Source: K. DeepMind Control Suite Experiments · p. 23
  • The illustrated DMC backgrounds include defaults without distractors, simple ideal-gas-video distractors, and natural Kinetics-video distractors.The examples show Finger Spin, Cheetah Run, and Walker Walk.Source: Figure 11 · p. 23
  • Figure 2 reports average return on DMC tasks with default backgrounds, with mean and standard deviation computed over five runs.The supplied caption does not state comparative outcomes.Source: Figure 12 · p. 23
  • UCB-DrAC outperforms PPO and RAD with the best augmentations on DMC tasks with simple distractor backgrounds.The figure reports average return over five seeds.Source: Figure 13 · p. 23, K. DeepMind Control Suite Experiments · p. 23

Related papers

Loading 2006.12862v2…