Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels
Authors: Denis Yarats
Ilya Kostrikov
Rob Fergus
Venue: ICLR 2021 Spotlight (2021)

Abstract: We propose a simple data augmentation technique that can be applied to standard model-free reinforcement learning algorithms, enabling robust learning directly from pixels without the need for auxiliary losses or pre-training.  The approach leverages input perturbations commonly used in computer vision tasks to transform input examples, as well as regularizing the value function and policy.  Existing model-free approaches, such as Soft Actor-Critic (SAC), are not able to train deep networks effectively from image pixels. However, the addition of our augmentation method dramatically improves SAC’s performance, enabling it to reach state-of-the-art performance on the DeepMind control suite, surpassing model-based (Hafner et al., 2019; Lee et al., 2019; Hafner et al., 2018) methods and recently proposed contrastive learning (Srinivas et al., 2020).  Our approach, which we dub DrQ: Data-regularized Q, can be combined with any model-free reinforcement learning algorithm. We further demonstrate this by applying it to DQN and significantly improve its data-efficiency on the Atari 100k benchmark.

Published as a conference paper at ICLR 2021
IMAGE AUGMENTATION Is ALL YOU NEED: REGULARIZING DEEP REINFORCEMENT LEARNING FROM PIXELS
Denis Yarats*
Iya Kostrikov
New York University & Facebook Al Research denisyarats@ca nyu. edu
New York University
kostrikovêcs nyu, edu
Rob Fergus
New York University
fergusêcs edu
ABSTRACT
Existing model-free reinforcement learning (RL) approaches are effective when trained on states but struggle to learn directly from image observations. We propose an augmentation technique that can be applied to standard model-free RL algorithms, enabling robust learning directly from pixels without the need for auxiliary losses pre-training. The approach leverages input perturbations commonly used in computer vision tasks to transform input examples, as well regularizing the value function and policy. Our approach reaches of-the-art performance on DeepMind control suite and Atari 100k benchmark, surpassing previous model-free (Haarnoja al. 20118) van Hasselt al. (019a). model-basee badel-based (Hafner et al. 2019 Lee et al. [2019] Halner et al. 2018 Kaiser et al. 2019) and contrastive learning (Srimivas et al. (020) approaches. also closes the gap between state-based and image- based RL trainin. Our method, which we dub DrQ: regulaegul Q, can be combined with any model-free RL algorithm.
state-
To the best of our knowledge, our approach is the first effective data augmentation method for RL on these benchmarks.
INTRODUCTION
Sample-efficient deep reinforcement learning (RL) algorithms capable of directly training from image pixels would open up many real-world applications in control and robotics. However, simultaneously training convolutional encoder alongside policy network is challenging when given limited environment interaction, strong correlation between samples and typically sparse reward signal. Limited supervision common problem across Al and two approaches are commonly taken: (i) training with an additional auxiliary losses, such as those based on self-supervised learning (SSL) and (ii) training with data augmentation.
A wide range of auxiliary loss functions have been proposed to augment supervised objectives, e.g. weight regularization, noise injection (Hinton et al. 20122). or various forms of auto-encoded (Kingma) et al. 2014). In RL, reconstruction losses (Jaderberg et al. 2017 Yarats et al. 2019) SSL objectives (Dwibedi et al. 2018 Sriiivas et al., 2020) are used. However, these objectives are unrelated to the task hand, thus have guarantee of inducing an appropriate representation for the policy network. SSL losses are highly effective in the large data regime, e.g. in domains such as vision (Chen et al. 2020 He al. 2019) and NLP (Collobert et al.) 201 Devlin et al. 2018) where large (unlabeled) datasets are readily available. However, in sample-efficient RI., training data is more limited due to restricted interaction between the agent and the environment, limiting their effectiveness.
Data augmentation methods are widely used in vision and speech domains, where output-invariant perturbations can easily be applied to the labeled input examples. Surprisingly, data augmentation
Equual contribution. Author ordering determined by coin flip. Both authors corresponding.

---

Published as a conference paper at ICLR 2021
has received little attention in the RL community. In this paper we propose augmentation approaches appropriate for sample-efficient RL and comprehensively evaluate them. The key idea of our approach is to use standard image transformations to perturb input observations, well regularizing the Q-function learned by the critic so that different transformations of the same input image have similar Q-function values. No further modifications standard actor-critic algorithms required. Our study is, to the best of our knowledge, the first careful examination of image augmentation in sample-efficient RL.
The main contributions of the paper are as follows: (i) the first to demonstrate that data augmentation greatly improves performance when training model-free RL algorithms from images; (ii) introducing a natural way to exploit MDP structure through two mechanisms for regularizing the value function, in a manner that generally applicable to model-free RL. and (iii) setting a new state-of-the-ard performance on the standard DeepMind control suite TTasaa et al. (2018), closing the gap between learning from states, and Atari 100k (Kaiser et al.) 2019) benchmarks.
2
RELATED WORK
Data Augmentation in Computer Vision Data augmentation via image transformations has been used to improve generalization since the inception of convolutional networks (Becker & Hinton 1992 Simard et al. [2003] LeCun et al. 1989; Ciresan et al. 2011 Ciregan et al. 2012). Following
AlexNet (Krizhevsky et al. (012), they have become a standard part of training pipelines. For object classification tasks, the transformations
selected to avoid changing the semantic category, i.e. translations, scales, color shifts, etc. While a similar set of transformations are potentially applicable control tasks, the RI. context does require modifications to be made to the underlying algorithm.
Data augmentation methods have also been used in the context of self-supervised learning. Dosovit- skiy et al. (2016) use per-exemplar perturbations in unsupervised classification framework. More recently, several approaches (Chen et al. 2020 He et al. 2019)] Misra & van der Maaten, 2019) have used invariance imposed image transformations in contrastive learning schemes, producing of-te-of-thh-arr results on downstream recognition tasks. By contrast, our scheme addresses control tasks, utilizing different types of invariance.
Data Augmentation in RL In contrast to computer vision, data augmentation is rarely used in RL. Certain approaches implicitly adopt it, for SaampeeLeeing et al. (2018); Kalashnikov et al. (2018)
use image augmentation as part of the AlexNet training pipeline without analysing the benefits occurring from it, thus being overlooked in subsequent work. HER (Andrychowicz et al.] 2017) exploits information about the observation space by goal and reward relabeling, which can be viewed as way to perform data augmentation. Other work uses data augmentation to improve generalization in domain transfer (Cobbe et al.) 2018). However, the classical image transformations used in vision have not previously been shown to definitively help on standard RL benchmarks. Concurren with work, RAD (Laskin al. 2020) performs exploration of different data augmentation approaches, but is limited to transformations of the image alone, without the additional augmentation of the Q-function used in our approach. Moreover, RAD can be regarded as special case of our algorithm. Multiple follow ups to our initial preprint appeared on ArXiv (Raileanu et al. 2020 Okada & Taniguchi 2020), using similar techniques on other tasks, thus supporting the effectiveness and generality of data augmentation in RL.
Continuous Control from Pixels There are a variety of methods addressing the sample-efficiency of RL algorithms that directly learn from pixels. The most prominent approaches for this can be classified into two groups, -adell-asedd and model-free methods. The model-based methods attempt to learn the system dynamics in order to acquire a compact latent representation of high-dimensional
observations to later perform policy search (Hafner et al. 2018 Lee et al.) 2019 Hafner et al. 2019). In contrast, the model-free methods either learn the latent representation indirectly by optimizing the RL objective (Barth-Maron Maron et al. 2018 Abdolmaleki et al. 2018) ) or by employing auxiliary losses
that provide additional supervision Yarats et al. 2019 Srinivas et al. 2020. Sermanet al. 2018 Dwibedi et al. 2018). Our approach is complementary to these methods and can be combined with them to improve performance.
2

---

