PERIL: Probabilistic Embeddings for hybrid Meta-Reinforcement and Imitation Learning

1 INTRODUCTION . Reinforcement Learning ( RL ) and Imitation Learning ( IL ) are two popular approaches for teaching an agent , such as a robot , a new task . However , in their standard form , both require a very large amount of data to learn : exploration in the case of RL , and demonstrations in the case of IL . In recent years , meta-RL and meta-IL have emerged as promising solutions to this , by leveraging a meta-training dataset of tasks to learn representations which can quickly adapt to this new data . However , both these methods have their own limitations . Meta-RL typically requires hand-crafted , shaped reward functions to describe each new task , which is tedious and not practical for nonexperts . A more natural way to describe a task is to provide demonstrations , as with meta-IL . But after adaptation , these methods can not continue to improve the policy in the way that RL methods can , and are restricted by the similarity between the new task and the meta-training dataset . A third limitation , which both methods can suffer from , is that defining a low-dimensional representation of the environment for efficient learning ( as opposed to learning directly from high-dimensional images ) , requires hand-crafting this representation . As with rewards , this is not practical for nonexperts , but more importantly , it does not allow generalisation across different task families , since each task family would require its own unique representation and dimensionality . In this work , we propose a new method , PERIL , which addresses all three of these limitations , in a hybrid framework that combines the merits of both RL and IL . Our method allows for tasks to be defined using demonstrations only as with IL , but upon adaptation to the demonstrations , it also allows for continual policy improvement through further exploration of the task . Furthermore , we define the state representation using only the agent ’ s internal sensors , such as position encoders of a robot arm , and through interaction we implicitly recover the state of external environment , such as the poses of objects which the robot is interacting with . Overall , this framework allows for learning of new tasks without requiring any expert knowledge in the human teacher . PERIL operates by implicitly representing a task by a latent vector , which is predicted during learning of a new task , through two means . First , during meta-training we encourage high mutual information between the demonstration data and the latent space , which during testing forms a prior for the latent space and represents the agent ’ s task belief from demonstrations alone . Second , we allow further exploration to continually update the latent space , by conditioning on the robot ’ s states and actions during exploration of this new task . We model the latent space via an encoder , from which posterior sampling can be done to encode the agent ’ s current task belief . In essence , the encoder aims to learn an embedding which can simultaneously ( i ) infer the task intent , and ( ii ) output a policy which can solve the inferred task . During meta-training , PERIL is optimised end-to-end , by simultaneously learning both a policy and the embedding function upon which the policy is conditioned . In our experiments we find PERIL achieves exceptional adaptation rates and is capable of exploiting demonstrations to efficiently explore unseen tasks . Through structured exploration , PERIL outperforms other Meta-RL and Meta-IL baselines and is capable of zero-shot learning . We show how our method is capable of multi-family meta-learning as well as out-of-family meta-learning by clustering distinct meta-trained latent space representations . As an extension , we also show how to use privileged information during training to create an auxiliary loss for training the embedding function , which helps to form a stronger relationship between the latent space and the true underlying state which defines the task . Supplementary videos are available at our anonymous webpage https : //sites.google.com/view/peril-iclr-2021 2 RELATED WORK . Meta-RL was conceptualised as an RNN-based task . Developed by Wang ( 2016 ) and Duan et al. , the authors use RNNs to feed a history of transitions into the model such that the policy can internalise the dynamics . On another line , Finn et al . ( 2017 ) developed a learning-to-learn strategy , model agnostic meta learning ( MAML ) , which meta-learns an initialisation that adapts the parameters of the policy network and fine tunes it during meta-testing . Although promising results in simple goalfinding tasks , MAML-based methods fail to produce efficient stochastic exploration policies and adapt to complex tasks ( Gupta et al. , 2018 ) . Meta-learning robust exploration strategies is key in order to improve sample efficiency and allow fast adaptation at test time . In light of this , contextbased RL was developed with the aim of reducing the uncertainty of newly explored tasks . These map transitions τ collected from an unseen task into a latent space z via an encoder qφ ( z|τ ) , such that the conditioned policy πθ ( τ |z ) can efficiently solve said task ( Rakelly et al . ; Wang & Zhou , 2020 ) . An underlining benefit of decoupling task encoding from the policy is that it disentangles task inference from reward maximisation , whilst gradient-based and RNN-based meta-RL policies do this internally . An important remark regarding training conditions of the discussed meta-RL methods is that they are typically meta-trained using dense reward functions ( Rakelly et al . ; Wang , 2016 ; Wang & Zhou , 2020 ) . These dense reward functions provide information-rich contexts of the unseen task . Considering that the ultimate goal is to allow agents to solve new tasks in the real world , adaptation during test time must be robust to sparse reward feedback ( Schoettler et al. , 2020 ) . In the context of RL , incorporating demonstrations has proven successful in aiding exploration strategies , stabilising learning and increasing sample efficiency ( Vecerik et al. , 2017 ) . Learning expressive policies from a set of demonstrations requires a vast amount of expert trajectories ( Ross et al . ) , particularly in high dimensional state-action spaces ( Rajeswaran et al. ) . In contrast , Meta-IL can be implemented in meta-RL by conditioning the agent with expert trajectories . On the another hand , Zhou et al . propose a MAML-based meta-IL approach which averages objective across demonstrations . The latter learns to adapt at test-time by receiving binary rewards . A similar strategy was also developed by and Mendonca et al .. The caveats of these approaches remain that of traditional IL : ( i ) Imitating expert trajectories hinders the policy from doing better than the demonstrations ; ( ii ) Cloning behaviours reduces flexibility and generalisation capacity . 3 HYBRID META-REINFORCEMENT AND IMITATION LEARNING . 3.1 PROBLEM STATEMENT . We assume access to a set of tasks T ∈ p ( T ) , with each task represented as a partially observable Markov decision process ( POMDP ) . Observations o are incomplete and only include measurements of the agent ’ s internal state . For example , a peg-in-hole task would include the robot ’ s ( and therefore the peg ’ s ) pose in the observation , but not the hole ’ s pose . The full state s contains both the robot ’ s internal state , and the state of the external environment . Our method involves an agent inferring z , a descriptor of a task which provides the information required to solve that task , such as a representation of s. Each task is then defined as T = { p ( o0|z ) , p ( o′ , z′|o , a , z ) , r ( o , a|z ) } , with an unknown initial observation distribution p ( o0|z ) , transition distribution p ( o′ , z′|o , a , z ) , and reward function r ( o , a|z ) , where a is the action taken by the agent . By leveraging task beliefs , PERIL exploits enriched observational spaces ( s ∼ { o ∩ z } ) with the objective of closing said POMDP into a stable MDP form . With the aim of supporting continuous task inference through interaction , we construct z as a probabilistic embedding conditioned on a set of recently collected transitions , referred to as the context c. We define a contexton cTt = ( ot , at , rt , ot+1 ) as a transition collected on task T at time-step t such that the context cT0 : t ( what we call c ) denotes the set of accumulated contextons . During meta-testing , the agent then attempts to find the true posterior p ( z|c ) over the task belief z ∼ p ( z|c ) , by conditioning on c. Demonstrations significantly reduce the search space in exploration whilst providing a natural means of communicating the task . Thus , access to expert trajectories provides information-rich context which can be exploited to pre-condition the policy via a prior over z . Subsequently , online adaptation through RL can further disambiguate z . In our approach , we leverage dual meta-learning objectives to perform hybrid adaptation to new tasks , with conditioning on both demonstration contexts and exploratory contexts , where we define a demonstration as a trajectory of T observations and actions d = { ( ot , at ) } Tt=1 , and an adaptation episode τe as a trajectory of N steps τe = { ( ot , at , rt , zt ) } Nt=1 : Primal Inference : The agent observes a set of k = 1 , ... , K demonstrations from an expert set Ddemo : = { dk } to form a context prior cdemo = { dk } Kk . Exploratory Adaptation : The agent explores with initialised context c ←− cdemo and adapts with sampled trajectories τ such that at time-step t , a hybrid context is formed c←− cdemo ∩ τ0 : t. Primal task inference aims to provide long-term inference on a task belief : meta-trained inference from previously solved tasks . Using posterior sampling , exploratory adaptation provides a proxy for short-term memory : as the context is updated with recent transitions , structured exploration can adjust primal inference by reasoning about the uncertainty of z , where c : = { cdemo∩ ( τTe , ... , τ∞Te ) } . 3.2 PROBABILISTIC EMBEDDINGS FOR META-RL . As we do not have access to the posterior p ( z|c ) , we use variational inference methods to produce an approximation qφ ( z|c ) . Through generative processes , we sample z ∼ qφ ( z|c ) and optimise φ by maximising a meta-objective conditioned on z. Variational Encoders ( VEs ) are capable of producing latent distributions which optimise an arbitrary task-dependent objective G ( T |· ) . We average the expectancy over a set of tasks from p ( T ) to formulate a meta-objective which rewards fast adaptation to unseen tasks ( Eq . 1 ) , where β controls the constraint for mutual information between the inferred variable z and context c. This information bottleneck filters down redundant information and compresses the z into a generalisable form . We set the prior p ( z ) to a unit Gaussian . φ∗ = argmax φ ET ∈p ( T ) [ Ez∼qφ ( z|cT ) [ G ( T |z ) + βDKL [ qφ ( z|cT ) ||p ( z ) ] ] ( 1 ) Traditional meta-RL methods leverage RNN-based inference systems to produce latent features from a history of recently collected transitions . The problem with this approach is that learning from entire trajectories leads to massive variances which hinder the learning process ( Duan et al . ; Humplik et al. , 2019 ) . In theory recurrence is not strictly required , since any MDP is defined by a distribution of sequentially invariant transitions . In light of this , we exploit task encoders as suggested by Rakelly et al. , where q ( z|c ) is modelled as a product of Gaussian Factors . Each factor ψφ ( z|cn ) parameterised by φ , is independently computed : ψφ ( z|cn ) = N ( fµψ ( cn ) , fσψ ( cn ) ) . We build on top of probabilistic meta-RL methods derived by Rakelly et al .. Decoupling task inference from task completion allows efficient off-policy objectives to be optimised . Moreover , conditioning a policy with diverse task descriptors aims to generalise behaviours amongst different dynamics and rewards . This is generally not the case when employing MAML-based or RNN-based meta-training kernels to ” learn how to learn ” new tasks , since policy gradients modify the policy parameters online which , in turn , can result in highly unstable and inefficient adaptation to new tasks , requiring up to hundreds of sample trajectories in unseen environments . To that end , we employ meta-adapted maximum entropy RL , where a proxy for the policy π ( a|s ) is defined by including the task belief z as part of the observational space πθ ( a|o , z ) , resulting in taskconditioned critic LTcritic and actor LTactor losses for the the Q-function Qθ ( · ) and the policy πθ ( · ) respectively ( see A.1 for more details ) . As VEs are a form of generative processes , G ( T |· ) can be defined with the aim of recovering transition dynamics or reconstructing state spaces ( Humplik et al. , 2019 ) . Instead , we optimise qφ ( z|c ) in a model-free manner as proposed in Rakelly et al .. Here , parameters φ can be optimised to maximise expected discounted rewards under a policy πθ ( ·|z ) through reconstruction of Qθ ( · , z ) . We fold LTcritic as a proxy for the task-dependent goal G ( T |z ) : G ( T |z ) ←− LTcritic ( z ) ( 2 )

Summary: This paper introduces PERIL, a meta RL method that combines demonstration trajectories and trajectories collected by the policy, in order to adapt to a new task. To this end, the authors combine ideas from metaRL (specifically from PEARL (Rakelly et al. 2019) and Humplik et al (2019)) where a set encoder is used to encode trajectories to a latent vector describing the task, with imitation learning techniques by (a) training this encoder also with demonstrations (b) initialising the latent vector at test time by feeding demonstrations through the encoder, and (c) having additional losses inspired by metaIL techniques. The motivation is that using demonstrations allows us to learn tasks that are difficult otherwise, for example because the rewards (at test time) are sparse.