Offline Meta Reinforcement Learning -- Identifiability Challenges and Effective Data Collection Strategies

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 training an agent on a large set of training environments [ 5 , 9 ] . 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 ‘ identify ’ the domain [ 5 , 28 , 18 ] , or by learning a parameter initialization that leads to good performance within a few gradient steps [ 9 ] . Another formulation of quick RL is Bayesian RL [ BRL , 11 ] . 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 . [ 36 ] 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 RL methods , by adding to the state the belief over the environment parameters . The VariBAD algorithm [ 36 ] implements this idea using a variational autoencoder ( VAE ) for belief estimation and deep neural network policies . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . 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 [ 7 , 21 ] . In this work we investigate the offline approach to meta-RL ( OMRL ) . Any offline RL approach is heavily influenced by the data collection policy . To ground our investigation , we focus on the following practical setting:1 we assume that data is collected by running standard RL agents on a set of environments from the environment distribution . While the data was not specifically collected for the meta-RL task , 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 method for 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 approach . 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 and practical algorithm . The offline setting , however , brings about another challenge – when the agent visits different parts of the state space in different environments , learning to identify the correct environment and obtain an accurate belief estimate becomes challenging , a problem we term MDP ambiguity . We formalize this problem , and discuss how it manifests in common scenarios such as sparse rewards or sparse differences in transitions . Based on our formalization , we propose a general data collection strategy that can mitigate the problem . Furthermore , when ambiguity is only due to reward differences , we show that a simple reward relabelling trick suffices , without changing data collection . We collectively term our data collection/relabelling and off-policy algorithm as Bayesian Offline Reinforcement Learning ( BOReL ) . In our experiments , we show that BOReL learns effective exploration policies from offline data on both discrete and continuous control problems . We demonstrate significantly better exploration than meta-RL methods based on Thompson sampling such as PEARL [ 28 ] , even when these methods are allowed to train online . Furthermore , we explore the issue of MDP ambiguity in practice , and demonstrate that , when applicable , our proposed solutions successfully mitigate it . An important implication of our study is that without suitably collected data , MDP ambiguity can make learning an effective offline BRL solution impossible . This stands in contrast to recent studies in offline ( non-Bayesian ) RL [ 30 , 10 , 19 ] , where the focus is on designing learning algorithms for arbitrarily collected data , typically by avoiding actions that were not explored enough . Outside offline RL , however , it is common to develop data collection methods based on formal limitations of the problem . For example , the use of randomized controlled trials in medical treatments is well motivated by the theory of causal inference [ 27 ] . The data collection methods we propose here are simple and practical to implement , and as we demonstrate , effectively handle MDP ambiguity . 1The theory and algorithms we develop , however , are not limited to any particular data collection protocol . Our main contributions are therefore 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 off-policy RL ; we formulate MDP ambiguity , which characterizes which problems are solvable under the offline BRL setting , and based on this formulation we propose several principled data collection strategies ; we show non-trivial empirical results that demonstrate significantly better exploration than meta-RL methods based on Thompson sampling ; finally , and of independent interest , our off-policy algorithm significantly improves the sample efficiency of conventional VariBAD in the online setting . 2 Background . We recapitulate meta-RL , BRL and the VariBAD algorithm . Meta-RL : In meta-RL , a distribution over tasks is assumed . A task Ti is described by a Markov Decision Process [ MDP , 2 ] Mi = ( S , A , Ri , Pi ) , where the state space S and the action space A are shared across tasks , and Ri 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 ) . 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 parametersR , P ∼ p ( R , P ) .2 Key here 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 , 3 ] , reducing the problem to solving a particular POMDP instance where the unobserved variables do not change in time . 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 , the state in BRL can be augmented with the belief to result in the Bayes-Adaptive MDP model [ BAMDP , 6 ] . 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 : ∀s+ ∈ S+ , a ∈ A we have that Q ( s+ , a ) = R+ ( s+ , a ) + γEs+′∼P+ [ max a′ Q ( s+ ′ , a′ ) ] . ( 2 ) Computing a Bayes-optimal agent amounts to solving the BAMDP , where the optimal policy is a function of the augmented state . 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 [ 36 ] 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 2For ease of presentation , we consider the infinite horizon discounted return . Our formulation easily extends to the episodic and finite horizon settings , as considered in our experiments . parametric generative model for the state-reward 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 . [ 36 ] 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 . While other neural belief representations could be used [ 15 ] , we chose to focus on VariBAD for concreteness . To approximately solve the BAMDP , [ 36 ] 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 ) . The policies are trained using policy gradients , optimizing J ( π ) = ER , PEπ [ H∑ t=0 γtr ( st , at ) ] . ( 3 ) 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 is doing quite a few things at once:  1. Proposing an algorithm for learning effective exploration policies for meta-RL from only offline data 2. Defining the offline meta-RL problem (not technically new but concurrent with other recent work) 3. Defining the MDP ambiguity problem, which proves troublesome for inference-based meta-RL algorithms, and providing two baseline approaches to solving it  The authors accomplish 1 by converting the VariBAD to work with off-policy data and 3 with intuitive strategies for additional data collection or augmentation, which provably address the problem. The experiments suggest that the proposed solutions are necessary to solve problems for which identifying the MDP is difficult.