summaries: (Auxiliary Review only -- the authors need not respond to this).  This work proposes issues with data augmentation in RL and proposes a couple of techniques to address those issues. It applies augmentation only on the current state and not future state to address erroneous bootstrapping. It also proposes a Q objective that uses both the augmented and un-augmented data. Also, for actor critic experiments, the actor is optimised only with un-augmented data and learn a  policy that is general via sharing parameters between the networks. These are simple techniques that help the paper demonstrate good results in a range of tasks.  This is a clear paper with good set of experiments. The problem has been described clearly and the work proposes simple refinements to the data augmentation strategies used in RL. While I am convinced of the quality of the method itself and the experiments, this work could have benefitted from more rigour with the baselines to make a stronger case for the method. Hence, I am offering a score of 6 noting that the paper is marginally above acceptance threshold. This paper studies the problem of high-variance Q-targets under data augmentation, and proposes a method to reduce the variance of the targets called SVEA. The key idea contains three parts: (1) Apply augmentation to the Q-value estimation of the current state and not the next state Q-targets for bootstrapping, and (2) modify the objective to include both the targets using augmented and original observations and (3) optimize the actor using only the unaugmented data for actor-critic algorithms, with parameter sharing.   The paper performed experiments in DeepMind control suite and its variants (with raw pixel inputs). Compared to other data augmentation methods (e.g. DrQ) their proposed SVEA method is robust to the type of data augmentation and is more sample efficient with a better final performance. Some ablations were also presented among the above components to indicate the usefulness of all components. Generalization performance was demonstrated in the DistractingCS and robotic manipulation environments. The method also appears to help when using a Vision Transformers architecture.   This work proposes a simple technique to improve the robustness and sample-efficiency of reinforcement learning. The technique can be applied to any algorithm via a simple modification of the input pipeline to said algorithm's Q-learning subroutine. The claim is that this modification reduces the variance of Q-targets and therefore stabilizes learning. The paper addresses the question of zero-shot or fast generalization of value-based deep reinforcement learning agents using visual inputs to image transformation which do not change the (conceptual) state of the environment. The long term goal is therefore training agents which learn efficient, high return control policies both which are robust to irrelevant variations in high dimensional input observations. The ultimate goal seems to be generalization of learned skills to novel environments, but as far as I can tell this particular evaluation philosophy is implemented in experiments.  The paper reviews previous attempts to learn robust state representations to visual variations and formulates clear hypotheses for explaining the causes of difficulties which prior methods faced and were hindered by. The said hypotheses suggest a simplified approach (SVEA) which is empirically developed and evaluated in the remainder of the paper. Control tasks from widely used benchmarks are used to evaluate the method, somewhat less usual is the focus on vision-based high-dimensional versions of such tasks, instead of standard domains such as Atari games and benchmarks for generalization such as Natural Atari [1], CoinRun/ProcGen [2].   The paper presents a way of improving training stability of DrQ (an image-based RL algorithm that uses data augmentation) by applying image augmentation more carefully. The authors propose to input a mix of augmented and un-augmented observations into the critic, and stop augmenting the actor's inputs. On the considered benchmarks the proposed method demonstrates better generalization results. 
abstractText: While agents trained by Reinforcement Learning (RL) can solve increasingly challenging tasks directly from visual observations, generalizing learned skills to novel environments remains very challenging. Extensive use of data augmentation is a promising technique for improving generalization in RL, but it is often found to decrease sample efficiency and can even lead to divergence. In this paper, we investigate causes of instability when using data augmentation in common off-policy RL algorithms. We identify two problems, both rooted in high-variance Q-targets. Based on our findings, we propose a simple yet effective technique for stabilizing this class of algorithms under augmentation. We perform extensive empirical evaluation of image-based RL using both ConvNets and Vision Transformers (ViT) on a family of benchmarks based on DeepMind Control Suite, as well as in robotic manipulation tasks. Our method greatly improves stability and sample efficiency of ConvNets under augmentation, and achieves generalization results competitive with state-of-the-art methods for image-based RL in environments with unseen visuals. We further show that our method scales to RL with ViT-based architectures, and that data augmentation may be especially important in this setting.†
authors.name: Nicklas Hansen
authors.name: Hao Su
authors.name: Xiaolong Wang
references.authors: Rishabh Agarwal Marlos C. Machado P.S. Castro Marc G. Bellemare
references.title: Contrastive behavioral similarity embeddings for generalization in reinforcement
references.venue: learning. ArXiv,
references.authors: Richard Bellman
references.title: A markovian decision process
references.venue: Journal of Mathematics and Mechanics,
references.authors: Emily M. Bender Timnit Gebru Angelina McMillan-Major Shmargaret Shmitchell
references.title: On the dangers of stochastic parrots: Can language models be too big
references.venue: In Proceedings of the 2021 ACM Conference on Fairness, Accountability, and Transparency. Association for Computing Machinery,
references.authors: Christopher Berner Greg Brockman Brooke Chan Vicki Cheung Przemyslaw Debiak Christy Dennison David Farhi Quirin Fischer
references.title: Dota 2 with large scale deep reinforcement
references.venue: learning. ArXiv,
references.authors: T. Brown Benjamin Mann Nick Ryder Melanie Subbiah J. Kaplan Prafulla Dhariwal Arvind Neelakantan Pranav Shyam
references.title: Language models are few-shot
references.venue: learners. arXiv,
references.authors: Yevgen Chebotar A. Handa Viktor Makoviychuk M. Macklin J. Issac Nathan D. Ratliff D. Fox
references.title: Closing the sim-to-real loop: Adapting simulation randomization with real world experience
references.venue: International Conference on Robotics and Automation (ICRA),
references.authors: Ting Chen Simon Kornblith Mohammad Norouzi Geoffrey Hinton
references.title: A simple framework for contrastive learning of visual representations, 2020
references.authors: K. Cobbe Oleg Klimov Christopher Hesse Taehoon Kim John Schulman
references.title: Quantifying generalization in reinforcement learning
references.venue: In Icml,
references.authors: Karl Cobbe Oleg Klimov Chris Hesse Taehoon Kim John Schulman
references.title: Quantifying generalization in reinforcement learning, 2018
references.authors: A. Dosovitskiy Lucas Beyer Alexander Kolesnikov Dirk Weissenborn Xiaohua Zhai Thomas Unterthiner M. Dehghani Matthias Minderer
references.title: An image is worth 16x16 words: Transformers for image recognition
references.authors: Frederik Ebert Chelsea Finn Sudeep Dasari Annie Xie Alex Lee Sergey Levine
references.title: Visual foresight: Model-based deep reinforcement learning for vision-based robotic control, 2018
references.authors: Lasse Espeholt Hubert Soyer R. Munos K. Simonyan V. Mnih Tom Ward Yotam Doron Vlad Firoiu
references.title: Impala: Scalable distributed deep-rl with importance weighted actor-learner
references.venue: architectures. ArXiv,
references.authors: Jesse Farebrother Marlos C. Machado Michael H. Bowling
references.title: Generalization and regularization in dqn
references.venue: ArXiv,
references.authors: Scott Fujimoto H.V. Hoof D. Meger
references.title: Addressing function approximation error in actor-critic methods
references.venue: ArXiv,
references.authors: Spyros Gidaris Praveer Singh Nikos Komodakis
references.title: Unsupervised representation learning by predicting image rotations, 2018
references.authors: R. Givan T. Dean M. Greig
references.title: Equivalence notions and model minimization in markov decision processes
references.venue: Artificial Intelligence,
references.authors: Tuomas Haarnoja Aurick Zhou Pieter Abbeel Sergey Levine
references.title: Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor
references.venue: arXiv preprint arXiv:1801.01290,
references.authors: Danijar Hafner T. Lillicrap Ian S. Fischer R. Villegas David R Ha H. Lee James Davidson
references.title: Learning latent dynamics for planning
references.venue: pixels. ArXiv,
references.authors: Danijar Hafner T. Lillicrap Jimmy Ba Mohammad Norouzi
references.title: Dream to control: Learning behaviors by latent imagination
references.venue: ArXiv, abs/1912.01603,
references.authors: Nicklas Hansen Xiaolong Wang
references.title: Generalization in reinforcement learning by soft data augmentation
references.venue: In International Conference on Robotics and Automation,
references.authors: Nicklas Hansen Rishabh Jangir Yu Sun Guillem Alenyà Pieter Abbeel Alexei A. Efros Lerrel Pinto Xiaolong Wang
references.title: Self-supervised policy adaptation during deployment
references.venue: In International Conference on Learning Representations,
references.authors: H.V. Hasselt A. Guez Matteo Hessel V. Mnih D. Silver
references.title: Learning values across many orders of magnitude
references.venue: Nips,
references.authors: H.V. Hasselt A. Guez D. Silver
references.title: Deep reinforcement learning with double q-learning
references.venue: In Aaai,
references.authors: M. Hausknecht P. Stone
references.title: Deep recurrent q-learning for partially observable mdps
references.venue: In AAAI Fall Symposia,
references.authors: Kaiming He X. Zhang Shaoqing Ren Jian Sun
references.title: Deep residual learning for image recognition
references.venue: IEEE Conference on Computer Vision and Pattern Recognition (CVPR),
references.authors: Kaiming He Haoqi Fan Yuxin Wu Saining Xie Ross B. Girshick
references.title: Momentum contrast for unsupervised visual representation learning
references.venue: IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR),
references.authors: Matteo Hessel Joseph Modayil H.V. Hasselt Tom Schaul Georg Ostrovski Will Dabney Dan Horgan Bilal Piot
references.title: Rainbow: Combining improvements in deep reinforcement learning
references.venue: Aaai,
references.authors: Max Jaderberg Volodymyr Mnih Wojciech Marian Czarnecki Tom Schaul Joel Z Leibo David Silver Koray Kavukcuoglu
references.title: Reinforcement learning with unsupervised auxiliary tasks, 2016
references.authors: Leslie Pack Kaelbling Michael L. Littman Anthony R. Cassandra
references.title: Planning and acting in partially observable stochastic domains
references.venue: Artificial Intelligence,
references.authors: D. Kalashnikov A. Irpan P. Pastor J. Ibarz A. Herzog Eric Jang Deirdre Quillen Ethan Holly
references.title: Qtopt: Scalable deep reinforcement learning for vision-based robotic
references.venue: manipulation. ArXiv,
references.authors: Diederik P. Kingma Jimmy Ba
references.title: Adam: A method for stochastic optimization
references.venue: CoRR, abs/1412.6980,
references.authors: Ilya Kostrikov Denis Yarats Rob Fergus
references.title: Image augmentation is all you need: Regularizing deep reinforcement learning from pixels
references.venue: International Conference on Learning Representations,
references.authors: K.G. Larsen A. Skou
references.title: Bisimulation through probabilistic testing (preliminary report)
references.venue: In Proceedings of the 16th ACM SIGPLAN-SIGACT Symposium on Principles of Programming Languages,
references.authors: Michael Laskin Kimin Lee Adam Stooke Lerrel Pinto Pieter Abbeel Aravind Srinivas
references.title: Reinforcement learning with augmented data
references.venue: arXiv preprint arXiv:2004.14990,
references.authors: Kimin Lee Kibok Lee Jinwoo Shin Honglak Lee
references.title: A simple randomization technique for generalization in deep reinforcement
references.venue: learning. ArXiv,
references.authors: Sergey Levine Chelsea Finn Trevor Darrell Pieter Abbeel
references.title: End-to-end training of deep visuomotor policies
references.venue: The Journal of Machine Learning Research,
references.authors: T. Lillicrap J. Hunt A. Pritzel N. Heess T. Erez Y. Tassa D. Silver Daan Wierstra
references.title: Continuous control with deep reinforcement learning
references.authors: L.J. Lin
references.title: Self-improving reactive agents based on reinforcement learning, planning and teaching
references.venue: Machine Learning,
references.authors: Xingyu Lin Harjatin Singh Baweja George Kantor David Held
references.title: Adaptive auxiliary task weighting for reinforcement learning
references.venue: Advances in neural information processing systems,
references.authors: Clare Lyle Mark Rowland Georg Ostrovski Will Dabney
references.title: On the effect of auxiliary tasks on representation dynamics
references.venue: In Aistats,
references.authors: P. Mirowski Razvan Pascanu Fabio Viola Hubert Soyer Andy Ballard Andrea Banino Misha Denil R. Goroshin
references.title: Learning to navigate in complex
references.venue: environments. ArXiv,
references.authors: Volodymyr Mnih Koray Kavukcuoglu David Silver Alex Graves Ioannis Antonoglou Daan Wierstra Martin Riedmiller
references.title: Playing atari with deep reinforcement learning
references.venue: arXiv preprint arXiv:1312.5602,
references.authors: Ashvin V Nair Vitchyr Pong Murtaza Dalal Shikhar Bahl Steven Lin Sergey Levine
references.title: Visual reinforcement learning with imagined goals
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Mehdi Noroozi Paolo Favaro
references.title: Unsupervised learning of visual representations by solving jigsaw puzzles
references.venue: In European Conference on Computer Vision,
references.authors: Deepak Pathak Philipp Krahenbuhl Jeff Donahue Trevor Darrell Alexei A Efros
references.title: Context encoders: Feature learning by inpainting
references.venue: In Proceedings of the IEEE conference on computer vision and pattern recognition,
references.authors: Deepak Pathak Dhiraj Gandhi A. Gupta
references.title: Self-supervised exploration via disagreement
references.venue: ArXiv, abs/1906.04161,
references.authors: Xue Bin Peng Marcin Andrychowicz Wojciech Zaremba Pieter Abbeel
references.title: Sim-to-real transfer of robotic control with dynamics randomization
references.venue: IEEE International Conference on Robotics and Automation (ICRA),
references.authors: Lerrel Pinto Abhinav Gupta
references.title: Supersizing self-supervision: Learning to grasp from 50k tries and 700 robot hours
references.venue: In 2016 IEEE international conference on robotics and automation (ICRA),
references.authors: Lerrel Pinto Marcin Andrychowicz Peter Welinder Wojciech Zaremba Pieter Abbeel
references.title: Asymmetric actor critic for image-based robot learning
references.venue: arXiv preprint arXiv:1710.06542,
references.authors: Roberta Raileanu M. Goldstein Denis Yarats Ilya Kostrikov R. Fergus
references.title: Automatic data augmentation for generalization in deep reinforcement learning
references.authors: Fabio Ramos Rafael Possas Dieter Fox
references.title: Bayessim: Adaptive domain randomization via probabilistic inference for robotics simulators
references.venue: Robotics: Science and Systems
references.authors: Max Schwarzer Ankesh Anand Rishab Goel R Devon Hjelm Aaron Courville Philip Bachman
references.title: Data-efficient reinforcement learning with self-predictive representations
references.venue: arXiv preprint arXiv:2007.05929,
references.authors: Ramanan Sekar Oleh Rybkin Kostas Daniilidis Pieter Abbeel Danijar Hafner Deepak Pathak
references.title: Planning to explore via self-supervised world models, 2020
references.authors: Evan Shelhamer Parsa Mahmoudieh Max Argus Trevor Darrell
references.title: Loss is its own reward: Selfsupervision for reinforcement
references.venue: learning. ArXiv,
references.authors: Connor Shorten T. Khoshgoftaar
references.title: A survey on image data augmentation for deep learning
references.venue: Journal of Big Data,
references.authors: Xingyou Song Yiding Jiang Stephen Tu Yilun Du Behnam Neyshabur
references.title: Observational overfitting in reinforcement learning
references.venue: ArXiv, abs/1912.02975,
references.authors: Aravind Srinivas Michael Laskin Pieter Abbeel
references.title: Curl: Contrastive unsupervised representations for reinforcement learning
references.venue: arXiv preprint arXiv:2004.04136,
references.authors: Austin Stone Oscar Ramirez K. Konolige Rico Jonschkowski
references.title: The distracting control suite - a challenging benchmark for reinforcement learning
references.venue: pixels. ArXiv,
references.authors: Adam Stooke Kimin Lee Pieter Abbeel Michael Laskin
references.title: Decoupling representation learning from reinforcement learning
references.venue: ArXiv, abs/2004.1499,
references.authors: R. Sutton
references.title: Learning to predict by the methods of temporal differences
references.venue: Machine Learning,
references.authors: Richard S. Sutton Andrew G. Barto
references.title: Reinforcement Learning: An Introduction
references.authors: Yuval Tassa Yotam Doron Alistair Muldal Tom Erez Yazhe Li Diego de Las Casas David Budden Abbas Abdolmaleki
references.title: Deepmind control suite
references.venue: Technical report,
references.authors: Yonglong Tian Dilip Krishnan Phillip Isola
references.title: Contrastive multiview coding
references.venue: arXiv preprint arXiv:1906.05849,
references.authors: Josh Tobin Rachel Fong Alex Ray Jonas Schneider Wojciech Zaremba Pieter Abbeel
references.title: Domain randomization for transferring deep neural networks from simulation to the real world
references.venue: IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS),
references.authors: Emanuel Todorov Tom Erez Yuval Tassa
references.title: Mujoco: A physics engine for model-based control
references.venue: In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems,
references.authors: Aaron van den Oord Yazhe Li Oriol Vinyals
references.title: Representation learning with contrastive predictive coding, 2018
references.authors: Ashish Vaswani Noam M. Shazeer Niki Parmar Jakob Uszkoreit Llion Jones Aidan N. Gomez Lukasz Kaiser Illia Polosukhin
references.title: Attention is all you
references.venue: need. arXiv,
references.authors: Oriol Vinyals I. Babuschkin Wojciech Czarnecki Micha "el Mathieu Andrew Dudzik J. Chung D. Choi Richard Powell
references.title: Grandmaster level in starcraft ii using multi-agent reinforcement
references.venue: learning. Nature,
references.authors: K. Wang Bingyi Kang Jie Shao Jiashi Feng
references.title: Improving generalization in reinforcement learning with mixture regularization
references.venue: ArXiv, abs/2010.10814,
references.authors: Xudong Wang Long Lian Stella X. Yu
references.title: Unsupervised visual attention and invariance for reinforcement learning
references.venue: ArXiv, abs/2104.02921,
references.authors: Ziyu Wang Tom Schaul Matteo Hessel H.V. Hasselt Marc Lanctot N.D. Freitas
references.title: Dueling network architectures for deep reinforcement
references.venue: learning. ArXiv,
references.authors: Zhirong Wu Yuanjun Xiong Stella X Yu Dahua Lin
references.title: Unsupervised feature learning via nonparametric instance discrimination
references.venue: In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition,
references.authors: Zhenlin Xu Deyi Liu Junlin Yang M. Niethammer
references.title: Robust and generalizable visual representation learning via random convolutions
references.venue: ArXiv, abs/2007.13003,
references.authors: Denis Yarats Amy Zhang Ilya Kostrikov Brandon Amos Joelle Pineau Rob Fergus
references.title: Improving sample efficiency in model-free reinforcement learning from images, 2019
references.authors: Denis Yarats Rob Fergus Alessandro Lazaric Lerrel Pinto
references.title: Reinforcement learning with prototypical representations
references.venue: arXiv preprint arXiv:2102.11271,
references.authors: Sarah Young Dhiraj Gandhi Shubham Tulsiani Abhinav Gupta Pieter Abbeel Lerrel Pinto
references.title: Visual imitation made easy
references.venue: CoRL,
references.authors: Tianhe Yu Saurabh Kumar A. Gupta Sergey Levine Karol Hausman Chelsea Finn
references.title: Gradient surgery for multi-task learning
references.venue: ArXiv, abs/2001.06782,
references.authors: A. Zhang Rowan McAllister R. Calandra Y. Gal Sergey Levine
references.title: Learning invariant representations for reinforcement learning without reconstruction
references.authors: C. Zhang Oriol Vinyals R. Munos S. Bengio
references.title: A study on overfitting in deep reinforcement
references.venue: learning. ArXiv,
references.authors: Richard Zhang Phillip Isola Alexei A Efros
references.title: Colorful image colorization
references.venue: In European conference on computer vision,
references.authors: Yuke Zhu Roozbeh Mottaghi Eric Kolve Joseph J Lim Abhinav Gupta Li Fei-Fei Ali Farhadi
references.title: Target-driven visual navigation in indoor scenes using deep reinforcement learning
references.venue: In Icra,
references.authors: Brian D Ziebart Andrew Maas J Andrew Bagnell Anind K Dey
references.title: Maximum entropy inverse reinforcement learning
references.venue: In Proceedings of the 23rd National Conference on Artificial Intelligence,
sections.heading: 1 Introduction
sections.text: Reinforcement Learning (RL) from visual observations has achieved tremendous success in various applications such as video-games [43, 4, 70], robotic manipulation [37], and autonomous navigation [42, 83]. However, it is still very challenging for current methods to generalize the learned skills to novel environments, and policies trained by RL can easily overfit to the training environment [81, 13], especially for high-dimensional observation spaces such as images [8, 58].
  Increasing the variability in training data via domain randomization [66, 50] and data augmentation [57, 35, 33, 51] has demonstrated encouraging results for learning policies invariant to changes in environment observations. Specifically, recent works on data augmentation [35, 33] both show improvements in sample efficiency from simple cropping and translation augmentations, but the studies also conclude that additional data augmentation in fact decrease sample efficiency and even cause divergence. While these augmentations have the potential to improve generalization, the increasingly varied data makes the optimization more challenging and risks instability. Unlike supervised learning, balancing the trade-off between stability and generalization in RL requires substantial trial and error.
  In this paper, we illuminate causes of instability when applying data augmentation to common off-policy RL algorithms [43, 38, 15, 18]. Based on our findings, we provide an intuitive method for stabilizing this class of algorithms under use of strong data augmentation. Specifically, we find two main causes of instability in previous work’s application of data augmentation: (i) indiscriminate application of data augmentation resulting in high-variance Q-targets; and (ii) that Q-value estimation strictly from augmented data results in over-regularization.
  †Website and code is available at: https://nicklashansen.github.io/SVEA.
  35th Conference on Neural Information Processing Systems (NeurIPS 2021)
  To address these problems, we propose SVEA: Stabilized Q-Value Estimation under Augmentation, a simple yet effective framework for data augmentation in off-policy RL that greatly improves stability of Q-value estimation. Our method consists of the following three components: Firstly, by only applying augmentation in Q-value estimation of the current state, without augmenting Q-targets used for bootstrapping, SVEA circumvents erroneous bootstrapping caused by data augmentation; Secondly, we formulate a modified Q-objective that optimizes Q-value estimation jointly over both augmented and unaugmented copies of the observations; Lastly, for SVEA implemented with an actorcritic algorithm, we optimize the actor strictly on unaugmented data, and instead learn a generalizable policy indirectly through parameter-sharing. Our framework can be implemented efficiently without additional forward passes nor introducing additional learnable parameters.
  We perform extensive empirical evaluation on the DeepMind Control Suite [64] and extensions of it, including the DMControl Generalization Benchmark [21] and the Distracting Control Suite [60], as well as a set of robotic manipulation tasks. Our method greatly improve Q-value estimation with ConvNets under a set of strong data augmentations, and achieves sample efficiency, asymptotic performance, and generalization that is competitive or better than previous state-of-the-art methods in all tasks considered, at a lower computational cost. Finally, we show that our method scales to RL with Vision Transformers (ViT) [10]. We find that ViT-based architectures are especially prone to overfitting, and data augmentation may therefore be a key component for large-scale RL.
