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 method to learn a three-leveled hierarchy of skills offline, from a dataset of demonstrations, that can then be applied to accelerate reinforcement learning. The three-level architecture is novel. It encodes a discrete selection, a continuous contextual variable (dependent on the discrete selection) and a low level policy dependant of the continuous contextual variable. The results show some improvements over baselines, especially in a sparse reward context where the presented method capitalizes from the offline learned strategies for exploration.