Learning transferable motor skills with hierarchical latent mixture policies

1 INTRODUCTION . Reinforcement learning is a powerful and flexible paradigm to train embodied agents , but relies on large amounts of agent experience , computation , and time , on each individual task . Learning each task from scratch is inefficient : it is desirable to learn a set of skills that can efficiently be reused and adapted to related downstream tasks . This is particularly pertinent for real-world robots , where interaction is expensive and data-efficiency is crucial . There are numerous existing approaches to learn transferable embodied skills , usually formulated as a two-level hierarchy with a high-level controller and low-level skills . These methods predominantly represent skills as being either continuous , such as goal-conditioned ( Lynch et al. , 2019 ; Pertsch et al. , 2020b ) or latent space policies ( Haarnoja et al. , 2018 ; Merel et al. , 2019 ; Singh et al. , 2021 ) ; or discrete , such as mixture or option-based methods ( Sutton et al. , 1999 ; Daniel et al. , 2012 ; Florensa et al. , 2017 ; Wulfmeier et al. , 2021 ) . Our goal is to combine these perspectives to leverage their complementary advantages . We propose an approach to learn a three-level skill hierarchy from an offline dataset , capturing both discrete and continuous variations at multiple levels of behavioural abstraction . The model comprises a low-level latent-conditioned controller that can learn motor primitives , a set of continuous latent mid-level skills , and a discrete high-level controller that can compose and select among these abstract mid-level behaviours . Since the mid- and high-level form a mixture , we call our method Hierarchical Latent Mixtures of Skills ( HeLMS ) . We demonstrate on challenging object manipulation tasks that our method can decompose a dataset into distinct , intuitive , and reusable behaviours . We show that these skills lead to improved sample efficiency and performance in numerous transfer scenarios : reusing skills for new tasks , generalising across unseen objects , and transferring from state to vision-based policies . Further analysis and ablations reveal that both continuous and discrete components are beneficial , and that the learned hierarchical skills are most useful in sparse-reward settings , as they encourage directed exploration of task-relevant parts of the state space . Our main contributions are as follows : • We propose a novel approach to learn skills at different levels of abstraction from an offline dataset . The method captures both discrete behavioural modes and continuous variation using a hierarchical mixture latent variable model . • We present two techniques to reuse and adapt the learned skill hierarchy via reinforcement learning in downstream tasks , and perform extensive evaluation and benchmarking in different transfer settings : to new tasks and objects , and from state to vision-based policies . • We present a detailed analysis to interpret the learned skills , understand when they are most beneficial , and evaluate the utility of both continuous and discrete skill representations . 2 RELATED WORK . A long-standing challenge in reinforcement learning is the ability to learn reusable motor skills that can be transferred efficiently to related settings . One way to learn such skills is via multi-task reinforcement learning ( Heess et al. , 2016 ; James et al. , 2018 ; Hausman et al. , 2018 ; Riedmiller et al. , 2018 ) , with the intuition that behaviors useful for a given task should aid the learning of related tasks . However , this often requires careful curation of the task set , where each skill represents a separate task . Some approaches avoid this by learning skills in an unsupervised manner using intrinsic objectives that often maximize the entropy of visited states while keeping skills distinguishable ( Gregor et al. , 2017 ; Eysenbach et al. , 2019 ; Sharma et al. , 2019 ; Zhang et al. , 2020 ) . A large body of work explores skills from the perspective of unsupervised segmentation of repeatable behaviours in temporal data ( Niekum & Barto , 2011 ; Ranchod et al. , 2015 ; Krüger et al. , 2016 ; Lioutikov et al. , 2017 ; Shiarlis et al. , 2018 ; Kipf et al. , 2019 ; Tanneberg et al. , 2021 ) . Other works investigate movement or motor primitives that can be selected or sequenced together to solve complex manipulation or locomotion tasks ( Mülling et al. , 2013 ; Rueckert et al. , 2015 ; Lioutikov et al. , 2015 ; Paraschos et al. , 2018 ; Merel et al. , 2020 ; Tosatto et al. , 2021 ; Dalal et al. , 2021 ) . Some of these methods also employ mixture models to jointly model low-level motion primitives and a high-level primitive controller ( Muelling et al. , 2010 ; Colomé & Torras , 2018 ; Pervez & Lee , 2018 ) . Several existing approaches employ architectures in which the policy is comprised of two ( or more ) levels of hierarchy . Typically , a low-level controller represents the learned set of skills , and a high-level policy instructs the low-level controller via a latent variable or goal . Such latent variables can be discrete ( Florensa et al. , 2017 ; Wulfmeier et al. , 2020 ) or continuous ( Nachum et al. , 2018 ; Haarnoja et al. , 2018 ) and regularization of the latent space is often crucial ( Tirumala et al. , 2019 ) . The latent variable can represent the behaviour for one timestep , for a fixed number of timesteps ( Ajay et al. , 2021 ) , or options with different durations ( Sutton et al. , 1999 ; Bacon et al. , 2017 ; Wulfmeier et al. , 2021 ) . One such approach that is particularly relevant ( Florensa et al. , 2017 ) learns a diverse set of skills , via a discrete latent variable that interacts multiplicatively with the state to enable continuous variation in a Stochastic Neural Network policy ; this skill space is then transferred to locomotion tasks by learning a new categorical controller . Our method differs in a few key aspects : our proposed three-level hierarchical architecture explicitly models abstract discrete skills while allowing for temporal dependence and lower-level latent variation in their execution , enabling diverse object-centric behaviours in challenging manipulation tasks . Our work is related to methods that learn robot policies from demonstrations ( LfD , e.g . ( Rajeswaran et al. , 2018 ; Shiarlis et al. , 2018 ; Strudel et al. , 2020 ) ) or more broadly from logged data ( offline RL , e.g . ( Wu et al. , 2019 ; Kumar et al. , 2020 ; Wang et al. , 2020 ) ) . While many of these focus on learning single-task policies , several approaches learn skills offline that can be transferred online to new tasks ( Merel et al. , 2019 ; Lynch et al. , 2019 ; Pertsch et al. , 2020a ; Ajay et al. , 2021 ; Singh et al. , 2021 ) . These all train a two-level hierarchical model , with a high-level encoder that maps to a continuous latent space , and a low-level latent-conditioned controller . The high-level encoder can encode a whole trajectory ( Pertsch et al. , 2020a ; 2021 ; Ajay et al. , 2021 ) ; a short look-ahead state sequence ( Merel et al. , 2019 ) ; the current and final goal state ( Lynch et al. , 2019 ) ; or can even be simple isotropic Gaussian noise ( Singh et al. , 2021 ) that can be flexibly transformed by a flow-based low-level controller . At transfer time , a new high-level policy is learned from scratch : this can be more efficient with skill priors ( Pertsch et al. , 2020a ) or temporal abstraction ( Ajay et al. , 2021 ) . HeLMS builds on this large body of work by explicitly modelling both discrete and continuous behavioural structure via a three-level skill hierarchy . We use similar information asymmetry to Neural Probabilistic Motor Primitives ( NPMP ) ( Merel et al. , 2019 ; 2020 ) , conditioning the highlevel encoder on a short look-ahead trajectory . However HeLMS explicitly captures discrete modes of behaviour via the high-level controller , and learns an additional mid-level which is able to transfer abstract skills to downstream tasks , rather than learning a continuous latent policy from scratch . 3 METHOD . This paper examines a two-stage problem setup : an offline stage where a hierarchical skill space is learned from a dataset , and an online stage where these skills are transferred to a reinforcement learning setting . The dataset D comprises a set of trajectories , each a sequence of state-action pairs { xt , at } Tt=0 . The model incorporates a discrete latent variable yt ∈ { 1 , . . . , K } as a high-level skill selector ( for a fixed number of skills K ) , and a mid-level continuous variable zt ∈ Rnz conditioned on yt which parameterises each skill . Marginally , zt is then a latent mixture distribution representing both a discrete set of skills and the variation in their execution . A sample of zt represents an abstract behaviour , which is then executed by a low-level controller p ( at | zt , xt ) . The learned skill space can then be transferred to a reinforcement learning agent π in a Markov Decision Process defined by tuple { S , A , T , R , γ } : these represent the state , action , and transition distributions , reward function , and discount factor respectively . When transferring , we train a new high-level controller that acts either at the level of discrete skills yt or continuous zt , and freeze lower levels of the policy . We explain our method in detail in the following sections . 3.1 LATENT MIXTURE SKILL SPACES FROM OFFLINE DATA . Our method employs the generative model in Figure 1a . As shown , the state inputs can be different for each level of the hierarchy , but to keep notation uncluttered , we refer to all state inputs as xt and the specific input can be inferred from context . The joint distribution of actions and latents over a trajectory is decomposed into a latent prior p ( y0 : T , z1 : T ) and a low-level controller p ( at | zt , xt ) : p ( a1 : T , y0 : T , z1 : T |x1 : T ) = p ( y0 : T , z1 : T ) T∏ t=1 p ( at | zt , xt ) p ( y0 : T , z1 : T ) = p ( y0 ) T∏ t=1 p ( yt |yt−1 ) p ( zt |yt ) . ( 1 ) Intuitively , the categorical variable yt can capture discrete modes of behaviour , and the continuous latent zt is conditioned on this to vary the execution of each behaviour . Thus , zt follows a mixture distribution , encoding all the relevant information on desired abstract behaviour for the low-level controller p ( at | zt , xt ) . Since each categorical latent yt is dependent on yt−1 , and zt is only dependent on yt , this prior can be thought of as a Hidden Markov model over the sequence of z1 : T . To perform inference over the latent variables , we introduce the following variational approximation : q ( y0 : T , z1 : T |x1 : T ) = p ( y0 ) T∏ t=1 q ( yt |yt−1 , xt ) q ( zt |yt , xt ) ( 2 ) Here , the selection of a skill yt ∼ q ( yt |yt−1 , xt ) is dependent on that of the previous timestep ( allowing for temporal consistency ) , as well as the input . The mid-level skill is then parameterised by zt ∼ q ( zt |yt , xt ) based on the chosen skill and current input . p ( y0 ) and p ( yt |yt−1 ) model a skill prior and skill transition prior respectively , while p ( zt |yt ) represents a skill parameterisation prior to regularise each mid-level skill . While all of these priors can be learned in practice , we only found it necessary to learn the transition prior , with a uniform categorical for the initial skill prior and a simple fixed N ( 0 , I ) prior for p ( zt |yt ) . Training via the Evidence Lower Bound The proposed model contains a number of components with trainable parameters : the prior parameters ψ = { ψa , ψy } for the low-level controller and categorical transition prior respectively ; and posterior parameters φ = { φy , φz } for the high-level controller and mid-level skills . For a trajectory { x1 : T , a1 : T } ∼ D , we can compute the Evidence Lower Bound for the state-conditional action distribution , ELBO ≤ p ( a1 : T |x1 : T ) , as follows : ELBO = Eqφ ( y0 : T , z1 : T |x1 : T ) [ log pψ ( a1 : T , y0 : T , z1 : T |x1 : T ) − log qφ ( y0 : T , z1 : T |x1 : T ) ] ≈ T∑ t=1 [ ∑ yt q ( yt |x1 : t ) ( per-component action recon︷ ︸︸ ︷ log pψa ( at | z̃ { yt } t , xt ) −βz per-component KL regulariser︷ ︸︸ ︷ KL ( qφz ( zt |yt , xt ) || p ( zt |yt ) ) ) ] −βy T∑ t=1 [ ∑ yt−1 q ( yt−1 |x1 : t−1 ) KL ( qφy ( yt |yt−1 , xt ) || pψy ( yt |yt−1 ) ) ︸ ︷︷ ︸ categorical regulariser ] ( 3 ) where z̃ { yt } t ∼ q ( zt |yt , xt ) . The coefficients βy and βz can be used to weight the KL terms , and the cumulative component probability q ( yt |x1 : t ) can be computed iteratively as q ( yt |x1 : t ) =∑ yt−1 qφy ( yt |yt−1 , xt ) q ( yt−1 |x1 : t−1 ) . In other words , for each timestep t and each mixture component , we compute the latent sample and the corresponding action log-probability , and the KL-divergence between the component posterior and prior . This is then marginalised over all yt , with an additional KL over the categorical transitions . For more details , see Appendix C. Information-asymmetry As noted in previous work ( Tirumala et al. , 2019 ; Galashov et al. , 2019 ) , hierarchical approaches often benefit from information-asymmetry , with higher levels seeing additional context or task-specific information . This ensures that the high-level remains responsible for abstract , task-related behaviours , while the low-level executes simpler motor primitives . We employ similar techniques in HeLMS : the low-level inputs xLL comprise the proprioceptive state of the embodied agent ; the mid-level inputs xML also include the poses of objects in the environment ; and the high-level xHL concatenates both object and proprioceptive state for a short number of lookahead timesteps . The high- and low-level are similar to ( Merel et al. , 2019 ) , with the low-level controller enabling motor primitives based on proprioceptive information , and the high-level using the lookahead information to provide additional context regarding behavioural intent when specifying which skill to use . The key difference is the categorical high-level and the additional mid-level , with which HeLMS can learn more object-centric skills and transfer these to downstream tasks . Network architectures The architecture and information flow in HeLMS are shown in Figure 1b . The high-level network contains a gated head , which uses the previous skill yt−1 to index into one of K categorical heads , each of which specify a distribution over yt . For a given yt , the corresponding mid-level skill network is selected and used to sample a latent action zt , which is then used as input for the latent-conditioned low-level controller , which parameterises the action distribution . The skill transition prior p ( yt |yt−1 ) is also learned , and is parameterised as a linear softmax layer which takes in a one-hot representation of yt−1 and outputs the distribution over yt . All components are trained end-to-end via the objective in Equation 3 .

