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: This paper proposes a regularization method for reinforcement learning that encourages the Q-value of the original image (i.e., original state) and the Q-value of the transformed image (new state) to be the same. This method enhances the robustness of RL methods against environment variation. This paper introduces the background and the motivation of the proposed method. Discussion and comparison of the difference between related works, such as SAC and DrQ, is also provided. Experiments show that the proposed method can improve the performance of image-based methods, even outperforming several state-based methods.