Published as a conference paper at ICLR 2021
3 BACKGROUND
Reinforcement Learning from Images We formulate image-based control as an infinite-horizo partially observable Markov decision process (POMDP) (Bellman 1957 Kaelbling et al. 1998). An POMDP can be described as the tuple (O, A, P, p,,,), where O is the high-dimensional observation space (image pixels), A is the action space, the transition dynamics p = Pr(o)]ost, a.) capture the probability distribution over the next observation of given the history of previous observations o and current action a₁, r O x A -> is the reward function that maps the current observation and action reward r2 r(oga, at), and [0, 1) is discount factor. Per common practice (Mnih) et al. 2013). throughout the paper the POMDP is converted into an MDP (Bellman 1957) by stacking several consecutive image observations into state 5, {or, 1t-1, 2t-2, ...). ... For simplicity we Pr(s(|s,, ae) and the reward function Te = r(sr. 0e). We then aim find a policy (()))ss)) that maximizes the cumulative discounted return E, 1Σ y'relag ~ I (-1ss), s ~ p(-(ss, ~ p())-
redefine the transition dynamics
Soft Actor-Critic The Soft Actor-Critic (SAC) (Haarnoja al., 2018 learns state-action value function Qo stochastic policy To and temperature to find an optimal policy for an MDP (S, A. p. r. 7) by optimizing 7-discuunted maximum-enee -nntropp objective (Ziebart et al. 2008). used generically denote the parameters updated through training in each part of the model.
Deep Q-learning DQN (Mnih et al., 2013) also learns a convolutional neural net to approximate Q-function over states and actions. The main difference is that DQN operates discrete actions spaces, thus the policy can be directly inferred from Q-values. In practice, the standard version of DQN is frequently combined with set of refinements that improve performance and training stability, commonly known as Rainbow (yan Hasselt et al.) 2013). For simplicity, the rest of the paper describes generic actor-critic algorithm rather than DQN or SAC in particular. Further background on DQN and SAC can be found in (Appendix A.
4 SAMPLE EFFICIENT REINFOR LEARNING FROM PIXELS
4.1 OPTIMALITY INVARIANT IMAGE TRANSFORMATIONS FOR Q FUNCTION
We first introduce general framework for regularizing the value function through transformations of the input state. For given task, we define an optimality invariant state transformation f S x S mapping that preserves the Q-values
Qss a) Qf v), for all E S, E and v e T.
where y are the parameters of f()), drawn from the set of all possible parameters One example of such transformations are the random image translations successfully applied in the previous section.
For every state, the transformations allow the generation of several surrogate states with the same Q-values, thus providing mechanism to reduce the variance of Q-function estimation. In particular, for an arbitrary distribution of states u(-) and policy т, instead of using a single sample s" p(-). a" ~ TT-Is*) estimation of the following expectation
E soru()) [Q(s, a)) v Q("",a") (-1((1))
generate K samples via random transformations and obtain estimate with lower variance
K Q(f(s" vk). ox) where vy E T and ag ~ π T--If(=", vx)).
E s [Q(s,a) (-)))
~ K
Σ
This suggests two distinct ways to regularize Q-function. First, we use the data augmentation to compute the target values for every transition tuple (st, as rs, s/) as
1 K
y1 ri
Qo(f(s', M(,)), dh,x) where dea ((115((),,/,,))))
(I)
3

---

Published as a conference paper at ICLR 2021
where v4,2 T corresponds to transformation parameter of s/. Then the Q-function is updated using these targets through an SGD update using learning rate X
Σ (o((((s v.)² 1 N t -- -
In tandem, we note that the same target from Equation (1) can be used for different augmentations of Si resulting in the second regularization approach
N.M
0
Σ 1im==
(Qo(S(s)) - y;)²
(3)
NM
When both regularization methods are used, Vi, and vL drawn independently.
4.2 PRACTICAL INSTANTIATION OF OPTIMALITY INVARIANT IMAGE TRANSFORMATION
A range of successful image augmentation techniques have been developed in computer vision (Ciregan et al., [012]] Ciresan et al., 2011 Simard et al. 2003; Krizhevsky et al. 2012 Chen et al. 2020). These apply transformations to the input image for which the task labels are invariant, e.g. for object recognition tasks, image flips and rotations do not alter the semantic label. However, tasks in RL differ significantly from those in vision and in many cases the reward would not be preserved by these transformations. We examine image transformations from Chen et al. (2020) (random shifts, random cutouts, horizontal/vertical flips, rotations and intensity shifts) in Appendix E. and conclude that random shifts strike a good balance between simplicity and performance, we therefore limit our choice of transformation function f(-) to random shifts.
We apply shifts to the images sampled from the replay buffer. For example, images from the DeepMind control suite used in our experiments are 84 84. We pad each side by pixels (by repeating boundary pixels) and then select a random 84 x 84 crop, yielding the original image shifted by #4 pixels. This procedure is repeated every time an image sampled from the replay buffer.
4.3 OUR APPROACH: DATA-REGULARIZED Q (DRQ)
Our approach, DrQ, is the union of the three separate regularization mechanisms introduced above:
1. transformations of the input image (Section 4.2).
2. averaging the Q target over K image transformations EEquation (I)
3. averaging the Q function itself over M image transformations (Equation (3)
Algorithm 1 details how they are incorporated into a generic pixel-based off-policy actor-critic algorithm. Note that if [M=I,M==] then DrQ reverts to image transformations alone, this makes applying DrQ to any model-free RL algorithm straightforward.
For the experiments in this paper, we pair DrQ with SAC (Haarnoja et al., 2018) and DQN (Mnih al.) 2013). popular model-free algorithms for control in continuous and discrete action spaces respectively. We select image shifts as the class of image transformations f, with v + 4, as explained in Section 4.2
EXPERIMENTS
5.1 ABLATION EXPERIMENT
Figure I shows the effect of image shift augmentation applied to three tasks from the DeepMind
control suite (Tassa et al. 2018 Figure la shows unmodified SAC (Haarnoja et al. 2018) parameter- ized with different image encoders, taken from: NatureDON (Mnih et al. 2013). Dreamer (Hafner
et al. 2019), Impala (Espeholt et al. 2018), SAC-AE (Yarats et al,] 2019), and D4PG (Barth-Madd
al. 2018). The encoders vary significantly in their architecture and capacity, with parameter
4

---

Published as a conference paper at ICLR 2021
Algorithm 1 DrQ: retalregularized Q applied to generic off-policy actor critic algorithm. Black: unmodified off-policy actor-critic.
Orange: image transformation.
Green: target Q augmentation.
Blue: Q augmentation.
Hy pyperparameter:: Total number of environment steps T, mini-batch size N. learning rate Ag. target network update rate т, image transformation f, number of target Q augmentations K, number of Q augmentations M. for each timestep & = 1..T'do
a₂ T(+1s))
5/ p(-(s) at)
+ DU r(sr, at),
Add transition to the replay buffer
UPDATECRITIC(D)
UPDATEACTOR(D) Data augmentation applied to the samples for actor training as well end for
procedure UPDATECRITIC(D)
((s.. ai, Te, ~ (/L M1x ~ -U(T),
Sample mini batch from the replay buffer
1.
Uniformly sample target augmentations
for each = I..N do
df (--ss)) or d/12 ~ f((, k= 1.K
Q = Q0 (8/, d)) or Q = k Ek= Qr (((,) x(
Vi t r(s;, a;) + +7②;
end for
{Vi.m Vi, ~ ((()))) i 1= N, m = 1..M}
b Uniformly sample Q augmentations E (Qo((s y1) or JQ NM Exm= (Qo (f Pt,m), o.) y1)
JQ(0)
0 0 Xn (0)
D Update the critic >> Update the critic target
0' + (1- +70
end procedure
counts ranging from 220k to 2.4M. None of these train satisfactorily, with performance decreasing for the larger capacity models. Figure 1b) shows SAC with the application our random shifts transformation of the input images (i.e. just Section 4.2] not Q augmentation also). The results for all encoder architectures improve dramatically, suggesting that our method is general and can assist many different encoder architectures. To the best of our knowledge, this is the first successful demonstration of applying image augmentation on the standard benchmarks for continuous control. Furthermore, Figuree 2 shows the full DrQ. with both image shifts and augmentation (Section 4. as well as ablated versions. Q augmentation provides additional consistent gain over image shift augmentation alone (full results are in Appendix
5.2 DEEPMIND CONTROL SUITE EXPERIMENTS
In this section we evaluate our algorithm (DrQ) on the two commonly used benchmarks based on the DeepMind control suite (Tassa et al. 2018), namely the PlaNet (Hafner et al., 2018) and Dreamer (Hafner et al.) 2019) setups. Throughout these experiments all hyper-parameters of the algorithm are kept fixed: the actor and critic neural networks are trained using the Adam opti- mizer Kingmaa & Ba 2014) with default parameters and mini-bate size of 512 For SAC, the soft target update rate 7 is 0.01. initial temperature is 0.1, and target network and the actor updates made every 2 critic updates (as in Yarats et al. (2019)) We use the image encoder architecture from SAC-AE Yaratt et al. 2019) and follow their training procedure. The full set of parameters can be found in Appepidx B Following Henderson et al. (2018), the models are trained using 10 different seeds; for every seed the mean episode returns are computed every 10000 environment steps, averaging over 10 episodes. All figures plot the mean performance over the 10 seeds, together with + 1 standard deviation shading. We compare our DrQ approach to leading model-free and model-based
Note that DrQ does not utilize additional information beyond transitions sampled from the replay buffer (i.c. does use more observations per mini-batch), thus the mini-batch size is the same as for unmodified SAC.
5

---

Published as a conference paper at ICLR 2021
Cartpole Swingup
Finger Spin
Walker Walk
Eptsode Tantern
822-
Eptende Rettrr
600
600
620
600
MAD wate mo 23M1
203
Bromar 2aaaig
200
200
BiG peld loc BLA
Dell Orc ean
SAL befo Enc S2444
Stapa (x:10))
Stups
(a) Unmodified SAC.
Cartpole Swingup
Finger Spin
Walker Walk
1.000
1020
1000
Spicode Decuum
Epionde Recur
600
620
600
Cell
te 19740)
5AT FQt by F.240
racad Spot
200
SACAACe 31M
200
5AC SSC trecl 5114
566 D4P0 (2440
6.0
Ensironment Slaas (:18f)
Enviroement (x200)
Envirumment Siaps (x:1
(b) SAC with random shifts augmentation.
The performance of SAC trained from pixels on the DeepMind control suite using image encoder networks of different capacity (network architectures taken from recent RL algorithms, with parameter count indicated). (a): unmodified SAC. Task performance can be seen to get worse as the capacity of the encoder increases. For Walker Walk (right), all architectures provide mediocre performance, demonstrating the inability of SAC to train directly from pixels on harder problems. (b): SAC combined with image augmentation in the form of random shifts. The task performance is now similar for all architectures, regardless of their capacity, which suggests the generality of our method. There is also clear performance improvement relative to (a), particularly for the more challenging Walker Walk task.
Figure I:
Run
Reacher Easy
Walker
1200
3000
6/00
Esoood Recurn
800
800
E240000
Estoode
410-
Aigmanssinn
Aupracctionn Miat
200-
0.58
275
2.20
Lis
150 135
3.00
0.25
1:00
125
1.00
LT5
202
0.00
0.25
0.00
0.75
Envisnment Stops (x20%)
Environeent Stass (SID')
Enviroement Staps (x229)
Figure 2: Different combinations of our three regularization techniques on tasks from Tassa et al. 2018) using SAC. Black: standard SAC. Blue: DrQ [K=1,M=1], SAC augmented with random shifts. Red: DrQ 2 random shifts Target Q augmentations. Purple: DrQ [K=,, random shifts + Target Q + Q augmentations. All three regularization methods correspond to Algorithm 1 with different K.M showing clear gains when both Target Q and Q augmentations used.
approaches: PlaNet (Hafner et al. CURL (Srinivas et al.) (020) and Dreamer (Hafner al. 2019). The comparisons use the results
2018), SAC-AE (Yarats et al. 2019). SLAC (Lee et al., 2019).
provided by the authors of the corresponding papers.
PlaNet Benchmark (Hatner al. 2018) consists of six challenging control tasks from (Tassa et al. 2018) with different traits. The benchmark specifies different action-repeat paraeeerr for each of the six task Following common practice (Hafner et al. 2018 Lee al. 2019 Yarats et al.
This means the number of training observations is fraction of the environment steps (e.g. an episode of 1000 steps with action-repeat 4 results in 250 training observations).

---

Published as a conference paper at ICLR 2021
Boll in Cup Catch
Cartpole Swinguup
Cheetai
10000
1000
970
Eplecde Benter
K76
Epleede Beruurn mm
400
Apert
See
CLR
Fatton
ORL
KAC
0.3
2.5
6.5
0.0
0.2
3.4
0.
0.50
025
030
11
2.00
Environement Stoee (x2f)
Envirommmnt Stopp bx3071
Environment Steos (xUF)
Finger Spin
Reacher Easy
Walker Walk
3000
Epitade Rutum
Enknnt
690
400
Agent
430
400
Agant
Fialet
300
Fille
LAGAC
SAGAC
Ero loua
Ong sound
DrG (oust
Steps
Figure 3: The PlaNet benchmark. Our algorithm (DrQ (K=2,M=2)) outperforms the other methods and demonstrates the state-of-the-art performance. Furthermore, on several tasks DrQ is able to match the upper-bound performance of SAC trained directly on internal state, rather than images. Finally, our algorithm not only shows improved sample-efficiency relative to other approaches, but is also faster in terms of wall clock time.
2019; Mnih al., 2 we report the performance using true environment steps, thus are invariant to the action-repeat typer-parameter Aside from action-repeat, all other hyper-parameters of our algorithm are fixed across the six tasks, using the values previously detailed.
Figure 3) compares DrQ 2= to PlaNet (Hafner et al. 2018), SAC-AE Yarats et al. 2019).
CURL (Srinivas et al. 2020). SLAC (Lee et al. 2019). and an upper bound performance provided by SAC HHarnooa et al. 2018) that directly learns from internal states. We use the version of SLAC that performs one gradient update per an environment step to ensure a fair comparison to other approaches. all the tasks, despite being much simpler than other methods. Furthermore, since DrQ does not learn a model (Hafner et al.] [2018] [eee al. 2019) any auxiliary tasks (Srinivas al. 2020), the wall clock time also compares favorably to the other methods.
DrQ achieves state-of ah-t performance this benchmark
In Table 1 we also compare performance given at fixed number of environment interactions (e.g. 100k and 500k). Furthermore, in Appendix G we demonstrate that DrQ is robust to significant changes in hyper-parameter settings.
Dreamer Benchmark is a more extensive testbed that was introduced in Dreamer (Hafner et al. 2019), featuring a diverse set of tasks from the DeepMind control suite. Tasks involving sparse reward were excluded (c.g. Acrobot and Quadruped) since they require modification of SAC to incorporate multi-step returns (Barth-Maron Maron et al. 2018), which is beyond the scope of this work. We evaluate on the remaining 15 tasks, fixing the action-repeat
to 2 as in Hafner et al. (2019).
We compare DrQ [K= to Dreamer (Hafner al., 2019) and the upper- performance of SAC (Haarnoja et al. 2018) from states") Again, we keep all the pyper-parameters of our algorithm fixed across all the tasks. Figure 4, DrQ demonstrates the state-of-the-art results by collectively outperforming Dreamer (Hafner al. (019), although Dreamer is superior on 3 of the 15 tasks (Walker Run, Cartpole Swingup Sparse and Pendulum Swingup). On many tasks DrQ approaches the upper-bound performance of SAC (Haarnoja et al. 2018) trained directly states.
Noo other publicly reported results are available for the other methods due to the recency of the Dreamer (Hafner al. 2019) benchmark.
7

---

Published as a conference paper at ICLR 2021
Table 1: The PlaNet benchmark at 100k and 500k environment steps. Our method (DrQ (K=2,M=2]) outperforms other approaches in both the data-efficient (100k) and asymptotic performance (500k) regimes. Random shifts only version (e.g. DrQ [M=],)=l]) has a competitive performance but is consistently inferior to DrQ [K=2,M=2], particularly for 100k steps. We emphasize, that both versions of DrQ use exactly the same number of interactions with both the environment and replay buffer. Note that DrQ [K=1,M=1] is almost identical to RAD (Laskin et al. 2020), modulo some differences.
500k step scores
DrQ [K=2
DrQ
CURL
PlaNet
SAC-AE
SLAC
SAC State
Finger Spin
938= 103
913E151
874=131
7IK=40
914=107
771
927=43
Caripole Swingup
868::10
845:399
861430
787-46
730=152
870447
Reacher Easy
942:171
857+120
904194
5881471
601=135
97545
Cheetah Run
660:196
460-559
5004991
568=21
5441:50
64
772-660
Walker Walk
921:145
897147
906:1:6
478::16
858:182
865
96418
Cup Catch
963-19
961412
9584:13
939-143
810:1121
100k step scores
Finger Spin
901=104
744+144
779+108
560
747-=130
680 30
672+76
Caripole Swingup
759--92
537-:119
592:1170
563-173
2764:38
812:445
Reacher Easy
601=2113
451+:210
517=113
82:174
225=1164
919::123
Cheetah Run
344:167
250-558
307148
252::173
240:138
391-477
228-955
Walker Walk
612=164
501-688
344=132
221=:43
395-1588
428=744
604-317
Ball In Cup Catch
913:153
6671114
7724241
7101-217
398-1196
607:1173
957
Caech
Emince
Canpoe
Carboc Seigga>-ppo
Cheese
Watter
Wileem Wark
100
Figure 4: The Dreamer benchmark. Our method (DrQ (K=2,M=2]) again demonstrates superior performance over Dreamer on 12 out 15 selected tasks. In many cases it also reaches the upper-bound performance of SAC that learns directly from states.
5.3 ATARI 100K EXPERIMENTS
We evaluate DrQ [K=I,M=]] on the Atari 100k benchmark (Kaiser et al., 2019) sample- constrained evaluation for discrete control algorithms. The underlying RL approach which DrQ is applied is a DQN, combined with double Q-learning (van Hasselt et al. 2015), n-step returns (Mnih al. 2016). and dueling critic architecture (Wang et al. 2015). As per common practice (Kaiser et al. 2019 van Hasselt et al.) 2019a). we evaluate our agent for 125k environment steps at the end of training and average performance over 5 random seeds. Figure 5 shows the median human- normalized episode returns performance (as in Mnih et al. (2013) of the underlying model, which we refer to as Efficient DQN, in pink. When DrQ is added there is significant increase in performance (cyan), surpassing OTRainbow (Kielak 2020) and Data Efficient Rainbow (van Hasselt et al. 2019a) 1a).. DrQ is also superior to CURL (Srinivas et al. 2020) that uses an auxiliary loss built on top hybrid between OTRainbow and Efficient rainbow. DrQ combined with Efficient DQN thus achieves state-of-the-ar performance, despite being significantly simpler than the other approaches. The experimental setup and full results detailed in Appendix Cand Appendix D respectively
8

---

Published as a conference paper at ICLR 2021
Median human normalized returns
0.25
0.20
0.15
0.10
0.05
0.00
Rainbow
SimPLe
OTRainbow
Eff. Rainbow OT/Eff. Rainbow +CURL
Eff. DON
DON +rQ (Ours)
Figure 5: The Atari 100k benchmark. Compared to set of leading baselines, our method (DrQ [K=1,M=1], combined with Efficient DQN) achieves the state-of-the-ard performance, despite being considerably simpler. Note the large improvement that results from adding DrQ to Efficient DQN (pink eyan). By contrast, the gains from CURL, that utilizes tricks from both Data Efficient Rainbow and OTRainbow, are more modest over the underlying RL methods.
CONCLUSION
We have introduced regularization technique, based on image shifts and Q-function augmentation, that significantly improves the performance of model-fre RL algorithms trained directly from images. In contrast to the concurrent work of Laskin et al. (2020), which is special case of DrQ, our method exploits the MDP structure of the problem, demonstrating gains over image augmentations alone. Our method is easy implement and adds negligible computational burden. We compared our method to state-of-the-art approaches on the DeepMind control suite, outperforming them on the majority of tasks and closing the gap with state-based training. On the Atari 100k benchmark DrQ outperforms other SOTA methods in the median metric. To the best of our knowledge, this is the first convincing demonstration of the utility of data augmentation on these standard benchmarks. Furthermore, we demonstrate the method to be robust to the choice of hyper--aramend parameters.

---

Published as a conference paper at ICLR 2021
REFERENCES
Abbas Abdolmaleki, Jost Tobias Springenberg, Yuval Tassa, Remi Munos, Nicolas Heess, and Martin Riedmiller. Maximum a posteriori policy optimisation. arXiv preprint / 1806.06666600 0096.06920, 2018.
Marcin Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong. Peter Welinder, Bob McGrew, Josh Tobin, OpenAl Pieter Abbeel, and Wojeiech Zaremba. Hindsight experience replay. In Advances in neural information processing systems, pp. 5048-5058, 2017.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization. arXiv e-prints, 2016.
Gabriel Barth-Maron, Matthew W. Hoffman, David Budden, Will Dabney, Dan Horgan, Dhruva TB, Alistair Muldal, Nicolas Heess, and Timothy Lillicrap. Distributional policy gradients. In International Conference on Learning Representations, 2018.
Becker and G. E. Hinton. Self-organizing neural network that discovers surfaces in random-dot Stereograms. Nature, 1992.
Richard Bellman. markovian decision process. Indiana Univ. Math. J., 1957.
Ting Chen, Simon Korbllihh, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. arXiv preprint arXiv:2002.05700 2020.
Dan Ciregan, Ueli Meier, and Jurgen Schmidhuber. Multi-column deep neural networks for image classification. In 2012 IEEE conference on computer vision and pattern recognition, pp. 3642-3649, 2012.
Dan C Ciresan, Ueli Meier, Jonathan Masci, Luca M Gambardella, and Jurgen Schmidhuber. High- performance neural networks for visual object classification. arXiv preprint arXiv: 102.0183, 2011.
Karl Cobbe, Oleg Klimov, Chris Hesse, Tachoon Kim, and John Schulman. Quantifying generalization in reinforcement learning. arXiv preprint 11Xiv:1812.02341, 2018.
Ronan Collobert, Jason Weston, Leon Bottou, Michael Karlen, Koray Kavukcuoglu, and Pavel Kuksa. Natural language processing (almost) from scratch. Journal of machine learning research, 2011 2.
Jacob Devlin, Ming-Woi Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint 0rXiv:1810.0480) 2018.
Terrance Vries and Graham W Taylor. Improved regularization of convolutional neural networks with cutout. arXiv preprint arXiv: 04555,, 2017.
Alexey Dosovitskiy, Philipp Fischer, Jost Tobias Springenberg, Martin Riedmiller, and Thomas Brox. Discriminative unsupervised feature learning with exemplar convolutional neural networks. TPAMI, 2016.
Debidatta Dwibedi, Jonathan Tompson, Corey Lynch, and Pierre Sermanet. Learning actionable representations from visual observations. CoRR, 2018.
Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Volodymir Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, al. Impala: Scalable distributed deep-rl with importance weighted leamnrr architectures. arXiv preprint arXiv: 1802:0166. 2018.
Scott Fujimoto, Herke van Hoof, and David Meger. Addressing function approximation error in actor-critic methods. In Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmussan. Stockholm, Sweden, July 10- 15, 2018, 2018
Tuomas Haarojj,, Aurick Zhou, Kristian Hartikainen, George Tucker, Sehoon Ha, Jie Tan, Vikash Kumar, Henry Zhu, Abhishek Gupta, Pieter Abbeel, et al. Soft actor-critic algorithms and applications. arXiv preprint arXiv: 18/2.059890 2018.
10

---

Published as a conference paper at ICLR 2021
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. arXiv preprint arXiv: 1.0455,, 2018.
Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv: 2019.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. arXiv preprint arXiv: 1911.0522, 2019.
Peter Henderson, Riashat Islam, Philip Bachman, Joelle Pineau, Doina Precup, and David Meger. Deep reinforcement learning that matters. Thirty-Second AAAI Conference On Artificial Intelli- gence (AAAI), 2018.
Geoffrey E Hinton, Nitish Srivastava, Alex Krizhevsky, llya Sutskever, and Ruslan R Salakhutdinov. Improving neural networks by preventing co-adaptation of feature detectors. arXiv preprint arXiv: 1207.00200.0 2012.
Max Jaderberg, Volodymyr Mnih, Wojciech Czarnecki, Tom Schaul, Joel Z. Leibo, David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. International Conference on Learning Representations, 2017.
Leslie Pack Kaclbling, Michael Littman, and Anthony R Cassandra. Planning and acting in partially observable stochastic domains. Artificial intelligence, 1998.
Lukasz Kaiser, Mohammad Babacizadch. Piotr Milos, Blazej Osinski, Roy H. Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, Ryan Sepassi, George Tucker, and Henryk Michalewski. Model-based reinforcement learning for atari. arXiv preprint 190::::9900000744 2019.
Dmitry Kalashnikov, Alex Irpan, Peter Pastor, Julian Ibarz, Alexander Herzog, Eric Jang. Deirdre Quillen, Ethan Holly, Mrinal Kalakrishnan, Vincent Vanhoucke, et al. Qt-opt: Scalable deep reinforcement learning for aasion-based robotic manipulation. arXiv preprint 18Xi 10293, 2018.
Kacper Piotr Kielak. Do recent advancements in model-based deep reinforcement learning really improve data efficiency? openreview, 2020. URL
IddBBeeulHHeg
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1rXiv:1412.69800 2014.
Durk P Kingma, Shakir Mohamed, Danilo Jimenez Rezende, and Max Welling. Semi-supervised learning with deep generative models. In Advances in neural information processing systems, pp. 3581-3589, 2014.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey Hinton. Imagenet classification with deep convolu- tional neural networks. In Advances in neural information processing systems, 2012.
Michael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Reinforcement learning with augmented data, 2020.
Yann LeCun, Bernhard Boser, John S Denker, Donnie Henderson, Richard Howard, Wayne Hubbard, and Lawrence D Jackel. Backpropagation applied to handwritten zip code recognition. Neural computation, 1989.
A. X. Lee, A. Nagabandi, P. Abbeel. and S. Levine. Stochastic latent actor-critic: Deep reinforcement learning with a latent variable model. arXiv e-prints, 2019.
Sergey Levine, Peter Pastor, Alex Krizhevsky, Julian Ibarz, and Deirdre Quillen. Learning hand- eye coordination for robotic grasping with deep learning and large-scale data collection. The International Journal of Robotics Research, 37(4-5):4214473 2018.
11

---

Published as a conference paper at ICLR 2021
Timothy P. Lillicrap. Jonathan J. Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. CoRR, 2015.
Ishan Misra and Laurens van der Maaten. Self-supervised learning of 192.09911,.01911, 2019.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, loannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv e-prints, 2013.
Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy P. Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. CoRR, 2016.
Masashi Okada and Tadahiro Taniguchi. Dreaming: Model-based reinforcement learning by latent imagination without reconstruction. arXiv preprint 140::20007.14555,, 2020.
Roberta Raileanu, Max Goldstein, Denis Yarats, Ilya Kostrikov, and Rob Fergus. Automatic data augmentation for generalization in deep reinforcement learning. 2020.
Edgar Riba, Dmytro Mishkin, Daniel Ponsa, Ethan Rublee, and Gary Bradski. Kornia: an open source differentiable computer vision library for pytorch. In The IEEE Winter Conference on Applications of Computer Vision, pp. 3674-3683, 2020.
Andrew M. Saxe, James L. McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv e-prints, 2013.
Pierre Sermanet, Corey Lynch, Yevgen Chebotar, Jasmine Hsu, Eric Jang, Stefan Schaal, Sergey Levine, and Google Brain. Time-contrastive networks: Sulf-super learning from video. In 2018 IEEE International Conference on Robotics and Automation (ICRA), pp. 1134-1141, IEEE, 2018.
Patrice Y Simard, David Steinkraus, John C Platt, et al. Best practices for convolutional neural networks applied to visual document analysis. In Iedar, 2003.
Aravind Srinivas, Michael Laskin, and Pieter Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. arXiv preprint arXiv:2004. 04/ 04136, 2020.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint arXiv: 1801.0000, 2018
Hado van Hasselt, Arthur Guez, and David Silver. Deep reinforcement learning with double q- learning. arXiv e-prints, 2015.
Hado Hasselt, Matteo Hessel, and John Aslanides. When use parametric models in reinforce- ment learning? arXiv preprint arXiv: 052433, 2019a.
Hado P van Hasselt, Mattco Hessel, and John Aslanides. When to use parametric models in reinforcement learning? In Advances in Neural Information Processing Systems, 2019b.
Ziyu Wang, Tom Schaul, Matteo Hessel, Hado Van Hasselt, Marc Lanctot, and Nando De Freitas. Dueling network architectures for deep reinforcement learning. arXiv preprint 1.1.6688, 2015.
Denis Yarats and Ilya Kostrikov. Soft actor-critic (sac) implementation in pytorch. https:// 2020.
Denis Yarats, Amy Zhang, Ilya Kostrikov, Brandon Amos, Joelle Pineau, and Rob Fergus. Im- proving sample efficiency in model-free reinforcement learning from images. arXiv preprint 191.v01900.0170 2019.
Brian D. Ziebart, Andrew Maas, J. Andrew Bagnell, and Anind K. Dey. Maximum entropy inverse reinforcement learning. In Proceedings of the 23rd National Conference Artificial Intelligence Volume 3, 2008
12

---

Published as a conference paper at ICLR 2021
APPENDIX
A EXTENDED BACKGROUND
Reinforcement Learning from Images We formulate image-based control an infinite-horizon partially observable Markov decision process (POMDP) (Bellman 1957 Kaelbling et al. 1998). An POMDP can be described as the tuple p,,, where O is the high-dimensional observation space (image pixels), A is the action space, the transition dynamics p = Pr(o/]osr, ax) capture the probability distribution over the next observation of given the history of previous observations ose and current action r O x A - IR is the reward function that maps the current observation and action reward F4 r(oge, at). and Y [0, 1) is a discount factor. Per common practice (Mnih et al. 2013). throughout the paper the POMDP is converted into an MDP (Bellman 1957) by stacking several consecutive image observations into state 8; {ot 0:-1, 0t-2; ... -). For simplicity we redefine the transition dynamics
Pr de) and the reward function r(st. ac). We then
aim to find a policy T())())) that maximizes the cumulative discounted return Ex, 12 7 ~ T( (-1ss),, s/ p(-|st, 4z), $1 p(-))-
Soft Actor-Critic The Soft Actor-Critic (SAC) (Haarnoja et al. 2018) learns a state-action value function Qo, stochastic policy To and temperature to find an optimal policy for MDP (S, A, P. 7.)) by optimizing 7-discounted maximum-entropy objective (Ziebart et al. 2008 is used generically denote the parameters updated through training in each part of the model. The actor policy (a (t) is a parametric tanh-Gaussian that given samples a, = tanh (s.) + (8))e)) where e NO, 1) and pe and dg are parametric mean and standard deviation.
The policy evaluation step learns the critic Qa(s;, ax) network by optimizing single-step of the soft Bellman residual
(D) = (1 (QQa at) - y
U/ = r(st, 0z) + (00 (6/es ob) a log (d)1)))))
where D a replay buffer of transitions, 0' an exponential moving average of the weights as done in (Lillicrap et al. 2015). SAC uses clipped double-Q learning (van Hasselt et al.] [0155] Fujimoto ctt al. 2018). which we omit from our notation for simplicity but employ in practice.
The policy improvement step then fits the actor policy Th (a(a)ss)) network by optimizing the objective
(D) EEG>D DDLL (To (s)]l exp( 1 a (ss.))))))
Finally, the temperature a is learned with the loss
Jo (D) =
s,nD I-a log no (arass))) - -aH),
a (-19e)
where H E IR is the target entropy hyper-parameter that the policy tries to match, which in practice is usually to H -IA].
Deep Q-learning DQN (Mnih et al. 2013) also learns a convolutional neural net to approximate Q-funetion over states and actions. The main difference is that DQN operates on discrete actions spaces, thus the policy can be directly inferred from Q-values. The parameters of DQN are updated by optimizing the squared residual error
(D)
ut) y() yt r(st, at) + max Qo (s),
o
In practice, the standard version of DQN is frequently combined with a set of tricks that improve performance and training stability, wildly known Rainbow (van Hasselt et al. 2015).
13

