Learning Subgoal Representations with Slow Dynamics

1 INTRODUCTION . Deep Reinforcement Learning ( RL ) has demonstrated increasing capabilities in a wide range of domains , including playing games ( Mnih et al. , 2015 ; Silver et al. , 2016 ) , controlling robots ( Schulman et al. , 2015 ; Gu et al. , 2017 ) and navigation in complex environments ( Mirowski et al. , 2016 ; Zhu et al. , 2017 ) . Solving temporally extended tasks with sparse or deceptive rewards is one of the major challenges for RL . Hierarchical Reinforcement Learning ( HRL ) , which enables control at multiple time scales via a hierarchical structure , provides a promising way to solve those challenging tasks . Goal-conditioned methods have long been recognized as an effective paradigm in HRL ( Dayan & Hinton , 1993 ; Schmidhuber & Wahnsiedler , 1993 ; Nachum et al. , 2019 ) . In goal-conditioned HRL , higher-level policies set subgoals for lower-level ones periodically , and lower-level policies are incentivized to reach these selected subgoals . A proper subgoal representation function , abstracting a state space to a latent subgoal space , is crucial for effective goal-conditioned HRL , because the abstract subgoal space , i.e. , high-level action space , simplifies the high-level policy learning , and explorative low-level behaviors can be induced by setting different subgoals in this compressed space as well . Recent works in goal-conditioned HRL have been concentrated on implicitly learning the subgoal representation in an end-to-end manner with hierarchical policies ( Vezhnevets et al. , 2017 ; Dilokthanakul et al. , 2019 ) , e.g. , using a variational autoencoder ( Péré et al. , 2018 ; Nair & Finn , 2019 ; Nasiriany et al. , 2019 ) , directly utilizing the state space ( Levy et al. , 2019 ) or a handcrafted space ( Nachum et al. , 2018 ) as a subgoal space . Sukhbaatar et al . ( 2018 ) proposed to learn subgoal embeddings via self-play , and Ghosh et al . ( 2018 ) designed a representation learning objective using an actionable distance metric , but both of the methods need a pretraining process . Near-Optimal ∗Denotes equal contribution 1Videos available at https : //sites.google.com/view/lesson-iclr 2Find open-source code at https : //github.com/SiyuanLee/LESSON Representation ( NOR ) for HRL ( Nachum et al. , 2019 ) learns an abstract space concurrently with hierarchical policies by bounding the sub-optimality . However , the NOR subgoal space could not support efficient exploration in challenging deceptive reward tasks . In this paper , we develop a novel method , which LEarns the Subgoal representation with SlOw dyNamics ( LESSON ) along with the hierarchical policies . Subgoal representation in HRL is not only a state space abstraction , but also a form of high-level action abstraction . Since the high-level agent makes decisions at a low temporal resolution , our method extracts features with slow dynamics from observations as the subgoal space to enable temporal coherence . LESSON minimizes feature changes between adjacent low-level timesteps , in order for the learned feature representation to have the slowness property . To capture dynamic features and prevent the collapse of the learned representation space , we also introduce an additional contrastive objective that maximizes feature changes between high-level temporal intervals . We provide a theoretical motivation for the slowness objective . That is , selecting slow features as the subgoal space can achieve the most efficient hierarchical exploration when the subgoal space dimension is low and fixed . We illustrate on a didactic example that our method LESSON accomplishes the most efficient state coverage among all the compared subgoal representation functions . We also compare LESSON with state-of-theart HRL and exploration methods on complex MuJoCo tasks ( Todorov et al. , 2012 ) . Experimental results demonstrate that ( 1 ) LESSON dramatically outperforms previous algorithms and learns hierarchical policies more efficiently ; ( 2 ) our learned representation with slow dynamics can provide interpretability for the hierarchical policy ; and ( 3 ) our subgoal representation and low-level policies can be transferred between different tasks . 2 PRELIMINARIES . In reinforcement learning , an agent interacts with an environment modeled as an MDP M = ( S , A , P , R , γ ) , where S is a state space , A is an action space . P : S × A × S → [ 0 , 1 ] is an unknown dynamics model , which specifies the probability P ( s′|s , a ) of transitioning to next state s′ from current state s by taking action a. R : S × A → R is a reward function , and γ ∈ [ 0 , 1 ) is a discount factor . We optimize a stochastic policy π ( a|s ) , which outputs a distribution over the action space for a given state s. The objective is to maximize the expected cumulative discounted reward Eπ [ ∑∞ t=0 γ trt ] under policy π . 3 METHOD . In this section , we present the proposed method for LEarning Subgoal representations with SlOw dyNamics ( LESSON ) . First , we describe a two-layered goal-conditioned HRL framework . We then introduce a core component of LESSON , the slowness objective for learning the subgoal representation of HRL . Finally , we summarize the whole learning procedure . 3.1 FRAMEWORK . Following previous work ( Nachum et al. , 2018 ; 2019 ) , we model a policy π ( a|s ) as a two-level hierarchical policy composed of a high-level policy πh ( g|s ) and a low-level policy πl ( a|s , g ) . The high-level policy πh ( g|s ) selects a subgoal g in state s every c timesteps . The subgoal g is in a low dimensional space abstracted by representation function φ ( s ) : S → Rk . The low-level policy πl ( a|s , g ) takes the high-level action g as input and interacts with the environment every timestep . Figure 1 depicts the execution process of the hierarchical policy . LESSON iteratively learns the subgoal representation function φ ( s ) with the hierarchical policy . To encourage policy πl to reach the subgoal g , we train πl with an intrinsic reward function based on the negative Euclidean distance in the latent space , rl ( st , at , st+1 , g ) = −||φ ( st+1 ) − g||2 . Policy πh is trained to optimize the expected extrinsic rewards renvt . We use the off-policy algorithm SAC ( Haarnoja et al. , 2018 ) as our base RL optimizer . In fact , our framework is compatible with any standard RL algorithm . Apparently , a proper subgoal representation φ ( s ) is critical not only for learning an effective lowlevel goal-conditioned policy but also for efficiently learning an optimal high-level policy to solve a given task . As the feature dimension k is low , φ ( s ) has a compression property , which is necessary to make the hierarchical policy learning easier . If φ ( s ) is exactly an identity function without any abstraction , the high-level policy πh still needs to explore in a large space and the complicated subgoal g for the low-level policy is hard to reach as well . In this circumstance , the hierarchical structure can not simplify the MDP and has no advantage over a flat structure . 3.2 LEARNING SUBGOAL REPRESENTATIONS . Inspired by physics-based priors , features with slow dynamics preserve higher temporal coherence and less noise ( Wiskott & Sejnowski , 2002 ) . As the high-level policy acts at a lower temporal resolution compared to the low-level policy , it is sensible to learn a subgoal representation function with a slowness objective . To solve large-scale problems , we parameterize the representation function φ ( s ) with a neural network to extract slow features . One natural way of learning φ ( s ) is to minimize the squared difference between feature values at times t and t+ 1 , min φ E ( st , st+1 ) ∼D [ ||φ ( st ) − φ ( st+1 ) ||2 ] , ( 1 ) where D is a replay buffer . This loss function eliminates fast features , but can be trivially optimized if we allow lossy representation function φ ( e.g. , if φ ( s ) = 0 for ∀s ∈ S ) . To avoid such trivial solutions and capture dynamic features , we propose a contrastive loss to maximize the distance between high-level state transitions in the latent subgoal space , i.e. , minφ E ( st , st+c ) ∼D [ −||φ ( st ) − φ ( st+c ) ||2 ] . To trade off these two loss functions , we adopt the technique of triplet loss ( Chopra et al. , 2005 ) , i.e. , imposing the latent distance between high-level transitions larger than a margin parameterm , as shown by Eq . 2 . If we remove the margin parameterm and themax operator , Eq . 2 will be dominated by the maximizing distance part . Margin m defines a unit of distance in the latent space , which prevents trivial solutions as well . min φ E ( st , st+1 , st+c ) ∼D [ ||φ ( st ) − φ ( st+1 ) ||2 +max ( 0 , m− ||φ ( st ) − φ ( st+c ) ||2 ) ] . ( 2 ) The above learning objective abstracts the state space to a latent subgoal space with slow dynamics . As Eq . 2 optimizes the squared difference between feature values , the learned representation can preserve the spatial locality property of the state space , so a subgoal g can be selected in the neighborhood of φ ( s ) . In the next section , we give a theoretical motivation for the slowness objective . That is , selecting slow features as the subgoal space can promote efficient exploration . Algorithm 1 shows the learning procedure of our method . We update φ ( s ) and πl at the same frequency so that the low-level reward function varies in a stationary way . The high-level policy is updated less frequently , as the high-level transitions are less . 4 EFFICIENT EXPLORATION WITH SLOW SUBGOAL REPRESENTATION . In this section , we provide a theoretical motivation for subgoal representation learning with slow dynamics from a statistical view . To support a formal analysis , we consider selecting a subset of features from the state space as a subgoal space . We prove that , given a fixed subgoal space dimension , selecting slow features as the subgoal space can achieve the most efficient hierarchical exploration . We first define a measure for exploration and describe assumptions of our analysis . Then , we present a theorem about the optimality property and corresponding implications . Algorithm 1 LESSON algorithm 1 : Input : Number of training steps N , margin m , replay buffer D. 2 : Initialize : Learnable parameters for πh ( g|s ) , πl ( a|s , g ) and φ ( s ) . 3 : for t = 1 .. N do 4 : Collect experience ( st , gt , at , st+1 , renvt ) under πh and πl . 5 : Compute low-level reward rlt = −||φ ( st+1 ) − gt||2 . 6 : Update the replay buffer D. 7 : Optimize πh by maximizing cumulative task rewards with D every c timesteps . 8 : Optimize πl by maximizing cumulative low-level rewards with D every timestep . 9 : Sample a batch of state transitions from D and update φ with Eq . 2 every timestep . 10 : end for 11 : Return : πh , πl and φ .

Summary: This paper proposes a new method for learning subgoal representations in HRL. The method learns a representation that emphasises features that change slowly, through a “slowness objective”. The slowness objective minimises changes in the subgoal representation between low level time steps, while maximising feature changes between the high-level temporal intervals. This objective allows for efficient exploration, which the paper justifies theoretically, and supports with some empirical experiments on challenging control domains.