sections.heading: 2 Related Work
sections.text: Representation Learning. Learning visual invariances using data augmentation and self-supervised objectives has proven highly successful in computer vision [46, 45, 82, 74, 68, 65, 75, 27, 7]. For example, Chen et al. [7] perform an extensive study on data augmentation (e.g. random cropping and image distortions) for contrastive learning, and show that representations pre-trained with such transformations transfer effectively to downstream tasks. While our work also uses data augmentation for learning visual invariances, we leverage the Q-objective of deep Q-learning algorithms instead of auxiliary representation learning tasks.
  Visual Learning for RL. Numerous methods have been proposed with the goal of improving sample efficiency [29, 56, 68, 76, 40, 59, 61, 54, 77] of image-based RL. Recently, using self-supervision to improve generalization in RL has also gained interest [80, 47, 55, 1, 22, 21, 72]. Notably, Zhang et al. [80] and Agarwal et al. [1] propose to learn behavioral similarity embeddings via auxiliary tasks (bisimulation metrics and contrastive learning, respectively), and Hansen et al. [21] learn visual invariances through an auxiliary prediction task. While these results are encouraging, it has also been shown in [29, 40, 22, 79, 41] that the best choice of auxiliary tasks depends on the particular RL task, and that joint optimization with sub-optimally chosen tasks can lead to gradient interference. We achieve competitive sample-efficiency and generalization results without the need for carefully chosen auxiliary tasks, and our method is therefore applicable to a larger variety of RL tasks.
  Data Augmentation and Randomization for RL. Our work is directly inspired by previous work on generalization in RL by domain randomization [66, 50, 48, 52, 6] and data augmentation [36, 9, 71, 35, 33, 51, 61, 21]. For example, Tobin et al. [66] show that a neural network trained for object localization in a simulation with randomized visual augmentations improves real world generalization. Similarly, Lee et al.[36] show that application of a random convolutional layer to observations during training improve generalization in 3D navigation tasks. More recently, extensive studies on data augmentation [35, 33] have been conducted with RL, and conclude that, while small random crops and translations can improve sample efficiency, most data augmentations decrease sample efficiency and cause divergence. We illuminate main causes of instability, and propose a framework for data augmentation in deep Q-learning algorithms that drastically improves stability and generalization.
  Improving Deep Q-Learning. While deep Q-learning algorithms such as Deep Q-Networks (DQN) [43] have achieved impressive results in image-based RL, the temporal difference objective is known to have inherent instabilities when used in conjunction with function approximation and off-policy data [63]. Therefore, a variety of algorithmic improvements have been proposed to improve convergence [24, 73, 25, 23, 53, 38, 15, 14, 28]. For example, Hasselt et al. [24] reduce overestimation of Q-values by decomposing the target Q-value estimation into action selection and action evaluation using separate networks. Lillicrap et al. [38] reduce target variance by defining the target Q-network as a slow-moving average of the online Q-network. Our method also improves Q-value estimation, but we specifically address the instability of deep Q-learning algorithms on augmented data.
