Latent Skill Planning for Exploration and Transfer

1 INTRODUCTION Humans can effortlessly compose skills , where skills are a sequence of temporally correlated actions , and quickly adapt skills learned from one task to another . In order to build re-usable knowledge about the environment , Model-based Reinforcement Learning ( MBRL ) ( Wang et al. , 2019 ) provides an intuitive framework which holds the promise of training agents that generalize to different situations , and are sample efficient with respect to number of environment interactions required for training . For temporally composing behaviors , hierarchical reinforcement learning ( HRL ) ( Barto & Mahadevan , 2003 ) seeks to learn behaviors at different levels of abstraction explicitly . A simple approach for learning the environment dynamics is to learn a world model either directly in the observation space ( Chua et al. , 2018 ; Sharma et al. , 2019 ; Wang & Ba , 2019 ) or in a latent space ( Hafner et al. , 2019 ; 2018 ) . World models summarize an agent ’ s experience in the form of learned transition dynamics , and reward models , which are used to learn either parametric policies by amortizing over the entire training experience ( Hafner et al. , 2019 ; Janner et al. , 2019 ) , or perform online planning as done in Planet ( Hafner et al. , 2018 ) , and PETS ( Chua et al. , 2018 ) . Amortization here refers to learning a parameterized policy , whose parameters are updated using samples during the training phase , and which can then be directly queried at each state to output an action , during evaluation . Fully online planning methods such as PETS ( Chua et al. , 2018 ) only learn the dynamics ( and reward ) model and rely on an online ∗Kevin and Homanga contributed equally to this work . search procedure such as Cross-Entropy Method ( CEM ; Rubinstein , 1997 ) on the learned models to determine which action to execute next . Since rollouts from the learned dynamics and reward models are not executed in the actual environment during training , these learned models are sometimes also referred to as imagination models ( Hafner et al. , 2018 ; 2019 ) . Fully amortized methods such as Dreamer ( Hafner et al. , 2019 ) , train a reactive policy with many rollouts from the imagination model . They then execute the resulting policy in the environment . The benefit of the amortized method is that it becomes better with experience . Amortized policies are also faster . An action is computed in one forward pass of the reactive policy as opposed to the potentially expensive search procedure used in CEM . Additionally , the performance of the amortized method is more consistent as CEM relies on drawing good samples from a random action distribution . On the other hand , the shortcoming of the amortized policy is generalization . When attempting novel tasks unseen during training , CEM will plan action sequences for the new task , as per the new reward function while a fully amortized method would be stuck with a behaviour optimized for the training tasks . Since it is intractable to perform fully online random shooting based planning in high-dimensional action spaces ( Bharadhwaj et al. , 2020 ; Amos & Yarats , 2019 ) , it motivates the question : can we combine online search with amortized policy learning in a meaningful way to learn useful and transferable skills for MBRL ? To this end , we propose a partially amortized planning algorithm that temporally composes high-level skills through the Cross-Entropy Method ( CEM ) ( Rubinstein , 1997 ) , and uses these skills to condition a low-level policy that is amortized over the agent ’ s experience . Our world model consists of a learned latent dynamics model , and a learned latent reward model . We have a mutual information ( MI ) based intrinsic reward objective , in addition to the predicted task rewards that are used to train the low level-policy , while the high level skills are planned through CEM using the learned task rewards . We term our approach Learning Skills for Planning ( LSP ) . The key idea of LSP is that the high-level skills are able to abstract out essential information necessary for solving a task , while being agnostic to irrelevant aspects of the environment , such that given a new task in a similar environment , the agent will be able to meaningfully compose the learned skills with very little fine-tuning . In addition , since the skill-space is low dimensional , we can leverage the benefits of online planning in skill space through CEM , without encountering intractability of using CEM for planning directly in the higher dimensional action space and especially for longer time horizons ( Figure 1 ) . In summary , our main contributions are developing a partially amortized planning approach for MBRL , demonstrating that high-level skills can be temporally composed using this scheme to condition low level policies , and experimentally demonstrating the benefit of LSP over challenging locomotion tasks that require composing different behaviors to solve the task , and benefit in terms of transfer from one quadruped locomotion task to another , with very little adaptation in the target task . 2 BACKGROUND . We discuss learning latent dynamics for MBRL , and mutual information skill discovery , that serve as the basic theoretical tools for our approach . 2.1 LEARNING LATENT DYNAMICS AND BEHAVIORS IN IMAGINATION . Latent dynamics models are special cases of world models used in MBRL , that project observations into a latent representation , amenable for planning ( Hafner et al. , 2019 ; 2018 ) . This framework is general as it can model both partially observed environments where sensory inputs can be pixel observations , and fully observable environments , where sensory inputs can be proprioceptive state features . The latent dynamics models we consider in this work , consist of four key components , a representation module pθ ( st|st−1 , at−1 , ot ) and an observation module qθ ( ot|sT ) that encode observations and actions to continuous vector-valued latent states st , a latent forward dynamics module qθ ( st|st−1 , at−1 ) that predicts future latent states given only the past states and actions , and a task reward module qθ ( rt|st ) , that predicts the reward from the environment given the current latent state . To learn this model , the agent interacts with the environment and maximizes the following expectation under the dataset of environment interactions D = { ( ot , at , rt ) } J .= ED ( ∑ t ( J tO + J tR + J tD ) ) + const J tO . = ln q ( ot | st ) J tR . = ln q ( rt | st ) J tD . = −βKL ( p ( st | st−1 , at−1 , ot ) ∥∥ q ( st | st−1 , at−1 ) ) . ( 1 ) For optimizing behavior under this latent dynamics model , the agent rolls out trajectories in imagination and estimates the value V ( · ) of the imagined trajectories { sτ , aτ , rτ } t+Hτ=t through TD ( λ ) estimates as described by Sutton & Barto ( 2018 ) ; Hafner et al . ( 2019 ) . The agent can either learn a fully amortized policy qφ ( a|s ) as done in Dreamer , by backpropagating through the learned value network vψ ( · ) or plan online through CEM , for example as in Planet . 2.2 MUTUAL INFORMATION SKILL DISCOVERY . Some methods for skill discovery have adopted a probabilistic approach that uses the mutual information between skills and future states as an objective ( Sharma et al. , 2019 ) . In this approach , skills are represented through a latent variable z upon which a low level policy π ( a|s , z ) is conditioned . Given the current state s0 , skills are sampled from some selection distribution p ( z|s0 ) . The skill conditioned policy is executed under the environment dynamics pd ( st+1|st , a ) resulting in a series of future states abbreviated s′ : = { s } . Mutual information is defined as : MI ( z , { s } |s0 ) = H ( z|s0 ) −H ( z| { s } , s0 ) = H ( { s } |s0 ) −H ( { s } |s0 , z ) It quantifies the reduction in uncertainty about the future states given the skill and vice versa . By maximizing the mutual information with respect to the low level policy , the skills are encouraged to produce discernible future states . 3 PARTIAL AMORTIZATION THROUGH HIERARCHY . Our aim is to learn behaviors suitable for solving complex control tasks , and amenable to transfer to different tasks , with minimal fine-tuning . To achieve this , we consider the setting of MBRL , where the agent builds up re-usable knowledge of the environment dynamics . For planning , we adopt a partial amortization strategy , such that some aspects of the behavior are re-used over the entire training experience , while other aspects are learned online . We achieve partial amortization by forming high level latent plans and learning a low level policy conditioned on the latent plan . The three different forms of amortization in planning are described visually through probabilistic graphical models in Figure 2 and Figure 3 . We first describe the different components of our model , motivate the mutual information based auxiliary objective , and finally discuss the complete algorithm . World model . Our world model is a latent dynamics model consisting of the components described in section 2 . Low level policy . The low-level policy qφ ( at|st , z ) is used to decide which action to execute given the current latent state st and the currently active skill z . Similar to Dreamer ( Hafner et al. , 2019 ) , we also train a value model vψ ( st ) to estimate the expected rewards the action model achieves from each state st. We estimate value the same way as in equation 6 of Dreamer , balancing bias and variance . The action model is trained to maximize the estimate of the value , while the value model is trained to fit the estimate of the value that alters as the action model is updated , as done in a typical actor-critic setup ( Konda & Tsitsiklis , 2000 ) . High level skills . In our framework high level skills are continuous random variables that are held for a fixed number K steps . The high-level skills z are sampled from a skill selection distribution p ( z1 : dH/Ke|ζ ) = N ( µ , Σ ) which is optimized for task performance through CEM . Here , H denotes the planning horizon . For the sake of notational convenience we denote z1 : dH/Ke as z . Let ( j ) denote the jth CEM iteration . We first sample G skills { z ( g ) } Gg=1 ∼ p ( z|ζ ( j ) ) , execute G parallel imaginary rollouts of horizon H in the learned model with the skill-conditioned policy qφ ( at|st , z ( g ) ) . Instead of evaluating rollouts based only on the sum of rewards , we utilize the value network and compute value estimates { Vg } Gg=1 . We sort { Vg } Gg=1 , choose the top M values , and use the corresponding skills to update the sampling distribution parameters as ζ ( j+1 ) = ( µ ( j+1 ) , Σ ( j+1 ) ) µ ( j+1 ) = Mean ( { z ( m ) } Mm=1 ) Σ ( j+1 ) = Variance ( { z ( m ) } Mm=1 )

Summary: The paper proposes combining model-based RL with high-level skill learning and composition through hierarchical RL, into a single reinforcement learning framework. More specifically, the proposed approach leverages planning and composing skills in the low-dimensional, high-level representation, and learn low-level skills conditioned on the high-level skills. Only the low-level policies are executed in the environment to generate experiences. A mutual information objective is used to learn low-level policies conditioned on high-level skills, and this was shown to improve sample efficiency as the low-level policies do not learn to ignore the high-level skills they are conditioned on.