---

Published as a conference paper at ICLR 2021
B THE DEEPMIND CONTROL SUITE EXPERIMENTS SETUP
Our PyTorch SAC (Haarnoja al. 2018) implementation is based off of Yarats & Kostrikov (2020).
B.1 ACTOR AND CRITIC NETWORKS
We employ clipped double Q-learning (van Hasselt et al. 2015 Fujimoto et al., 2018 for the critic, where each Q-function is parametrized except of the last. The actor is also 3-layer MLP with ReLUs that outputs mean and covariance for the diagonal Gaussian that represents the policy. The hidden dimension is set 1024 for both the critic and actor.
3-layer MLP with ReLU activations after each layer
B.2 ENCODER NETWORK
We employ encoder architecture from Yarats et al. 2019). This encoder consists of four convo- lutional layers with 3 x 3 kernels and 32 channels. The ReLU activation is applied after each cony layer. We use stride to 1 everywhere, except of the first conv layer, which has stride 2. The output of the convnet is feed into a single fully-connected layer normalized by LayerNorm rm (Ba et al. 2016). Finally,
apply tanh nonlinearity the 50 dimensional output of the fully-connected layer. We initialize the weight matrix of fully-connected and convolutional layers with the orthogonal initialization (Saxe et al., 2013) and set the bias be zero.
The actor and critic networks both have separate encoders, although we share the weights of the conv layers between them. Furthermore, only the critic optimizer allowed update these weights (e.g. we stop the gradients from the actor before they propagate to the shared conv layers).
B.3 TRAINING AND EVALUATION SETUP
Our agent first collects 1000 seed observations using random policy. The further training obser- vations are collected by sampling actions from the current policy. We perform one training update every time receive new observation. In cases where observations is only a fraction of the environment steps (e.g. a 1000 steps episode at action repeat will only results into 250 training observations). We evaluate our agent every 10000 true environment steps by computing the average episode return over 10 evaluation episodes. During evaluation we take the mean policy action instead of sampling.
action repeat, the number of training
B.4 PLANET AND DREAMER BENCHMARKS
We consider two evaluation setups that were introduced PlaNet (Hafner et al., 2018) and Dreamer (Hafner al. 2019), both using tasks from the DeepMind control suite (Tassa et al. 2018).
The PlaNet benchmark consists six tasks various traits. Importantly, the benchmark proposed to use different action repeat syper-parameter for each task, which we summarize in Table 2
The Dreamer benchmark considers an extended set of tasks, which makes it more difficult that the the same set hyper-parameters for each
PlaNet setup. Additionally, this benchmark requires task, including action repeat (set to 2). which further increases the difficulty.
Table 2: The action repeat hyper- parameter used for each task in the PlaNet benchmark.
Task name
Action repeat
Cartpole Swingup
8
Reacher Fasy
4
Cheetah Run
4
Finger Spin
2
Ball Cup Catch
4
Walker Walk
2
14

