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 studies the problem of Offline Meta-Reinforcement Learning (OMRL). In this setting, one observes M trajectories coming from N different environments sampled from the same underlying distribution. The goal of the agent is to perform well in expectation on a new unseen MDP drawn from the same distribution. To do so, the authors propose relabelling strategies to mitigate the problem of MDP ambiguity they introduce/formalize and to enable off-policy learning for the VariBad algorithm. The resulting algorithm was tested and compared with PEARL (online) and MBML (offline) on a series of environments with discrete or continuous action spaces, and sparse or dense rewards.