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

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. We evaluate our method 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 approach also outperforms methods specifically designed to improve generalization in RL, thus setting a new state-of-the-art on Procgen. 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. 

Under review as a conference paper at ICLR 2021
AUTOMATIC DATA AUGMENTATION FOR GENERALIZATION IN REINFORCEMENT LEARNING
Anonymous authors Paper under double-bling review
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, we 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 use of data augmen- tation theoretically sound for actor-critic algorithms. We evaluate our method on the Procgen benchmark which consists of 16 procedurally generated environments and show that it improves test performance by 40% relative to standard RL algo- rithms. Our approach also outperforms methods specifically designed to improve generalization in RI,, thus setting a new state-of-the-arr on Procgen. In addition, our agent learns policies and representations which are more robust to changes the environment that are irrelevant for solving the task, such as the background.
INTRODUCTION
Generalization to new environments remains major challenge in deep reinforcement learning (RL). Current methods fail to generalize unseen environments even when trained on similar
settings (Farebrother et al. 2018 Packer et al. 2018 Zhang et al. [0188| Cobbe et al. 2018 Gamrian & Goldberg] [219 Cobbe et al., 2019 (Song et al., 2020). 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 of regularization (Farebrother al. 2018 /Zhang
et al.
2018a] Cobbe et al. 2018 [gl et al., 2019). data augmentation (Cobbe et al.) 2018 Lee et al. Ye et al. 2020 Kostrikov et al. 2020: Laskin et al., 2020), or representation learning (Zhang 2020a c). In this work, we focus
2020
the use of data augmentation in RL. We identily key differences between supervised learning and reinforcement learning which need to be taken into
et al.
account when using data augmentation in RL.
More specifically, we show that a naive application of data augmentation can lead to both thcoretical and practical problems with standard RI. 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 theoretically sound way. Specifically, we introduce two regularization terms which constrain the agent's is 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 in which a naive use of data augmentation completely fails or converges to a sub-optimal policy. While we use Proximal Policy Optimization (PPO, [Schulman et al. (2017)) to describe and validate our approach, the method can easily integrated with any actor-critic algorithm with discrete stochastic policy such as A3C (Mnih et al. 2013), SAC (Haarnoja et al. 2018). or IMPALA (Espeholt et al. 2018).
The current use of data augmentation in RL either relies on expert knowledge to pick an appropriate augmentation (Cobbe et al. 2018 Lee et al. 2020 Kostrikov et al. 2020) or separately evaluates a large number of transformations to find the best one (Ye et al.) 20200 Laskin et al., 2020). In this paper, propose three methods for automatically finding useful augmentation for given RL task. The first two learn to select the best augmentation from a fixed set, using either a variant of the upper confidence bound algorithm (UCB, Auer (2002)) or meta-leaming (RI.?. Waan et al. (016))).

---