---

Published as a conference paper at ICLR 2021
B.5 PIXELS PREPROCESSING
We construct an observational input as an 3-stack of consecutive frames (Mnih et al. 2013). where each frame is RGB rendering of size 84 84 from the Othh camera. We then divide each pixel by 255 to scale down to [0, 1] range.
B.6 OTHER HYPER PARAMETERS
Due to computational constraints for the continuous control ablation experiments the main paper and appendix we use minibatch size of 128, while for the main results we use minibatch of size 512. In Table 3 we provide comprehensive overview of all the other nyper-parameters
Table 3: An overview of used hyper- paper-prraeeters in the DeepMind control suite experiments.
Parameter
Setting 100000
Replay buffer capacity
Seed steps
1000
Ablations minibatch size
128
Main results minibatch size
512
Discount y
0.99
Optimizer
Adam
Learning rate
10-3 3
Critic target update frequency
2
Critic Q-function soft-update rate т
0.01
Actor update frequency
2
Actor log stddev bounds
|-10, 2]
Init temperature
0.1
15

---

Published as a conference paper at ICLR 2021
C THE ATARI 100K EXPERIMENTS SETUP
For ease reproducibility in Table we report the hyper-paramen settings used in the Atari 100k experiments. We largely reuse the hyper-parameten from OTRainbow (Kielak) 2020), but adapt them for DQN (Mih et al., 2013). Per common practise, we average performance of our agent over random seeds. The evaluation is done for 125k environment steps the end of training for 100k environment steps.
Table 4: A complete overview of hyper parameters used in the Atari 100k experiments.
Parameter
Setting
Data augmentation
Random shifts and Intensity
Grey-scaling
True
Observation down-sampling
84 x
Frames stacked
4
Action repetitions
Reward clipping
[-1, 1]
Terminal on loss life
True
Max frames per episode
108k
Update
Double Q True
Dueling
Target network: update period
1
Discount factor
0.99
Minibatch size
32
Optimizer
Adam
Optimizer: learning rate
0.0001
Optimizer: $1
0.9
Optimizer: $2
0.999
Optimizer:
0.00015
Max gradient norm
10
Training steps
100k
Evaluation steps
125k
Min replay size for sampling
1600
Memory size
Unbounded 1 step
Replay period every
Multi-step return length
10
Q network: channels
32, 64, 64
Q network: filter size
x 8, 4 x 4, x 4,2, 1
x
Q network: stride
Q network: hidden units
512
Non- Mon-linearity
ReLU e-greedy 5000
Exploration
e-deay
16

