Near-Optimal Representation Learning for Hierarchical Reinforcement Learning

1 INTRODUCTION . Hierarchical reinforcement learning has long held the promise of extending the successes of existing reinforcement learning ( RL ) methods ( Gu et al. , 2017 ; Schulman et al. , 2015 ; Lillicrap et al. , 2015 ) to more complex , difficult , and temporally extended tasks ( Parr & Russell , 1998 ; Sutton et al. , 1999 ; Barto & Mahadevan , 2003 ) . Recently , goal-conditioned hierarchical designs , in which higher-level policies communicate goals to lower-levels and lower-level policies are rewarded for reaching states ( i.e . observations ) which are close to these desired goals , have emerged as an effective paradigm for hierarchical RL ( Nachum et al . ( 2018 ) ; Levy et al . ( 2017 ) ; Vezhnevets et al . ( 2017 ) , inspired by earlier work Dayan & Hinton ( 1993 ) ; Schmidhuber & Wahnsiedler ( 1993 ) ) . In this hierarchical design , representation learning is crucial ; that is , a representation function must be chosen mapping state observations to an abstract space . Goals ( desired states ) are then specified by the choice of a point in this abstract space . Previous works have largely studied two ways to choose the representation : learning the representation end-to-end together with the higher- and lower-level policies ( Vezhnevets et al. , 2017 ) , or using the state space as-is for the goal space ( i.e. , the goal space is a subspace of the state space ) ( Nachum et al. , 2018 ; Levy et al. , 2017 ) . The former approach is appealing , but in practice often produces poor results ( see Nachum et al . ( 2018 ) and our own experiments ) , since the resulting representation is under-defined ; i.e. , not all possible sub-tasks are expressible as goals in the space . On the other hand , fixing the representation to be the full state means that no information is lost , but this choice is difficult to scale to higher dimensions . For example , if the state observations are entire images , the higher-level must output target images for the lower-level , which can be very difficult . We instead study how unsupervised objectives can be used to train a representation that is more concise than the full state , but also not as under-determined as in the end-to-end approach . In order to do so in a principled manner , we propose a measure of sub-optimality of a given representation . This measure aims to answer the question : How much does using the learned representation in place of the full representation cause us to lose , in terms of expected reward , against the optimal policy ? This question is important , because a useful representation will compress the state , hopefully making ∗Also at UC Berkeley . 1See videos at https : //sites.google.com/view/representation-hrl 2Find open-source code at https : //github.com/tensorflow/models/tree/master/ research/efficient-hrl the learning problem easier . At the same time , the compression might cause the representation to lose information , making the optimal policy impossible to express . It is therefore critical to understand how lossy a learned representation is , not in terms of reconstruction , but in terms of the ability to represent near-optimal policies on top of this representation . Our main theoretical result shows that , for a particular choice of representation learning objective , we can learn representations for which the return of the hierarchical policy approaches the return of the optimal policy within a bounded error . This suggests that , if the representation is learned with a principled objective , the ‘ lossy-ness ’ in the resulting representation should not cause a decrease in overall task performance . We then formulate a representation learning approach that optimizes this bound . We further extend our result to the case of temporal abstraction , where the higher-level controller only chooses new goals at fixed time intervals . To our knowledge , this is the first result showing that hierarchical goal-setting policies with learned representations and temporal abstraction can achieve bounded sub-optimality against the optimal policy . We further observe that the representation learning objective suggested by our theoretical result closely resembles several other recently proposed objectives based on mutual information ( van den Oord et al. , 2018 ; Ishmael Belghazi et al. , 2018 ; Hjelm et al. , 2018 ) , suggesting an intriguing connection between mutual information and goal representations for hierarchical RL . Results on a number of difficult continuous-control navigation tasks show that our principled representation learning objective yields good qualitative and quantitative performance compared to existing methods . 2 FRAMEWORK . Following previous work ( Nachum et al. , 2018 ) , we consider a two-level hierarchical policy on an MDP M = ( S , A , R , T ) , in which the higher-level policy modulates the behavior of a lowerlevel policy by choosing a desired goal state and rewarding the lower-level policy for reaching this state . While prior work has used a sub-space of the state space as goals ( Nachum et al. , 2018 ) , in more general settings , some type of state representation is necessary . That is , consider a state representation function f : S → Rd . A two-level hierarchical policy on M is composed of a higher-level policy πhi ( g|s ) , where g ∈ G = Rd is the goal space , that samples a high-level action ( or goal ) gt ∼ πhi ( g|st ) every c steps , for fixed c. A non-stationary , goal-conditioned , lower-level policy πlo ( a|st , gt , st+k , k ) then translates these high-level actions into low-level actions at+k ∈ A for k ∈ [ 0 , c − 1 ] . The process is then repeated , beginning with the higher-level policy selecting another goal according to st+c . The policy πlo is trained using a goal-conditioned reward ; e.g . the reward of a transition g , s , s′ is −D ( f ( s′ ) , g ) , where D is a distance function . In this work we adopt a slightly different interpretation of the lower-level policy and its relation to πhi . Every c steps , the higher-level policy chooses a goal gt based on a state st. We interpret this state-goal pair as being mapped to a nonstationary policy π ( a|st+k , k ) , π ∈ Π , where Π denotes the set of all possible c-step policies acting on M. We use Ψ to denote this mapping from S × G to Π . In other words , on every cth step , we encounter some state st ∈ S. We use the higher-level policy to sample a goal gt ∼ πhi ( g|st ) and translate this to a policy πt = Ψ ( st , gt ) . We then use πt to sample actions at+k ∼ πt ( a|st+k , k ) for k ∈ [ 0 , c− 1 ] . The process is then repeated from st+c . Although the difference in this interpretation is subtle , the introduction of Ψ is crucial for our subsequent analysis . The communication of gt is no longer as a goal which πhi desires to reach , but rather more precisely , as an identifier to a low-level behavior which πhi desires to induce or activate . The mapping Ψ is usually expressed as the result of an RL optimization over Π ; e.g. , Ψ ( st , g ) = arg max π∈Π c∑ k=1 γk−1EPπ ( st+k|st ) [ −D ( f ( st+k ) , g ) ] , ( 1 ) where we use Pπ ( st+k|st ) to denote the probability of being in state st+k after following π for k steps starting from st. We will consider variations on this low-level objective in later sections . From Equation 1 it is clear how the choice of representation f affects Ψ ( albeit indirectly ) . We will restrict the environment reward function R to be defined only on states . We use Rmax to denote the maximal absolute reward : Rmax = supS |R ( s ) | . 3 HIERARCHICAL POLICY SUB-OPTIMALITY . In the previous section , we introduced two-level policies where a higher-level policy πhi chooses goals g , which are translated to lower-level behaviors via Ψ . The introduction of this hierarchy leads to a natural question : How much do we lose by learning πhi which is only able to act on M via Ψ ? The choice of Ψ restricts the type and number of lower-level behaviors that the higher-level policy can induce . Thus , the optimal policy onM is potentially not expressible by πhi . Despite the potential lossy-ness of Ψ , can one still learn a hierarchical policy which is near-optimal ? To approach this question , we introduce a notion of sub-optimality with respect to the form of Ψ : Let π∗hi ( g|s , Ψ ) be the optimal higher-level policy acting onG and using Ψ as the mapping fromG to low-level behaviors . Let π∗hier be the corresponding full hierarchical policy onM . We will compare π∗hier to an optimal hierarchical policy π ∗ agnostic to Ψ . To define π∗ we begin by introducing an optimal higher-level policy π∗∗hi ( π|s ) agnostic to Ψ ; i.e . every c steps , π∗∗hi samples a low-level behavior π ∈ Π which is applied toM for the following c steps . In this way , π∗∗hi may express all possible low-level behaviors . We then denote π∗ as the full hierarchical policy resulting from π∗∗hi . We would like to compare π∗hier to π ∗ . A natural and common way to do so is in terms of state values . Let V π ( s ) be the future value achieved by a policy π starting at state s. We define the sub-optimality of Ψ as SubOpt ( Ψ ) = sup s∈S V π ∗ ( s ) − V π ∗ hier ( s ) . ( 2 ) The state values V π ∗ hier ( s ) are determined by the form of Ψ , which is in turn determined by the choice of representation f . However , none of these relationships are direct . It is unclear how a change in f will result in a change to the sub-optimality . In the following section , we derive a series of bounds which establish a more direct relationship between SubOpt ( Ψ ) and f . Our main result will show that if one defines Ψ as a slight modification of the traditional objective given in Equation 1 , then one may translate sub-optimality of Ψ to a practical representation learning objective for f . 4 GOOD REPRESENTATIONS LEAD TO BOUNDED SUB-OPTIMALITY . In this section , we provide proxy expressions that bound the sub-optimality induced by a specific choice of Ψ . Our main result is Claim 4 , which connects the sub-optimality of Ψ to both goalconditioned policy objectives ( i.e. , the objective in 1 ) and representation learning ( i.e. , an objective for the function f ) . 4.1 SINGLE-STEPS ( c = 1 ) AND DETERMINISTIC POLICIES For ease of presentation , we begin by presenting our results in the restricted case of c = 1 and deterministic lower-level policies . In this setting , the class of low-level policies Π may be taken to be simply A , where a ∈ Π corresponds to a policy which always chooses action a . There is no temporal abstraction : The higher-level policy chooses a high-level action g ∈ G at every step , which is translated via Ψ to a low-level action a ∈ A . Our claims are based on quantifying how many of the possible low-level behaviors ( i.e. , all possible state to state transitions ) can be produced by Ψ for different choices of g. To quantify this , we make use of an auxiliary inverse goal model ϕ ( s , a ) , which aims to predict which goal g will cause Ψ to yield an action ã = Ψ ( s , g ) that induces a next state distribution P ( s′|s , ã ) similar to P ( s′|s , a ) .3 We have the following theorem , which bounds the sub-optimality in terms of total variation divergences between P ( s′|s , a ) and P ( s′|s , ã ) : Theorem 1 . If there exists ϕ : S ×A→ G such that , sup s∈S , a∈A DTV ( P ( s ′|s , a ) ||P ( s′|s , Ψ ( s , ϕ ( s , a ) ) ) ) ≤ , ( 3 ) then SubOpt ( Ψ ) ≤ C , where C = 2γ ( 1−γ ) 2Rmax . 3In a deterministic , c = 1 setting , ϕmay be seen as a state-conditioned action abstraction mappingA→ G. Proof . See Appendices A and B for all proofs . Theorem 1 allows us to bound the sub-optimality of Ψ in terms of how recoverable the effect of any action in A is , in terms of transition to the next state . One way to ensure that effects of actions in A are recoverable is to have an invertible Ψ . That is , if there exists ϕ : S × A → G such that Ψ ( s , ϕ ( s , a ) ) = a for all s , a , then the sub-optimality of Ψ is 0 . However , in many cases it may not be desirable or feasible to have an invertible Ψ . Looking back at Theorem 1 , we emphasize that its statement requires only the effect of any action to be recoverable . That is , for any s , ∈ S , a ∈ A , we require only that there exist some g ∈ G ( given by ϕ ( s , a ) ) which yields a similar next-state distribution . To this end , we have the following claim , which connects the sub-optimality of Ψ to both representation learning and the form of the low-level objective . Claim 2 . Let ρ ( s ) be a prior and f , ϕ be so that , for K ( s′|s , a ) ∝ ρ ( s′ ) exp ( −D ( f ( s′ ) , ϕ ( s , a ) ) ) ,4 sup s∈S , a∈A DKL ( P ( s ′|s , a ) ||K ( s′|s , a ) ) ≤ 2/8 . ( 4 )

Summary: The paper studies the problem of representation learning in the context of hierarchical reinforcement learning by building on the framework of  HIRO (Nachum et al. (2018)). The papers propose a way to handle sub-optimality in the context of learning representations which basically refers to the overall sub-optimality of the entire hierarchical polity with respect to the task reward. And hence, the only practical different from the HIRO paper is that the proposed method considers representation learning for the goals, while HIRO was directly using the state space.