Temporal abstractions-augmented temporally contrastive learning: an alternative to the Laplacian in RL

1 INTRODUCTION . Representation learning has been at the core of many recent machine learning advances ( c.f . Bengio et al. , 2013 ) . With the advent of deep reinforcement learning ( Mnih et al. , 2015 ) , representation learning has also become one of the main topics of interest in reinforcement learning ( RL ) . For example , in the goal-conditioned hierarchical setting ( Vezhnevets et al. , 2017 ; Nachum et al. , 2019a ) , one learns a representation that maps observations to an abstract space , the representation space , in which the higher-level policy defines the desired behavior of the lower-level policy . Distance in the representation space can then be used to reward and guide the lower-level policy towards specific goal states . Moreover , environments with rich observations and complex dynamics ( e.g. , Bellemare et al. , 2020 ) have motivated recent works on learning representations as controllable or contingent features ( Bengio et al. , 2017 ; Choi et al. , 2019 ) , on top of which one can potentially learn latent models in the perspective of planning ( Hafner et al. , 2019b ; Nasiriany et al. , 2019 ; Schrittwieser et al. , 2020 ) and control ( Watter et al. , 2015 ; Banijamali et al. , 2018 ; Hafner et al. , 2019a ) . In this work , we are interested in the reward-agnostic setting in which an RL agent first interacts with the environment to build a representation , φ , of the state space , S , without relying on any task-specific reward signal . This representation can later be used to solve a task posed in the environment in the form of a reward function . In this setting , the environment dynamics are the only informative interaction channel available to the agent . This has naturally motivated graph Laplacianbased methods to address the task-agnostic phase ; where the graph vertices correspond to the states and its edges to the transitions probabilities . The Laplacian ’ s eigenvectors can been leveraged as a holistic state representation , termed the Laplacian representation , which captures the environment ’ s dynamics structure and geometry ( Mahadevan , 2005 ; Mahadevan & Maggioni , 2007 ) . Wu et al . ( 2019 ) recently proposed an efficient approximation of the Laplacian representation ( LAPREP ) by framing the graph drawing objective as a temporally-contrastive loss ( see Section 2.2 ) . While this formulation works around potentially prohibitive eigendecompositions , which extends the representation ’ s applicability to large and continuous state spaces , it assumes access to a uniform sampling prior over S. In practice , this translates in the ability to reset the agent to a uniform random starting state in the environment , which artificially alleviates the exploration problem . As we will show in Section 4 , the uniformity of that prior is crucial for the quality of the learned representation . However , such sampling is not trivial in the absence of the uniform prior privilege , since the agent has to learn to explore the state space to be able to access arbitrary states . In effect , one must handle the exploration along the representation learning in order to preserve the representation ’ s quality . In this work , we propose a representation learning framework that conciliate a similar temporallycontrastive approach with exploration in the task-agnostic setting . In practice , the representation is trained on data collected with a uniformly random policy , πµ ( random walk trajectories ) . Without a uniform access to the state space , the collected data is concentrated around accessible starting states . To achieve a better data collection , we tie the representation learning problem to that of learning a covering strategy . Briefly , our method consists in using the available representation to learn a skill-based ( hierarchical ) covering policy that is in turn used to discover yet unseen parts of the state space , providing novel data to refine and expand the representation . Our approach , illustrated in Figure 1 , is inspired by the cyclic option discovery framework ( Machado , 2019 ) , which motivated several related methods ( Machado et al. , 2017 ; 2018 ; Jinnai et al. , 2020 ) . In addition , we propose to integrate the temporal abstractions learned by the skills in the contrastive representation learning objective to encourage temporally-extended exploration and enforce the representation ’ s dynamics-awareness , i.e . how representative the φ-induced euclidean distance is of distances in the state space . We empirically show our agent ’ s ability to progressively explore the state space and to consistently extend the representation covered domain in a non-uniform prior setting . We show that our representation leads to better value predictions than LAP-REP , and that it recovers the representation quality expected from a uniform prior . We also evaluate our representation in shaping rewards for goal-achieving tasks , and we show it outperforms LAP-REP , confirming both its superior ability in capturing dynamics and in scaling to larger environments . Finally , the skills learned in our framework also prove to be successful at difficult continuous navigation tasks with sparse rewards , where other standard skill discovery methods are limited . 2 PRELIMINARIES . 2.1 TASK-AGNOSTIC REINFORCEMENT LEARNING . We describe a task-agnostic RL environment as a task-agnostic Markov decision process ( MDP ) M = ( S , A , P , γ , d0 ) where S is the state space , A the action space , P : S × A → ∆ ( S ) the transition dynamics defining the next state distribution given current state and action taken , γ ∈ [ 0 , 1 ) the discount factor , and d0 the initial state distribution . A policy π : S → ∆ ( A ) maps states s ∈ S to distributions over actions . Knowledge acquired from task-agnostic interactions with the environment ( e.g. , a representation or a policy ) can then be leveraged for specific tasks . A task is instantiated with a reward function , R : S → R , which is combined with the task-agnostic MDP . The task objective is to find the optimal policy maximizing the expected discounted return , Eπ , d0 [ ∑ t γ tR ( st , at ) ] , starting from state s0 ∼ d0 and acting according to at ∼ π ( ·|st ) . 2.2 THE LAPLACIAN REPRESENTATION . The Laplacian representation ( LAP-REP ) , as proposed by Wu et al . ( 2019 ) , can be learned with the following contrastive objective : LLap ( φ ; Dπµ ) = E ( u , v ) ∼Dπµ [ ‖φ ( u ) − φ ( v ) ‖22 ] +β E u∼Dπµ v∼Dπµ [ ( φ ( u ) > φ ( v ) ) 2 − ‖φ ( u ) ‖22 − ‖φ ( v ) ‖22 ] , ( 1 ) where β is a hyperparameter , πµ is the uniformly random policy , Dπµ a set of trajectories from πµ ( random walks ) . We use ( u , v ) ∼ Dπµ to denote the sampling of a random transition from Dπµ , and similarly u ∼ Dπµ for a random state . Wu et al . ( 2019 ) showed the competitiveness of the Laplacian representation when provided with a uniform prior over S during the collection of Dπµ . Their objective ( Eq . 1 ) is a temporally-contrastive loss : it is comprised of an attractive term that forces temporally close states to have similar representations and a repulsive term that keeps temporally far states ’ representations far apart . Here , the repulsive term was specifically derived from the orthonormality constraint of the Laplacian eigenvectors . 2.3 THE NON-UNIFORM PRIOR SETTING . In RL , representation learning is deeply coupled to the problem of exploration . Indeed , the induced state distribution defines the representation ’ s training distribution . However , LAP-REP ( Wu et al. , 2019 ) has been learned in the specific uniform prior setting that alleviates the exploration challenge . In this setting , Dπµ , from Eq . 1 , is a collection of random walks with uniformly random starting states , which provides a uniform training distribution to the representation learning objective . In the case of a non-uniform prior , the induced visitation distribution can be quite concentrated around the start state distribution when solely relying on random walks , hence the need for an exploration strategy for a better covering distribution . To study the problem described above , we investigate the setting in which the environment has a fixed predefined state s0 to which it resets with a probability pr every K steps ; with K of the order of diameter of S. With a uniformly random behavior policy , this setting is equivalent to a initial state distribution that is concentrated around s0 and whose density decays exponentially away from it . We will refer to this setting as the non-uniform prior ( non-µ ) setting , as opposed to the uniform prior ( µ ) setting where the agent has access to the uniform state distribution . 3 TEMPORAL ABSTRACTIONS-AUGMENTED REPRESENTATION LEARNING . In this section , we present Temporal Abstractions-augmented Temporally-Contrastive learning ( TATC ) , a representation learning approach in which the representation works in tandem with a skill-based covering policy for a better representation learning in the non-uniform prior setting . We first propose an alternative objective to Eq . 1 that suits this setting , then describe the exploratory policy training . Finally , we introduce an augmentation of the proposed objective based on the learned temporal abstractions to improve exploration and enforce the representation ’ s dynamics-awareness . 3.1 TEMPORALLY-CONTRASTIVE REPRESENTATION OBJECTIVE . As mentioned in Section 2.2 , the repulsive term in LAP-REP ’ s objective ( Eq . 1 ) derives from the eigenvectors ’ orthonormality constraint . However , because the environment is expected to be progressively covered in the non-uniform prior setting , the orthonormality constraint can make online representation learning highly non-stationary.1 For this reason , we adopt the following objective with a generic repulsive term that is more amenable to online learning : Lcont ( φ ; Dπµ ) : = E ( u , v ) ∼Dπµ [ ‖φ ( u ) − φ ( v ) ‖22 ] + β E u∼Dπµ v∼Dπµ [ exp ( −‖φ ( u ) − φ ( v ) ‖2 ) ] . ( 2 ) 1In general , even within a given matrix ’ s perturbation neighborhood , its eigenvectors can show a highly nonlinear sensitivity ( Trefethen & Bau , 1997 ) . 3.2 REPRESENTATION-BASED COVERING POLICY . In the non-uniform prior setting , exploration is required to provide the representation with a better training distribution . To this purpose , we adopt a hierarchical RL approach to leverage the exploratory efficiency of options ( Sutton et al. , 1999 ; Nachum et al. , 2019b ) , or skills . The agent acts according to a bi-level policy ( πhi , πlow ) . The high-level policy πhi : S → ∆ ( Ω ) defines , at each state s , a distribution over a set Ω of directions ( unit vectors ) in the representation space ( Ω = { δ | δ ∈ Rd , ‖δ‖2=1 } ) . Each direction corresponds to a fixed length skill encoded by the low-level policy πlow : S × Ω → ∆ ( A ) . These skills are expected to travel in the representation space along the directions instructed by πhi . In short , given a sampled direction πhi ( ·|s ) ∼ δ ∈ Ω , the low-level policy executes the directional skill πlow ( ·|s , δ ) for a fixed number of steps c before a new direction is sampled . Now , we describe the intrinsic rewards used to train the policies πhi and πlow . Low-level Policy . πlow is simply trained to follow directions defined by πhi in the representation space . For a given δ ∈ Ω ⊂ Rd , the corresponding skill πlow ( ·|s , δ ) is trained to maximize the intrinsic reward function : rδ ( s , s′ ) : = cos ( δ , φ ( s′ ) − φ ( s ) ) = δ > ( φ ( s′ ) − φ ( s ) ) ‖φ ( s′ ) − φ ( s ) ‖ ( 3 ) where ( s , s′ ) is an observed state transition , and φ the representation being learned . We use the cosine similarity as a way to encourage learning diverse directional skills . Indeed , skills cospecialization is avoided by rewarding the agent for the steps induced along the instructed direction δ regardless of their magnitudes . High-level Policy . The high-level policy is expected to guide the covering strategy . It should do so by sampling the skills of the most promising directions in terms of exploration : affording new discoveries while avoiding to spend more time than needed in previously explored areas . For this purpose , we design a reward function defined over a sequence of L consecutive skills . Let { shik } Lk=1 be the sequence of their initial states and their respective sampled directions δk ∼ πhi ( ·|shik ) . Since φ is trained to capture the dynamics , the travelled distance in φ ’ s space is a good proxy of how far the choices made by πhi eventually brought the agent in the environment . Therefore , for a given high-level trajectory , τ hi = ( shi1 , s hi 2 , ... , s hi L , s hi f ) , with s hi f the final state reached by the last skill , the high-level policy is trained to maximize the following quantities : ∀k ∈ { 1 , ... , L } , Rhi ( shik , δk ) : = ‖φ ( shi1 ) − φ ( shif ) ‖2 , ( 4 ) where δk ∼ πhi ( ·|shik ) is the direction sampled at shik . From the policy optimization perspective , each of these quantities plays the role of the return cumulated along the sampled high-level trajectory and not just a single ( high-level ) step reward . This term looks at reaching shif as the result of a sequential collaboration of L skills , rewarding them equally . It values how far this sequence of skills has eventually brought the agent . These policy training choices are closely related to how the representation is trained . Indeed , the exploratory behavior emerges from the interaction between the policy and the representation while training . In the following section , we describe how the representation benefits from the temporal abstractions learned by the covering policy ( πhi , πlow ) .

Summary: This paper introduces a representation learning framework for reward-agnostic RL, which is based on temporal contrastive learning. The proposed method (TATC) is motivated by the Laplacian representation approximation approach (Wu et al., 2019) and overcomes the limitation of requiring a uniform prior. The experiments demonstrate the effectiveness of TATC.