---

Published as a conference paper at ICLR 2021
D FULL ATARI 100K RESULTS
Besides reporting in Figure 5 median human-normalized episode returns over the 26 Atari games used in (Kaiser et al., 2019), we also provide the mean episode return for each individual game in Table 5
Table 5: Mean episode returns each of 26 Atari games from the setup in Kaiser et al. (2019). The results are recorded at the end training and averaged across random seeds (the CURL's results averaged over 3 seeds reported in Srinivas et al. (2020). On each game mark bold the highest score. Our method demonstrates better overall performance (as reported in Figure
Game
Rainbow
SimPLe
OTRainbow
Efl. Rainbow
OT/EIL Rainbow
EN. DQN
ED. DQN
+CURL
+D:Q (Ours)
Alien
318.7
616.9
824.7
739.9
1148.2
558.1
702.5
32.5
88.0
232.3
63.7
100.2
Assault
231.0
527.2
351.9
431.2
543.7
589.5
490.3
Asterix
243.6
1128.3
628.5
470.8
$24.3
341.9
577.9
BankHeist
15.6
34.2
182.1
51.0
193.7
74.0
205.3
BattleZone
2360.0
5184.4
4060.6
10124.6
11208.0
4760.8
6240.0
Boxing
-24.8
9.1
4.8
Breakout
1.2
16.4
9.8
1.9
18.2
7.3
14.3
ChopperCommen Command
120.0
1246.9
1033.3
861.8
1198.0
624.4
CrazyClimber
2254.5
62583.6
21327.8
16185.3
27805.6
$430.6
20072.2
DemcknAtacc
163.6
208.1
711.8
508.0
834.0
403.5
1086.0
Freeway
0.0
20.3
25.0
27.9
27.9
3.7
20.0
Frostbite
60.2
254.7
231.6
866.8
924.0
202.9
889
Gopher
431.2
8014
Hero
487.0
2656.6
6458.8
6857.0
6235.1
2200.1
4083.7
Jamesbond
47.4
125.3
112.3
301.6
400.1
133.2
330.3
Kangaroo
0.0
323.1
605.4
779.3
345,3
448.6
1282.6
Krall
1468.0
4539.9
3277.9
2851.5
3833.6
2999.0
4163.0
KungFuMaster
0.0
17257.2
5722.2
14346.1
14280.0
2020.9
7649.0
MsPacman
67.0
1480.0
941.9
1204.1
1492.8
872.0
1015.9
Pong
-20.6
PrivateFiye
0.0
58.3
100.0
97.8
105.2
351_3
-50.4
Qbert
123.5
1288,8
509.3
1152.9
1225.6
627.5
769.1
RoadRunmer
1588.5
$640.6
2696.7
9600.0
6786.7
1491.9
8296.3
Seaquest
131.7
683.3
286.9
354.1
408.0
240.1
299
UpNDown
504.6
3350.3
2847.6
2877.4
2735.2
2901.7
3134.8
Median human-nommen normalised
episode returss
0.020
0.135
0.208
0.147
0.240
0.094
0.270
E IMAGE AUGMENTATIONS ABLATION
Following (Chen et al. 2020). evaluate popular image augmentation techniques, namely random shifts, cutouts, vertical and horizontal flips, random rotations and imagewise intensity jittering. Below, provide comprehensive overview of each augmentation. Furthermore, examine effectiveness of these techniques in Figure
We bring our attention to random shifts that are commonly used to regularize neural networks trained on small images (Becker & Hinton| 1992, Simard et al., 20033 LeCun al. 1989 Ciresan et al., 2011 Ciregan et al. 2012). In our implementation of this method images of size 84 84 are padded each side by pixels (by repeating boundary pixels) and then randomly cropped back to the original 84 x 84 size.
Random Shift
Cutout Cutouts introduced in DeVries & Taylor (207)) represent generalization of Dropout (Hin- ton et al., 2012). Instead of masking individual pixels cutouts mask square regions. Since image pixels can be highly correlated, this technique is proven to improve training of neural networks.
Horizontal/Vertical Flip with probability 0.1.
This technique simply flips an image either horizontally or vertically
Rotate Here, image is rotated by degrees, where is uniformly sampled from [-5, 5).
17