Under review as a conference paper at ICLR 2021
Augmented Observations
Environment Observations
Augmentation
UCB
DrAC
random-sonv
Figure I: Overview of UCB-DrAC. A UCB bandit selects image transformation (e.g. random- conv) and applies it to the observations. The augmented and original observations are passed to a regularized actor-critic agent (i.e. DrAC) which uses them learn a policy and value function which are invariant to this transformation.
We refer to these methods as UCB-DrAC and RL22-rrC,, respectively. The third method, Meta- ArAC, directly meta-learns the weights of a convolutional network, without access to predefined transformations (MAML, Finn et al. (2017)) Figure I gives an overview of UCB-DrAC.
We evaluate these approaches on the Procgen generalization benchmark (Cobbe al. 2019) which consists of 16 procedurally generated environments with visual observations. Our results show that UCB-DrAC is the most effective among these finding good augmentation, and is comparable
better than using DrAC with the best augmentation from a given set. UCB UCB-DrAC also outperforms
baselines specifically designed to improve generalization in RL (Igl et al., 2019 Lee et al. 2020 Laskin et al.) 2020) on both train and test. In addition, we show that our agent learns policies and representations that more invariant 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, our 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 RL settings, (iii) we show that the use of data augmentation leads to policies and representations that better capture task invariances, and (iv) we demonstrate state-of-the-an th-ttt-artttt results on the Proegen benchmark.
2 BACKGROUND
We consider a distribution ((m) of Markov decision processes (MDPs, Bellman (1957] m E M. with defined by the tuple A,,, Tm, Rm, Pm, 7), where S., is the state space, A is the action P
space, Tm (n,('") (6'[,, a) the transition function, Rm, (s, a) the reward funetion, and pm (80) the initial state distribution. During training. restrict access a fixed set of MDPs, Mirain {mi..... m,). where m, ~ q, Vi n. The goal is to find a policy πe which maximizes the expected discounted reward over the entire distribution of MDPs, J (TT)) = E.g,,, Tm. Et= 2 (82,a:)]
In practice,
the Procgen benchmark which contains 16 procedurally generated games. Each game corresponds to a distribution MDPs g(m), and each level of a game corresponds to an MDP determined by the seed (i.e. integer)
sampled from that game's 's distribution m ~ q. The MDP used to generate the corresponding level. Following the setup from Cobbe et al. (2019), agents are trained on a 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. (2017)) an actor-critic algorithm that learns policy and a value function Vo with the goal of finding an optimal policy for given MDP. PPO alternates between sampling data through interaction with the environment and maximizing clipped surrogate objective function JPPo using stochastic gradient ascent. See Appendix[A]ic for full description of PPO. One component of the PPO objective is the policy gradient term JPG. which is estimated using importance sampling:
no (a(a))) Te (als) Ao (s,
JrG(0) = Σ Ta(ass)Ao.u (s. a) Ea
(1)
aEA
where A(-) is an estimate of the advantage function, Me is the behavior policy used to collect trajectories (i.e. that generates the training distribution of states and actions), and To is the policy we want optimize (i.e. that generates the true distribution of states and actions).
2

---

Under review as a conference paper at ICLR 2021
3 AUTOMATIC DATA AUGMENTATION FOR RL
3.1 DATA AUGMENTATION IN RL
Image augmentation has been successfully applied in computer vision for improving generaliza- tion on object classification tasks (Simard et al.) 2003; Cireşan et al.) 2011 Ciregan et al. 2012 Krizhevsky al. 2012). As noted by Kostrikov al. (020), those tasks are invariant to certain image transformations such as rotations or 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 RL settings without other algorithmic changes (Cobbe et al. 2018 Ye et al. 2020
Laskin
al. 2020), argue that this approach is not theoretically sound.
transformations are naively applied to observations in PPO's buffer, as done in Laskin et al. 2020 the PPO objective changes and equation (I) is replaced by
Re (a)(((s)) Aaass (s. To (als)
JPG (0)
Σ To(aa)) Ansa (5,0) BawnToald GEA
where f S X H + S is the image transformation. However, the right hand side of the above equation not a sound estimate of the left hand side because fs)) & The (a(a)s), since nothing constrains No (al (s)) to be close (a(as)) Moreover, can define certain transformations f()) that result in an arbitrarily large ratio of (s))/T((a(s))
Figure 2] shows examples where naive of data augmentation prevents PPO from learning good policy in practice, suggesting that this is not just a theoretical concern. In the following section, we propose 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. (020), we propose two novel regularization terms for the policy and value functions that enable the proper use of data augmentation for actor-critic algorithms. Our algorithmic contribution differs from that of Kostrikov et al.] (2020) in that constrains both the actor and the critic, opposed to only regularizing the Q-function.
Following Kostrikov al. (220), we define an optimality-invarintt state transformation f:S+ H + S as mapping that preserves both the agent's 's policy and its value function V such that V(s) V(f(s, and π (als)) = т f(,)f(s))) Ys E S, v E H, where v are the parameters of f()), drawn from the set of all possible parameters H.
To ensure that the policy and value functions are invariant to such transformation of the input state, propose additional loss term for regularizing the policy,
G KL (a)s) fa,
(3)
as well as an extra loss term for regularizing the value functionn
Gv = Vv (f(s
(4)
Thus, our dagu-regd actor-critic method, or DrAC, maximizes the following objective:
JD:AC JpPo (G, + Gy),
where a,, 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 can be used to impose certain inductive biases relevant for the task (e.g. invariance with respect colors or translations). In addition, G2 and Gy can be added to the objective of any actor-critic algorithm with a discrete stochastic policy (e.g. A3C, TRPO, ACER, SAC, or IMPALA) without any other changes.
Note that when using DrAC, as opposed to the method proposed by Laskin et al. (2020), we still use the correct importance sampling estimate of the left hand side objective in equation (1) (instead of
3

---

Under review as a conference paper at ICLR 2021
wrong estimate as in equation 2 This is because the transformed observations f(s) 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 B.11 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.
Algorithm 1 DrAC: Gata-regularized Actor-Critic applied PPO Black: unmodified actor-critic algorithm.
Cyan: image transformation.
Red: policy regularization.
Blue: value function regularization.
Hyperparameters: image transformation f. regularization loss coefficient Cy, minibatch size M, replay buffer size T, number of updates K.
2: for k k=1, l...... K do
3: Collect new set of transitions D = {(*;, ai ri, 8141)) using wp-
4: for = 1,..., T
do
Lad
5:
{(s, a, ray 8 )=
J==1 ~ D
Sample minibatch of transitions
6:
for i = 1..... M do
7:
V H
Sample the augmentation parameters & Compute the policy targets
8:
f 4 (-11))
9:
1/2 + Vg((s)
Compute the value function targets
10:
end for
11:
G. (0) M M KL h Te(-)f(s,,ui)
Regularize the policy
12:
(Ø) = M Σ
(v: V (f(s;,
Regularize the value function
13:
JD Φ) Jrro (0, Φ) - (G.((()) Gv)
> Compute the total loss function b Update the policy D- Update the value function
arg max. JDIAC
15:
① - arg max& JDAC
end for
17:
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 a technique would significantly reduce the computational requirements for applying data augmentation in RL. In this section, we describe three approaches for doing this. trained at the same time 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.
of them, the augmentation learner
Upper Confidence Bound. The problem of selecting data augmentation from given set can be formulated as a multi-armed bandit problem, where the action space is the set of available transformations F = "i......'.. A popular algorithm for such settings is the upper confidence bound or UCB Auer 2002), which selects actions according to the following policy:
log(t)
ft argmax fEF (i(f) + t N.(f)
(6)
where ft the transformation selected time step t, .N) is the number of times transformation has been selected before time step f. and e 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) update the agent's 's policy and value function. We also update the counter: N₂(f) = (f) Next, we collect rollouts with the new policy and update the Q-function: Qe(f) = x R(fi f). which is computed as a sliding window average of the past K mean returns obtained by the agent after being updated using augmentation We refer to this algorithm as UCB-DrAC (Algorithm 4).
4

---

Under review as a conference paper at ICLR 2021
Note that rewards from the entire history. However, the choice of estimating Q(f) using only more recent rewards is crucial due to the nonstationarity of the problem.
estimation of Q(f) differs from that of typical UCB algorithm which uses
Meta-Learning the Selection of an Augmentation. Alternatively, the problem of selecting a data augmentation from given set can be formulated meta-learner like the one proposed by Wang et al. (2016). Before each DrAC update, the meta-learner selects an augmentation, which is then used to update the agent using equation ( We then collect rollouts using the new policy and update the meta-learner using the mean return of these trajectories. We refer this approach as RL2-DrAC (Algorithm 4.
meta-learning problem. Here, consider
Meta-Learning 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 an augmentation from a given set. In this work, we focus on 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 jet al. (2017). 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 sets. We refer to this approach as Meta-DrAC (Algorithm More details about the implementation of these methods can be found in Appendix D
EXPERIMENTS
In this section, we evaluate our methods the Procgen benchmark (Cobbe et al., 2019) which consists of 16 procedurally generated games (see Figure 6 in Appendix G). Procgen has number of attributes that make it a good testbed for generalization in RL: (i) it has a diverse set of games in a similar spirit with the ALE benchmark (Bellemare et al., 2013), (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 64 3 RGB observa- tions. We use Procgen's 's easy setup, so for each game, agents are trained on 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 can be found in Appendix E
Data Augmentation. In our experiments, we use set of eight transformations: crop, grayscale, cutout, cutout-color, fiip rotate, random convolution and color-jitter (Krizhevsky et al. 2012 Vries & Taylor, 2017). We use RAD's (Laskin et al. (2020) implementation of these transformations, except for crop, which we pad the image with 12 (boundary) pixels 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, we will refer to our implementation as RAD. DrAC uses the same set of transformations as RAD. but is trained with additional regularization losses for the actor and the critic, described in Section 3.2
Automatic Selection of Data Augmentation. We compare three different approaches for automat- ically finding an effective transformation: UCB-DrAC which uses UCB (Auer 2002) to select an augmentation from a given set, RL2-DrAC which uses RL (Wang et al. 2016) to do the same, and Meta-DrAC which uses MAML (Finn et al. 2017) to meta-learn the weights of a convolutional network. Meta-DrAC is implemented using the higher library (Grefenstette et al. 2019). Note that we do not expect these approaches to 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 uses the best augmentation during the entire training process.
Ablations. Rand-DrAC uses uniform distribution to select an augmentation each time. Crop- DrAC uses crop for all games (which is the most effective augmentation on half of the Proegen games). UCB-RAD combines UCB with RAD (i.e. it does not use the regularization terms).
Baselines. We also compare with Rand-FM Lee et al. 2020) and IBAC-SNI [Ig et al. 2019), two methods specifically designed for improving generalization in RL and previously tested on CoinRun, one of the Procgen games. Rand-FM uses a random convolutional networks to regularize the learned representations, while IBAC-SNI uses an information bottleneck with selective noise injection.
5

---

Under review as a conference paper at ICLR 2021
Table 1: Train and test performance for the Procgen benchmark (aggregated over all 16 tasks, 10 seeds), (a) compares PPO with two 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 on both train and test. (c) compares different approaches for automatically finding an augmentation for each task, namely using UCB or RL for selecting the best transformation from given set, or meta- veeta-leanning the weights of convolutional network (Meta-DrAC). (d) shows additional ablations: Rand-DrAC selects an augmentation using uniform distribution, Crop-DrÁC uses image crops for all tasks, and UCB-RAD is an ablation that does not use 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).
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
(a)
Rand-FM
93.4
87.6
8.9
91.6
78.0
9.0
IBAC-SNI
91.9
103.4
8.5
86.2
102.9
8.6
DrAC (Best)
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)
R1.2-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
Rand- DrAC
100.4
99.5
8.4
102.4
103.4
7.0
(d) 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
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 Tb. For a per-game breakdown, see Tables aand] Appendix
4.1 GENERALIZATION ABILITY
Table I shows train and test performance on Procgen. UCB-DrAC significantly outperforms PPO, Rand-FM, and IBAC-SNI. Regularizing the policy and value function leads to 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). 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 meta- learning algorithms. See Figures 7 and in Appendix comparison of these three approaches each game. Moreover, automatically selecting the augmentation from a given set using UCB-DrAC performs similarly well throughout the entire training process. UCB-DrAC also achieves higher returns than an ablation that uses uniform distribution 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 Tables and from Appendix H).
better than method that the best augmentation for each task
4.2 REGULARIZATION EFFECT
In Section 3.1 argued that additional regularization terms needed in order to make the of data augmentation 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

---

Under review as a conference paper at ICLR 2021
Miner
StarPilot
PPO
10
RAD (wndom--coro)
RAD (graysca
(Burs) (rendom-curr)
25
DYAC (Ours) igrayscalel
Score
15
5
15 Step [1c6]
10
20
25
10 15 Stop [1e6]
20
25
10 15 Step [1c6]
20
25
Figure 2: Comparison between RAD and DrAC with the same augmentations, grayscale and random convolution, on the test environments of Chaser (left), Miner (center), and StarPilot (right). While DrAC's performance is comparable or better than PPO's, not using the regularization terms, i.e. using RAD, significantly hurts performance relative to PPO. This is because, in contrast to DrAC, RAD does not use a principled (importance sampling) estimate of PPO's objective.
function. For this purpose, we compare the performance of RAD and DrAC with grayscale and random convolution augmentations on Chaser, Miner, and StarPilot.
Figure 2 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 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.3 AUTOMATIC AUGMENTATION
Our experiments indicate there is not a single augmentation that works best across all Procgen games (see Tables 4 and in Appendix HH. Moreover, intuitions regarding the best transformation for each game might be misleading. For example, at a first sight, Ninja appears to be somewhat similar Jumper, but the augmentation that performs best Ninja is color-jitter, while for Jumper is random-conv (see Tables 4 and [5). In contrast, Miner seems like a different type of game than Climber 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 I aalong with Figures and in in the Appendix compare different approaches for automatically finding an augmentation, showing that UCB-DrAC performs best and reaches the asymptotic per- formance obtained when the most effective transformation for each game is used throughout the entire training process. Figure FigureBJillustraaee 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. Figure in Appendix Fillustrates how UCB's s behavior varies with its exploration coefficient.
4.4 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-Sihe -hannon divergence (JSD) between the agent's policy for an observation from training level and a modified version of that observation with a different background theme (i.e. color and pattern). Note that the JSD also represents a lower bound for the joint empirical risk across train and test [Isse et al. 2020). 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 semanties the original one (with respect to the reward function), so the agent should have the same policy in both cases. Note that many of the backgrounds uniform and can contain items such trees or planets which can be easily
7

---

Under review as a conference paper at ICLR 2021
Ninja
Ninga
Dodgeball
Dodgebbll
ctp tandem-ccam CaTN
PPO
Liac tnckepptid
13000
500
SFapscalc
500
400
200
Trank
200
10
15
15
25
1.
15
20
(a) UCB Selection
(b) Performance
(c) UCB Selection
(d) Performance
Figure 3: 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.
Table 2: JSD and 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-DrAC learns more robust policies and representations that invariant changes in the observation that irrelevant for the task.
CyceeCoosissency (%)
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
misled for objects the agent can interact with. As seen in Table UCB-DrAC has a lower JSD than PPO, indicating that it learns a policy that is more robust changes in the background.
To quantitatively evaluate the quality of the learned representation, we use the sycle-consistency metric proposed by Aytar et al. (2018) and also used by Lee et al. (2020). See Appendix [orr more details about this metric. Table reports the percentage of input observations in the seen environment that are cycle-conse consssssttt with trajectories in modified unseen environments, which have different background but the same layout. UCB-DrAC has higher cycle-consistency than PPO, suggesting that learns representations that better capture relevant task invariances.
4.5 DEEPMIND CONTROL SUITE EXPERIMENTS
In this section, we evaluate our approach on the DeepMind Control Suite from pixels (DMC. Tassa jt al. (2088). 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 (Kay et al. 2017). as introduced in Zhang et al. (2020b). See Figure 9 for few examples. Note that in the simple and natural settings, the background is sampled from a list of videos at the beginning each episode, which creates spurious correlations between the backgrounds and the rewards. In the simple and natural distractor settings, shown in Figure [1] UCB-DrAC outperforms PPO and RAD with the best augmentation on all these environments in the most challenging setting with natural distractors. See Appendix for results DMC with default and simple distractor backgrounds, where our method also outperforms the baselines.
RELATED WORK
Generalization in Deep RL. recent body of work has pointed out the problem of overfitting in
deep RL (Rajeswaran al. 2017]] Machado al. [2018] Packer al. 2018 Zhang al. [018aa bi
Cobbe et al.] 20188 2019: Yarats et al. 2019; Raileanu Rocktäschel] 2020). A promising approach
prevent overfitting is to apply regularization techniques originally developed for supervised learning
8

---

Under review as a conference paper at ICLR 2021
Caripole Balance- - Natural
Finger Spin Nabural
Walker Wilk Natioall
Cheelah Bun Nat Ratural
430
250
PP0
500
UC3CmeC
420
NA bount
400
FAB
120
EAD tren
350
320
to
200
50
10
250
100
150
04
0.6
0.2
0.2
0.4
0.5
02
0.5
(100)
Enufromment Siee
(1.91
Figure 4: Average return on DMC tasks with natural video backgrounds with mean and standard de- viation computed over 5 seeds. UCB-DrAC outperforms PPO and RAD with the best augmentations.
such dropout (Srivastava et al. 2014) or batch normalization (loffe & Szegedy 2013). Farebrother
jet al. (2018) and Cobbe et al. (2018) show that such regularization methods can improve the gen-
eralization ability of RL agents in Atari (Machado al. 2018 and CoinRun (Cobbe al. 2018). respectively. Similarly, |gl et al. (2019) use selective noise injection with a variational information bottleneck, while Lee et al. 2020) regularize the agent's s representation with respect to random convolutional transformations. More recently, Sonar et al. (2020) learn invariant policies. ZZhngg al. (2020a) and Zhang et al. (2020c) learn state abstractions using bisimulation, Roy & Konidaris (2020) align the features of two domains using Wasserstein distance, while Igl et al.. (2020) reduce non-stationarity using policy distillation, and |azoure al. (2020) maximize the mutual information between the agent's internal representation of successive time steps. More similar to our work, [Cobbe] al. (2018), Ye et al. (2200) and Laskin et al. (2020) add augmented observations to the training buffer of RL agent. However,
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.
Data Augmentation has been extensively used in computer vision for both supervised (LeCun et al.
1989 Becker & Hinton| 1992 LeCun et al., 1998; (Siard et al.) [2003] Ciresan et al. [201 Cirresan
et al.
2011; Krizhevsk et al. 2012) and self-supervised (Dosovitskiy et al. 2016 Misra & van der
Maaten 2019) learning. More recent work uses data augmentation for contrastive learning, leading
to ofate-of-the-art results on downstream tasks (Ye et al. 2019; Hénaff et al. 2019 He et al.] 2019
Chen et al. 2020). Domain randomization can also be considered a type of data augmentation, which has proven useful for transferring RL policies from simulation to the real world (Tobin et al. (017))
However, domain randomization requires access to a physics simulator, which is not always available. Recently, few papers propose the use of data augmentation in RL (Cobbe et al. 2018: [ee et al., 2020, Srinivas et al.) (2200] Kostrikov et al., 20200] Laskin et al. 2020). 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. (2020), who propose to regularize the Q-function in Soft Actor-Critic (SAC) (Haarnoja al.) 2018) using random shifts of the input image. Our work differs from theirs in that automatically selects an augmentation from a given set, regularizes both the actor and the critic, and focuses the problem of generalization rather than sample efficiency. While there is body of work on the automatic use of data augmentation (Cubuk et al. (2019) a Fang et al., 2019 Shi et al., 2019 Li et al., 2020), these approaches were designed for supervised learning and, as we explain here, cannot be applied to RL without further algorithmic changes.
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 a new oftt--of-the-rt on the Procgen benchmark. In addition, the learned policies and representations are more invariant to spurious correlations between observations and rewards. A promising avenue for future research is to use a more expressive function class for meta-learning the augmentations.

---

Under review as a conference paper at ICLR 2021
REFERENCES
Peter Auer. Using confidence bounds for exploitationnexxloraiion trade-offs. Journal of Machine Learning Research, 3(Nov):397-47) 422, 2002.
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, pp. 2930 2941, 2018.
Suzanna Becker and Geoffrey E. Hinton. Self-organizing neural network that discovers surfaces in random-dot stereograms. Nature, 355:161-163, 1992.
Mare G Bellemare, Yavar Naddaf. Joel Veness, and Michael Bowling. The arcade learning environ- ment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47: 253-279, 2013.
Richard Bellman. markovian decision process. Journal of mathematics and mechanics, pp- 679-684, 1957.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey E. Hinton. A simple framework for contrastive learning of visual representations. ArXiv, abs/2002.05709, 2020.
Dan Ciregan, Ueli Meier, and Jürgen Schmidhuber. Multi-column deep neural networks for image classification. In 2012 IEEF conference computer vision and pattern recognition, pp. 3642-3649, IEEE, 2012.
Dan C Ciresan, Ueli Meier, Jonathan Masci, Luca M Gambardella, and Jürgen Schmidhuber. High- performance neural networks for visual object classification. arXiv preprint arXiv: 102.0183, 2011.
Dan C. Ciresan, Ueli Meier, Jonathan Masci, Luca Maria Gambardella, and Jürgen Schmidhuber. High-performance neural networks for visual object classification. ArXiv, abs/ 102.0188, 2011.
Karl Cobbe, Oleg Klimov, Chris Hesse, Taehoon Kim, and John Schulman. Quantifying generalization in reinforcement learning. arXiv preprint arXiv: 1..02344, 2018.
Karl Cobbe, Christopher Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation to benchmark reinforcement learning. arXiv preprint arXiv: 1912.0158,, 2019.
Ekin D. Cubuk, Barret Zoph, Jonathon Shlens, and Quoc V. Le. Randaugment: Practical automated data augmentation with a reduced search space. arXiv: Computer Vision and Pattern Recognition, 2019a.
Ekin Dogus Cubuk, Barret Zoph, Dandelion Mané, Vasudevan, and Quoe Le. Autoaugment: Learning augmentation strategies from data. 2019 IEEE/CVF Conference Computer Vision and Pattern Recognition (CVPR), pp. 113-123, 2019b.
Terrance DeVries and Graham W Taylor. Improved regularization of convolutional neural networks with cutout. arXiv preprint arXiv: /708. 04552, 2017.
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:17334111 17444,177,,11 2016.
Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Volodymir Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, Iain Dunning, et al. Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures. arXiv preprint 1rXiv:1802.0166) 2018.
Boli Fang, Miao Jiang, and Jerry J. Shen. Paganda : An adaptive task-independent automatic data augmentation. 2019.
Jesse Farebrother, Marlos C. Machado, and Michael H. Bowling. Generalization and regularization in dqn. ArXiv, abs/18 8 10.00123, 2018.
10

---

Under review as a conference paper at ICLR 2021
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-leaming for fast adaptation of deep networks. In Proceedings of the 34th International Conference on Machine Learning-Ver Volume 70, pp. 1 1226-1135 JMLR. org, 2017.
Shani Gamrian and Yoav Goldberg. Transfer learning for related reinforcement learning tasks via image-to-iman translation. ArXiv, abs/1800.07337., 2019.
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 arXiv: 2019.
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.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross B. Girshick. Momentum contrast for unsupervised visual representation learning. ArXiv, abs/ 0tss/9911.05722 2019.
Olivier J. 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, abs/1905.09272. 2019.
S. Hochreiter and Schmidhuber. Long short-term memory. Neural Computation, 9:1735-1788 1997.
Maximilian Igl, Kamil Ciosek, Yingzhen Li, Sebastian Tschiatschek, Cheng Zhang, Sam Devlin, and Katja Hofmann. Generalization reinforcement learning with selective noise injection and information bottleneck. In Advances in Neural Information Processing Systems, pp. 13956-13968, 2019.
Maximilian Igl, Gregory Farquhar, Jelena Luketina, Wendelin Böhmer, and Shimon Whiteson. The impact of non-stationarity generalisation in deep reinforcement learning. ArXiv, abs/2006.05826. 2020.
Maximilian llse, Jakub M Tomezak, and Patrick Forré. Designing data augmentation for simulating interventions. arXiv preprini arXiv:2005.01856, 2020.
Sergey Toffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. ArXiv, abs/ 100
W. Kay, J. Carreira, K. Simonyan, Brian Zhang, Chloe Hillier, Sudheendra Vijayanarasimhan, F. Viola, T. Green, T. Back, A. Natsev, Mustafa Suleyman, and Andrew Zisserman. The kinetics human action video dataset. ArXiv, abs/1 705.06950, 2017.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. CoRR, abs/1442.66980, 2015.
Ilya Kostrikov. Pytorch implementationn of reinforcement learning algorithms. https://githu>
2018.
Ilya Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. arXiv preprint arXiv:2004. / 13649, 2020.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolu- tional neural networks. In Advances in neural information processing systems, pp. 1097-1105, 2012.
Michael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Reinforcement learning with augmented data. arXiv preprint arXiv:2004. 14990, 2020.
Yann LeCun, Bernhard E. Boser, John 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.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient- based learning applied to document recognition. 1998.
11

---

Under review as a conference paper at ICLR 2021
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://opeereview. netfforum, 2020.
Yonggang Li, Guosheng Hu, Yongtao Wang, Timothy M. Hospedales, Neil M. Robertson, and Yongxing Yang. Dada: Differentiable automatic data augmentation. ArXiv, abs/2003.03780, 2020.
Marlos C. Machado, Mare G. Bellemare, Erik Talvitie, Joel Veness, Matthew Hausknecht, and Michael H. Bowling. Revisiting the arcade learning environment: Evaluation protocols and open problems for general agents. In IJCAI, 2018.
Bogdan Mazoure, R R'emi Tachet des Combes, Thang Doan, Philip Bachman, and R. Devon Hjelm. Deep reinforcement and infomax learning. ArXiv, abs/2006.07217, 2020.
Ishan Misra and Laurens van der Maaten. Self-supervised learning of pretext-invariant representations. ArXiv, abs/1912.01991, 2019.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Joannis Antonoglou, Daan Wierstra, and Martin A. Riedmiller. Playing atari with deep reinforcement learning. ArXiv, abs/1312.5602, 2013.
Charles Packer, Katelyn Gao, Jernej Kos, Philipp Krähenbühl, Vladlen Koltun, and Dawn Xiaodong Song. Assessing generalization in deep reinforcement learning. ArXiv,
2018.
Roberta Raileanu and Tim Rocktäschel. Ride: Rewarding impact-driven exploration for procedurally- generated environments. ArXiv, abs/2002. 12292, 2020.
Aravind Rajeswaran, Kendall Lowrey, Emanuel Todorov, and Sham M. Kakade. Towards generaliza- tion and simplicity in continuous control. ArXiv, abs/ htss1703.00666) 2017.
Josh Roy and George Konidaris. Visual transfer for reinforcement learning via wasserstein domain confusion. arXiv preprint arXiv:2006.03440 2020.
John Schulman, Sergey Levine, Pieter Abbeel, Michael I. Jordan, and Philipp Moritz. Trust region policy optimization. In ICML, 2015.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347 2017.
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.
Patrice Y Simard, David Steinkraus, John Platt, al. Best practices for convolutional neural networks applied to visual document analysis. In Icdar, volume 3, 2003.
Anoopkumar Sonar, Vincent Pacelli, and Anirudha Majumdar. Invariant policy optimization: Towards stronger generalization in reinforcement learning. ArXiv, 0tss2200601199 2020.
Xingyou Song, Yiding Jiang, Stephen Tu, Yilun Du and Behnam Neyshabur. Observational overfitting in reinforcement learning. ArXiv, abs/ 191200297, 2020.
Aravind Srinivas, Michael Laskin, and Pieter Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. ArXiv, abs/2004.04136, 2020.
Nitish Srivastava, Geoffrey E. Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: simple way to prevent neural networks from overfitting. J. Mach. Learn. Res., 15: 1929-1958, 2014.
Y. Tassa, Yotam Doron, Alistair Muldal, T. Erez, Y. Li, D. Casas, D. Budden, Abbas Abdolmaleki, J. Merel, Andrew Lefrancq, T. Lillicrap, and Martin A. Riedmiller. Deepmind control suite. ArXiv, abs/ Abs/1801.00690) 2018.
Joshua Tobin, Rachel H Fong, Alex Ray, Jonas Schneider, Wojciech 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), pp. 23-30, 2017.
12