sections.heading: 3 Preliminaries
sections.text: Problem formulation. We formulate the interaction between environment and policy as a Markov Decision Process (MDP) [2] M = hS,A,P, r, i, where S is the state space, A is the action space, P : S ⇥ A 7! S is the state transition function that defines a conditional probability distribution P (·|st,at) over all possible next states given a state st 2 S and action at 2 A taken at time t, r : S⇥A 7! R is a reward function, and 2 [0, 1) is the discount factor. Because image observations only offer partial state observability [30], we define a state st as a sequence of k + 1 consecutive frames (ot,ot 1, . . . ,ot k), o 2 O, where O is the high-dimensional image space, as proposed in Mnih et al. [43]. The goal is then to learn a policy ⇡ : S 7! A that maximizes discounted return Rt = E ⇠⇡[ P T t=1 t r(st,at)] along a trajectory = (s0, s1, . . . , sT ) obtained by following policy ⇡ from an initial state s0 2 S to a state sT with state transitions sampled from P , and ⇡ is parameterized by a collection of learnable parameters ✓. For clarity, we hereon generically denote parameterization with subscript, e.g. ⇡✓. We further aim to learn parameters ✓ s.t. ⇡✓ generalizes well (i.e., obtains high discounted return) to unseen MDPs, which is generally unfeasible without further assumptions about the structure of the space of MDPs. In this work, we focus on generalization to MDPs M = hS,A,P, r, i, where states st 2 S are constructed from observations ot 2 O, O ✓ O of a perturbed observation space O (e.g. unseen visuals), and M ⇠M for a space of MDPs M. Deep Q-Learning. Common model-free off-policy RL algorithms aim to estimate an optimal stateaction value function Q⇤ : S ⇥ A 7! R as Q✓(s,a) ⇡ Q⇤(s,a) = max⇡✓ E [Rt|st = s,at = a] using function approximation. In practice, this is achieved by means of the single-step Bellman residual ⇣ r(st,at) + maxa0t Q tgt (st+1,a0t) ⌘ Q✓(st,at) [62], where parameterizes a target state-action value function Qtgt. We can choose to minimize this residual (also known as the temporal difference error) directly wrt ✓ using a mean squared error loss, which gives us the objective
  LQ(✓, ) = Est,at,st+1⇠B
  " 1
  2
  ✓ r(st,at) + max
  a0t Q
  tgt (st+1,a 0 t ) ◆ Q✓(st,at) 2# , (1)
  where B is a replay buffer with transitions collected by a behavioral policy [39]. From here, we can derive a greedy policy directly by selecting actions at = argmaxat Q✓(st,at). While Qtgt = Q and periodically setting ✓ exactly recovers the objective of DQN [43], several improvements have been proposed to improve stability of Eq. 1, such as Double Q-learning [24], Dueling Q-networks [73], updating target parameters using a slow-moving average of the online Q-network [38]:
  n+1 (1 ⇣) n + ⇣✓n (2) for an iteration step n and a momentum coefficient ⇣ 2 (0, 1], and others [25, 23, 53, 14, 28]. As computing maxa0t Q tgt (st+1,a0t) in Eq. 1 is intractable for large and continuous action spaces, a number of prominent actor-critic algorithms that additionally learn a policy ⇡✓(st) ⇡ argmaxat Q✓(st,at) have therefore been proposed [38, 15, 18].
  Soft Actor-Critic (SAC) [18] is an off-policy actor-critic algorithm that learns a state-action value function Q✓ and a stochastic policy ⇡✓ (and optionally a temperature parameter), where Q✓ is optimized using a variant of the objective in Eq. 1 and ⇡✓ is optimized using a -discounted maximum-entropy objective [84]. To improve stability, SAC is also commonly implemented using Double Q-learning and the slow-moving target parameters from Eq. 2. We will in the remainder of this work describe our method in the context of a generic off-policy RL algorithm that learns a parameterized state-action value function Q✓, while we in our experiments discussed in Section 6 evaluate of our method using SAC as base algorithm.
  4 Pitfalls of Data Augmentation in Deep Q-Learning In this section, we aim to illuminate the main causes of instability from naïve application of data augmentation in Q-value estimation. Our goal is to learn a Q-function Q✓ for an MDP M that generalizes to novel MDPs M ⇠M with unseen visuals, and we leverage data augmentation as an optimality-invariant state transformation ⌧ to induce a bisimulation relation [34, 17] between a state s and its transformed (augmented) counterpart saug = ⌧(s, ⌫) with parameters ⌫ ⇠ V . Definition 1 (Optimality-Invariant State Transformation [33]). Given an MDP M, a state transformation ⌧ : S ⇥ V 7! S is an optimality-invariant state transformation if Q(s,a) = Q(⌧(s, ⌫),a) 8 s 2 S, a 2 A, ⌫ 2 V , where ⌫ 2 V parameterizes the transformation ⌧ .
  Following our definitions of M,M from Section 3, we can further extend the concept of optimalityinvariant transformations to MDPs, noting that a change of state space (e.g. perturbed visuals) itself can be described as a transformation ⌧ : S ⇥ V 7! S with unknown parameters ⌫ 2 V . If we choose the set of parameters V of a state transformation ⌧ to be sufficiently large such that it intersects with V with high probability, we can therefore expect to improve generalization to state and observation spaces not seen during training. However, while naïve application of data augmentation as in previous work [35, 33, 61, 54] may potentially improve generalization, it can be harmful to Q-value estimation. We hypothesize that this is primarily because it dramatically increases the size of the observed state space, and consequently also increases variance Var [Q(⌧(s, ⌫))] Var [Q(s)] , ⌫ ⇠ V when V is large. Concretely, we identify the following two issues:
  Pitfall 1: Non-deterministic Q-target. For deep Q-learning algorithms, previous work [35, 33, 61, 54] applies augmentation to both state saug
  t , ⌧(st, ⌫) and successor state saugt+1 , ⌧(st+1, ⌫0) where
  ⌫, ⌫ 0 ⇠ V . Compared with DQN [43] that uses a deterministic (more precisely, periodically updated) Q-target, this practice introduces a non-deterministic Q-target r(st,at) + maxa0t Q tgt (saug t+1,a 0 t ) depending on the augmentation parameters ⌫0. As observed in the original DQN paper, high-variance target values are detrimental to Q-learning algorithms, and may cause divergence due to the “deadly triad” of function approximation, bootstrapping, and off-policy learning [63]. This motivates the work to introduce a slowly changing target network, and several other works have refined the Q-target update rule [38, 15] to further reduce volatility. However, because data augmentation is inherently non-deterministic, it can greatly increase variance in Q-target estimation and exacerbates the issue of volatility, as shown in Figure 1 (top). This is particularly troubling in actor-critic algorithms such as DDPG [38] and SAC [18], where the Q-target is estimated from (st+1,a0), a0 ⇠ ⇡(·|st+1), which introduces an additional source of error from ⇡ that is non-negligible especially when st+1 is augmented.
  Pitfall 2: Over-regularization. Data augmentation was originally introduced in the supervised learning regime as a regularizer to prevent overfitting of high-capacity models. However, for RL, even learning a policy in the training environment is hard. While data augmentation may improve generalization, it greatly increases the difficulty of policy learning, i.e., optimizing ✓ for Q✓ and potentially a behavior network ⇡✓. Particularly, when the temporal difference loss from Eq. 1 cannot be well minimized, the large amount of augmented states dominate the gradient, which significantly impacts Q-value estimation of both augmented and unaugmented states. We refer to this issue as over-regularization by data augmentation. Figure 1 (bottom) shows the mean difference in Q-predictions made with augmented vs. unaugmented data in fully converged DrQ [33] agents trained with shift augmentation. Augmentations such as affine-jitter, random convolution, and random overlay incur large differences in estimated Q-values. While such difference can be reduced by regularizing the optimization with each individual augmentation, we emphasize that even the minimal shift augmentation used throughout training incurs non-zero difference. Since is commonly chosen to be a
  moving average of ✓ as in Eq. 2, such differences caused by over-regularization affect Q✓ and Q tgt equally, and optimization may therefore still diverge depending on the choice of data augmentation. As such, there is an inherent trade-off between accurate Q-value estimation and generalization when using data augmentation. In the following section, we address these pitfalls.