---

Published as a conference paper at ICLR 2021
Intensity Each N C. 84 x 84 image tensor is multiplied by single scalar s, which is computed as $ + a clip(r,-2) 2), where r ~ (o, 1). For our experiments we use μt = 1.0 and ø = 0.1.
Ball In Cup Catch
Cartpole Swingup
Cheatah Run
1.00
870
600
Epicode Decuur
600
Eplioda dun-
Farcer
6000
refora
lone
vartea
Vettall 05
Hacce fls
Steps
Finger Spin
Easy
Walker
1000
1.003
-003
Exiorde Betum
500
400
600
600
200
mandy
63
6.1
0.5
0.0
0.2
6.3
64
3.5
02
0.6
0.8
Enstromment Steas Exlif)
Environennt Steps (xi1d))
Environment Steps Ix Ltf)
Figure 6: Various image augmentations have different effect on the agent's performance. Overall, we conclude that using image augmentations helps fight overfitting. Moreover, we notice that random shifts proven to be the most effective technique for tasks from the DeepMind control suite.
Implementation Finally, we provide Python-like implementation for the aforementioned augmen- tations powered by Kornia (Riba al. 2020).
import torch import torch.nn as nn import kernia. augmenttttonn as aug
random_ahts ft
Rnn. nn. nn. Se.Sunuettalll (nn. htw.ppeeeniall...e.elaaaaaaaaaaa (4) aug. RandomC (84, 84)))
cutout aug. (an.doonnrrsiinp-0.555)
h flip = agg.aannoollooi1ooottall1pppppp
v_flip = aug. agg.Banddom/eerr.caallllggp....
rotate =
aug. www.aandoonttaiionndeeeeeessssee
intensity = Intensity (scale==..)
class Intensity (nn Module); def Anit__see:: scale): super() ()_ern)t_t____ self.scale scale
def forward (self, x):
torch.randi ((), 1, 1, 1), device=x.devien noise = 1.0 + (self. r.clamp (-2.0, 2.0))
return x noise
18

