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 work seeks to efficiently learn new tasks by combining meta-RL and imitation learning (IL). Such a combination is a natural thing to try, as both lines of work improve sample complexity of learning a new task: meta-RL by leveraging experience on prior related tasks, and IL by leveraging demonstrations. Demonstrations also form a natural way of specifying a new task to the agent.