Mastering Visual Continuous Control: Improved Data-Augmented Reinforcement Learning

We present DrQ-v2 , a model-free reinforcement learning ( RL ) algorithm for visual continuous control . DrQ-v2 builds on DrQ , an off-policy actor-critic approach that uses data augmentation to learn directly from pixels . We introduce several improvements that yield state-of-the-art results on the DeepMind Control Suite . Notably , DrQ-v2 is able to solve complex humanoid locomotion tasks directly from pixel observations , previously unattained by model-free RL . DrQ-v2 is conceptually simple , easy to implement , and provides significantly better computational footprint compared to prior work , with the majority of tasks taking just 8 hours to train on a single GPU . Finally , DrQ-v2 ’ s implementation is publicly released to provide RL practitioners with a strong and computationally efficient baseline . 1 INTRODUCTION . Creating sample-efficient continuous control methods that observe high-dimensional images has been a long standing challenge in reinforcement learning ( RL ) . Over the last three years , the RL community has made significant headway on this problem , improving sample-efficiency significantly . The key insight to solving visual control is the learning of better low-dimensional representations , either through autoencoders ( Yarats et al. , 2019 ; Finn et al. , 2015 ) , variational inference ( Hafner et al. , 2018 ; 2019 ; Lee et al. , 2019 ) , contrastive learning ( Srinivas et al. , 2020 ; Yarats et al. , 2021a ) , self-prediction ( Schwarzer et al. , 2020b ) , or data augmentations ( Yarats et al. , 2021b ; Laskin et al. , 2020 ) . However , current state-of-the-art model-free methods are still limited in three ways . First , they are unable to solve the more challenging visual control problems such as quadruped and humanoid locomotion . Second , they often require significant computational resources , i.e . lengthy training times using distributed multi-GPU infrastructure . Lastly , it is often unclear how different design choices affect overall system performance . In this paper we present DrQ-v2 , a simple model-free algorithm that builds on the idea of using data augmentations ( Yarats et al. , 2021b ; Laskin et al. , 2020 ) to solve hard visual control problems . Most notably , it is the first model-free method that solves complex humanoid tasks directly from pixels . Compared to previous state-of-the-art model-free methods , DrQ-v2 provides significant improvements in sample efficiency across tasks from the DeepMind Control Suite ( Tassa et al. , 2018 ) . Conceptually simple , DrQ-v2 is also computationally efficient , which allows solving most tasks in DeepMind Control Suite in just 8 hours on a single GPU ( see Figure 1 ) . Recently , a model-based method , DreamerV2 ( Hafner et al. , 2020 ) was also shown to solve visual continuous control problems and it was first to solve the humanoid locomotion problem from pixels . While our model-free DrQ-v2 matches DreamerV2 in terms sample efficiency and performance , it does so 4⇥ faster in terms of wall-clock time to train . We believe this makes DrQ-v2 a more accessible approach to support research in visual continuous control and it reinforces the question on whether model-free or model-based is the more suitable approach to solve this type of tasks . DrQ-v2 , which is detailed in Section 3 , improves upon DrQ ( Yarats et al. , 2021b ) by making several algorithmic changes : ( i ) switching the base RL algorithm from SAC ( Haarnoja et al. , 2018a ) to DDPG ( Lillicrap et al. , 2015a ) with clipped double Q-learning from TD3 ( Fujimoto et al. , 2018 ) , ( ii ) this allows us straightforwardly incorporating multi-step return , ( iii ) adding bilinear interpolation to the random shift image augmentation , ( iv ) introducing an exploration schedule , ( v ) selecting better hyper-parameters including a larger capacity of the replay buffer . A careful ablation study of these design choices is presented in Section 4.4 . Furthermore , we re-examine the original implementation of DrQ and identify several computational bottlenecks such as replay buffer management , data augmentation processing , batch size , and frequency of learning updates ( see Section 3.2 ) . To remedy these , we have developed a new implementation that both achieves better performance and trains around 3.5 times faster with respect to wall-clock time than the previous implementation on the same hardware with an increase in environment frame throughput ( FPS ) from 28 to 96 ( i.e. , it takes 106/96/3600 ⇡ 2.9 hours to train for 1M environment steps ) . DrQ-v2 ’ s implementation is available at https : //anonymous.4open.science/r/drqv2 . 2 BACKGROUND . 2.1 REINFORCEMENT LEARNING FROM IMAGES . We formulate image-based control as an infinite-horizon Markov Decision Process ( MDP ) ( Bellman , 1957 ) . Generally , in such a setting , an image rendering of the system is not sufficient to perfectly describe the system ’ s underlying state . To this end and per common practice ( Mnih et al. , 2013 ) , we approximate the current state of the system by stacking three consecutive prior observations . With this in mind , such MDP can be described as a tuple ( X , A , P , R , , d0 ) , where X is the state space ( a three-stack of image observations ) , A is the action space , P : X ⇥A ! ( X ) is the transition function1 that defines a probability distribution over the next state given the current state and action , R : X ⇥ A ! [ 0 , 1 ] is the reward function , 2 [ 0 , 1 ) is a discount factor , and d0 2 ( X ) is the distribution of the initial state x0 . The goal is to find a policy ⇡ : X ! ( A ) that maximizes the expected discounted sum of rewards E⇡ [ P1 t=0 trt ] , where x0 ⇠ d0 , and 8t we have at ⇠ ⇡ ( ·|xt ) , xt+1 ⇠ P ( ·|xt , at ) , and rt = R ( xt , at ) . 2.2 DEEP DETERMINISTIC POLICY GRADIENT . Deep Deterministic Policy Gradient ( DDPG ) ( Lillicrap et al. , 2015a ) is an actor-critic algorithm for continuous control that concurrently learns a Q-function Q✓ and a deterministic policy ⇡ . For this , DDPG uses Q-learning ( Watkins and Dayan , 1992 ) to learn Q✓ by minimizing the one-step Bellman residual J✓ ( D ) = E ( xt , at , rt , xt+1 ) ⇠D [ ( Q✓ ( xt , at ) rt Q✓̄ ( xt+1 , ⇡ ( xt+1 ) ) 2 ] . The policy ⇡ is learned by employing Deterministic Policy Gradient ( DPG ) ( Silver et al. , 2014 ) and maximizing J ( D ) = Ext⇠D [ Q✓ ( xt , ⇡ ( xt ) ) ] , so ⇡ ( xt ) approximates argmaxaQ✓ ( xt , a ) . Here , D is a replay buffer of environment transitions and ✓̄ is an exponential moving average of the weights . DDPG is amenable to incorporate n-step returns ( Watkins , 1989 ; eng and Williams , 1996 ) when estimating TD error beyond a single step ( Barth-Maron et al. , 2018 ) . In practice , n-step returns allow for faster 1Here , ( X ) denotes a distribution over the state space X . reward propagation and has been previously used in policy gradient and Q-learning methods ( Mnih et al. , 2016b ; Barth-Maron et al. , 2018 ; Hessel et al. , 2017 ) . 2.3 DATA AUGMENTATION IN REINFORCEMENT LEARNING . Recently , it has been shown that data augmentation techniques , commonplace in Computer Vision , are also important for achieving the state-of-the-art performance in image-based RL ( Yarats et al. , 2021b ; Laskin et al. , 2020 ) . For example , the state-of-the-art algorithm for visual RL , DrQ ( Yarats et al. , 2021b ) builds on top of Soft Actor-Critic ( Haarnoja et al. , 2018a ) , a model-free actor-critic algorithm , by adding a convolutional encoder and data augmentation in the form of random shifts . The use of such data augmentations now forms an essential component of several recent visual RL algorithms ( Srinivas et al. , 2020 ; Raileanu et al. , 2020 ; Yarats et al. , 2021a ; Stooke et al. , 2020 ; Hansen and Wang , 2021 ; Schwarzer et al. , 2020b ) . 3 DRQ-V2 : IMPROVED DATA-AUGMENTED REINFORCEMENT LEARNING . In this section , we describe DrQ-v2 , a simple model-free actor-critic RL algorithm for image-based continuous control , that builds upon DrQ . 3.1 ALGORITHMIC DETAILS . Image Augmentation As in DrQ we apply random shifts image augmentation to pixel observations of the environment . In the settings of visual continuous control by DMC , this augmentation can be instantiated by first padding each side of 84 ⇥ 84 observation rendering by 4 pixels ( by repeating boundary pixels ) , and then selecting a random 84⇥ 84 crop , yielding the original image shifted by ±4 pixels . We also find it useful to apply bilinear interpolation on top of the shifted image ( i.e , we replace each pixel value with the average of the four nearest pixel values ) . In our experiments , this modification provides an additional performance boost across the board . Image Encoder The augmented image observation is then embedded into a low-dimensional latent vector by applying a convolutional encoder . We use the same encoder architecture as in DrQ , which first was introduced introduced in SAC-AE ( Yarats et al. , 2019 ) . This process can be succinctly summarized as h = f⇠ ( aug ( x ) ) , where f⇠ is the encoder , aug is the random shifts augmentation , and x is the original image observation . Actor-Critic Algorithm We use DDPG ( Lillicrap et al. , 2015a ) as a backbone actor-critic RL algorithm and , similarly to Barth-Maron et al . ( 2018 ) , augment it with n-step returns to estimate TD error . This results into faster reward propagation and overall learning progress ( Mnih et al. , 2016a ) . While some methods ( Hafner et al. , 2020 ) employ more sophisticated techniques such as TD ( ) or Retrace ( ) ( Munos et al. , 2016 ) , they are often computationally demanding when n is large . We find that using simple n-step returns , without an importance sampling correction , strikes a good balance between performance and efficiency . We also employ clipped double Q-learning ( Fujimoto et al. , 2018 ) to reduce overestimation bias in the target value . Practically , this requires training two Qfunctions Q✓1 and Q✓2 . For this , we sample a mini-batch of transitions ⌧ = ( xt , at , rt : t+n 1 , xt+n ) from the replay buffer D and compute the following two losses : L✓k , ⇠ ( D ) = E⌧⇠D ⇥ ( Q✓k ( ht , at ) y ) 2 ⇤ 8k 2 { 1 , 2 } , ( 1 ) with the TD target y defined as : y = n 1X i=0 irt+i + n min k=1,2 Q✓̄k ( ht+n , at+n ) , where ht = f⇠ ( aug ( xt ) ) , ht+n = f⇠ ( aug ( xt+n ) ) , at+n = ⇡ ( ht+n ) + ✏ , ✓̄1 and ✓̄2 are the slowmoving weights for the Q target networks . We note , that in contrast to DrQ , we do not employ a target network for the encoder f⇠ and always use the most recent weights ⇠ to embed xt and xt+n . The exploration noise ✏ is sampled from clip ( N ( 0 , 2 ) , c , c ) similar to TD3 ( Fujimoto et al. , 2018 ) , with the exception of decaying , which we describe below . Finally , we train the deterministic actor ⇡ using DPG with the following loss : L ( D ) = Ext⇠D ⇥ min k=1,2 Q✓k ( ht , at ) ⇤ , ( 2 ) where ht = f⇠ ( aug ( xt ) ) , at = ⇡ ( ht ) + ✏ , and ✏ ⇠ clip ( N ( 0 , 2 ) , c , c ) . Similar to DrQ , we do not use actor ’ s gradients to update the encoder ’ s parameters ⇠ . Scheduled Exploration Noise Empirically , we observe that it is helpful to have different levels of exploration at different stages of learning . At the beginning of training we want the agent to be more stochastic and explore the environment more effectively , while at the later stages of training , when the agent has already identified promising behaviors , it is better to be more deterministic and master those behaviors . Similar to Amos et al . ( 2020 ) , we instantiate this idea by using linear decay ( t ) for the variance 2 of the exploration noise defined as : ( t ) = init + ( 1 min ( t T , 1 ) ) ( final init ) , ( 3 ) where init and final are the initial and final values for standard deviation , and T is the decay horizon . Key Hyper-Parameters We conduct an extensive hyper-parameter search and identify several hyper-parameter changes compared to DrQ . The three most important hyper-parameters are : ( i ) the size of the replay buffer , ( ii ) mini-batch size , and ( iii ) learning rate . Specifically , we use a 10 times larger replay buffer than DrQ . We also use a smaller mini-batch size of 256 without any noticeable performance degradation . This is in contrast to CURL ( Srinivas et al. , 2020 ) and DrQ ( Yarats et al. , 2021b ) that both use a larger batch size of 512 to attain more stable training in the expense of computational efficiency . Finally , we find that using smaller learning rate of 1⇥ 10 4 , rather than DrQ ’ s learning rate of 1⇥ 10 3 , results into more stable training without any loss in learning speed .

Summary: This paper introduces DrQ-v2, an improvement over DrQ, a model free off-policy actor-critic approach which uses SAC+data augmentation to learn directly from pixels. DrQ-v2, on the other hand, switches SAC with DDPG and proposes a series of algorithmic, hyperparameter choices and implementation improvements. These improvements make it both faster computationally and better empirically as it is able to solve the Humanoid environment from pixel.