---

Published as a conference paper at ICLR 2021
F AND M HYPER-PARAMETERE ABLATION
We further ablate the K,M hyper-parameters from Algorithm to understand their effect on perfor- mance. In Figure we observe that increase values of K,M improves the agent's performance. We choose to use the (K=2,M=2) parametrization as it strikes good balance between performance and computational demands.
Cup
Cartpole Swingup
1000
000
Eppondd Setum
800
Eplocds Rutuum
600
800
600
Apert
Org 210001
200
200
6.3
6.2
33
0.5
20
03
022
0.3
0.4
3.00
0.5
5.50
6.75
3.00
3.25
150
1.75
2.00
Environemnt Stepp (:10f)
Envirommntt Stapa (x109)
Emvtronment Staps (4
Finger Spii
Reacher Easy
Walker Walk
1230
1.000
1890
600
Ephodde treum
Epicadl Behum
470
Spri
400
Aprit
200
200
Org 22 -4
(r) X--La--
Envfronment
Environment
Envtronment
Figure 7: Increasing values of K.M hyper-parameters generally correlates positively with the agent's performance, especially on the harder tasks, such as Cheetah Run.
G ROBUSTNESS INVESTIGATION
To demonstrate the robustness our approach (Henderson et al. 2018), we perform a comprehensive study on the effect different hyper-parameter choices have on performance. A review of prior work (Hafner et al. 2018 2019 Lee et al. 2019; Srinnvas et al. 2020) shows consistent values for discount 0.01 parameters, but variability mini-batch sizes, learning rates. Since our method is based on SAC (Haarnoj et al. 2018). we also check whether the initial value of the temperature is important, as it plays a crucial role in the initial phase of exploration. We omit search over network architectures since Figure 1b shows our method to be robust to the exact choice. We thus focus on three hyper- parame-rrameters: mini-batch size, learning rate, and initial temperature.
y = 0.99 and target update rate
network architectures,
Due to computational demands, experiments are restricted to a subset of tasks from Tassa et al. (0088): Walker Walk, Cartpole Swingup. and Finger Spin. These were selected to be diverse, requiring different behaviors including locomotion and goal reaching. A grid search is performed over mini- batch sizes (128, 256, 512), learning rates (0.0001, 0.0005, 0.001, 0.005), ), and initial temperatures (0.005, 0.01, 0.05, 0.13. We follow the experimental setup from from//rpeentt Appenddix B except that only 3 seeds used due the computation limitations, but since variance low the results representative.
19

---

Published as a conference paper at ICLR 2021
Iuc
Cenonn 224
(a) Walker Walk.
Stea 2011 poalt
1el Ingenbbeed
(b) Cartpole Swingup.
Kape
Delon Stee
(c) Finger Spin.
Figure 8: A robustness study of our algorithm (DrQ) to changes in mini- batch size, learning rate, and initial temperature hyper-parameter on three different tasks from (Tassa et al.) 2018). Each row corresponds to different mini-batch size. The low variance of the curves and heat-maps shows DrQ to be generally robust to exact hyper-parameter settings.
20

