Automatic Data Augmentation for Generalization in Reinforcement Learning
Authors: Roberta Raileanu
Maxwell Goldstein
Denis Yarats
Ilya Kostrikov
Rob Fergus
Venue: NeurIPS 2021 Poster (2021)
Keywords: reinforcement learning
generalization
data augmentation
automatic machine learning

Abstract: Deep reinforcement learning (RL) agents often fail to generalize beyond their training environments. To alleviate this problem, recent work has proposed the use of data augmentation. However, different tasks tend to benefit from different types of augmentations and selecting the right one typically requires expert knowledge. In this paper, we introduce three approaches for automatically finding an effective augmentation for any RL task. 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 actor-critic algorithms. Our method achieves a new state-of-the-art on the Procgen benchmark and outperforms popular RL algorithms on DeepMind Control tasks with distractors. In addition, our agent learns policies and representations which are more robust to changes in the environment that are irrelevant for solving the task, such as the background. 

Automatic Data Augmentation for Generalization in Reinforcement Learning
Roberta Raileanu New York University
Max Goldstein New York University mag10380nyu. edu
Denis Yarats
New York University
raileanuêes. nyu...n..u..... nyu.edu
Facebook Al Research
demisyarats@en. nyu edu
Ilya Kostrikov
Rob Fergus
New York University kostrikovêcs nyu. edu
New York University
fergusecs myu. edu
Abstract
Deep reinforcement learning (RL) agents often fail to generalize beyond their training environments. To alleviate this problem, recent work has proposed the use of data augmentation. However, different tasks tend to benefit from different types augmentations and selecting the right one typically requires expert knowledge. In this paper, introduce three approaches for automatically finding effective augmentation for any RL task. These are combined with two novel regularization terms for the policy and value function, required to make the of data augmen- tation theoretically sound for actor-critic algorithms. Our method achieves a new slate-of-the-aal arte-rr(] pn the Procgen benchmark and outperforms popular RL algorithms on DeepMind Control tasks with distractors. In addition, our agent learns policies and representations which more robust to changes in the environment that are irrelevant for solving the task, such as the background. Our code is available at attps://github wttp:/cco/rrraaaaaeeeeaeeeeae http://raaalaaaaautdrrreeee
Introduction
Generalization to new environments remains major challenge in deep reinforcement learning (RL). Current methods fail to generalize to unseen environments even when trained similar settings 19 51 71 ITI 21 12 60]. This indicates that standard RL agents memorize specific trajectories rather than learning transferable skills. Several strategies have been proposed to alleviate
this problem, such as the use regularization [9 71
11] 28 data augmentation [1] 69 38 41
representation learning 772 74]. In this work, we focus on the use of data augmentation in RL. We identify key differences between supervised learning and reinforcement learning which need to be taken into account when using data augmentation in RL.
More specifically, we show that a naive application of data augmentation can lead both theoretical and practical problems with standard RL algorithms, such unprincipled objective estimates and poor performance. As a solution, we propose Data-regularized Actor-Critic or DrAC, a new algorithm that enables the use of data augmentation with actor-critic algorithms in theoretically sound way. Specifically, we introduce two regularization terms which constrain the agent's policy and value function to be invariant to various state transformations. Empirically, this approach allows the agent to learn useful behaviors (outperforming strong RL baselines) in settings which a naive use of data augmentation completely fails converges to a sub-optimal policy. While Policy Optimization (PPO, Schulman et al. (56) to describe and validate our approach, the method
Proximal
in June 2020, at the time of making this work publicly available on arXiv. it has since been surpassed.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

can be easily integrated with any actor-critic algorithm with discrete stochastic policy such as A3C
49
SAC [2
IMPALA
Augmented Observations
Environment Observations
Augmentation
UCB (f)
fi random-cony
DrAC
Figure 1: Overview of UCB-DrAC. A UCB bandit selects an image transformation (e.g. random- conv) and applies the observations. The augmented and original observations are passed to regularized actor-critic agent (i.e. DrAC) which uses them to learn a policy and value function which are invariant to this transformation.
The current use of data augmentation in RL either relies expert knowledge pick appropriate augmentation [1] 44 38) or separately evaluates a large number of transformations to find the best one [69] 41. In this paper, we propose three methods for automatically finding useful augmentation for a given RL task. The first two Iearn to select the best augmentation from a fixed set, using either variant of the upper confidence bound algorithm (UCB, Auer (2) or meta-learning (RL Wang et al. (6)). We refer to these methods as UCB-DrAC and RL22DrAC, respectively. The third method, Meta-DrAC, directly meta-learns the weights of a convolutional network, without access to predefined transformations (MAML, Finn et al. (0) Figure gives an overview of UCB-DrAC.
We evaluate these approaches the Procgen generalization benchmark [12] which consists of 16 procedurally generated environments with visual observations. Our results show that UCB-DrAC the most effective among these finding good augmentation, and is comparable
better than using DrAC with the best augmentation from given set. UCB-DrAC also outperforms baselines specifically designed to improve generalization RL (28] 44 41] on both train and test. addition, we show that our agent learns policies and representations that are more invariant to changes in the environment which do not alter the reward transition function (i.e. they are inconsequential for control), such as the background theme.
To summarize, work makes the following contributions: (i) introduce principled way of using data augmentation with actor-critic algorithms, (ii) we propose a practical approach for automatically selecting an effective augmentation in RL settings, (iii) show that the of data augmentation leads to policies and representations that better capture task invariances, and (iv) we demonstrate state-of-the-art results on the Procgen benchmark and outperform popular RL methods on four DeepMind Control tasks with natural and synthetic distractors.
2 Background
We consider distribution g(m) of Partially Observable Markov Decision Processes (POMDPs, 6 m M, with m defined by the tuple (Sm, Om, A, Tmy Ry, 7), where Sm is the state space, Om is the observation space, A is the action space, (m (s'ls, [s, a) is the transition function, Rm (s, a) the reward function, and 7 is the discount factor. During training, we restrict access to a fixed set of POMDPs, Mtrain {mi...., m. J. where m,
Vi n. The goal is to find policy Ty which maximizes the expected discounted reward over the entire distribution POMDPs, [(=))= = =Fggr.TT...p Et=0" 7R Rm az)]
In practice, game corresponds to distribution POMDPs (mm), and each level of a game corresponds to a POMDP sampled from that game's distribution
the Procgen benchmark which contains 16 procedurally generated games. Each
q. The POMDP m is determined by the seed (i.e. integer) used to generate the corresponding level. Following the setup from Cobbe et al. [12], agents trained fixed set of 200 levels (generated using seeds from 200) and tested on the full distribution of levels (generated by sampling seeds uniformly at random from all computer integers).
Proximal Policy Optimization (PPO, Schulman al. (56) is an actor-critic algorithm that learns policy Te and a value function Va with the goal of finding optimal policy for given POMDP.

---

PPO alternates between sampling data through interaction with the environment and maximizing a clipped surrogate objective function using stochastic gradient ascent.
One component of the PPO objective is the policy gradient term JPG, which estimated using importance sampling:
To (a(as) Telia (als) Abmia A
(PG = Σ
To(als)Aoond (s. a) Eu aEA
(I)
where A(-) is an estimate of the advantage function, Ta is the behavior policy used to collect trajectories (i.e. that generates the training distribution of states and actions), and Ty the policy we want to optimize (i.e. that generates the true distribution of states and actions). See Appendix A for full description of PPO.
3
Automatic Data Augmentation for RL
3.1 Data Augmentation in RL
Image augmentation has been successfully applied in computer vision for improving generalization on object classification tasks [58] 9 8 39]. As noted by Kostrikov et al. (38), those tasks are invariant certain image transformations such as rotations flips, which is not always the case in RL. For example, if your observation is flipped, the corresponding reward will be reversed for the left and right actions and will not provide an accurate signal to the agent. While data augmentation has been previously used in RL settings without other algorithmic changes approach is not theoretically sound.
[ 69 41], we argue that this
transformations are naively applied to observations in PPO's buffer, as done in Laskin et al. 41 the PPO objective changes and equation (T) is replaced by
To (a) (()) Teaaa (a|s)
ABasa (8,
(f) = Ea
where f S x 7- -> S is the image transformation. However, the right hand side of the above equation not sound estimate of the left hand side because (a f()) o Th (a(a)s) since nothing constrains Te (al fs)) to be close to To(a) (a(a)s). Note that in the on-policy case when 0 = Oold. the ratio used estimate the advantage should be equal one, which is not necessarily the case when using equation (p. In fact, one can define certain transformations f(-) that result in an arbitrarily large ratio Ty (s)) (als))
Figure 3 shows examples where a naive use of data augmentation prevents PPO from learning a good policy practice, suggesting that this is not just theoretical concern. In the following section, propose an algorithmic change that enables the use of data augmentation with actor-critic algorithms in principled way.
3.2 Policy and Value Function Regularization
Inspired by the recent work of Kostrikov et al. 38 we propose two novel regularization terms for the policy and value functions that enable the proper use of data augmentation for actor-critic al algorithms. Our algorithmic contribution differs from that of Kostrikov et al. (38) in that it constrains both the actor and the critic, as opposed only regularizing the Q-function. This allows method to be used with a different (and arguably larger) class of RL algorithms, namely those that learn a policy and value function.
Following Kostrikov et al. as a mapping that preserves both the agent's 's policy π and its value function V such that V(s) = (ff v)) and π (a|s) = x (a)f(s, (s,)(. v)), Vs E S, v E H. where v are the parameters f-). drawn from the set of all possible parameters H.
38 we define an aptimality-invarian state transformation f S +
To ensure that the policy and value functions are invariant such transformation of the input state, we propose an additional loss term for regularizing the policy,
G. = KL (a)s) TH ma(aff(s))))
(3)
3

