Offline Meta Learning of Exploration

1 INTRODUCTION . A central question in reinforcement learning ( RL ) is how to learn quickly ( i.e. , with few samples ) in a new environment . Meta-RL addresses this issue by assuming a distribution over possible environments , and having access to a large set of environments from this distribution during training ( Duan et al. , 2016 ; Finn et al. , 2017 ) . Intuitively , the meta-RL agent can learn regularities in the environments , which allow quick learning in any environment that shares a similar structure . Indeed , recent work demonstrated this by training memory-based controllers that can ‘ identify ’ the domain ( Duan et al. , 2016 ; Rakelly et al. , 2019 ; Humplik et al. , 2019 ) , or by learning a parameter initialization that can lead to good performance with only a few gradient updates ( Finn et al. , 2017 ) . Another approach to quick RL is Bayesian RL ( Ghavamzadeh et al. , 2016 , BRL ) . In BRL , the environment parameters are treated as unobserved variables , with a known prior distribution . Consequentially , the standard problem of maximizing expected returns ( taken with respect to the posterior distribution ) explicitly accounts for the environment uncertainty , and its solution is a Bayes-optimal policy , wherein actions optimally balance exploration and exploitation . Recently , Zintgraf et al . ( 2020 ) showed that meta-RL is in fact an instance of BRL , where the meta-RL environment distribution is simply the BRL prior . Furthermore , a Bayes-optimal policy can be trained using standard policy gradient methods , simply by adding to the state the posterior belief over the environment parameters . The VariBAD algorithm ( Zintgraf et al. , 2020 ) is an implementation of this approach that uses a variational autoencoder ( VAE ) for parameter estimation and deep neural network policies . Most meta-RL studies , including VariBAD , have focused on the online setting , where , during training , the meta-RL policy is continually updated using data collected from running it in the training environments . In domains where data collection is expensive , such as robotics and healthcare to name a few , online training is a limiting factor . For standard RL , offline ( a.k.a . batch ) RL mitigates this problem by learning from data collected beforehand by an arbitrary policy ( Ernst et al. , 2005 ; Levine et al. , 2020 ) . In this work we investigate the offline approach to meta-RL ( OMRL ) . It is well known that any offline RL approach is heavily influenced by the data collection policy . To ground our investigation , we focus on the following practical setting : we assume that data has been collected by running standard RL agents on a set of environments from the environment distribution . Importantly - we do not allow any modification to the RL algorithms used for data collection , and the meta-RL learner must make use of data that was not specifically collected for the meta-RL task . Nevertheless , we hypothesize that regularities between the training domains can still be learned , to provide faster learning in new environments . Figure 1 illustrates our problem : in this navigation task , each RL agent in the data learned to find its own goal , and converged to a behavior that quickly navigates toward it . The meta-RL agent , on the other hand , needs to learn a completely different behavior that effectively searches for the unknown goal position . Our key idea to solving OMRL is an off-policy variant of the VariBAD algorithm , based on replacing the on-policy policy gradient optimization in VariBAD with an off-policy Q-learning based method . This , however , requires some care , as Q-learning applies to states of fully observed systems . We show that the VariBAD approach of augmenting states with the belief in the data applies to the off-policy setting as well , leading to an effective algorithm we term Off-Policy VariBAD . The offline setting , however , brings about another challenge – when the agent visits different parts of the state space in different environments , it becomes challenging to obtain an accurate belief estimate , a problem we term MDP ambiguity . When the ambiguity is due to differences in the reward between the environments , we propose a simple solution , based on a reward relabelling trick that significantly improves the performance of the VariBAD VAE trained on offline data . Our experimental results show that our method can learn effective exploration policies from offline data on both discrete and continuous control problems . Our main contributions are as follows . To our knowledge , this is the first study of meta learning exploration in the offline setting . We provide the necessary theory to extend VariBAD to offpolicy RL . We show that a key difficulty in OMRL is MDP ambiguity , and propose an effective solution for the case where tasks differ in their rewards . We show non-trivial empirical results that demonstrate significantly better exploration than meta-RL methods based on Thompson sampling such as PEARL ( Rakelly et al. , 2019 ) , even when these methods are allowed to train online . Finally , our method can also be applied in the online setting , and demonstrates significantly improved sample efficiency compared to conventional VariBAD . 2 BACKGROUND . Our work leverages ideas from meta-RL , BRL and the VariBAD algorithm , as we now recapitulate . Meta-RL : In meta-RL , a distribution over tasks is assumed . A task Ti is described by a Markov Decision Process ( MDP , Bertsekas , 1995 ) Mi = ( S , A , Ri , Pi ) , where the state space S and the action spaceA are shared across tasks , andRi and Pi are task specific reward and transition functions . Thus , we write the task distribution as p ( R , P ) . For simplicity , we assume throughout that the initial state distribution Pinit ( s0 ) is the same for all MDPs . The goal in meta-RL is to train an agent that can quickly maximize reward in new , unseen tasks , drawn from p ( R , P ) . To do so , the agent must leverage any shared structure among tasks , which can typically be learned from a set of training tasks . Bayesian Reinforcement Learning : The goal in BRL is to find the optimal policy π in an MDP , when the transitions and rewards are not known in advance . Similar to meta-RL , we assume a prior over the MDP parameters p ( R , P ) , and seek to maximize the expected discounted return , Eπ [ ∞∑ t=0 γtr ( st , at ) ] , ( 1 ) where the expectation is taken with respect to both the uncertainty in state-action transitions st+1 ∼ P ( ·|st , at ) , at ∼ π , and the uncertainty in the MDP parameters R , P ∼ p ( R , P ) .1 A key observation is that this formulation naturally accounts for the exploration/exploitation tradeoff – an optimal agent must plan its actions to reduce uncertainty in the MDP parameters , if such leads to higher rewards . One way to approach the BRL problem is to modelR , P as unobserved state variables in a partially observed MDP ( POMDP , Cassandra et al. , 1994 ) , reducing the problem to solving a particular POMDP instance where the unobserved variables can not change in time ( P andR are assumed to be stationary ) . The belief at time t , bt , denotes the posterior probability overR , P given the history of state transitions and rewards observed until this time bt = P ( R , P|h : t ) , where h : t = { s0 , a0 , r1 , s1 . . . , rt , st } ( note that we denote the reward after observing the state and action at time t as rt+1 = r ( st , at ) ) . The belief can be updated iteratively according to Bayes rule , where b0 ( R , P ) = p ( R , P ) , and : bt+1 ( R , P ) = P ( R , P|h : t+1 ) ∝ P ( st+1 , rt+1|h : t , R , P ) bt ( R , P ) . Similar to the idea of solving a POMDP by representing it as an MDP over belief states ( Cassandra et al. , 1994 ) , the state in BRL can be augmented with the belief to result in the Bayes-Adaptive MDP ( BAMDP ) model ( Duff & Barto , 2002 ) . Denote the augmented state s+t = ( st , bt ) and the augmented state space S+ = S×B , where B denotes the belief space . The transitions in the BAMDP are given by : P+ ( s+t+1|s + t , at ) = Ebt [ P ( st+1|st , at ) ] δ ( bt+1 = P ( R , P|h : t+1 ) ) , and the reward in the BAMDP is the expected reward with respect to the belief : R+ ( s+t , at ) = Ebt [ R ( st , at ) ] . The Bayes-optimal agent seeks to maximize the expected discounted return in the BAMDP , and the optimal solution of the BAMDP gives the optimal BRL policy . As in standard MDPs , the optimal action-value function in the BAMDP satisfies the Bellman equation : Q ( s+ , a ) = R+ ( s+ , a ) + γ ∑ s+ ′∈S+ P+ ( s+ ′ |s+ , a ) max a′ Q ( s+ ′ , a′ ) , ∀s+ ∈ S+ , a ∈ A . ( 2 ) Computing a Bayes-optimal agent amounts to solving the BAMDP , where the optimal policy is a function of the augmented state . However , for most problems this is intractable , as the augmented state space is continuous and high-dimensional , and the posterior update is also intractable in general . The VariBAD Algorithm : VariBAD ( Zintgraf et al. , 2020 ) approximates the Bayes-optimal solution by combining a model for the MDP parameter uncertainty , and an optimization method for the corresponding BAMDP . The MDP parameters are represented by a vector m ∈ Rd , corresponding to the latent variables in a parametric generative model for the statereward trajectory distribution conditioned on the actions P ( s0 , r1 , s1 . . . , rH , sH |a0 , . . . , aH−1 ) =∫ pθ ( m ) pθ ( s0 , r1 , s1 . . . , rH , sH |m , a0 , . . . , aH−1 ) dm . The model parameters θ are learned by a variational approximation to the maximum likelihood objective , where the variational approximation to the posterior P ( m|s0 , r1 , s1 . . . , rH , sH , a0 , . . . , aH−1 ) is chosen to have the structure qφ ( m|s0 , a0 , r1 , s1 . . . , rt , st ) = qφ ( m|h : t ) . That is , the approximate posterior is conditioned on the history up to time t. The evidence lower bound ( ELBO ) in this case is ELBOt = Em∼qφ ( ·|h : t ) [ log pθ ( s0 , r1 , s1 . . . , rH , sH |m , a0 , . . . , aH−1 ) ] −DKL ( qφ ( m|h : t ) ||pθ ( m ) ) . The main claim of Zintgraf et al . ( 2020 ) is that qφ ( m|h : t ) can be taken as an approximation of the belief bt . In practice , qφ ( m|h : t ) is represented as a Gaussian distribution q ( m|h : t ) = N ( µ ( h : t ) , Σ ( h : t ) ) , where µ and Σ are learned recurrent neural networks . To approximately solve the BAMDP , Zintgraf et al . ( 2020 ) exploit the fact that an optimal BAMDP policy is a function of the state and belief , and therefore consider neural network policies that take the augmented BAMDP state as input π ( at|st , qφ ( m|h : t ) ) , where the posterior is practically represented by the distribution parameters µ ( h : t ) , Σ ( h : t ) . To train such policies , Zintgraf et al . ( 2020 ) maximize the BRL objective , J ( π ) = ER , PEπ [ H∑ t=0 γtr ( st , at ) ] , ( 3 ) 1For ease of presentation , we consider the infinite horizon discounted return . Our formulation easily extends to the episodic and finite horizon settings . using policy gradient based methods . The expectation over MDP parameters in ( 3 ) is approximated by averaging over training environments , and the RL agent is trained online , alongside the VAE .

Summary: The paper studies the problem of offline Meta-Reinforcement Learning (RL). In this problem, N separate RL environments are considered which are drawn from a specific underlying distribution. For each such environment, M trajectories each of horizon H are provided beforehand. The task is to train an RL agent that performs well in expectation on a new RL environment drawn from the same distribution.  The authors adapt the online VariBAD algorithm (Zintgraf et al., 2020) by the use of a techniques called state relabeling and reward relabeling.