---

Published as a conference paper at ICLR 2021
Figure 8 shows performance curves for each configuration well heat map the mean performance of the final evaluation episodes, similar to Mnih et al. (2016). Our method demonstrates good stability and is largely invariant the studied simplicity the experiments in SSection 5 juse the default learning rate of Adam (Kingma & Ba 2014) (0.001), even though it is not always optimal.
We emphasize that for
H IMPROVED DAFA-EFFICIENT REINFORCEMENT LEARNING FROM PIXELS
Our method allows to generate many various transformations from training observation due to the data augmentation strategy. Thus, we further investigate whether performing more training updates per an environment step can lead even better sample-efficiency. Following van Hassell et al. (019b) we compare single update with mini-batch of 512 transitions with 4 updates with 4 different mini-batches of size 128 samples each. Performing more updates per an environment step leads to even worse over-fitting on some tasks without data augmentation (see Figure 9a), while our method DrQ, that takes advantage of data augmentation, demonstrates improved sample-efficiency (see Figure 9b
Ball
Cartpola Swingup
Cheetah
1000
3200
502
902
Edccdde
-00
K00
202
Apmr
Pick with bedl
20
Envirenmnnt Staps (x10P)
Environment Stapa (x201)
Enstronment Steps (xtP))
(a) Unmodified SAC.
Ball In Cup Catch
Cartpole Swingup
Cheotah Run
1.00
Eroodd Rletim
600
600
800
200
6amd
23
60
20
300
20
60
30
100
20
60
90
107
Environment Stapa (c10P)
Environment Steps (x119)
Environment Staps (xCP")
(b) Our method DrQ.
Figure 9: In the data-efficient regime, where measure performance 100k environment steps, DrQ is able to enhance its efficiency by performing more training iterations per an environment step. This is because DrQ allows generate various transformations for training observation.
21


Reviews:
[{"review_id": "vtA0ftF8mQC", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 5, "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: To enable robust policy learning with image observations, the paper proposes a simple data augmentation technique that can be used with existing model-free reinforcement learning algorithms. It defines a notion of optimality invariant state transformation which preserves the Q function. An example of such transformations can be random image translations. It uses these transformations to (i) transform the input images, (ii) average the target Q values, and (iii) average the Q function themselves. Using this simple technique, they are able to get SOTA on DM control tasks and Atari 100k benchmark. On DM control tasks, it’s able to outperform SAC trained on state representations. Additionally, the paper provides ablation studies on different image transformations and robustness analysis with respect to hyperparameter settings.\n\nNovelty: While data augmentation techniques are common in computer vision, this was the first work (concurrently with RAD) to apply the technique in the context of reinforcement learning.\n\nReasons for score: Overall, I vote for accepting. The reasons are as follows. (i) It’s a simple technique which can be used with any RL algorithm to improve the performance of the algorithm (ii) Good and detailed evaluation  (iii) additional ablation studies and robustness analysis present\n\nQuestions: Will using different random image transformations in sequence help? (My hunch is it won’t as shown in RAD but still wanted your view given this method also does the average of Q function and target Q values)", "extra_scores": {}, "extra_text": {}}, {"review_id": "T6gJRTrzKGc", "reviewer": "AnonReviewer4", "summary": "", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "PROS\n-------------------------------------------\n- Finding that image augmentation helps in learning a good policy is indeed a substantial contributions. \n- Other two aspects to regularize Q-values are also helping to learn a better policy. \n\nCONS\n--------------------------------------------\nThe paper presents many experiments but there are a few crucial ones which are missing. \nFor example, \n- what is the impact on training time of DrQ as compared conventional DQN? \n- How much augmentation is good? etc.\n\nQuestion that needs justification:\n---------------------------------------------\n- Overall, the paper is well written and explain various things but the algo-1 needs explanation of all parameters. For example, \\math{D} (reader has to think that it's a replay buffer), \\math{u(T)} is not at all clear. Many notation are not explained even in text and needs a clear explanation for reader from RL and non-RL domain. \n\n- The primary claim of the paper is that image augmentation improves the performance. Sec5.1 shows significant improvement when image augmentation is used with different methods but it is very strange to see the improvement is just by adding 4 pixels on image boundary. Is there implication when we do more augmentation by increasing the size of random shift? How does the augmented image compare visually to the original image? The figure-6 in appendix shows all the results but which one is the random shift? \n\n- In figure4, the SAC state is significantly better than the DrQ and no explanation is provided. Why one should not used SAC compared to DrQ?\n\n- Paper claims that the proposed method can work with any model-free RL algorithm. Any justification or experiments to support the claim? If not, the contribution needs a re-writing. \n", "extra_scores": {}, "extra_text": {}}, {"review_id": "P3wlhlWEkat", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 5, "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:\n \nThis paper investigates data augmentation in the context of RL and proposes a novel augmentation algorithm to enabling robust learning directly from pixels without the need for auxiliary losses or pre-training. The authors propose to average both the Q function and its target over multiple image transformations. The experiments on DeepMind control suite and Atari 100k benchmark show that their method outperforms previous model-free, model-based and contrastive learning approaches.\n\n##########################################################################\n\nPros: \n \n1. This paper tackles a valuable problem of improving RL by data augmentation. It will have a broad impact on the area of both representation learning and reinforcement learning.  \n\n2. The idea of averaging both the Q function and its target over multiple image transformations is interesting and promising. This approach is easy to use and can be combined with any model-free RL algorithm.\n \n3. The paper is well written and the results section is well structured. They outperform baseline methods on two popular benchmarks and conduct ablation studies to verify the contribution of each component.\n \n##########################################################################\n\nCons: \n \n1. The proposed idea is very similar to RAD, a concurrent work by Laskin et al. The performance is also similar between these two approaches. More discussions and comparisons will help the readers better understand the difference.\n\n2. They claim data augmentation is all you need. To support this strong claim. I think the authors should conduct more experiments on more base algorithms. Can we get the same conclusion if we add the data augmentation techniques to other model-free RL algorithms (e.g., PPO or TD3) or a model-based RL algorithm (e.g., PlaNet or Dreamer)?\n\n##########################################################################\n\nTaken both pros and cons in to consideration, I vote for an acceptance because of the novelty of the proposed idea and large-scale comparisons to previous model-free, model-based and contrastive approaches. However, the experiments in the paper are not sufficient to support their claims “data augmentation is all you need” and I do suggest the authors to include more experiments to make it clear.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "5wvwZS1Zv37", "reviewer": "AnonReviewer3", "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 tackle the effectiveness of data augmentation in reinforcement learning. Authors have introduced a regularization technique, based on image shifts and Q-function augmentation,(DrQ) that significantly improves the performance of model-free RL algorithms trained directly from images.\n\nHere are two positive points in the paper:\n1. The authors have combined a model-free RL and regularized method, which shows a promising direction with a training strategy with augmented data.\n2. The authors have solid improve in continuous control by image pixel.\n\nThe negative point is that their DrQ method is very rough and very intuitive, authors can tackle more on the transformations of the images.  \n\nGenerally the paper is acceptable and they have shown a promising direction. ", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper describes a new data augmentation approach for image based RL.  The approach is both simple and effective.  It improves significantly the performance of several algorithms across a number of tasks.  The reviewers were unanimous about the benefits of the proposed technique.  This represents an important advance for RL.
Decision: Accept (Spotlight)