---

as well as an extra loss term for regularizing the value function,
Gv = (Va(s) Ve (J))
(4)
Thus, data-regularized actor-critic method, DrAC, maximizes the following objective:
JDAAC = JPPo a, + Gv),
(5)
where or, is the weight of the regularization term (see Algorithm D
The use of G. and Gv ensures that the agent's policy and value function are invariant to the transformations induced by various augmentations. Particular transformations certain inductive biases relevant for the task (e.g. invariance with respect to colors or translations). In addition, G and Gv can be added to the objective of any actor-critic algorithm with discrete stochastic policy (e.g. A3C, TRPO, ACER, SAC, or IMPALA) without any other changes.
be used to impose
Note that when using DrAC, as opposed to the method proposed by Laskin et al. [41] still use the correct importance sampling estimate the left hand side objective in equation (instead of a wrong estimate as in equation ) This is because the transformed observations f(s are only used to compute the regularization losses G and Gy, and thus are not used for the main PPO objective. Without these extra terms, the only way to data augmentation is explained in Section [3.1 which leads to inaccurate estimates of the PPO objective. Hence, DrAC benefits from the regularizing effect of using data augmentation, while mitigating adverse consequences on the RI. objective. See Appendix B for more detailed explanation of why naive application of data augmentation with certain policy gradient algorithms is theoretically unsound.
Data- regullagge Actor-Critic applied to PPO Black: unmodified actor-critic algorithm. Cyan: image transformation.
Algorithm 1 DrAC:
Red: policy regularization.
Blue: value function regularization.
1: Hyperparameters: image transformation f. regularization loss coefficient 0,,, minibatch size M, replay buffer size T, number of updates K.
2: for k k= L...., K do
3: Collect new set of transitions = {(86, as Tis 8141) using Te.
4: for = 1,..., 17 do
5:
{(s, a; Fi, 8;-1) 1) ~ DD
b. Sample a minibatch of transitions
6:
for i 1 1..... M do
7:
P ~ H
b Sample the augmentation parameters D Compute the policy targets
8:
t; 4 m
9:
V/2 + Vs(ss)
Compute the value function targets
10:
end for
11:
Gn (0) = 1/4 EM
KL |
Regularize the policy
Gv( M ΣM (v, V (f(s
Regularize the value function
13:
JD,AC(0,4) (0, Jpro (0, の) or (G+ Gv)
Compute the total loss function D Update the policy Update the value function
14:
- arg maxg JDIAC
15:
< arg max, JDAAC
16: end for
end for
3.3 Automatic Data Augmentation
Since different tasks benefit from different types of transformations, we would like to design method that can automatically find an effective transformation for any given task. Such technique would significantly reduce the computational requirements for applying data augmentation in RL. this section, we describe three approaches for doing this. In all of them, the augmentation learner is trained at the same time as the agent learns to solve the task using DrAC. Hence, the distribution of rewards varies significantly as the agent improves, making the problem highly nonstationary.

---

Upper Confidence Bound. The problem of selecting a data augmentation from a given set can be formulated as a multi-armes bandit problem, where the action space is the set of available transformations F == [f" ,...., J. A popular algorithm for such settings is the upper confidence bound or UCB 2 which selects actions according to the following policy:
log(t)
ft = argmax JEF
FEF Qo(f)
(6)
Ni(f)
where ft the transformation selected at time step I, N&(.)) is the number of times transformation has been selected before time step I. and c is UCB's exploration coefficient. Before the t-th DrAC update, we use equation (6) to select an augmentation f. Then, we use equation (5) to update the agent's policy and value function. We also update the counter: Ni(f) NN-1(f)++ + Next, collect rollouts with the new policy and update the Q-function: Qi(f) & -i=i-K R(f, f). which is computed as a sliding window average of the past K mean returns obtained by the agent after being updated using augmentation f. We refer to this algorithm as UCB-DrAC (Algorithm Note that UCB-DrAC's estimation of Q()) differs from that of a typical UCB algorithm which uses rewards from the entire history. However, the choice of estimating Q()) using only more recent rewards is crucial due to the nonstationarity of the problem.
Meta-Learning the Selection of an Augmentation. Alternatively, the problem of selecting data augmentation from a given set can be formulated as a meta-learning problem. Here, we consider a meta-learner like the one proposed by Wang al. [66]. Before each DrAC update, the meta- leea-ee selects an augmentation, which is then used to update the agent using equation (5). We then collect rollouts using the new policy and update the meta-learner using the mean return of these trajectories. We refer to this approach as R122-rAC (A (Algorithm 4)
Meta-I Leranng the Weights of an Augmentation. Another approach for automatically finding an appropriate augmentation to directly learn the weights of a certain transformation rather than selecting augmentation from given set. In this work, we focus meta-learning the weights of a convolutional network which can be applied to the observations to obtain a perturbed image. We meta-learn the weights of this network using approach similar to the one proposed by Finn et al. [20). For each agent update, we also perform a meta-update of the transformation function by splitting DrAC's buffer into meta-train and meta-test test sets. We refer to this approach as Meta-DrAC (Algorithm 4). More details about the implementation of these methods can be found in Appendix D
Experiments
In this section, we evaluate our methods on four DeepMind Control environments with natural and synthetic distractors and the full Procgen benchmark [12) which consists of 16 procedurally generated games (see Figure 6 in Appendix G). Procgen has number of attributes that make it good testbed for generalization in RL: (i) it has diverse set of games in a similar spirit with the ALE benchmark (5). (ii) each of these games has procedurally generated levels which present agents with meaningful generalization challenges, (iii) agents have to learn motor control directly from images, and (iv) it has a clear protocol for testing generalization.
All environments use a discrete 15 dimensional action space and produce 64 x 64 x 3 RGB observa- tions. We use Procgen's 's easy setup, so for each game, agents are trained 200 levels and tested on the full distribution of levels. We use PPO as a base for all our methods. More details about our experimental setup and hyperparameters
be found in Appendix B
Data Augmentation. In our experiments, we use set of eight transformations: crop, grayscale, cutout, cutout-color, flip, rotate, random convolution and color-jitter [39 15). We use RAD's [4] implementation of these transformations, except for crop, in which we pad the image with 12 (boundary) pixels on each side and select random crops of 64 x 64. We found this implementation of crop to be significantly better on Procgen, and thus it can be considered an empirical upper bound of RAD in this case. For simplicity, set of transformations as RAD, but is trained with additional regularization losses for the actor and the critic, as described in Section 3.2
will refer to our implementation as RAD. DrAC uses the same
Automatic Selection of Data Augmentation. We compare three different approaches for automati- cally finding effective transformation: UCB-DrAC which uses UCB [2] to select an augmentation
5

---

Table 1: Train and test performance for the Procgen benchmark (aggregated over all 16 tasks, 10 seeds). (a) compares PPO with four baselines specifically designed to improve generalization in RL and shows that they do not significantly help. (b) compares using the best augmentation from our set with and without regularization, corresponding to DrAC and RAD respectively, and shows that regularization improves performance both train and test. (c) compares different approaches for automatically finding augmentation for each task, namely using UCB or RL for selecting the best transformation from given set, or meta-learning the weights (Meta-DrAC). (d) shows additional ablations: DrÅ regularizes only the actor, DrC regularizes only the critic, Rand-DrAC selects an augmentation using a uniform distribution, Crop-DrAC uses image crops for all tasks, and UCB-RAD is an ablation that does not the regularization losses. UCB-DrAC performs best on both train and test, and achieves a return comparable with or better than DrAC (which uses the best augmentation).
convolutional network
PPO-Normalized Return (%)
Train
Test
Method
Median
Mean
Std
Median
Mean
Std
PPO
100.0
100.0
7.2
100.0
100.0
8.5
Rand-FM
93.4
87.6
8.9
91.6
78.0
9.0
IBAC-SNI
91.9
103.4
8.5
102.9
Mixreg
95.8
104.2
3.1
105.9
114.6
3.3
PLR
101.5
106.7
5.6
107.1
128.3
5.8
DrAC (Best) (Ours)
114.0
119.6
9.4
118.5
138.1
10.5
(b)
RAD (Best)
103.7
109.1
9.6
114.2
131.3
9.4
UCB-DrAC (Ours)
102.3
118.9
8.8
118.5
139.7
8.4
(c)
RL2-DrAC
96.3
95.0
8.8
99.1
105.3
7.1
Meta-DrAC
101.3
100.1
8.5
101.7
101.2
7.3
DrA (Best)
102.6
117.7
11.1
110.8
126.6
9.0
DrC (Best)
103.3
108.2
110.6
115.4
8.5
(b) Rand-DrAC
100.4
99.5
8.4
102.4
103.4
7.0
Crop-DrAC
97.4
112.8
9.8
114.0
132.7
11.0
UCB-RAD
100.4
104.8
8.4
103.0
125.9
9.5
from given set, RL2-DrAC which uses RL3 67 do the same, and Meta-DrAC which MAML (20 to meta-learn the weights of a convolutional network. Meta-DrAC is implemented using the higher library 222 Note that we do not expect these approaches be better than DrAC with the best augmentation. In fact, DrAC with the best augmentation can be considered to be an upper bound for these automatic approaches since it the best augmentation during the entire training process.
Ablations. DrC and DrA are ablations to DrAC that use only the value or only the policy regulariza- tion terms, respectively. DrC be thought of an analogue of DrQ (38 applied to PPO rather than SAC. Rand-DrAC uses uniform distribution to select an augmentation each time. Crop-DrAC crop for all games (which is the most effective augmentation half of the Procgen games). UCB-RAD combines UCB with RAD (i.e. does not use the regularization terms).
Baselines. We also compare with Rand-FM 44 IBAC-SNI 28 Mixreg (66), and PLR 32 four methods specifically designed for improving generalization in RL and previously tested on Procgen environments, Rand-FM uses random convolutional networks to regularize the learned representations, while IBAC-SNI uses an information bottleneck with selective noise injection. Mixreg uses mixtures of observations to impose lincarity constraints between the agent's inputs and the outputs, while PLR samples levels according to their learning potential.
Evaluation Metrics. At the end of training, for each method and each game, we compute the average score over 100 episodes and 10 different seeds. The scores are then normalized using the corresponding PPO score on the same game. We aggregate the normalized scores over all 16 Procgen games and report the resulting mean, median, and standard deviation (Table TT. For a per-game breakdown, Tables 6 and Appendix!!

---

Carlpole Belance- - Natural
Finger Spit Natural
Chealah Bun- Natural
450
PPO
PP0
300
FPO
540
FPO
HBD jasd)
120
400
220
320
1:0
300
80
2:0
250
200
50
210
100
Tndrermerr (105
Endfrommnt She (107)
Erairemment Shep (101)
Enstrnmentt Spep (137)
Figure 2: Average return DMC tasks with natural video backgrounds with mean and standard de- viation computed over seeds. UCB-DrAC outperforms PPO and RAD with the best augmentations.
4.1 Generalization Performance on Procgen
Table T shows train and test performance on Procgen. UCB-DrAC significantly outperforms PPO, Rand-FM, IBAC-SNI, PLR, and Mixreg, As shown in Jiang et al. B22, combining PLR with UCB- DrAC achieves new ttate-of-the-an art on Procgen leading to 76% gain over PPO. Regularizing the policy and value function leads improvements over merely using data augmentation, and thus the performance of DrAC is better than that of RAD (both using the best augmentation for each game). In addition, we demonstrate the importance of regularizing both the policy and the value function rather than either one of them by showing that DrAC is superior to both DrA and DrC (see Figures 9 and i0 in Appendix K.. Our experiments show that the most effective way of automatically finding an augmentation is UCB-DrAC, As expected, meta-learning the weights of CNN using Meta-DrAC performs reasonably well on the games in which the random convolution augmentation helps. But overall, Meta-DrAC and RL2-DrAC are worse than UCB-DrAC. In addition, UCB is generally more stable, easier to implement, and requires less fine-tuning compared to meta-learning algorithms. See Figures [ and in Appendix K for a comparison of these three approaches on each game. Moreover, automatically selecting the augmentation from a given set using UCB-DrAC performs similarly well even better than method that uses the best augmentation for each task throughout the entire training process. UCB-DrAC also achieves higher returns than an ablation that uses a uniform distribution to select augmentation each time, Rand-DrAC. Nevertheless, UCB-DrAC is better than Crop-DrAC, which uses crop for all the games (which is the best augmentation for eight of the Procgen games as shown in Tables and from Appendix H
4.2 DeepMind Control with Distractors
In this section, we evaluate our approach on the DeepMind Control Suite from pixels (DMC, Tassa al. (64). We use four tasks, namely Cartpole Balance, Finger Spin, Walker Walk, and Cheetah Run, in three settings with different types of backgrounds, namely the default, simple distractors, and natural videos from the Kinetics dataset [35], as introduced in Zhang et al. (73]. See Figure [ for a few examples. Note that in the simple and natural settings, the background is sampled from list of videos the beginning of each episode, which creates spurious correlations between the backgrounds and the rewards. In the simple and natural distractor settings, as shown in Figure [3] UCB-DrAC outperforms PPO and RAD on all these environments in the most challenging setting with natural distractors. See Appendix [ results on DMC with default and simple distractor backgrounds, where our method also outperforms the baselines.
4.3 Regularization Effect
In Section 3.1
we argued that additional regularization terms needed in order to make the use of data augmentation in RL theoretically sound. However, one might wonder if this problem actually appears in practice. Thus, empirically investigate the effect of regularizing the policy and value function. For this purpose, we compare the performance of RAD and DrAC with grayscale and random convolution augmentations on Chaser, Miner, and StarPilot.
Figure 3 shows that not regularizing the policy and value function with respect to the transformations used can lead to drastically worse performance than vanilla RL methods, further emphasizing the
7

---

Chaser
Miner
Starpilot
PPO
30
(raysca
25-
(Ourst conv)
(Bura) (grayacale)
Score
Score
5-
15 Step[1a6) 66]
10
20
25
10 15 Step [1o6]
20
25
15 Step [1
20
25
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 comparable or better than PPO's, not using the regularization terms, iee. using RAD, significantly hurts performance relative to PPO. This is because, in contrast to DrAC, RAD does not use principled (importance sampling) estimate of PPO's objective.
Nirja
Nirja
Dodgebal
Dodgebaal
PFD
500
300
200
200
coler-ttter
Tran
200
10
155
30 15 20
20
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.
importance of these loss terms. In contrast, using the regularization terms as part of the RL objective (as DrAC does) results in an agent that is comparable or, in some cases, significantly better than PPO.
4.4 Automatic Augmentation
Our experiments indicate there is not single augmentation that works best across all Procgen games (see Tables 4 and [ in Appendix [H). Moreover, our intuitions regarding the best transformation for each game might be misleading. For example, at first sight, Ninja appears to be somewhat similar to Jumper,
the augmentation that performs best on Ninja color-jitter, while for Jumper is random-conv (see Tables and [5) In contrast, Miner seems like a different type of game than Climber or Ninja, but they all have the same best performing augmentation, namely color-jitter. These observations further underline the need for a method that can automatically find the right augmentation for each task.
Table D allnn with Figures X and 7 in the Appendix compare different approaches for automatically finding an augmentation, showing that UCB-DrAS performs best and reaches the asymptotic per- formance obtained when the most effective transformation for each game is used throughout the entire training process. Tigure4iillustattee an example of UCB's policy during training on Ninja and Dodgeball, showing that it converges to always selecting the most effective augmentation, namely color-jitter for Ninja and crop for Dodgeball. Finally, Figure in Appendix illustrates how UCB's behavior varies with its exploration coefficient.

---

Table 2: JSD and Cycle- Sycle-Consistency (%) (aggregated across all Procgen tasks) for PPO, RAD and UCB-DrAC, measured between observations that vary only in their background themes (i.e. colors and patterns that do not interact with the agent). UCB- JCB-DrAC learns more robust policies and representations that are more invariant to changes in the observation that are irrelevant for the task.
Sycle-Conssstency (%)
JSD
2-way
3-way
Method
Mean
Median
Mean
Median
Mean
Median
PPO
0.25
0.23
20.50
18.70
12.70
5.60
RAD
0.19
0.18
24.40
22.20
15.90
8,50
UCB-DrAC
0.16
0.15
27.70
24.80
17.30
10.30
4.5
Robustness Analysis
To further investigate the generalizing ability of these agents, we analyze whether the learned policies and state representations are invariant changes in the observations which are irrelevant for solving the task.
We first measure the Jensen-Shannon divergence (JSD) between the agent's policy for observation from a training level and a modified version of that observation with a different background theme (iee color and pattern). Note that the JSD also represents lower bound for the joint empirical risk across train and test [30] The background theme is randomly selected from the set of backgrounds available for all other Procgen environments, except for the one of the original training level. Note that the modified observation has the same semantics as the original one (with respect to the reward function), so the agent should have the same policy in both cases. Moreover, many of the backgrounds are not uniform and can contain items such as trees or planets which can be easily misled for objects the agent can interact with. As seen in Table 2 UCB-DrAC has a lower JSD than PPO, indicating that learns policy that is more robust to changes in the background.
To quantitatively evaluate the quality of the learned representation, metric proposed by Aytar et al. [3] and also used by Lee et al. [44]. See Appendix [C for more details about this metric. Table 2 reports the percentage of input observations in the seen environment that are cycle-consistent with trajectories in modified unseen environments, which have a different background but the same layout. UCB-DrAC has higher cycle-consistency than PPO, suggesting that it learns representations that better capture relevant task invariances.
the cycle- coneiseenss
5 Related Work
Generalization in Deep RL. recent body of work has pointed out the problem of overfitting
in deep RL [53 46 34 71 775, (50 [1] [2 33
[2]
40 233. A promising approach to prevent
overfitting is to apply regularization techniques originally developed for supervised learning such
as dropout [62 28 or batch normalization [3] 19 28). For example, Igl et al. (28) use selective
noise injection with variational information bottleneck, while Lee et al. 44 regularize the agent's
representation with respect to random convolutional transformations. The use of state abstractions has also been proposed for improving generalization in RL [72
74 I. Similarly, Roy and Konidaris and Sonar et al. (59 learn domain-invariant policies via feature alignment, while Stooke et al. decouple representation from policy learning. Igl et al.
54
29) reduce non-stationarity using policy distillation, while Mazoure et al. [47 maximize the mutual information between the agent's representation of successive time steps. Jiang et al. [32] improve efficiency and generalization by sampling levels according to their lcarning potential, while Wang et al. [67] use mixtures of observations to impose linearity constraints between the agent's inputs and the outputs. More similar to our work, Cobbe et al. ΑЛ Ye et al. (69 and Laskin et al. 41 add augmented observations the training buffer of an RL agent. However, as we show here, naively applying data augmentation in RL can lead to both theoretical and practical issues. Our algorithmic contributions alleviate these problems while still benefitting from the regularization effect of data augmentation.
63
Data Augmentation has been extensively used in computer vision for both supervised 42 4%
58 9 TO 39) and self-supervised [16 48 learning. More recent work uses data augmentation for contrastive learning, leading to sfate-of-the-an results downstream tasks [70] 26 25 7 Domain

---

randomization can also be considered type of data augmentation, which has proven useful for transferring RL policies from simulation the real world [65]. However, requires access to physics simulator, which is not always available. Recently, few papers propose the use of data augmentation in RI. [1T
41,, but all of them use a fixed (set of) augmentation(s) rather than automatically finding the most effective one. The most similar work to ours is that of Kostrikov et al. (38), who propose to regularize the Q-function in Soft Actor-Critic (SAC) [24] using random shifts of the input image. Our work differs from theirs in that it automatically selects an augmentation from given set, regularizes both the actor and the critic, and focuses on the problem of generalization rather than sample efficiency. While there is body of work the automatic use of data augmentation [14 13 18 57 45), these approaches were designed for supervised learning and, as we explain here, cannot be applied to RL without further algorithmic changes.
6 Discussion
In this work, we propose UCB-DrAC, a method for automatically finding an effective data augmenta- tion for RL tasks. Our approach enables the principled use of data augmentation with actor-critic algorithms by regularizing the policy and value functions with respect to state transformations. We show that UCB-DrAC avoids the theoretical and empirical pitfalls typical in naive applications of data augmentation in RL. Our approach improves training performance by 19% and test performance by 40% on the Procgen benchmark, and sets new cate-of-the-art on the Procgen benchmark. In addition, the learned policies and representations are more invariant to spurious correlations between observations and rewards. One limitation of our work the assumption that using a single type of data augmentation throughout the entire training process is optimal. In future work, we plan to study the effect of using multiple augmentations at different stages during the training of the agent's policy. Another promising avenue for future research is to use more expressive function class for meta-learning the transformations in order to capture a wider range of inductive biases.
Acknowledgments and Disclosure of Funding
We would like to thank our NeurlPS reviewers for their valuable feedback on this work. Roberta and Max were supported by the DARPA Machine Commonsense program.
References
[1] Rishabh Agarwal, Marlos C. Machado, P. S. Castro, and Marc G. Bellemare. Contrastive behavioral similarity embeddings for generalization in reinforcement learning. 2021.
[2] Peter Auer. Using confidence bounds for exploitation-exxplration trade-offs. of Machine Learning Research, 3(Nov):397 422, 2002.
[3] Yusuf Aytar, Tobias Pfaff, David Budden, Thomas Paine, Ziyu Wang, and Nando de Freitas. Playing hard exploration games by watching youtube. In Advances in Neural Information Processing Systems, pages 2930-2941, 2018.
[4] Suzanna Becker and Geoffrey E. Hinton. Self-organizing neural network that discovers surfaces in random-dot stereograms. Nature, 355: 1992.
[5] Mare G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47:253-279, 2013.
[6] Richard Bellman. A markovian decision process. Journal of mathematics and mechanics, pages 679-684, 1957.
[7] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. ArXiv, abs/2002.05709, 2020.
[8] Dan Ciregan, Ueli Meier, and Jürgen Schmidhuber. Multi-column deep neural networks for image classification. In 2012 IEEE conference on computer vision and pattern recognition, pages 3642-3649, IEEE, 2012.
10

---

[9] Dan C Cireşan, Ueli Meier, Jonathan Masci, Luca M Gambardella, and Jürgen Schmid- huber. High-peeformance neural networks for visual object classification. arXiv preprint arXiv: NoKiv:/102.00 100.0183, 2011.
[10] Dan C. Ciresan, Ueli Meier, Jonathan Masci, Luca Maria Gambardella, and Jürgen Schmidhuber. High-performan perfforrremee neural networks for visual object classification, ArXiv, abs/1102.00 102.0188, 2011.
[11] Karl Cobbe, Oleg Klimov, Chris Hesse, Taehoon Kim, and John Schulman. Quantifying generalization in reinforcement learning. arXiv preprint arXiv: /1/2.82/22..1111 2018.
[12] Karl Cobbe, Christopher Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation to benchmark reinforcement learning. arXiv preprint arXiv: 2019.
[13] Ekin D. Cubuk, Barret Zoph, Jonathon Shlens, and Quoe V. Le. Randaugment: Practical automated data augmentation with reduced search space. arXiv: Computer Vision and Pattern Recognition, 2019.
[14] Ekin Dogus Cubuk, Barret Zoph, Dandelion Mané, V. Vasudevan, and Quoe V. Le. Autoaugment: Learning augmentation strategies from data. 2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 113-123, 2019.
[15] Terrance DeVries and Graham W Taylor. Improved regularization of convolutional neural networks with cutout. arXiv preprint arXiv: 04552, 2017.
[16] Alexey Dosovitskiy, Philipp Fischer, Jost Tobias Springenberg, Martin A. Riedmiller, and Thomas Brox. Discriminative unsupervised feature learning with exemplar convolutional neural networks. IEEE Transactions on Pattern Analysis and Machine Intelligence, 38:1 38:1734-1747 2016.
[17] Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Volodymir Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, et al. Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures. arXiv preprint arXiv: 1802.0 062.01161, 2018.
[18] Boli Fang, Miao Jiang, and Jerry J. Shen. Paganda: An adaptive task-independent automatic data augmentation. 2019.
[19] Jesse Farebrother, Marlos C. Machado, and Michael H. Bowling, Generalization and regulariza- tion in dqn. ArXiv, 11s00100,,,,,,, 2018.
[20] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adap- tation of deep networks. In Proceedings of the 34th International Conference on Machine Learning-Volunn Volume 70, pages 1126-1135. JMLR. org, 2017.
[21] Shani Gamrian and Yoav Goldberg. Transfer learning for related reinforcement learning tasks via image-to-image translation. ArXiv, has/1806600377,, 2019.
[22] Edward Grefenstette, Brandon Amos, Denis Yarats, Phu Mon Htut, Artem Molchanov, Franziska Meier, Douwe Kiela, Kyunghyun Cho, and Soumith Chintala. Generalized inner loop meta- learning, arXiv preprint 1rXiv:1910.01127, 2019.
[23] Jake Grigsby and Yanjun Qi. Measuring visual generalization in continuous control from pixels. ArXiv, abs/2010.06740, 2020.
[24] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In ICMI,, 2018.
[25] Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross B. Girshick. Momentum contrast for unsupervised visual representation learning. ArXiv, abs/191 1111911002222. 2019.
[26] Olivier Hénaff Aravind Srinivas, Jeffrey De Fauw, Ali Razavi, Carl Doersch, S. M. Ali Eslami, and Aäron van den Oord. Data-efficient image recognition with contrastive predictive coding. ArXiv, 0es//995509272.. 2019.
[27] Hochreiter and J. Schmidhuber. Long short-term memory. Neural Computation, 9:1735-1780. 1997.
11

---

[28] Maximilian Igl, Kamil Ciosek, Yingzhen Li, Sebastian Tschiatschek, Cheng Zhang, Sam Devlin, and Katja Hofmamn. Generalization in reinforcement learning with selective noise injection and information bottleneck. In Advances in Neural Information Processing Systems, pages 13956-13968, 2019.
[29] Maximilian Igl, Gregory Farquhar, Jelena Luketina, Wendelin Böhmer, and Shimon White- son. The impact of non-stationarity generalisation in deep reinforcement learning. ArXiv, abs/2006.05826, 2020.
[30] Maximilian lsse, Jakub M Tomczak, and Patrick Forré. Designing data augmentation for simulating interventions. arXiv preprint arXiv:2005.01856, 2020.
[31] Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. ArXiv, abs/1502.03316, 2015.
[32] Minqi Jiang, Edward Grefenstette, and Tim Rocktäschel. Prioritized level replay. ArXiv, abs/2010.03934, 2020.
[33] Arthur Juliani, Ahmed Khalifa, Vincent- Pierre Berges, J. Harper, Hunter Henry, Adam Crespi, J. Togelius, and D. Lange. Obstacle tower: A generalization challenge in vision, control, and planning. ArXiv, thss/192.00378. 2019.
[34] Niels Justesen, R. Torrado, Philip Bontrager, A. Khalifa, J. Togelius, and S. Risi. Illuminating generalization in deep reinforcement learning through procedural level generation. arXiv: Learning, 2018.
[35] W. Kay, J. Carreira, K. Simonyan, Brian Zhang, Chloe Hillier, Sudheendra Vijaynarrasimhan, Viola, T. Green, Back, A. Natsev, Mustafa Suleyman, and Andrew Zisserman. The kinetics human action video dataset. ArXiv, 705.06950)) 2017.
[36] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. CoRR, abs/1412.6990, 2015.
[37] Ilya Kostrikov. Pytorch implementations of reinforcement learning algorithms. https:// github. hatbbbccoo..tttttr.ooo//oooooo-.
http::/hwwwopppoooaa-----------o -ttpoorrcc-----p----cc---------e ackttrpp-acktr-rre acktr- gaall 2018.
[38] llya Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation all you need: Regularizing deep reinforcement learning from pixels. arXiv preprint arXiv:2004. /3649, 2020.
[39] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolutional neural networks. In Advances in neural information processing systems, pages 1097-1105, 2012.
[40] Heinrich Kuttler, Nantas Nardelli, Alexander H. Miller, Roberta Raileanu, Marco Selvatici, Edward Grefenstette, and Tim Rocktäschel. The nethack learning environment. ArXiv, abs/2006.1 13760, 2020.
[41] Michael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Reinforcement learning with augmented data. arXiv preprint arXiv:2004. 14990, 2020.
[42] Yann LeCun, Bernhard E. Boser, John S. Denker, Donnie Henderson, Richard E. Howard, Wayne E. Hubbard, and Lawrence D. Jackel. Backpropagation applied to handwritten zip code recognition. Neural Computation, 1:541-551, 1989.
[43] Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. 1998.
[44] Kimin Lee, Kibok Lee, Jinwoo Shin, and Honglak Lee. Network randomization: A simple technique for generalization in deep reinforcement learning. In International Conference on Learning Representations. https://openreview net/forum, 2020.
[45] Yonggang Li, Guosheng Hu, Yongtao Wang, Timothy M. Hospedales, Neil M. Robertson, and Yongxing Yang. Dada: Differentiable automatic data augmentation, ArXiy, abs/2003.03766 2020.
12

---

[46] Marlos C. Machado, Mare G. Bellemare, Erik Talvitie, Joel Veness, Matthew J. Hausknecht, and Michael H. Bowling. Revisiting the arcade learning environment: Evaluation protocols and open problems for general agents. In IJCAI, 2018.
[47] Bogdan Mazoure, R' emm Tachet des Combes, Thang Doan, Philip Bachman, and R. Devon Hjelm. Deep reinforcement and infomax learning. ArXiv, abs/2006.07217, 2020.
[48] Ishan Mísra and Laurens van der Maaten, Self-supervised learning of pretext-invariant repre- sentations. ArXiv, bbss/192..0191, 2019.
[49] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wienstra, and Martin A. Riedmiller. Playing atari with deep reinforcement learning. ArXiv, abs/1312.5602, 2013.
[50] Alex Nichol, V. Pfau, Christopher Hesse, O. Klimov, and John Schulman. Gotta learn fast: new benchmark for generalization in rl. ArXiv, abs/ 804.03320,, 2018.
[51] Charles Packer, Katelyn Gao, Jernej Kos, Philipp Krähenbühl, Vladlen Koltun, and Dawn Xi- aodong Song. Assessing generalization in deep reinforcement learning. ArXiv, abs/1810.122555 2018.
[52] Roberta Raileanu and Tim Rocktäschel. Ride: Rewarding impact-driven exploration for vmocedurally-genenaae environments. ArXiv, abs/2002. 122220 2020.
[53] Aravind Rajeswaran, Kendall Lowrey, Emanuel Todorov, and Sham M. Kakade. Towards generalization and simplicity in continuous control. ArXiv,
2017.
[54] Josh Roy and George Konidaris. Visual transfer for reinforcement learning via wasserstein domain confusion. arXiv preprint 20ti:200000036633, 2020.
[55] John Schulman, Sergey Levine, Pieter Abbeel, Michael I. Jordan, and Philipp Moritz. Trust region policy optimization. In ICML, 2015.
[56] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint 177.v:/70..063/ 2017.
[57] Yinghuan Shi, Tiexin Qin, Yong Liu, Jiwen Lu, Yang Gao, and Dinggang Shen. Automatic data augmentation by learning the deterministic policy. ArXiv, abs/1910.08343, 2019.
[58] Patrice Y Simard, David Steinkraus, John C Platt, et al. Best practices for convolutional neural networks applied visual document analysis. In ledar, volume 2003.
[59] Anoopkumar Sonar, Vincent Pacelli, and Anirudha Majumdar. Invariant policy optimization:
Towards stronger generalization in reinforcement learning. ArXiv, 0196,0006.00066, 2020.
[60] Xingyou Song, Yiding Jiang, Stephen Tu, Yilun Du, and Behnam Neyshabur. Observational overfitting in reinforcement learning. ArXiv, abs/1912.02975, 2020.
[61] Aravind Srinivas, Michael Laskin, and Pieter Abbeel. Curl: Contrastive unsupervised represen- tations for reinforcement learning. ArXiv, abs/2004.001333 2020.
[62] Nitish Srivastava, Geoffrey E. Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdi- nov. Dropout: a simple way to prevent neural networks from overfitting. J. Mach. Learn. Res., 15:1929-1958, 2014.
[63] Adam Stooke, Kimin Lee, P. Abbeel, and M. Laskin. Decoupling representation learning from reinforcement learning. ArXiv, abs/2009.082 008319, 2020.
[64] Y. Tassa, Yotam Doron, Alistair Muldal, T. Erez, Y. Li, D. Casas, D. Budden, Abbas Abdol- maleki, J. Merel, Andrew Lefrancq, T. Lillicrap, and Martin A. Riedmiller. Deepmind control suite. ArXiv,
2018.
[65] Joshua Tobin, Rachel H Fong, Alex Ray, Jonas Schneider, Wojeiech Zaremba, and Pieter Abbeel. Domain randomization for transferring deep neural networks from simulation to the real world. 2017 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pages 23-30, 2017.
13

---

[66] Jane X Wang. Zeb Kurth-Nelson, Dhruva Tirumala, Hubert Soyer, Joel Z Leibo, Remi Munos, Charles Blundell, Dharshan Kumaran, and Matt Botvinick. Learning to reinforcement learn. arXiv preprini /61
2016.
[67] K. Wang, Bingyi Kang, Jie Shao, and Jiashi Feng. Improving generalization in reinforcement learning with mixture regularization. ArXiv, abs/2010. 108 14, 2020.
[68] R. J. Williams. Simple statistical wradient-following algorithms for connectionist reinforcement learning. Machine Learning. 8:229-2555 256, 2004.
[69] Chang Ye, Ahmed Khalifa, Philip Bontrager, and Julian Togelius. Rotation, translation, and cropping for zero-shot generalization. arXiv preprint w0X009000.09990,, 09908, 2020.
[70] Mang Ye, Xu Zhang, Pong C. Yuen, and Shih-Fu Chang. Unsupervised embedding learning via invariant and spreading instance feature. 2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 6203-6212, 2019.
[71] Amy Zhang, Nicolas Ballas, and Joelle Pineau. A dissection of overfitting and generalization in continuous reinforcement learning. ArXiv, 0oos/1800.07933 2018.
[72] Amy Zhang, Clare Lyle, Shagun Sodhani, Angelos Filos, Marta Kwiatkowska, Joelle Pineau, Yarin Gal, and Doina Precup. Invariant causal prediction for block mdps. arXiv preprini 2rXiv:2003.06016 2020.
[73] Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. 2020.
[74] Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. arXiv preprint arXiv:2006.10742 2020.
[75] Chiyuan Zhang, Oriol Vinyals, Rémi Munos, and Samy Bengio. study on overfitting in deep reinforcement learning, ArXiv, abs/1804.06899 2018.
14


Reviews:
[{"review_id": "Hpr4TD2sgHm", "reviewer": "Reviewer_FoA8", "summary": "The paper proposes a framework for automatic data augmentation in deep RL. It proposes three algorithms based on 1) UCB 2) RL2 3) meta-learning, combining PPO and regularization of both the actor and critic to make sure that the policy and value function are invariant to the transformation corresponding to the augmentation. Experiments on ProcGen and the DMC with distractors show that the UCB-based algorithm (UCB-DrAC) outperforms vanilla deep RL and previous data augmentation baselines and is competitive with the best augmentation type chosen post hoc.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "# Originality #\n\nThe paper combines existing ideas to create a novel framework for automatic data augmentation in RL: 1) UCB, RL2, MAML 2) regularizing the value function. Thus, in my opinion, this work has sufficient novelty.\n\n# Quality #\n\nThe paper does not contain theory, aside from some intuition for why we should include regularization. I do not consider this a major weakness since the paper is primarily an applied work, but I do think there may be some theoretical support behind regularization, as TRPO does something very similar.\n\nWith respect to the empirical results, I found the results to be fairly convincing. The authors use two benchmarks, ProcGen and DCM,  and compare to recent work for generalization in deep RL, both based on data augmentation and other techniques. Ablations are done on the proposed algorithm showing that both the regularization and UCB are necessary for improvements over baselines. The performance of UCB-DrAC is comparable to that of choosing the best augmentation post hoc, with lower sample efficiency. \n\nI would've liked to see more discussion comparing the three proposed algorithms, in particular in relation to the assumption that there is a single best augmentation. If there is a single best augmentation (that happens to be in the set of eight transformations), then it is not surprising that UCB-DrAC works best, as it has the fewest parameters. Do the authors think that that is the case, or could it be that the higher expected sample complexity of RL2-DrAC and Meta-DrAC require training for more timesteps than was done in the paper?\n\n# Clarity #\n\nOverall, the paper is clear, and the results were presented well. However, there are a few places where edits are recommended:\n- The description of RL2-DrAC is a little incomplete, it is not clear what the input to the meta-learner is from the main text. After looking at the appendix, I see that it is the sequence of augmentations, not the sequence of states as in the original RL2. I found this to be a bit confusing, so I would suggest the authors to clarify in the final draft.\n- The second paragraph of section 4.5 seems a little out of place. Should it have been earlier?\n\n# Significance #\n\nThe significance of this work is that it is the first attempt at an automatic data augmentation framework for deep RL. It shows competitive performance to choosing the best augmentation after running all of them, thus leading to improved sample efficiency. As UCB-DrAC seems fairly simple to implement, I think it would be useful for practitioners. \n\n# Updated review #\n\nAfter reading the responses and reviews, I have decided to keep my score, with the understanding that the authors include more discussion of the assumptions of each algorithm in the final draft. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "# Limitations #\n\nThe authors mention one limitation of the work in the conclusion, that only one augmentation is chosen. I agree that this is a good direction for future work, as in practice many augmentations are often used. Another possible limitation is that the choice of the augmentation is currently based on only the reward signal, but intuitively the state (image) would also be helpful for picking the augmentation.\n\n# Societal Impact #\n\nThe authors do not discuss any potential societal impact. I think this is fine since the paper is methodological, and I'd like to mention the following paper on the societal impacts of deep RL.\n- Whittlestone, Jess and Arulkumaran, Kai and Crosby, Matthew. The Societal Implications of Deep Reinforcement Learning. JAIR 70, 2021.\n", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "FmPRZIRZXzz", "reviewer": "Reviewer_jCPN", "summary": "Summary: Reinforcement learning has been shown to overfit to training environments. Some recent works (laskin et al, kostrikov et al) have shown data-augmentation to be an effective technique in curbing the overfitting issue. However, different games have different biases and hence require different data augmentation. This paper proposes to use UCB algorithm to automatically select a good data augmentation (from a list of data augmentations) for a game. However, naively applying data augmentation to RL algorithms like PPO isn’t theoretically sound. To address this problem, the paper introduces novel policy and value regularization. Overall, the paper contains extensive experiments and ablation studies.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Originality: The paper is a mix of original ideas as well as novel combinations of well known techniques. The paper introduces novel policy and value regularization to combine data augmentation with PPO in a theoretically sound manner. In addition, it uses UCB algorithm to select appropriate data augmentation for a game to be used during the policy training.  \n\nQuality: The claims are empirically well supported, the experiments and the analysis are extensive and proper ablations are considered. \n\nClarity:  The paper is clearly written and well organized.\n\nSignificance: The results are important in a way as the paper suggests a theoretically sound and a stable way to use data augmentation with RL. \n \nMinor point: I wouldn’t call the method SOTA as DAAC/iDAAC gets better results (Raileanu et al). But that didn’t affect my final score as I feel this method is complementary to DAAC/iDAAC and it fixes (or atleast makes a good attempt at fixing) the issue of stability wrt data-augmentation in RL.\n \nReferences:\nMichael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, Aravind Srinivas. Reinforcement Learning with Augmented Data. NeurIPS 2020.\nIlya Kostrikov, Denis Yarats, Rob Fergus. Image Augmentation is all you need. ICLR 2021.\nRoberta Raileanu, Rob Fergus. Decoupling Value and Policy for Generalization in Reinforcement Learning. ICML 2021.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "It's not discussed in the paper.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "1 hr"}, "extra_text": {}}, {"review_id": "5adg6jUT03W", "reviewer": "Reviewer_Toxj", "summary": "This paper claims that the current data augmentation methods in RL relies on expert knowledge and large number of transformation methods, and proposes three methods that can automatically find a useful augmentation for a given RL task.\nThe paper introduce a way of using data augmentation with actor-critic algorithms under settings where the input observation has images.\n", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Originality: This paper deals with the data augmentation problem in RL. While the problem setting is not new, the method proposed in this paper do have some merits in automatically learning the augmentation method to reduce the computational requirements. This paper also provides different perspectives on data augmentation, and provided three methods to augment data correspondingly. In the related work, it also discussed adequate paper and differs itself from them. \n\nQuality: This paper technically sound under the scope that the input observations are images.  This paper conducted a bunch of experiments under different environments to validate the effectiveness of their proposed method. However, it seems like in the experiments the input have to be images so that the transformations (candidate data augmentation methods) are applicable. It would be more interesting to see how this method could generate to scenarios where the inputs are vectors (as in many real world cases).  However, I did not find any discussions on this issue in this paper.\nClarity: This paper is clearly written and well organized.\n\nSignificance: The results seems to be less important considering the input have to be images.  But I do appreciate that in the experiment part, this paper has investigated the robustness, which is corresponding to their motivations mentioned in the introduction.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "1"}, "extra_text": {}}, {"review_id": "06tP7L8WSL", "reviewer": "Reviewer_CGaL", "summary": "This paper presents a method for improving generalization in reinforcement learning using data augmentation. It identifies a key theoretical limitation of previous work in the naive application of data augmentation to actor critic methods and proposes a solution that avoids biased estimation of the policy gradient objective while still allowing for regularization of the policy and value function. The resulting loss is paired with three distinct approaches for identifying an optimal data augmentation strategy. The best of these approaches, a bandit algorithm, obtains SOTA performance (relative to published works) on the Procgen environment. \n", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n\n- The paper provides a well-motivated and justified means of applying data augmentation to improve generalization in reinforcement learning.\n\n- The use of data augmentation to promote invariance in the value function and policy directly, rather than indirectly via applying the data augmentation to the inputs used to compute the policy gradient, is a nice way of avoiding the issue of biased gradients in policy gradient learners.\n\n- The UCB method to select data augmentation transformations is a nice addition to promote more effective augmentation policies. \n\n- The paper is well-written and clear, with the experimental settings easy to follow and the theoretical motivation compelling.\n\n- The empirical analysis of different approaches to automatically learning data augmentation is reasonably complete and provides some insight into the robustness of the proposed regularization strategy to possibly-suboptimal augmentation sets.\n\nWeaknesses:\n\n- The proposed method isn’t particularly novel as data augmentation approaches have been widely used in both supervised and unsupervised learning to regularize the representation. Applying data augmentation only to enforce invariance in a representation, as opposed to using it to increase the size of the training dataset, has been deployed previously in papers such as BYOL.\n\n- The meta-learning baseline is unclear: does the meta-learner only have the capacity to apply a single augmentation to the inputs? If so, this seems like a fairly severe limitation. A more fair comparison might have been against a learned Spatial Transformer Network, which has the capacity to apply more complex sets of augmentations. \n\n- The empirical performance indeed outperforms the baselines, but in many instances the performance improvements are relatively modest. \nUCB assumes a fixed reward distribution, but as mentioned by the authors the policy improvements of the reinforcement learner mean that the rewards are likely to be highly nonstationary. Additionally, the use of the “training” environment to select the data augmentation, rather than a “validation” environment, seems like it might encourage overfitting.\n\nQuestions and Points of Improvement\n\n- There are additional approaches to RL that seek to throw away irrelevant information to improve generalization, such as DBC and PBL. I would be interested in seeing how the method compares to these other approaches that use methods other than data augmentation to improve robustness and generalization.\n\n- The theoretical setup should clarify whether the goal is truly to maximize the expected return over all environments or to maximize some sort of normalized return. For example, if one environment provides significantly more rewards than another, it might make sense to focus policy learning on this environment. This presumably isn’t an issue for the environments under consideration in the empirical analysis because of the shared structure, but may be worth noting for correctness. \n\n- I would have liked to see more justification for the use of training-set performance improvement in selecting a transformation rather than a validation-environment performance improvement as the reward signal for UCB, as it seems the principal goal of this method is not to increase the performance on the training environments, but to improve test-set performance.\n\n- The best distribution of augmentations to maximize generalization performance may not be to greedily select the augmentation that improves the policy improvement step the most. I would be interested in seeing if a similar approach that can identify stochastic optimal policies over augmentations might yield better performance.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "While I don't believe this work is likely to lead to any direct negative societal impacts, the authors should include a statement saying as much in the text.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper presents a solid contribution to improve generalization in reinforcement learning through data augmentation. Reviewers appreciated the theoretical motivations provided for the method and the addition of UCB to select transformations. I encourage the authors to make revisions promised in their response and address the issues raised by the reviewers. 
Decision: Accept (Poster)