Summary: This paper presents a hierarchical model for skills extraction from an online dataset. In particular, the model can be seen as a 3-level hierarchy. The highest level of the hierarchy selects a discrete variable that corresponds to the index of a skill. The mid-level selects the continuous parameter that defines the skill. The lower level executes the skill.  The considered problem is crucial in reinforcement learning. In particular, the definition of skills to be reused across different tasks is essential to achieve better sample efficiency, and crucial to allow reinforcement learning to be applied directly on challenging real robotic tasks.   When dealing with an off-line dataset, one wants potentially to first extract a set of skills (say, low-level policies) and then to learn via RL the selector of such policies. In this way, one both simplifies the problem (increasing sample efficiency) and allows the skills to be reused in different tasks.  The authors propose a graphical model that consists, per time step, of four variables: $x_t$, $y_t$ $z_t$ and $a_t$. $x_t$ represents the MDP's state, $y_t$ is the discrete selector of the skill, $z_t$ the (continuous) parameters of the skill, and $a_t$ is the output at time $t$ of the selected skill (and can represent, for example, the joint velocities of the robot).  After learning via variational inference the probabilistic model, the authors propose to refine the learned behavior via MPO, learning the skill selector $p(y_t | y_{t-1},x_t)$, the skill parameters $p(z_t | y_t, x_t)$, and the skill p(a_t | z_t, y_t).  Further, the authors propose 1) to use an asymmetric information approach, where each level of the described hierarchy sees a different definition of the state. In particular, the lower layer sees the robotic state (proprioception), the mid-layer sees the position of objects in the scene, and the highest level sees both the information 2) the models use "gated"-heads: the categorical variable $y_t$ is fed to the last part of the neural network.   Once the hierarchy has been learned, one can tune it on the same or different tasks. In particular, the lowest layer is kept fixed, while only the highest, or the highest, and the mid-layer are learned.   The authors propose to use MPO as a reinforcement learning algorithm. They perform a few ablation studies and experiments to empirically prove the efficacy of their method.