---

Under review as a conference paper at ICLR 2021
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 preprint arXiv:1611.05763, 2016.
R. Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine Learning, 8:229-256, 2004.
Denis Yarats, Amy Zhang, llya Kostrikov, Brandon Amos, Joelle Pineau, and Rob Fergus. Improving sample efficiency in model-free reinforcement learning from images. ArXiv, abs/1910.01741, 2019.
Chang Ye, Ahmed Khalifa, Philip Bontrager, and Julian Togelius. Rotation, translation, and cropping for zero-shot generalization. arXiv preprint arXiv:2001 099008, 2020.
200000,0000908,
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), pp. 6203-6212, 2019.
Amy Zhang, Nicolas Ballas, and Joelle Pineau. A dissection of overfitting and generalization in continuous reinforcement learning. ArXiv, abs/aabs./n/eeee 806.07937, 201 8a.
Amy Zhang, Clare Lyle, Shagun Sodhani, Angelos Filos, Marta Kwiatkowska, Joelle Pineau, Yarin Gal, and Doina Precup. Invariant causal prediction for block mdps. arXiv preprint arXiv:2003.06016, 2020a.
Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. 2020b.
Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. arXiv preprint arXiv:2006. 10742, 2020c.
Chiyuan Zhang, Oriol Vinyals, Rémi Munos, and Samy Bengio. A study on overfitting in deep reinforcement learning. ArXiv, abs/1 808.00444999 06893, 2018b.
13

