Return-Based Contrastive Representation Learning for Reinforcement  Learning

1 INTRODUCTION . Deep reinforcement learning ( RL ) algorithms can learn representations from high-dimensional inputs , as well as learn policies based on such representations to maximize long-term returns simultaneously . However , deep RL algorithms typically require large numbers of samples , which can be quite expensive to obtain ( Mnih et al. , 2015 ) . In contrast , it is usually much more sample efficient to learn policies with learned representations/extracted features ( Srinivas et al. , 2020 ) . To this end , various auxiliary tasks have been proposed to accelerate representation learning in aid of the main RL task ( Suddarth and Kergosien , 1990 ; Sutton et al. , 2011 ; Gelada et al. , 2019 ; Bellemare et al. , 2019 ; François-Lavet et al. , 2019 ; Shen et al. , 2020 ; Zhang et al. , 2020 ; Dabney et al. , 2020 ; Srinivas et al. , 2020 ) . Representative examples of auxiliary tasks include predicting the future in either the pixel space or the latent space with reconstruction-based losses ( e.g. , Jaderberg et al. , 2016 ; Hafner et al. , 2019a ; b ) . Recently , contrastive learning has been introduced to construct auxiliary tasks and achieves better performance compared to reconstruction based methods in accelerating RL algorithms ( Oord et al. , 2018 ; Srinivas et al. , 2020 ) . Without the need to reconstruct inputs such as raw pixels , contrastive learning based methods can ignore irrelevant features such as static background in games and learn more compact representations . Oord et al . ( 2018 ) propose a contrastive representation learning method based on the temporal structure of state sequence . Srinivas et al . ( 2020 ) propose to leverage the prior knowledge from computer vision , learning representations that are invariant to image augmentation . However , existing works mainly construct contrastive auxiliary losses in an unsupervised manner , without considering feedback signals in RL problems as supervision . In this paper , we take a further step to leverage the return feedback to design a contrastive auxiliary loss to accelerate RL algorithms . Specifically , we propose a novel method , called Return-based ∗This work is conducted at Microsoft Research Asia . The first two authors contributed equally to this work . Contrastive representation learning for Reinforcement Learning ( RCRL ) . In our method , given an anchor state-action pair , we choose a state-action pair with the same or similar return as the positive sample , and a state-action pair with different return as the negative sample . Then , we train a discriminator to classify between positive and negative samples given the anchor based on their representations as the auxiliary task . The intuition here is to learn state-action representations that capture return-relevant features while ignoring return-irrelevant features . From a theoretical perspective , RCRL is supported by a novel state-action abstraction , called Zπirrelevance . Zπ-irrelevance abstraction aggregates state-action pairs with similar return distributions under certain policy π . We show that Zπ-irrelevance abstraction can reduce the size of the stateaction space ( cf . Appendix A ) as well as approximate the Q values arbitrarily accurately ( cf . Section 4.1 ) . We further propose a method called Z-learning that can calculate Zπ-irrelevance abstraction with sampled returns rather than the return distribution , which is hardly available in practice . Zlearning can learn Zπ-irrelevance abstraction provably efficiently . Our algorithm RCRL can be seen as the empirical version of Z-learning by making a few approximations such as integrating with deep RL algorithms , and collecting positive pairs within a consecutive segment in a trajectory of the anchors . We conduct experiments on Atari games ( Bellemare et al. , 2013 ) and DeepMind Control suite ( Tassa et al. , 2018 ) in low data regime . The experiment results show that our auxiliary task combined with Rainbow ( Hessel et al. , 2017 ) for discrete control tasks or SAC ( Haarnoja et al. , 2018 ) for continuous control tasks achieves superior performance over other state-of-the-art baselines for this regime . Our method can be further combined with existing unsupervised contrastive learning methods to achieve even better performance . We also perform a detailed analysis on how the representation changes during training with/without our auxiliary loss . We find that a good embedding network assigns similar/dissimilar representations to state-action pairs with similar/dissimilar return distributions , and our algorithm can boost such generalization and speed up training . Our contributions are summarized as follows : • We introduce a novel contrastive loss based on return , to learn return-relevant representations and speed up deep RL algorithms . • We theoretically build the connection between the contrastive loss and a new form of stateaction abstraction , which can reduce the size of the state-action space as well as approximate the Q values arbitrarily accurately . • Our algorithm achieves superior performance against strong baselines in Atari games and DeepMind Control suite in low data regime . Besides , the performance can be further enhanced when combined with existing auxiliary tasks . 2 RELATED WORK . 2.1 AUXILIARY TASK . In reinforcement learning , the auxiliary task can be used for both the model-based setting and the model-free setting . In the model-based settings , world models can be used as auxiliary tasks and lead to better performance , such as CRAR ( François-Lavet et al. , 2019 ) , Dreamer ( Hafner et al. , 2019a ) , and PlaNet ( Hafner et al. , 2019b ) . Due to the complex components ( e.g. , the latent transition or reward module ) in the world model , such methods are empirically unstable to train and relies on different regularizations to converge . In the model-free settings , many algorithms construct various auxiliary tasks to improve performance , such as predicting the future ( Jaderberg et al. , 2016 ; Shelhamer et al. , 2016 ; Guo et al. , 2020 ; Lee et al. , 2020 ; Mazoure et al. , 2020 ) , learning value functions with different rewards or under different policies ( Veeriah et al. , 2019 ; Schaul et al. , 2015 ; Borsa et al. , 2018 ; Bellemare et al. , 2019 ; Dabney et al. , 2020 ) , learning from many-goals ( Veeriah et al. , 2018 ) , or the combination of different auxiliary objectives ( de Bruin et al. , 2018 ) . Moreover , auxiliary tasks can be designed based on the prior knowledge about the environment ( Mirowski et al. , 2016 ; Shen et al. , 2020 ; van der Pol et al. , 2020 ) or the raw state representation ( Srinivas et al. , 2020 ) . Hessel et al . ( 2019 ) also apply auxiliary task to the multi-task RL setting . Contrastive learning has seen dramatic progress recently , and been introduced to learn state representation ( Oord et al. , 2018 ; Sermanet et al. , 2018 ; Dwibedi et al. , 2018 ; Aytar et al. , 2018 ; Anand et al. , 2019 ; Srinivas et al. , 2020 ) . Temporal structure ( Sermanet et al. , 2018 ; Aytar et al. , 2018 ) and local spatial structure ( Anand et al. , 2019 ) has been leveraged for state representation learning via contrastive losses . CPC ( Oord et al. , 2018 ) and CURL ( Srinivas et al. , 2020 ) adopt a contrastive auxiliary tasks to accelerate representation learning and speed up main RL tasks , by leveraging the temporal structure and image augmentation respectively . To the best of our knowledge , we are the first to leverage return to construct a contrastive auxiliary task for speeding up the main RL task . 2.2 ABSTRACTION . State abstraction ( or state aggregation ) aggregates states by ignoring irrelevant state information . By reducing state space , state abstraction can enable efficient policy learning . Different types of abstraction are proposed in literature , ranging from fine-grained to coarse-grained abstraction , each reducing state space to a different extent . Bisimulation or model irrelevance ( Dean and Givan , 1997 ; Givan et al. , 2003 ) define state abstraction under which both transition and reward function are kept invariant . By contrast , other types of state abstraction that are coarser than bisimulation such as Qπ irrelevance orQ∗ irrelevance ( see e.g. , Li et al. , 2006 ) , which keep the Q function invariant under any policy π or the optimal policy respectively . There are also some works on state-action abstractions , e.g. , MDP homomorphism ( Ravindran , 2003 ; Ravindran and Barto , 2004a ) and approximate MDP homomorphism ( Ravindran and Barto , 2004b ; Taylor et al. , 2009 ) , which are similar to bisimulation in keeping reward and transition invariant , but extending bisimulation from state abstraction to stateaction abstraction . In this paper , we consider a new form of state-action abstraction Zπ-irrelevance , which aggregates state-action pairs with the same return distribution and is coarser than bisimulation or homomorphism which are frequently used as auxiliary tasks ( e.g. , Biza and Platt , 2018 ; Gelada et al. , 2019 ; Zhang et al. , 2020 ) . However , it is worth noting that Zπ-irrelevance is only used to build the theoretical foundation of our algorithm , and show that our proposed auxiliary task is well-aligned with the main RL task . Representation learning in deep RL is in general very different from aggregating states in tabular case , though the latter may build nice theoretical foundation for the former . Here we focus on how to design auxiliary tasks to accelerate representation learning using contrastive learning techniques , and we propose a novel return-based contrastive method based on our proposed Zπ-irrelevance abstraction . 3 PRELIMINARY . We consider a Markov Decision Process ( MDP ) which is a tuple ( S , A , P , R , µ , γ ) specifying the state space S , the action space A , the state transition probability P ( st+1|st , at ) , the reward R ( rt|st , at ) , the initial state distribution µ ∈ ∆S and the discount factor γ . Also , we denote x : = ( s , a ) ∈ X : = S × A to be the state-action pair . A ( stationary ) policy π : S → ∆A specifies the action selection probability on each state . Following the policy π , the discounted sum of future rewards ( or return ) is denoted by the random variable Zπ ( s , a ) = ∑∞ t=0 γ tR ( st , at ) , where s0 = s , a0 = a , st ∼ P ( ·|st−1 , at−1 ) , and at ∼ π ( ·|st ) . We divide the range of return into K equal bins { R0 = Rmin , R1 , · · · , RK = Rmax } such that Rk −Rk−1 = ( Rmax−Rmin ) /K , ∀k ∈ [ K ] , where Rmin ( resp . Rmax ) is the minimum ( reps. maximum ) possible return , and [ K ] : = { 1 , 2 , · · · , K } . We use b ( R ) = k ∈ [ K ] to denote the event that R falls into the kth bin , i.e. , Rk−1 < R ≤ Rk . Hence , b ( R ) can be viewed as the discretized version of the return , and the distribution of discretized return can be represented by a K-dimensional vector Zπ ( x ) ∈ ∆K , where the k-th element equals to Pr [ Rk−1 < Zπ ( x ) ≤ Rk ] . The Q function is defined as Qπ ( x ) = E [ Zπ ( x ) ] , and the state value function is defined as V π ( s ) = Ea∼π ( ·|s ) [ Zπ ( s , a ) ] . The objective for RL is to find a policy π that maximizes the expected cumulative reward J ( π ) = Es∼µ [ V π ( s ) ] . We denote the optimal policy as π∗ and the corresponding optimal Q function as Q∗ : = Qπ ∗ .

Summary: The authors propose the inclusion of an auxiliary task for training an RL model, where the auxiliary task objective is to learn an abstraction of the state-action space that clusters (s,a) pairs according to their expected return.  The authors first describe a basic abstraction learning framework (Z-learning) followed by the extension to Deep RL as an auxiliary task (RCRL). The authors present results in Atari (discrete action) building on Rainbow, showing an improvement compared to baselines on median HNS in the low-data regime, and results on DMControl (continuous action) building on SAC, showing similar or improved performance compared to baselines.