sections.heading: 5 Method
sections.text: We propose SVEA: Stabilized Q-Value Estimation under Augmentation, a general framework for visual generalization in RL by use of data augmentation. SVEA applies data augmentation in a novel learning framework leveraging two data streams – with and without augmented data, respectively. Our method is compatible with any standard off-policy RL algorithm without changes to the underlying neural network that parameterizes the policy, and it requires no additional forward passes, auxiliary tasks, nor learnable parameters. While SVEA in principle does not make any assumptions about the structure of states st 2 S , we here describe our method in the context of image-based RL.
sections.heading: 5.1 Architectural Overview
sections.text: An overview of the SVEA architecture is provided in Figure 2. Our method leverages properties of common neural network architectures used in off-policy RL without introducing additional learnable parameters. We subdivide the neural network layers and corresponding learnable parameters of a state-action value function into sub-networks f✓ (denoted the state encoder) and Q✓ (denoted the Q-function) s.t qt , Q✓(f✓(st),at) is the predicted Q-value corresponding to a given stateaction pair (st,at). We similarly define the target state-action value function s.t. q tgt t
  , r(st,at) + maxa0t Q tgt (f tgt (st+1),a0) is the target Q-value for (st,at), and we define parameters as an exponential moving average of ✓ as in Eq. 2. Depending on the choice of underlying algorithm, we may choose to additionally learn a parameterized policy ⇡✓ that shares encoder parameters with Q✓ and selects actions at ⇠ ⇡✓(·|f✓(st)).
  To circumvent erroneous bootstrapping from augmented data (as discussed in Section 4), we strictly apply data augmentation in Q-value estimation of the current state st, without applying data augmentation to the successor state st+1 used in Eq. 1 for bootstrapping with Q tgt
  (and ⇡✓ if applicable), which addresses Pitfall 1. If ⇡✓ is learned (i.e., SVEA is implemented with an actor-critic algorithm), we also optimize it strictly from unaugmented data. To mitigate over-regularization in optimization of f✓ and Q✓ (Pitfall 2), we further employ a modified Q-objective that leverages both augmented and unaugmented data, which we introduce in the following section.