---

Under review as a conference paper at ICLR 2021
A
PPO
Proximal Policy Optimization (PPO) SShhulan et al. (2017) an actor-critic RI. algorithm that learns policy нe and value function Ve with the goal of finding an optimal policy for given MDP. PPO alternates between sampling data through interaction with the environment and optimizing an objective function using stochastic gradient ascent. each iteration, PPO maximizes the following objective:
JPPO Jn Jy + 022
where or,, 02 are weights for the different loss terms, SH, is the entropy bonus for aiding exploration, Jy is the value function loss defined as
Jv (Vo(s) - Vttrget)2
The policy objective term J= is based on the policy gradient objective which be estimated using importance sampling in off- off-policy settings (i.e. when the policy used for collecting data different from the policy we want to optimize):
(P) - To(a)s) (8, Eawtoosa aEA
TTQ.aa (als) Aoaa (s, (s.00)
To(aas))
(8)
where A(-) is estimate of the advantage function, Ootd the policy parameters before the update, Teas is the behavior policy used to collect trajectories (f.e. that generates the training distribution of states and actions), and ve is the policy want optimize (i.e. that generates the true distribution of states and actions).
This objective also be written
(O) Ban (r(0) (s, (s)
where
xe (a(a)s
Aada (a|s)
is the importance weight for estimating the advantage function.
PPO is inspired by TRPO (Schulman et al. 2013), which constrains the update so that the policy does not change too much in one step. This significantly improves training stability and leads to better results than vanilla policy gradient algorithms. TRPO achieves this by minimizing the KL divergence between the old (i.e. before update) and the new (i.e. after an update) policy. PPO implements the constraint a simpler way by using a clipped surrogate objective instead of the more complicated TRPO objective. More specifically, PPO imposes the constraint by forcing r(f) to stay within a small interval around precisely [1 - e, 1 + e]. where is a hyperparameter. The policy objective term from equation becomes
J. = E, min (rad, clip (ro, 1 A)]
where A Agua (s, a) for brevity. The function clip(r(0), - clips the ratio to be no more than and no less than e. The objective function of PPO takes the minimum one between the original value and the clipped version so that agents are discouraged from increasing the policy update to extremes for better rewards.
Note that the use of the Adam optimizer (Kingma & Ba 2015) allows loss components of different magnitudes so can G₂ and Gy from equations (3) and 44) be used as part of the DrAC ob- jective in equation (5) with the same loss coefficient or- This alleviates the burden of hyperparameter search and means that DrAC only introduces single extra parameter a,..
14

---

Under review as a conference paper at ICLR 2021
B NAIVE APPLICATION OF DATA AUGMENTATION IN RL
In this section, we further clarify why a naive application of data augmentation with certain RL algorithms is theoretically unsound. This argument applies for all algorithms that use importance sampling for estimating the policy gradient loss, including TRPO, PPO, IMPALA, or ACER. The use of importance sampling is typically employed when the algorithm performs more than a single policy update using the same data order to correct for the off off-policy nature of the updates. For brevity, we will use PPO to explain this problem.
The correct estimate of the policy gradient objective used in PPO is the one in equation (I (or equivalently, equation (8) which does not use the augmented observations at all since we are estimating advantages for the actual observations, A(s, a). The probability distribution used to sample advantages is (als) (rather than Tota(al (s)) since we can only interact with the envi- ronment via the true observations and not the augmented ones (because the reward and transition functions are not defined for augmented observations). Hence, the correct importance sampling instead would be incorrect for the reasons mentioned above. What we argue is that, in the case of RAD, the only way use the augmented observations ((() is in the policy gradient objective, whether by π T [())))/=a((((l))) or ( (s)) (sa)) (alss) depending on the exact implementation, but both of these incorrect. In contrast, DrAC does not change the policy gradient objective at all which remains the one in equation (I) and instead uses the augmented observations in the additional regularization losses, as shown in equations
estimate uses π (a]ss) Tata(a)s). Using T
and
C CYCLE-CONSISTEE
Here is a description of the cycle-consistency metric proposed by Aytar et al. (2018) and also used in [ee et al. (2020) for analyzing the learned representations of RL agents. Given two trajectories V and U, E V first locates its nearest neighbor in the other trajectory uj = argmin,.g [[h(v=) hwwlll where h(-) denotes the output of the penultimate layer of trained agents. Then, the nearest neighbor of uy E V is located, i.e., 122 = argmin, [/((,,) h(u))]l 2 and ny is defined as cycle-consistent if İ k| 1, i.e., it can return to the original point. Note that this cycle-consistency implies that two trajectories are accurately aligned in the hidden space. Similar to Aytar et al. (2018), we also evaluate the three-way cycle-consistency by measuring whether vi remains cycle-consistent along both paths, V + U + + V and V + I + - V, where J is the third trajectory.
D AUTOMATIC DATA AUGMENTATION ALGORITHMS
In this section, we provide more details about the automatic augmentation approaches, as well as pseudocodes for all three methods we propose.
RL2-DrAC uses an LSTM (Hochreiter & Schmidhuber 1997 network select an effective augmen- lation from a given set, which is used to update the agent's 's policy and value function according to the DrAC objective from Equation (5). We will refer to this network (recurrent) selection policy. The LSTM network takes as inputs the previously selected augmentation and the average return obtained after performing one update of the DrAC agent with this augmentation (using Algorithm with K=1). The LSTM outputs an augmentation from the given set and is rewarded using the average return obtained by the agent after one update with the selected augmentation. The L.TTM trained using REINFORCE (Williams) 2004). See Algorithm 3 for pseudocode of RL2-DrAC.
Meta-DrAC meta-learns the weights of a convolutional neural network (CNN) which used to augment the observations in order to update the agent's 's policy and value function according to the DrAC objective from Equation (5). For each DrAC update of the agent, we split the trajectories the replay buffer into meta-train and meta-test using a 9 to I ratio. The CNN's weights are updated using MAML (Finn et al. 2017) where the objective function is maximizing the average return obtained by DrAC (atter update using Algorithm 1 with K = 1 and the CNN the transformation f).
15

---

Under review as a conference paper at ICLR 2021
Algorithm UCB-DrAC
1: Hyperparameters: Set of image transformations F = .f....., f"). exploration coefficient window for estimating the Q-functions W, number of updates K, initial policy parameters meg initial value function Va-
2: N()) 1, FFFF
Initialize the number of times each augmentation was selected Initialize the Q-functions for all augmentations > Initialize the lists of returns for all augmentations
3: Q(f) = O, YEEF
4: R(f) = FIFO(W), VIEF
5: for I..... K do
6: f JEF -
[og(k) NUT
Use UCB to select augmentation
7: Update the policy and value function according to Algorithm with f = fx and K 1: D Update the policy
8:
< arg max, JDAAC
9:
Φ +- arg max JDAAC
Update the value function
Compute the mean return obtained by the new policy rk-
11:
Add to the R(fa) list using the firs--n-ffrst-out oir rule.
Q(fa) RUN) 2reR(fe) r
13:
((a)) + (f +1
end for
Algorithm 3 RL2-DrAC
I: Hyperparameters: Set of image transformations F { )"), number of updates K, initial policy ne. initial value function Va-
2: Initialize the selection poicy an LSTM network y with parameters vb.
3: fo F
b- RRndomly initialize the augmentation Initialize the average return
4: ro 4-0
5: for k= L do
6:
f fu-k((fk-.Ptt1
> Select augmentation according to the recurrent policy Update the policy and value function according to Algorithm with f = fx and K 1: > Update the policy Update the value function
7:
8:
+- arg maxg JDAAC
9:
Φ + arg max; JDAAC
Compute the mean return obtained by the new policy rk-
11: Reward Is with 7g-
Update Iv. using REINFORCE.
13: end for
Algorithm 4 Meta-DrAC
1: Hyperparameters: Distribution over tasks (or levels) g(m), number of updates step size parameters a and 3, initial policy ma. initial value function Vg
2: Initialize the set of all training levels D = (
Initialize the augmentation
CNN g with parameters vb.
4: for k = l,..,,.44 K do a batch of tasks m, ~ ((m)
5: for all m, do
6: Collect trajectories on task using the current policy.
7:
Update the policy and value function according to Algorithm with f= f& and =
8:
D Update the policy Update the value function Compute the return of the new agent on task m, after being updated with gos Twn (ge)
+ arg maxg JDrac
9:
t arg maxg. JDAAC
10:
11: vi t- s/ + a viTm,, (gv)
12:
end for
13:
V 2 + VV En Em.rg(m) Tim (got)
14:
end for
16

---

Under review as a conference paper at ICLR 2021
Table 3: List of hyperparameters used to obtain the results in this paper.
Hyperparameter
Value
0.999
0.95
timesteps per rollout
256
# epochs per rollout
# minibatches per epoch
entropy bonus clip range reward normalization
0.01
0.2
yes
learning rate workers # environments per worker
5c-4
# total timesteps optimizer
25M
Adam
LSTM frame stack
no
ax
0.1
0.1
E HYPERPARAMETERS
We use Kostrikov (2018)'s implementation of PPO (Schulman et al. 2017), on top of which all our
methods are build. The agent is parameterized by the ResNet architecture from (Espeholt et al. 2018)
which was used to obtain the best results in Cobbe et al. (2019). Following Cobbe et al. (2019).
also share parameters between the policy value networks. To improve stability when training
with DrAC, we only backpropagate gradients through π f (s, v)) and (ff v/) in equations (3)
and (4), respectively. Unless otherwise noted, (2019) for the easy mode of Procgen (iee same experimental setup as the one used here), namely:
the best hyperparameter' found Cobbe et al.
did a grid search for the regularization coefficient Or [0.0001, 0.01, 0.05, 0.1, 0.5, 1.0] used in equation (5) and found that the best value is ar, = 0.1, which used produce all the results in this paper.
For DrAC,
For UCB-DrAC, we did grid searches for the exploration coefficient E [0.0, 0..00.1,..55 1.0,5.0 5.0) and the size of the sliding window used compute the Q-values K [10, 50, 100), We found that the best values are c = 0.1 and K = 10, which were used to obtain the results shown here.
For RL2-DrAC, we performed hyperparameter search for the dimension of recurrent hidden state h E [16, 32, 64], for the learning rate E [3e 4, 5e - 4, 7e - 4]. and for the entropy coefficient e E [1e - 4, 1e 3, le 2] and found 32, Adam with = 1e 5 as the optimizer.
4, and le work best. We used
For Meta-DrAC. the convolutional network whose weights we meta-learn consists of single convolutional layer with 3 input and 3 output channels, kernel size 3, stride and O padding. At each epoch, we perform one meta-update where we unroll the inner optimizer using the training set and compute the meta-test return the validation set. We did the same hyperparameter grid searches as for RL2-DrAC and found that the best values were 7e - and le 2 in this case. The buffer of experience (collected before each PPO update) was split into 90% for meta-training and 10% for meta-testing.
For Rand-FM Lee et al. (2020) implementation, which were the best values for CoinRun (Cobbe et al. 2018), one of the Procgen games used for evaluation in (Lee et al. 2020).
the recommended hyperparameters in the authors' released
17

---

Under review as a conference paper at ICLR 2021
Dodgeball
Dodgeball
Dodgeball
1000
randam-com
200
Number Selected
Number Sellectd
Number Selected
150
800
600
100
400
200
50
0
5
10
15
20
25
10
15
20
25
10
20
25
Step [1e6]
Step [1e6]
Step [1e6]
(a) 0.0
(b) 0.1
c=)
Figure 5: Behavior of UCB for different values of its exploration coefficient on Dodgeball. When too small, UCB might converge to a suboptimal augmentation. On the other hand, when c is too large, UCB might take too long to converge.
For IBAC-SNI [gl et al. (2019) we also use the authors' open sourced implementation. We use the parameters corresponding to IBAC-SNI data augmentation turned on, and value of 3 = .0001 which turns on the variational information bottleneck, and selective noise injection turned on. This corresponds to the best version of this approach, as found by the authors after evaluating it on CoinRun (Cobbe et al. (018). While IBAC- SNI outperforms the other methods on maze-like games such as heist, maze. and miner, still significantly worse than our approach the entire Procgen benchmark.
.5. We use weight regularization with 12 .0001,
For both baselines, Rand-FM and IBAC-SNI, we use the same experimental setup for training and testing the used for methods. Hence, train them for 25M frames the easy mode of each Procgen game, using (the same) 200 levels for training and the rest of the levels for testing.
We use the Adam (Kingma & Ba 2015) optimizer for all experiments. Note that by using Adam, we do not need separate coefficients for the policy and value regularization terms (since Adam rescales gradients for each loss component accordingly).
F ANALYSIS OF UCB's BEHAVIOR
In Figure show the behavior of UCB during training, along with train and test performance on the respective environments. In the case of Ninja, UCB converges to always selecting the best augmentation only after 15M training steps. This is because the augmentations have similar effects on the agent early in training, so it takes longer to find the best augmentation from the given set. In contrast, on Dodgeball, UCB finds the most effective augmentation much earlier in training because there is a significant difference between the effect of various augmentations. Early discovery of an effective augmentation leads to significant improvements over PPO, for both train and test environments.
Another important factor is the exploration coefficient used by UCB (see equation (6 to balance the exploration and exploitation of different augmentations. 5 compares UCB 's behavior for different values of the exploration coefficient. Note that the coefficient augmentation with the largest Q-value. This can sometimes lead to UCB converging augmentation due to the lack of exploration. However, the exploration term of equation ( is too large relative to the differences in the Q-values among various augmentations, UCB might take too long to converge. In our experiments, we found that an exploration coefficient of 0.1 results in good exploration-explootaiinn balance and works well across all Procgen games.
UCB always selects
suboptimal
18

---

Under review as a conference paper at ICLR 2021
G PROCGEN BENCHMARK
Figure 6: Screenshots of multiple cnttlurrall-generrttee levels from 15 Procgen environments: StarPilot, CaveFlyer, Dodgeball, FruitBot, Chaser, Miner, Jumper, Leaper, Maze, BigFish, Heist, Climber, Plunder, Ninja, BossFight (from left to right, top to bottom).
H BEST AUGMENTATIONS
Table 4: Best augmentation type for each game, evaluated on the test environments.
Game
BigFish
StarPilot
FruitBot
BossFight
Ninja
Plunder
CaveFlyer
CoinRun
Best Augmentation
crop
crop
crop
flip
color-jitter
crop
rotate
random-cony
Table 5: Best augmentation type for each game, evaluated on the test environments.
Game
Jumper
Chaser
Climber
Dodgeball
Heist
Leaper
Maze
Miner
Best Augmentation
random-cony
crop
color-jitter
crop
crop
crop
crop
color-jitter
19

---

Under review as a conference paper at ICLR 2021
1 BREAKDOWN OF PROCGEN SCORES
Table 6: Procgen scores on train levels after training on 25M environment steps. The mean and standard deviation computed using 10 runs. The best augmentation for each game is used when computing the results for DrAC and RAD.
Game
PPO
Rand FM
IBAC-SNI
DrAC
RAD
UCB-DrAC
RI.2-DrAC
Meta-DrAC
BigFish
8.9 1.5
6.0 0.9
19.1
13.1 + 2.2
13.2 + 2.8
13.2 2.2
10.1 + 1.9
9.28 + 1.9
StarPilot
29.8 2.3
0.8
38.0 0.7
36.5
2.2
30.6 2.6
30.5 3.9
FruitBot
2.1
2.72
0.8
2.0
26.1
1.2
29.4 1.0
BossFight
8.7
5.7
7.7
8.2
8.1
8.2
8.4
0.5
Ninja
7.7
7.2
8.8
8.8
0.9
8.3
8.1
7.8
Plunder
6.5
5.7
6.0
1.3
8.1
11.1
0.5
6.5
CaveFlyer
0.6
6.5
0.5
8.7
0.8
5.7
0.8
6.7
CoinRun
9.3
9.6
9.6
9.7
9.6
9.5
0.3
9.4
Jumper
8.3
8.9
0.6
9.1
8.6
8.1
0.6
8.5
Chaser
0.5
2.7
3.1
7.1
6.4
7.6
4.5
5.5
Climber
8.4
7.5
7.1
9.8
9.3
9.0
0.9
8.5
Dodgeball
0.5
0.3
9.6
1.0
5.7
8.9
1.1
0.6
Heist
0.5
0.5
4.7
6.8
6.9
6.9
5.6 0.8
6.6
Leaper
0.4
3.7
2.7
5.7
4.9
5.5
2.7
3.6
Maze
0.3
8.9
8.2
8.7
8.4
8.7
7.0
9.2
Miner
0.3
1..7
8.7
1.3
1..6
12.5
10.9 1.5
1.34
Table 7: Procgen scores on test levels after training on 25M environment steps. The mean and standard deviation are computed using 10 runs. The best augmentation for each game is used when computing the results for DrAC and RAD.
Game
PPO
Rand FM
IBAC-SNI
DrAC
RAD
UCB-DrAC
RI.2-DrAC
Meta-DrAC
BigFish
4.0 1.2
0.6 0.8
0.6
0.8 0.9
8.7 + 1.4
9.9
9.7
6.0 - 0.5
3.6
StarPilot
24.7 3.4
8.8 0.7
4.8
29.5 5.4
3.1
3..2
29.4.20000
26.6
FruitBot
2.8
0.7
2.87
2.8
1.8
2.9
1.6
2..4
BossFight
1.0
1.7
1.0
0.8
0.9
8.8
7.6
7.7
Ninja
5.9
6.1
9.2
0.0
6.9
6.9
6.2
5.9
Plunder
5.0
3.0
2.1
9.5
8.5
8.9
1.=
5.64
CaveFlyer
5.1
5.4
8.0
0.8
5.1
0.9
4.1
5.5
CoinRun
8.5
9.3
8.7
8.
9.0
8.5
8.3
8.6
Jumper
5.5
5.310.6
3.6
6.4
6.6
6.4
0.5
5.7
Chaser
5.0
1.4
1.505
5.7
1.9
6.7
0.5
5.1
Climber
5.7
5.3
3.6
7.1
6.9
6.5
0.5
6.6
Dodgeball
1..7
0.5
1.4
4.8
2.8
4.7
3.0=
1.9
Heist
0.5
0.6
9.8
4.0
4.1
4.0
2.4=0.4
2.6
Leaper
4.9
6.2
6.8
5.3+111
1.3
5.3
0.7
3.
Maze
5.7
8.0
10.0
6.6
6.1
6.3 0.6
0.3
5.2
Miner
0.5
7.7
8.0
9.8
9.2
9.7
8.0
9.2 0.7
20

---

Under review as a conference paper at ICLR 2021
J
PROCGEN LEARNING CURVES
BigFish
StarPilot
FruitBot
BossFight
30
(best)
UCB-DrAC
10
20
Score
RL2-DrAC
4-
Meta-DrAC
10
2
0
Ninja
Plunder
CaveFlyer
CoinRun
12
8
9
7
Score
8
6-
8-
5-
6-
3
Jumper
Chaser
Climber
Dodgeball
8
8
8
6
Score
4-
4-
2
Heist
Leaper
Maze
Miner
12.5
9
10.0
8
7
7.5
6
5.0
5
10
Step [1e6]
Step [1e6]
Step [1e6]
Step [1e6]
Figure 7: Train performance various approaches that automatically select an augmentation, namely UCB-DrAC, RL2-DrAC, and Meta-DrAC. The mean and standard deviation are computed using 10 runs.
21

---

Under review as a conference paper at ICLR 2021
BigFish
StarPilot
FruitBot
BossFight
8
PPO
DrAC (best)
20
6
UCB-DrAC
Score
RL2-DrAC
20
15
4
Meta-DrAC.
4
10
2
0:
Ninja
Plunder
CaveFlyer
CoinRun
7
6
Score
3-
6
Jumper
Chaser
Climber
Dodgeball
5
6.5
6.0
Score
5.5
3
5.0
2
4.5
Heist
Leaper
Maze
Miner
10
6.5
3.5
6.0
8-
3.0
5
5.5
6
5.0
4-
2.0
2
4.0
20
10
20
Step [1e6]
Step [1e6]
Step [1e6]
Step [1e6]
Figure 8: Test performance of various approaches that automatically select an augmentation, namely UCB-DrAC, R122-DrAC, and Meta-DrAC. The mean and standard deviation are computed using 10
runs.
22

---

Under review as a conference paper at ICLR 2021
K
DEEPMIND CONTROL SUITE EXPERIMENTS
Figure 9: DMC environment examples. Top row: default backgrounds without any distractors. Middle row: simple distractor backgrounds with ideal gas vidcos. Bottom row: natural distractor backgrounds with Kinetics videos. Tasks from left right: Finger Spin, Cheetah Run, Walker Walk.
Cartpoll Balance- Defaut
Finger Spin Defauk
Walber Walk Decault
Cheetsh Run De Defacit
buct
600
HRB
tuat
7:0
500
200
9:0
420
220
300
100
200
200
220
0.2
0.1
05
07
0.5
0.6
52
04 0.5 0s
Step
Figure 10: Average return on DMC tasks with default (i.e. no distractor) backgrounds with mean and standard deviation computed over 5 seeds. UCB-DrAC outperforms PPO and RAD with the best augmentations.
Caripele Balance Simple
Finger Soin Simple
Simple
400
450
RAP bunn
ND bett
3C0
RAB Tour
KAD bzeth
4:0
300
440
350
2200
200
Anccoge
:00
320
2:0
250
50
50
250
Imntrsmmen Step (105)
Invtrommnt Stes (101)
Ensterment Stop (131)
Instrotment Stop (127)
Figure 11: Average return on DMC tasks with simple (i.e. synthetic) distractor backgrounds with mean and standard deviation computed over seeds. UCB-DrAC outperforms PPO and RAD with the best augmentations.
For our DMC experiments, we ran a grid search over the learning rate in [le [le-44 4, 3e 7, 1, the number of minibatches in [32, 8, 16, 64]. the entropy coefficient in [0.0, 0.01, 0.001, 0.0001], and the number of ppo epochs per update in [3. 5, 10, 20). For Walker Walk and Finger Spin we use use 4. We also use 3 stacked frames observations. For Finger
action repeats and for the others Spin, we found 10 ppo epochs, 0.0 entropy coefficient, 16 minibatches, and 0.0001 learning rate to be the best. For Cartpole Balance, we used the same except for 0.001 learning rate. For Walker Walk, used 5 ppo epochs, 0.001 entropy coefficient, 32 minibatches, and 0.001 learning rate. For Cheetah Run, we used ppo epochs, 0.0001 entropy coefficient, 64 minibatches, and 0.001 learning rate. We use = 0.99. 0.95 for the generalized advantage estimates, 2-48 steps, 1 process, value loss coefficient 0.5, and linear rate decay over I million environment steps. We also found crop be the best augmentation from our set of eight transformations. Any other hyperaparameters not mentioned here were set the same values as the ones used for Procgen as described above.
23


Reviews:
[{"review_id": "oeLdbdVTzeX", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: \n\nThis paper proposes an automatic data augmentation approach for RL tasks. Specifically, it takes UCB for data selection and introduces two regularization terms for actor-critic algorithms' policy and value function. Then this paper evaluated the approach based on the Procgen benchmark and demonstrated that it outperforms existing methods. It is also shown that the learned policies and representations are robust to irrelevant factors.\n\n\nReasons for score:\n\nOn the one hand, I feel that the proposed approach is incremental (UCB for data selection with the actor-critic algorithm as the RL algorithm plus additional regularization terms). It would also be ideal if more experiments can be included to prove that the proposed approach is effective in most RL tasks. However, on the other hand, based on the current experimentation results, the proposed method seems promising and can be useful for generalization in reinforcement learning.\n\n\nPros\n\n1. The proposal of the automatic data augmentation/selection approach is useful, given that the existing primary methods either rely on expert knowledge or separately evaluate a large number of transformations to find the best one, both of which are expensive. \n\n2. The proposed approach achieves SOTA performance in the Procgen benchmark.\n\n3. It shows that the proposed method is robust to irrelevant factors. \n\t\n\nCons\n\n1. I think the primary difference between the proposed approach and existing approaches is the additional regularization losses. This would make the new approach seem incremental. \n\n2. There is a lack of performance comparison between the automatic data augmentation method and non-automatic approaches.\n\n3. The data augmentation approach is only evaluated in the Procgen benchmark. It would be great if this paper includes more experiments to demonstrate the performance, especially since the proposed method is intended for any RL task.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "QZlDOtzdNda", "reviewer": "AnonReviewer1", "summary": "", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper presents a method that utilizes data augmentation for image-based reinforcement learning. The data augmentation is used to regularize the policy and function approximation in the proposed method. In addition, a method for automatically identifying effective ways of data augmentation is proposed. The experimental results show that the proposed method outperforms the baseline methods.\n\nThe study shows that the regularization of policy and function approximation using the transformed images is more effective than training a policy by using the transformed image as a state. Regarding identification of effective ways of data augmentation,  the proposed method seems improve the performance of image-based RL methods, although the proposed approach is simple.\n\nI have some suggestions for improving the paper.\n- The term for regularizing the policy is given in Eq. (3). I think that the first \\pi(a|s) in the KL divergence should also have the subscript \\theta because it seems that the two policies are based on the same model. Likewise, the first term in Eq. (4) should have the same subscript as the second term. \n\n- It would be good to show the Jensen-Shannon divergence and the cycle consistency of RAD in Table 2. \n\n\n== comments after discussions and the paper update ==\n\nI appreciate the authors' efforts to improve the clarity and provide additional results. I believe that the proposed method is now clearly presented and the claims are properly supported by experiments. I raise the score to \"accept\". \n", "extra_scores": {}, "extra_text": {}}, {"review_id": "ODn9f8VMfOk", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 4, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary after Discussion Period:\n-----------------------------------------------\nAfter reading the other reviewer's comments and corresponding with the authors, I have become convinced that the author's proposed regularization method is novel and effective, and would recommend this avenue of research be further explored. Yet it has also become clear to me that the author's claims on why their method works are not yet supported by evidence. Further, I don't believe the author's proposed further ablation studies would fix the theory, since such experiments don't address whether their method works by fixing problems with Laskin’s work (as the author's claim) or because it provides a more direct way of enforcing invariance to transformation (as I claim).\n\nSo we're left with a difficult situation, the method and the experiments are good while the theory is lacking. In such a situation both acceptance or rejection seem reasonable. Yet, as per ICLR reviewer guidelines, one should answer three questions for oneself:\n\n - What is the specific question and/or problem tackled by the paper?\n - is the approach well motivated, including being well-placed in the literature?\n - Does the paper support the claims? This includes determining if results, whether theoretical or empirical, are correct and if they are scientifically rigorous.\n\nSince the theory is lacking and the approach is not well motivated, and since the theoretical claims haven't been rigorously supported, I feel as per ICLR guidelines the paper is not yet ready for acceptance.\n\nInitial Review:\n-------------------\n\nSummary In this paper, the authors introduce three approaches for automatically finding an augmentations for any RL task, and a novel regularization scheme to make such augmentations work effectively.\n\nPositive aspects:\n-----------------------\nThe paper’s language is clear and the authors provide a good overview of the problem of data augmentation for reinforcement learning. Furthermore, they nicely explain why data augmentation for RL isn’t as straightforward as augmenting data for supervised learning learning. I believe that data augmentation could be a nice tool in the Reinforcement Learner’s toolbox, and I’m glad to see a paper advancing the idea.\n\n\nMajor Concerns:\n-----------------------\n\nThis paper has not provided sufficient evidence that the author’s proposed way of doing data augmentation is effective. In this paper, there are two main novel methods for doing data augmentation / insights in RL. I will discuss my concerns with both methods separately.\n\nPolicy and Value function Regularization.\n-----------------------------------------------------\n\nThe authors criticize the naive application of transformers in the PPO’s buffer, as done in Laskin et al. (2020), saying that this changes the PPO objective. While I agree that such a naive transformation as in Eq. 2 is problematic, I fail to understand why application of transformation to states in the buffer would result in the Eq. 2, as the transformation would happen to the states being fed into both $\\pi_\\theta$ and $\\pi_{\\theta_\\text{old}}$, resulting in an equation different from Eq. 2. One just needs to save the old policies (and old Advantage function), so that $\\pi_{\\theta_\\text{old}}$ can be applied to transformed states, and not just use the actions from the buffer. This would seem like a straightforward fix.\n\nYet the authors have proposed a different regularization fix which judging the experiments does seem to work, as shown in Figure 2. I suspect it works for another reason: since the regularization forces $V(s) = V(f(s))$ and $\\pi(\\cdot | s) = \\pi(\\cdot | f(s))$ I wonder if this isn’t a method to allow prior knowledge to flow into the policy and value estimation. If the transformation(s) $f_i$ have been chosen such that one can be reasonably sure that true value and policy functions should be invariant to said transformations, then by enforcing $V(s) = V(f(s))$ and $\\pi(\\cdot | s) = \\pi(\\cdot | f(s))$  one is constraining V and \\pi to fit the prior knowledge contained in $f_i$.\n\nSo now we’re comparing apples and oranges, since your method (DrAC) gets to incorporate this prior knowledge, while PPO and RAD don’t, and DrAC’s good performance isn’t surprising.\n\nAutomatic Data Augmentation\n----------------------------------------\n\nHere, given some candidates for data augmentation, the authors propose three methods to discover which candidate work well. Unfortunately, I don’t fully understand the approach. \n\nThe authors are examining a Meta-Reinforcement Learning setting, where one wishes to find a policy which performs well not just on one MDP, but on a whole distribution of MDPs. This leads to an inner-and-outer for-loop like setting, in the inner for loop, the agent does multiple episodes with a single environment, in the outer loop the agent gets new environments.\n\nI had expected this inner-outer for-loop structure to pe present in the meta learning of augmentations, and for the authors to clearly describe how knowledge about the effectiveness of augmentations is transferred from past episodes on one environment to future episodes in the same environment, and how it’s transferred from past environments to future environments.\n\nThe only support given for these methods is the experimental results. Yet we see that the performance of DrAC and UCB-DrAC lie within a standard deviation of one another. So the evidence of the effectiveness of UCB-DrAC over the simpler DrAC is weak at best.\n\nMinor Comments:\n-------------------------\n$T_m(s’|s,a)$ is the transition function, $R_m(s,a)$ is the reward function: Here, $T$ and $R$ are generally distributions and not functions\n\nEq. 6 is confusing, since $f_*$ can refer to both $f_t$ (function at timestep $t$) and $f_i$ (the $i$-th transformation function). \n\nDo the update with something like $N_t(f) \\leftarrow N_{t-1}(f) + 1$ is a bit clearer, since then it’s the number of times $f$ has been pulled before timestep $t$\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "6t-9HCFwPM", "reviewer": "AnonReviewer2", "summary": "", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "#######################################\n\nSummary:\nThis paper tackles the problem of generalization in deep RL via data augmentation. It provides a framework for automatic data augmentation based on UCB, RL^2, or MAML. When UCB is combined with regularization of the policy and value function so that their outputs are invariant to transformations (such as rotation, crop, etc.), it shows improvements over similar algorithms on the ProcGen benchmark. \n\n#######################################\n\nPros:\n1. The paper is well written and the proposed algorithms are straightforward to understand.\n2. UCB-DrAC is shown to be statistically superior to the tested baselines. Surprisingly, when all ProcGen games are taken into account, some existing methods such as Rand-FM are actually worse than PPO.\n3. Ablation studies are able to show that both ingredients of UCB-DrAC are important. UCB is shown to find the best augmentation asymptotically, and the DrAC is shown to be better than PPO and RAD. \n\nCons:\n1. The experiments do not compare the proposed algorithms to DrQ, the algorithm proposed in Kostrikov et al. (2020). Since it also tackles generalization in deep RL through data augmentation, it seems that it should be included as a baseline. Can the authors explain why it was not included?\n2. The proposed algorithms are only combined with PPO. It would be good to have some results for other actor-critic algorithms, such as SAC, to verify if the SOTA behavior holds.\n\n#######################################\n\nOverall:\nI would vote for accepting this paper, due to the strength of its experimental results. The proposed approaches are novel in the data augmentation for generalization in deep RL subfield, although AutoAugment (cited in the paper) also uses RL for choosing data augmentations in the supervised learning case.\n\n#######################################\n\nFurther comments and questions:\n1. It may be helpful to rewrite some of the algorithms in Section 3.3 in pseudocode, as text is harder to read.\n2. Is there intuition for why only UCB-DrAC leads to statistically significant improvements over PPO, and not RL2-DrAC and Meta-DrAC? Can it be shown that the former has lower sample complexity than the latter two?\n3. Why is the cycle consistency percentage in Table 2 always low? Is it some idiosyncrasy of the metric or due to the inherent variance in the trajectories?\n\n#######################################\n\nUpdate after reading other reviews and author responses:\n\nI am happy to keep my score and support accepting this paper. I agree with Reviewer 4 that conducting a more detailed ablation study of the impact of regularizing both the policy and value function (i.e. a comparison with an algorithm like that of Kostrikov et al. (2020)) would improve the paper and hope that it will be included in the final paper. ", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: As of now, automatic data augmentation methods have mostly been proposed for supervised learning tasks, especially classification. This paper introduces automatic data augmentation to deep (image-based) reinforcement learning agents, aiming to make the agents generalize better to new environments. A new algorithm called data-regularized actor-critic (DrAC) is proposed, with three variants that correspond to different methods for automatically finding a useful augmentation: UCB-DrAC, RL2-DrAC, and Meta-DrAC. Promising results are reported on OpenAI’s Procgen generalization benchmark which consists of 16 procedurally generated environments (games) with visual observations. Further experiments have been added in the revised version.

**Strengths:**
  * This work is among the first attempts that propose an automatic data augmentation scheme for reinforcement learning.
  * The paper articulates well the problem of data augmentation for reinforcement learning.
  * The experiment results are generally promising.

**Weaknesses:**
  * Although the experiment results reported seem promising, there are missing pieces in order to help the readers gain a deeper understanding to justify more thoroughly why the proposed regularization-based scheme works.
  * Theoretical justification is lacking.

This is a borderline paper. While it presents some interesting ideas supported empirically by experiment results, the paper in its current form is premature for acceptance since a more thorough, scientific treatment is lacking before drawing conclusions. Moreover, considering that there are many competitive submissions to ICLR, I do not recommend that the paper be accepted. Nevertheless, the authors are encouraged to address the concerns raised to fill the gaps when revising their paper for resubmission in the future.

Decision: Reject