rQdia: Regularizing Q-Value Distributions With Image Augmentation

1 INTRODUCTION . Human perception is invariant to and remarkably robust against many perturbations , like discoloration , obfuscation , and low exposure . On the other hand , artificial neural networks do not intrinsically carry these invariance properties , not without regularizers or hand-crafted inductive biases like convolution , kernel rotation , dilation , attention , and recurrence . In deep reinforcement learning ( RL ) from pixels , an agent must learn to visually interpret a scene in order to decide on an action . Thus , recent approaches in RL have turned to the self-supervision and data augmentation techniques found in computer vision . Indeed , such contrastive learning auxiliary losses ( Srinivas et al. , 2020 ) or data augmentation regularizers ( Yarats et al. , 2021b ) have afforded greater sample efficiency and final scores in both the DeepMind Continuous Control Suite ( Tassa et al. , 2018 ) from pixels and Atari Arcade Learning Environments ( Bellemare et al. , 2013 ) . Nevertheless , pixel-based approaches continue to lag behind models that learn directly from state embeddings , not just in terms of sample efficiency , but also in longer-term asymptotic performance . For example , the recent DrQ ( Yarats et al. , 2021b ) , an image augmentation-based regularizer added to SAC ( Haarnoja et al. , 2018 ) , reports falling 14.5 % short of its state embedding-based SAC counterpart on the Cheetah Run task . Such discrepancies indicate that visual representations are not yet up to par with state embeddings , at least not for locomotive continuous control . State embeddings have many properties that facilitate generalization , such as location invariance , and to a degree , invariance between morphological relations . If a subset of the dimensions of a state embedding always indicates feet position , then the relation “ one foot in front of the other ” will be represented by those dimensions invariant to the placement of the robot ’ s arms , head , or other body parts . Parametric visual encodings are not guaranteed to learn such invariances with respect to the robot ’ s morphology . What other signals in deep RL can guide visual representation learning toward more invariant encodings ? We propose Q-value distributions , sets of cumulative discounted rewards , as such a signal . For state s , Q-function Q , and actions a ( 0 ) , . . . , a ( m ) ∼ D ( A ) from some statistical distribution D over action space A , we define Q-value distribution simply as : Q ( s , a ( 0 ) ) , . . . , Q ( s , a ( m ) ) . Since Q-values are optimally proportional to action probabilities , this “ distribution ” is representative of the actual policy distribution when D = π . Furthermore , it is a measure of the current and future value of each action for that state . We review the MDP framework in Section 2.1 . This signal is amenable to many of the same invariances afforded by state embeddings , if not additional ones . For example , if the optimal action is “ put one foot in front of the other , ” then the Q-value distribution reflects this action ’ s relation to other actions regardless of ( invariant to ) where the agent is located . While recent work shows individual Q-values benefit from regularizing across augmented images ( Yarats et al. , 2021b ) , we consider that Q-value distributions are also important , as they contain information not just about one action in isolation , but multiple actions in relation to one another . 2 BACKGROUND . 2.1 DEEP RL FROM PIXELS . A Markov Decision Process ( MDP ) consists of an action a ∈ IRda , state s ∈ IRds , and reward r ∈ IR . “ From pixels ” assumes that state s is an image frame or multiple image frames stacked together . The action is sampled from a policy at any t time step at ∼ π ( st ) . Taking such actions yields a trajectory τ = ( s0 , a0 , s1 , a1 , ... , sT ) via the dynamics model st+1 ∼ f ( st , at ) of the environment and its rewards rt+1 = R ( st , at ) . The goal is to maximize cumulative discounted reward R ( τ ) = ∑T t=0 rtγ t where γ is the temporal discount factor . The optimal action for a state a∗ ( s ) = argmaxaQ ∗ ( s , a ) thus depends on the state-action value function , also known as the Qvalue , Qπ ( s , a ) = Eτ∼π [ R ( τ ) |s0 = s , a0 = a ] . 2.2 SOFT ACTOR-CRITIC & DRQ . Soft-Actor Critic ( SAC ) ( Haarnoja et al. , 2018 ) is an RL algorithm which learns a min-reduced ensemble of Q-value functions Qφ ( s , a ) = mini=1,2Qφi ( s , a ) optimized with one-step Bellman error , and a stochastic Gaussian policy πθ ( s , a ) optimized by maximizing Qφ ( s , a ) and entropy H ( s ) = −πθ ( s ) log ( πθ ( s ) ) , made differentiable via the reparameterization trick . To further encourage exploration and avoid premature policy collapse , entropy H ( s ) is also added as part of the agent ’ s reward . In visual domains , Qφ and πθ are typically equipped with a shared convolutional neural net encoder . DrQ ( Yarats et al. , 2021b ) sets the Bellman target as the average of the augmented and non-augmented next-state targets , and minimizes Bellman error for both Qφ ( s , a ) and Qφ ( aug ( s ) , a ) where aug is the random augmentation . 2.3 DATA-EFFICIENT RAINBOW ( DER ) . Rainbow ( Hessel et al. , 2018 ) maps directly to Q-value estimates for discrete actions . Compared to vanilla DQN ( Mnih et al. , 2013 ) , several refinements are used : Q-values are sampled from a multivariate Gaussian probabilistically ( Dabney et al. , 2018 ) , noise is injected into network parameters ( Plappert et al. , 2017 ) , double Q networks ( Van Hasselt et al. , 2016 ) , dueling DQNs ( Wang et al. , 2016 ) , n-step returns ( Watkins , 1989 ) , and mini-batches are sampled from a prioritized experience replay ( Schaul et al. , 2015 ) . “ Data-efficient ” refers to the Atari sample limit of 100k environmental interactions , a much more challenging setting for notoriously inefficient RL . 3 RELATED WORK . 3.1 DATA-EFFICIENT RL . Data-efficient RL is a paramount concern to the practicality of RL in real-world use cases . Image augmentation has proven an extremely effective regularizer for improving the sample efficiency of model-free off-policy RL algorithms ( Yarats et al. , 2019a ; Srinivas et al. , 2020 ; Sermanet et al. , 2018 ; Dwibedi et al. , 2018 ) , so much so that basic augmentation techniques suffice to rival or surpass model-based RL algorithms ( Hafner et al. , 2019c ; Lee et al. , 2019b ; Hafner et al. , 2019b ) in the sample efficiency metric . Curiously , these methods have become progressively simpler . CURL ( Srinivas et al. , 2020 ) employed contrastive learning , using positive and negative samples extracted from the mini-batch , requiring a computation of quadratic complexity w.r.t . the mini-batch size . RAD ( Laskin et al. , 2020 ) and DrQ ( Yarats et al. , 2021b ) showed that simpler methods work just as well or better , either by augmenting images naively , or augmenting and averaging their Q-values , respectively . Even more recently ( contemporaneously ) , the as-yet unpublished DrQv2 ( Yarats et al. , 2021a ) found basic augmentation alone suffices under a DDPG-based algorithm , with significant efficiency improvements over prior methods despite the exceptional simplicity . While these recent methods have traded CURL ’ s mini-batch statistics for mere augmentation , rQdia marks the first combination of the two that uses mini-batch statistics to enforce consistency across Q-value distributions , in a manner both simple and complementary to the above implements . 3.2 MINI-BATCH REGULARIZATION IN RL . Mini-batch-based regularization has not commonly been used in RL . For example , Batch Norm ( Ioffe & Szegedy , 2015 ) , a common regularizer in computer vision , is not as notably employed in RL . This type of regularization remained unexploited in RL until CURL ( Srinivas et al. , 2020 ) showed that contrastive learning via image augmentation ( Chen et al. , 2020 ; He et al. , 2020 ; Misra & Maaten , 2020 ; Henaff , 2020 ) greatly improved RL data efficiency . CURL contrasts a state ’ s “ positive ” augmentation sample with the rest of the mini-batch ’ s “ negative ” augmentation samples . These negative samples can be thought of as sampled from a Uniform distribution over the agent ’ s experience replay , inspiring rQdia . However , contrastive learning en- forces a non-guaranteed ground truth , disassociating negative samples regardless of actual similarity . See Figure 2 for an example of this negative samples problem , where similar states are contrasted to have dissimilar encodings as an inadvertent byproduct of the uniform randomness . rQdia bypasses this flaw by only enforcing a guaranteed constraint : that the Q-value for any sampled action , regardless of statistical distribution , be consistent across the same states invariant to augmentation . This indeed should always be the case , thus yielding gains over CURL while remaining complementary to methods like DrQ and DrQv2 . 3.3 IMAGE AUGMENTATION . Image augmentation is commonly used to counter over-fitting in computer vision . Techniques include color shift , affinity translation , scale , etc . ( Ciregan et al. , 2012 ; Cireşan et al. , 2011 ; Simard et al. , 2003 ; Krizhevsky et al. , 2012 ; Chen et al. , 2020 ) . In Yarats et al . ( 2021b ) , the authors investigated several common image transformations and concluded that random shifts strike a good balance between simplicity and performance for the MuJoCo environments . A variety of different augmentations are useful for different games in the Procgen benchmark ( Raileanu et al. , 2020 ) , and Yarats et al . ( 2021b ) used Intensity variation for the Atari environments . These techniques have proven critical to MuJoCo from pixels . 4 METHODS . rQdia is the first RL Q-value regularizer that does not necessarily depend on either the on-policy or the off-policy states and actions . n , m such states and actions are instead sampled from arbitrary distributions , let ’ s call them D1 and D2 , of state space S and action space A , respectively : s ( 0 ) , ... , s ( n−1 ) ∼ D1 ( S ) a ( 0 ) , ... , a ( m−1 ) ∼ D2 ( A ) . ( 1 ) Then , the following constraint is enforced for Qφ ( · ) , the neural network that models Qπ ( · ) the Q-value function : Qφ ( s ( i ) , a ( j ) ) = Qφ ( aug ( s ( i ) ) , a ( j ) ) ∀i , j , ( 2 ) where aug ( · ) is an arbitrary augmentation transform . aug ( · ) and D1 , D2 could vary . For aug ( · ) , in line with Yarats et al . ( 2021b ) , we use translation for MuJoCo and intensity jittering for Atari . Specifically , to translate , we pad by 4 pixels , then crop randomly inward by 4 pixels ; to intensity jitter , each image is multiplied by some random noise s = 1.1×clip ( r , −2 , 2 ) , r ∼ N ( 0 , 1 ) . For D1 , D2 in MuJoCo continuous action spaces , we implement a simple approach analogous to CURL ’ s negative-sample sampling , except sampling both states and actions rather than just states . That is , states and actions are directly lifted from the mini-batchB ( n = m = |B| ) , in effect sampled from a Uniform distribution over the agent ’ s experience replay . By using historical states and actions as opposed to random noise , we compute Q-value distributions over state-action pairs that could more feasibly be encountered in a deployed roll-out . In discrete Atari , actions may be lifted from the action space directly , that is , a ( 0 ) , ... , a ( m−1 ) = A . Thus , given mini-batch states s ( 0 ) , ... , s ( n−1 ) and mini-batch ( or action space ) actions a ( 0 ) , ... , a ( m−1 ) , the following auxiliary loss is proposed to constitute a basic implement of rQdia : LrQdia = 1 nm ∑ i < n , j < m ( Qφ ( s ( i ) , a ( j ) ) −Qφ ( aug ( s ( i ) ) , a ( j ) ) ) 2 . ( 3 ) Then this auxiliary loss is simply added to the RL agent ’ s standard loss term . Voila , rQdia ( visualized “ in a nutshell ” in Figure 1 ) . This is applied in parallel for each s ( i ) , a ( j ) pair . If mini-batches or actions spaces are very large , it is possible to convolve a smaller subset of n states and m actions . In MuJoCo , we note that while D1 , D2 are treated as Uniform distributions over an agent ’ s history similar to CURL ’ s negative sampling , D1 , D2 could be more sophisticated . For example , the probability of sampling an action could be proportional to state similarities . Or , like MPO ( Abdolmaleki et al. , 2018 ) , actions could be sampled directly from the policy itself . Algorithm 1 provides pseudo-code for rQdia in continuous control algorithms like SAC-AE and DrQ . All code for rQdia will be released open-source .

Summary: Image augmentations have recently become a standard component of deep RL algorithms.  Previous work has enforced consistencies at a sample-level.  This paper proposes to look at the distribution of statistics at a minibatch-level in order to enforce consistencies.  Paper shows results on standard benchmarks in discrete and continuous control (Atari and DMC).