sections.heading: 5.2 Learning Objective
sections.text: Our method redefines the temporal difference objective from Eq. 1 to better leverage data augmentation. First, recall that qtgt
  t = r(st,at) + maxa0t Q
  tgt (f tgt (st+1),a0). Instead of learning to predict
  q tgt t
  only from state st, we propose to minimize a nonnegative linear combination of LQ over two individual data streams, st and s aug t = ⌧(st, ⌫), ⌫ ⇠ V , which we define as the objective
  L SVEA Q (✓, ) , ↵LQ st, q tgt t ; ✓,
  + LQ saug t , q tgt t ; ✓,
  (3)
  = Est,at,st+1⇠B h ↵ Q✓(f✓(st),at) qtgtt 2 2 +
  Q✓(f✓(saugt ),at) q tgt t 2 2
  i , (4)
  where ↵, are constant coefficients that balance the ratio of the unaugmented and augmented data streams, respectively, and qtgt
  t is computed strictly from unaugmented data. LSVEA Q (✓, ) serves as a
  data-mixing strategy that oversamples unaugmented data as an implicit variance reduction technique.
  As we will verify empirically in Section 6, data-mixing is a simple and effective technique for variance reduction that works well in tandem with our proposed modifications to bootstrapping. For ↵ = , the objective in Eq. 4 can be evaluated in a single, batched forward-pass by rewriting it as:
  gt = [st, ⌧(st, ⌫)]N (5) ht = ⇥ q tgt t , q tgt t ⇤ N (6)
  L SVEA Q (✓, ) = Est,at,st+1⇠B, ⌫⇠V h (↵+ ) kQ✓(f✓(gt),at) htk 2 2 i , (7)
  where [·]N is a concatenation operator along the batch dimension N for st, s aug t 2 RN⇥C⇥H⇥W and qtgt t 2 RN⇥1, which is illustrated as ⌦ in Figure 2. Empirically, we find ↵ = 0.5, = 0.5 to be both effective and practical to implement, which we adopt in the majority of our experiments. However, more sophisticated schemes for selecting ↵, and/or varying them as training progresses could be interesting directions for future research. If the base algorithm of choice learns a policy ⇡✓, its objective L⇡(✓) is optimized solely on unaugmented states st without changes to the objective, and a stop-grad operation is applied after f✓ to prevent non-stationary gradients of L⇡(✓) from interfering with Q-value estimation, i.e., only the objective from Eq. 4 or optionally Eq. 7 updates f✓ using stochastic gradient descent. As described in Section 5.1, parameters are updated using an exponential moving average of ✓ and a stop-grad operation is therefore similarly applied after Qtgt .
  We summarize our method for ↵ = applied to a generic off-policy algorithm in Algorithm 1.
  Algorithm 1 Generic SVEA off-policy algorithm (I naïve augmentation, I our modifications) ✓, ✓⇡, : randomly initialized network parameters, ✓ B Initialize to be equal to ✓ ⌘, ⇣: learning rate and momentum coefficient ↵, : loss coefficients, default: (↵ = 0.5, = 0.5)
  1: for timestep t = 1...T do act: 2: at ⇠ ⇡✓ (·|f✓(st)) B Sample action from policy 3: s0
  t ⇠ P(·|st,at) B Sample transition from environment
  4: B B [ (st,at, r(st,at), s0t) B Add transition to replay buffer update: 5: {si,ai, r(si,ai), s0i | i = 1...N} ⇠ B B Sample batch of transitions 6: si = ⌧(si, ⌫i), s0i = ⌧(s 0 i , ⌫ 0 i ), ⌫i, ⌫0i ⇠ V I Naïve application of data augmentation 7: for transition i = 1..N do 8: ✓⇡ ✓⇡ ⌘r✓⇡L⇡ (si; ✓⇡) (if applicable) B Optimize ⇡✓ with SGD 9: qtgt
  i = r(si,ai) + maxa0i Q
  tgt (f tgt (s0 i ),a0 i ) B Compute Q-target
  10: saug i = ⌧(si, ⌫i), ⌫i ⇠ V I Apply stochastic data augmentation 11: gi = ⇥ si, s aug i ⇤ N , hi = ⇥ q tgt i , q tgt i ⇤ N I Pack data streams 12: ✓ ✓ ⌘r✓LSVEAQ (gi, hi; ✓, ) I Optimize f✓ and Q✓ with SGD 13: (1 ⇣) + ⇣✓ B Update using EMA of ✓
  6 Experiments
  We evaluate both sample efficiency, asymptotic performance, and generalization of our method and a set of strong baselines using both ConvNets and Vision Transformers (ViT) [10] in tasks from DeepMind Control Suite (DMControl) [64] as well as a set of robotic manipulation tasks. DMControl offers challenging and diverse continuous control tasks and is widely used as a benchmark for image-based
  RL [19, 20, 76, 59, 35, 33]. To evaluate generalization of our method and baselines, we test methods under challenging distribution shifts (as illustrated in Figure 3) from the DMControl Generalization Benchmark (DMControl-GB) [21], the Distracting Control Suite (DistractingCS) [60], as well as distribution shifts unique to the robotic manipulation environment. Code is available at https://github.com/nicklashansen/dmcontrol-generalization-benchmark.
  Setup. We implement our method and baselines using SAC [18] as base algorithm, and we apply random shift augmentation to all methods by default. This makes our base algorithm equivalent to DrQ [33] when K=1,M=1; we refer to the base algorithm as unaugmented and consider stability under additional data augmentation. We use the same network architecture and hyperparameters for all methods (whenever applicable), and adopt the setup from Hansen and Wang [21]. Observations are stacks of 3 RGB frames of size 84⇥ 84⇥ 3 (and 96⇥ 96⇥ 3 in ViT experiments). In the DMControlGB and DistractingCS benchmarks, all methods are trained for 500k frames and evaluated on all 5 tasks from DMControl-GB used in prior work, and we adopt the same experimental setup for robotic manipulation. See Appendix H for hyperparameters and further details on our experimental setup.
  Baselines and data augmentations. We benchmark our method against the following strong baselines: (1) CURL [59], a contrastive learning method for RL; (2) RAD that applies a random crop; (3) DrQ that applies a random shift; (4) PAD [22] that adapts to test environments using self-supervision; (5) SODA [21] that applies data augmentation in auxiliary learning; as well as a number of ablations. We compare to the K=1,M=1 setting of DrQ by default, but also provide comparison to varying K,M . We experiment with a diverse set of data augmentations proposed in previous work on RL and computer vision, namely random shift [33], random convolution (denoted conv) [36], random overlay [21], random cutout [9], Gaussian blur, random affine-jitter, and random rotation [35, 16]. We provide samples for all data augmentations in Appendix C and test environments in Appendix E.
sections.heading: 6.1 Stability and Generalization on DMControl
sections.text: Stability. We evaluate the stability of SVEA and DrQ under 6 common data augmentations; results are shown in Figure 4. While the sample efficiency of DrQ degrades substantially for most augmentations, SVEA is relatively unaffected by the choice of data augmentation and improves sample efficiency in 27 out of 30 instances. While the sample efficiency of DrQ can be improved by increasing its K,M parameters, we find that DrQ requires approx. 6x wall-time to match the sample efficiency of SVEA; see Figure 6 (left). We further ablate each component of SVEA and report both training and test curves in Figure 5; we find that both components are key to SVEA’s success. Because we empirically find the conv augmentation to be particularly difficult to optimize, we provide additional stability experiments in Section 6.2 and 6.3 using this augmentation. See Appendix A for additional ablations.
  Generalization. We compare the test performance of SVEA to 5 recent state-of-the-art methods for image-based RL on the color_hard and video_easy benchmarks from DMControl-GB (results in Table 1), as well as the extremely challenging DistractingCS benchmark, where camera pose, background, and colors are continually changing throughout an episode (results in Figure 6 (right)). We here use conv and overlay augmentations for fair comparison to SODA, and we report additional results on the video_hard benchmark in Appendix F. SVEA outperforms all methods considered in 12 out of 15 instances on DMControl-GB, and at a lower computational cost than CURL, PAD, and SODA that all learn auxiliary tasks. On DistractingCS, we observe that SVEA improves generalization by 42% at low intensity, and its generalization degrades significantly slower than DrQ for high intensities. While generalization depends on the particular choice of data augmentation and test environments, this is an encouraging result considering that SVEA enables efficient policy learning with stronger augmentations than previous methods.
  6.2 RL with Vision Transformers
  Vision Transformers (ViT) [10] have recently achieved impressive results on downstream tasks in computer vision. We replace all convolutional layers from the previous experiments with a 4-layer ViT encoder that operates on raw pixels in 8⇥ 8 space-time patches, and evaluate our method using data augmentation in conjunction with ViT encoders. Importantly, we design the ViT architecture such that it roughly matches our CNN encoder in terms of learnable parameters. The ViT encoder is trained from scratch using RL, and we use the same experimental setup as in our ConvNet experiments. In particular, it is worth emphasizing that both our ViT and CNN encoders are trained using Adam [32] as optimizer and without weight decay. See Figure 7 (top) for an architectural overview, and refer to Appendix H for additional implementation details.
  Our training and test results are shown in Figure 7 (bottom). We are, to the best of our knowledge, the first to successfully solve image-based RL tasks without CNNs. We observe that DrQ overfits significantly to the training environment compared to its CNN counterpart (94 test return on color_hard for DrQ with ViT vs. 569 with a ConvNet on the Walker, walk task). SVEA achieves comparable sample efficiency and improves generalization by 706% and 233% on Walker, walk and Cartpole, swingup, respectively, over DrQ, while DrQ + conv remains unstable. Interestingly, we observe that our ViT-based implementation of SVEA achieves a mean episode return of 877 on the color_hard benchmark of the challenging Walker, walk task (vs. 760 using CNNs). SVEA might therefore be a promising technique for fu-
  ture research on RL with CNN-free architectures, where data augmentation appears to be especially important for generalization. We provide additional experiments with ViT encoders in Section 6.3 and make further comparison to ConvNet encoders in Appendix A.
  6.3 Robotic Manipulation
  We additionally consider a set of goal-conditioned robotic manipulation tasks using a simulated Kinova Gen3 arm: (i) reach, a task in which the robot needs to position its gripper above a goal indicated by a red mark; (ii) reach moving target, a task similar to (i) but where the robot needs to follow a red mark moving continuously in a zig-zag pattern at a random velocity; and (iii) push, a task in which the robot needs to push a cube to a red mark. The initial configuration of gripper, object, and goal is randomized, the agent uses 2D positional control, and policies are trained using dense rewards. Observations are stacks of RGB frames with no access to state information. Training and test environments are shown in Figure 8. See Appendix G for further details and environment samples.
  Results are shown in Figure 9 and Figure 10. For both CNN and ViT encoders, SVEA trained with conv augmentation has similar sample efficiency and training performance as DrQ trained without augmentation, while DrQ + conv exhibits poor sample efficiency and fails to solve the push task. Generalization results are shown in Table 2. We find that naïve application of data augmentation has a higher success rate in test environments than DrQ, despite being less
  successful in the training environment, which we conjecture is because it is optimized only from augmented data. Conversely, SVEA achieves high success rates during both training and testing.
  Conclusion. SVEA is found to greatly improve both stability and sample efficiency under augmentation, while achieving competitive generalization results. Our experiments indicate that our method scales to ViT-based architectures, and it may therefore be a promising technique for large-scale RL experiments where data augmentation is expected to play an increasingly important role.
  Broader Impact. While our contribution aims to reduce computational cost of image-based RL, we remain concerned about the growing ecological and economical footprint of deep learning – and RL in particular – with increasingly large models such as ViT; see Appendix J for further discussion.
  Acknowledgments and Funding Transparency. This work was supported by grants from DARPA LwLL, NSF CCF-2112665 (TILOS), NSF 1730158 CI-New: Cognitive Hardware and Software Ecosystem Community Infrastructure (CHASE-CI), NSF ACI-1541349 CC*DNI Pacific Research Platform, NSF grant IIS-1763278, NSF CCF-2112665 (TILOS), as well as gifts from Qualcomm, TuSimple and Picsart.
title: Stabilizing Deep Q-Learning with ConvNets and Vision Transformers under Data Augmentation