Offline Meta Learning of Exploration
Authors: Ron Dorfman
Aviv Tamar
Venue:  (2021)
Keywords: Meta-RL
Offline RL
Bayesian RL

Abstract: Consider the following problem: given the complete training histories of $N$ conventional RL agents, trained on $N$ different tasks, design a meta-agent that can quickly maximize reward in a new, unseen task from the same task distribution. In particular, while each conventional RL agent explored and exploited its own different task, the meta-agent must identify regularities in the data that lead to effective exploration/exploitation in the unseen task. This meta-learning problem is an instance of a setting we term Offline Meta Reinforcement Learning (OMRL). To solve our challenge, we take a Bayesian RL (BRL) view, and seek to learn a Bayes-optimal policy from the offline data. We extend the recently proposed VariBAD BRL algorithm to the off-policy setting, and demonstrate learning of approximately Bayes-optimal exploration strategies from offline data using deep neural networks. For the particular problem described above, our method learns effective exploration behavior that is qualitatively different from the exploration used by any RL agent in the data. Furthermore, we find that when applied to the online meta-RL setting (agent simultaneously collects data and improves its meta-RL policy), our method is significantly more sample efficient than the state-of-the-art VariBAD.

Under review as a conference paper at ICLR 2021
OFFLINE META LEARNING OF EXPLORATION
Anonymous authors Paper under double-blind review
ABSTRACT
Consider the following problem: given the complete training histories of N con- ventional RL agents, trained on N different tasks, design meta-agent that quickly maximize reward in a new, unseen task from the same task distribution. In particular, while each conventional RL agent explored and exploited its own different task, the meta-agent must identify regularities in the data that lead to effective exploration/exploitatinn in the unseen task. This meta- neta-learning problem is an instance of a setting we term Offline Meta Reinforcement Learning (OMRL). To solve our challenge, we take a Bayesian RL (BRL) view, and seek to learn a Bayes- optimal policy from the offline data. We extend the recently proposed VariBAD BRL. algorithm to the off-policy setting, and demonstrate learning of approximately Bayes-optimal exploration strategies from offline data using deep neural networks. For the particular problem described above, our method learns effective exploration behavior that is qualitatively different from the exploration used by any RL agent in the data. Furthermore, we find that when applied to the online meta-RL. setting (agent simultaneously collects data and improves its meta-RL policy), our method is significantly more sample efficient than the state-of-the-art VariBAD.
INTRODUCTION
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 distribution over possible environ- ments, 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 environ- ments, which allow quick learning any environment that shares a similar structure. Indeed, recent work demonstrated this by training memory-basee controllers that can 'identify' the domain (Duan) ctt al. 2016; Rakelly et al. 2019) Humplik et al.] 2019). or by learning a parameter initialization that can lead to good performance with only few gradient updates (Finn al. 2017).
Another approach to quick RL is Bayesian RL (Ghavamzadeh et al. 2016 BRL). In BRL, the environment parameters are treated unobserved variables, with known prior distribution. Conse- quentially, the standard problem of maximizing expected returns (taken with respect to the posterior distribution) explicitly accounts for the environment uncertainty, and its solution is Bayes-optimal policy, wherein ac 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, 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.) (020) is an implementation of this approach that uses a variational autoencoder (VAE) for parameter estimation and deep neural network policies.
Most RI. studies, including VariBAD, have focused on the online setting, where, during training, the meta-RI. policy is continually updated using data collected from running it the training environments. In domains where data collection is expensive, such robotics and healthcare 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 arbitrary policy (Ernst et al. 2005 Levine et al. 2020). In this work we investigate the offline approach to meta-RL 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 set of environments from the environment distribution.

---

Under review as a conference paper at ICLR 2021
Offinee Data
Mata-RL AgennR
1
Agent22
Fgiooda 2
Agent
Figure 1: Illustration of offline meta-RL: the task is navigate to goal position that can be anywhere on the semi-cirele. The reward is sparse (light-blue), and the offline data (left) contains training histories of conventional RL agents trained to find individual goals. The meta-RL. agent (right) needs to find a policy that quickly finds the unknown goal, here, by searching across the semi-circle. Note that this search behavior completely different than the dominant behaviors in the data.
Importantly 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 [
illlustrates our problem: in this navigation task, each RL agent in the data learned to find its own goal, and converged 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 well, leading to effective algorithm setting, however, brings about another challenge - when the agent visits different parts of the state space in different environments, it becomes challenging to obtain accurate belief estimate, problem we term MDP ambiguity. When the ambiguity is due to differences in the reward between the environments, we propose a simple solution, based reward relabelling trick that significantly improves the performance of the VariBAD VAE trained on offline data. Our experimental results show that method can learn effective exploration policies from offline data on both discrete and continuous control problems.
term Off-Policy VariBAD. The offline
Our main contributions exploration in the offline setting. We provide the necessary theory to extend VariBAD to off- policy RL. We show that 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.
follows. To knowledge, this is the first study of meta learning
2 BACKGROUND
Our work leverages ideas from meta-RL RL, BRL and the VariBAD algorithm, now recapitulate.
Meta-RL: meta-RL. RL, a distribution over tasks is assumed. A task T, is described by Markov Decision Process (MDP, Bertsekas 1995) M, (S, A, R, Pi). where the state space S and the action space A are shared across tasks, and R, and P. are task specific reward and transition functions. Thus, we write the task distribution as pRR P). For simplicity, we assume throughout that the initial state distribution (oo) 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 pRR, 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 not known in advance. Similar to meta-RI, we assume prior
2

---

Under review as a conference paper at ICLR 2021
over the MDP parameters p(R, P), and seek to maximize the expected discounted return,
Σ Er(r(s,.a))
(I)
where the expectation is taken with respect to both the uncertainty in state-action transitions key observation is that this formulation naturally accounts for the exploration//xploitatio tradeoff - an optimal agent must plan its actions reduce uncertainty in the MDP parameters, if such leads to higher rewards.
S141 P(-ls,, dz), dg n, and the uncertainty in the MDP parameters R, P) P)
One way to approach the BRL problem is to model R, P unobserved state variables in partially observed MDP (POMDP, Cassandra et al.) 1994). reducing the problem to solving particular POMDP instance where the unobserved variables cannot change in time (P and R are assumed to be stationary). The belief at time f, b,, denotes the posterior probability over R. P given the history of state transitions and rewards observed until this time b, P(R, Plh Plh:)), where h- {su, ag. F1, 81 ..... re, St} (note that we denote the reward after observing the state and action at time as T++1 = rss,, ax) )t)). The belief can be updated iteratively according to Bayes rule, where (R,, PR) = p(R, P), and: bo+1 (R. P) P(R. P Pllco-1) x 7( I1,, R, (R,,
[rossa/Caasaadra
Similar to the idea of solving a POMDP by representing it as an MDP over belief states (Cassandra) eet al. 1994). the state in BRL can be augmented with the belief to result in the Bayes- Adapsved MDP (BAMDP) model (Duff & Barto 2002). Denote the augmented state $t = (st, b,) and the augmented state space S+ = 8 x B, where B denotes the belief space. The transitions in the BAMDP are given by: (* ae) = Es, [P(s at)] 5 (bl+1 = P(R. P Plha+1)) and the reward in the BAMDP is the expected reward with respect to the belief: R+ .a) E,, [R(s,, at)]- The Bayes-optimal agent seeks to maximize the expected discounted return 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) +7
Σ P 1st, St'ES-
max Q(s ,a),, ss ESt est,aE E.
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.
Computing a
The VariBAD Algorithm: VariBAD (Zintgraf et al., 2020) approximates the Bayes-optimal solution by combining model for the MDP parameter uncertainty, and an optimization represented by vector m E Hdd corresponding to the latent variables in a parametric generative model for the state- reward trajectory distribution conditioned on the actions P(su, .1..., SH (a....., QH-1) J Po (m) (mm)pp((ss (Su, T1, 81 ... TH. &H |m, ay...., aH dm. The model parameters are learned by a variational approximation to the maximum likelihood objective, where the variational approxi- mation to the posterior Pmm] (s, r1, 81..,, TH, 8H, 00,..., (0,...,(H-1) 1)) is chosen to have the structure go(m|sp (a, do /1, ...,, r't; 8.)
method for the corresponding BAMDP. The MDP parameters
qo ()))))))). That is, the approximate posterior is conditioned on the history up to time t. The evidence lower bound (ELBO) in this case is ELBO, = Em (-1b,)) [log po (s00 T1, 31 .... rH, SH |m, ag. o D k (go (m (lee (m)). The main (9.
claim of Zintgraf et al. (2020) that go ( (h.) can be taken as an approximation of the belief b. In practice, (m (e.)) is represented as Gaussian distribution q (m(haa) = N ((uhha) E(ha)), where jl and Σ are learned recurrent neural networks.
To approximately solve the BAMDP, Zintgral et al. (020) exploit the fact that an optimal BAMDP policy is function of the state and belief, and therefore consider neural network policies that take the augmented BAMDP state as input π (ax (a.|st, (m/ (lah))) where the posterior is practically represented by the distribution parameters ()h.),, Ehh.). To train such policies, Zintgraf et al. (000 maximize the BRL objective,
H rr a.)
(iT) = E
Σ
(3)
For ease of presentation, we consider the infinite horizon discounted return. Our formulation easily extends the episodic and finite horizon settings.
3

---

Under review as a conference paper at ICLR 2021
using policy gradient based methods. The expectation MDP parameters in (3) is approximated by averaging over training environments, and the RL agent is trained online, alongside the VAE.
3 OMRL WITH OFF-POLICY VARIBAD
In this section, we derive off-policy variant of the VariBAD algorithm, and apply it the OMRL problem. We begin by describing our OMRL problem setting, and then present our algorithm.
3.1 OMRL PROBLEM DEFINITION
We follow the Meta-RL and BRL formulation described above, with prior distribution over MDP parameters pRR, P). We are provided training data of an agent interacting with N different MDPs, {R,, , sampled from the prior. We assume that each interaction is organized as M trajectories 1,...,,N, E I, where the rewards satisfy r (s, a) the transitions satisfy st. p:((ss) aa")). and the actions are chosen from an arbitrary data collection policy. To ground our work in specific context, we further assume that the trajectories are obtained from running conventional RL agent (i.e., the complete RL training history) in each one of the MDPs, which implicitly specifies the data collection policy. our experiments, we will later investigate the implications of this assumption. Our goal is to use the data for learning Bayes-optimal policy, i.e., a policy that maximizes Eq.
sampled of length H, 74. 50 i.j (J id S. i. i.j
4g
81
...,. , H
3.2 OFF-POLICY VARIBAD
The online VariBAD algorithm updates the policy using trajectories sampled from the current policy, and thus cannot be applied to our offline setting. Our first step is to modify VariBAD to work off-policy. We start with observation about the use of the BAMDP formulation in VariBAD, which will motivate our subsequent development.
Does VariBAD really optimize the BAMDP? Recall that a BAMDP is in fact a reduction of a POMDP to an MDP over augmented states s = (s. b). and with the rewards and transitions given by Ht and pt Thus, an optimal Markov policy for the BAMDP exists in the form of (Ts+), The VariBAD policy, as described above, similarly takes as input the augmented state, and is thus capable of representing an optimal BAMDP policy. However, VariBAD's policy optimization in Eq. 3 does not make use of the BAMDP parameters R+ and p+! While at first this may seem counterintuitive, Eq. (3) is in fact a sound objective for the BAMDP, as we now show 2
Proposition 1. Lei T 80, do: T1, .1... ..... HH,, 8H denote a random trajectory from fixed history dependent policy n, generated according to the following process. First, MDP parameters R,P drawn from the prior p(R, P). Then, the state trajectory is generated according so Pinit- a, ~ πi [-1 so, ag, FL, F:,).., St+1 ~ ((-(st, at) and T(+1 ~ R(se, Let by denote the posterior belief at time b. P(R, P|s0 0o, To,-...... .t). Then
P(stt1/50; ao; ",,,.,, St, a:) ER, P( (8t, at), and,
P(rei 180, 4o, Tl, Fl,..., 8t, 0+) = ER, R(re+ at).
For on-policy VariBAD, Proposition 1 shows that the rewards and transitions in each trajectory can be seen as sampled from a distribution that in expectation is equal to R+ and p+, and therefore maximizing Eq. 3 is valid However, off-policy RL does not take input trajectories, but tuples of the form (state, action, reward, next state), where states and actions can be sampled from any distribution. For an arbitrary distribution of augmented states, must replace the rewards and transitions in our data with R. and P+, which can be difficult to compute. Fortunately, Proposition | shows that when collecting data as described above (by sampling complete trajectories), this is not necessary, as in expectation, the rewards and transitions are correctly sampled from the BAMDP.
This result is closely related to the discussion in Ortega et al. (2019). here applied to our particular setting. To further clarify, if we could calculate Rt, replacing all rewards in the trajectories with R* will result in lower variance policy update, similar expected SARSA (Van Seijen 2009).
4

---

Under review as a conference paper at ICLR 2021
R1
R₂
Figure 2: Reward ambiguity: from the two trajectories, is impossible to know if there are two MDPs with different re- wards (blue and yellow circles), or one MDP with rewards at both locations.
Based on Proposition 1 we can use state augmentation method similar to VariBAD, which we refer to as state relabelling. Consider each trajectory in our data rdd s/6 dg defined above. Recall that the VariBAD VAE encoder provides an estimate of the belief given the state history h1.) (w(ha),, (h..))) Thus, we can run the encoder every partial t-length history Ftt" to obtain the belief at each time step. Following the BAMDP formulation, we define the augmented state si ttd (s/², ,b(""). where b = p((T E(T We next replace each state in our data s/ with St" .... effeetively transforming the data to as coming from a BAMDP.
3.J i. n ;..... , as
After applying state relabelling, any off- off-policy RL algorithm can be applied to the modified data, for learning a Bayes-optimal policy. In our experiments we used DQN (Mnih et al., 2015) for discrete action domains, and soft actor critic (SAC, Haarnoja et al. 2018) for continuous control.
3.3 OFFLINE META-RL RL AND MDP AMBIGUITY
We next consider the OMRL problem. While in principle, it is possible to simply run off-policy VariBAD on the offline data, we claim that in many problems this may not work well.
The reason is that the VariBAD VAE should reason about the uncertainty of the MDP parameters, which requires to effectively distinguish between the different possible MDPs. However, if trajectories from the different MDPs visit different parts of the state space, it becomes impossible to identify two trajectories come from different MDPs, or actually come from different parts of the same MDP. This problem, which we term MDP ambiguity, is illustrated in Figure 2 there are two MDPs, one with rewards in the blue circle, and the other with rewards in the yellow circle. If the data contains trajectories similar to the ones in the figure, it is impossible to distinguish between having two different MDPs with the indicated rewards, or single MDP with rewards at both the blue and yellow circles.
Note that MDP ambiguity is special to the offline meta-RI. setting, as in online meta-RI., the agent may be driven by the online adapting policy (or guided explicitly) to explore states that reduce its ambiguity. We also emphasize that this problem is not encountered in standard (non-meta) offline RL, as the problem here concerns the identificatinn of the MDP, which in standard RL is unique.
We next characterize MDP ambiguity in more detail. Different MDPs can differ from each other either in the reward function, the transition function, or both. The differences can also take place over the whole state space, or only in parts of it i..., in some states the MDP rewards or transitions are similar, while on others they are different. Let us denote the set of states where there are such differences the 'identifying states' of an MDP- if the agent has data such states obtained from different MDPs, it has the capability to identify which data samples belong to which MDP.
The MDP ambiguity problem arises when for each MDP in the data, the samples that belong to identifying states do not overlap. Since the agent does not know the MDPs in advance, it does not have information identify from which MDP the data came from (cf. Figure 2
Problems with sparse rewards, for example, are prone to MDP ambiguity the identifying states are the sparse set of rewarding states, and by collecting data as described in Section each agent mainly visits the rewarding states in its own MDP. It is important to note that there are problems where MDP ambiguity is not an issue, For example, MDPs with different transitions, where the set of rewarding states are similar. In this case, all agents will mainly visit the same states (though by applying different policies), and therefore the MDPs will be identifiable. Another example is when the optimal policy in each MDP visits some set of overlapping states, and these states are identifying. We provide examples of both problems in our experiments. In the following, we separate discussion to ambiguity due to different rewards, and ambiguity due to different transitions.
MDP ambiguity due to rewards: One way to solve MDP ambiguity is to increase the variability of the trajectories in the data, however, this may not be possible if the data has already been collected.
5

---

Under review as a conference paper at ICLR 2021
Instead, propose simple and effective alternative in case where the MDPs differ only in their reward function. We make the following assumptions:
1. The transition and reward uncertainties are decoupled, ie., the prior can be written as p(R, P) = Pn (R) x (P(().
2. For each MDP in the training data, we know the reward function R₁-
These assumptions are largely satisfied in most meta-RL studies date (e.g., Zintgraf et al. 2020 Li et 2020 Finn et al. 2017] Duan et al. 2016). The decoupled uncertainty assumption can be hard to verify in practice. However, we remark that assuming a decoupled prior when the true prior is coupled can only enlarge the support of the prior distribution, since if p(R, P) > 0, then the marginal distributions satisfy pn (n(R) > 0, pr (P) >> 0. Thus, this assumption may lead to less effective estimation (as there more MDP 'possibílities).. but should not prohibit the agent from estimating the true posterior with enough data.
We introduce reward relabelling, simple solution to the MDP ambiguity problem. We propose to make the state distribution in the offline data approximately uniform across all MDPs, by replacing the rewards in trajectory from some MDP in the data with rewards from another randomly chosen MDP + i. That is, for each E 1,..... N, we add to the data M trajectories {h,, J E
(.).. where the relabelled rewards & satisfy 1+1 (8. ). Note that our relabelling effectively samples data from an MDP with transitions P. and reward Ry, which has non-zero prior probability mass under the decoupled prior assumption. We only use the relabelled data for training the VariBAD VAE. While in principle it could also he used for training the off off-policy RI algorithm, we did not find it useful in practice.
1...., M, where gad = (#0 90 3.j pid 2- A 8H aij
MDP ambiguity due to transitions: when ambiguity is due to different transitions, the problem becomes significantly more complicated, as we cannot simply 'relabel' trajectories in the data F One solution is add trajectories that run the policy for agent on task v' (this is generalization of the reward relabelling idea). This, of course, requires collecting more data, and is not suitable for the offline setting. It does, however, provide a guideline for effective data collection for OMRL. In this work, we focus on ambiguity due to rewards, and leave the question of MDP ambiguity due to transitions to future research.
4 RELATED WORK
Meta-learning considers training agents that quickly solve a new learning problem. by exploiting structure in the problem distribution (Thrun & Prati 1998 [Hochreiter et al. 200T]. In this work we focus on meta-RI. quickly learning to solve RL problems. Gradient based approaches to meta- RL seek policy parameters that can be updated to the current task with a few gradient steps (Finn) ett al. 2017; Grant et al., 2018 Rothfuss et al. 2018 Clavera et al. 2018). These are essentially online methods, and several studies investigated learning of structured exploration strategies in this setting (Gupta al. 2018 Rothfuss et al., 2018 Stadie et al.) 2018). Memory-based meta-RL. on the other hand, map the observed history in task he to an action (Duan et al. 2016 Wang et al. 2016). These methods effectively treat the problem as a POMDP, and learn a memory based controller for
The connection between meta-learning and Bayesian methods, and between meta-RL. and Bayesian RL in particular, has been investigated in series of recent papers (Lee al., 2018 Humplik let al. 2019 Ortega et al., 2019 Zintgraf et al. 2020), and our work closely follows these ideas. based strategies,
In particular, these works elucidate the difference between hhompson-ssmpliing such as in PEARL (Rakelly et al., 2019). and Bayes-optimal policies, such as in VariBAD ZZingraa al. (020), and suggest to estimate the BAMDP belief using the latent state of deep generative models. Our contribution is an extension of these ideas to the offline RL setting, which to the best of our knowledge is novel. Technically, the VariBAD algorithm in Zintgraf et al. (2020) is limited to on-policy RI., and the off-policy method in Humplik et al. (2019) requires specific task descriptors during learning, while VariBAD, which our work is based on, does not.
WW remark that to our knowledge, previous work on online and offline meta-RL did not investigate such problems. For example, in the Walker environment of Zintgraf et al. (2020), the shape of the walker is varied, which manifests in almost every transition, and a successful agent must walk forward, thus many overlapping states visited.

---

Under review as a conference paper at ICLR 2021
Concurrently and independently with our work, Li et al. (2020) proposed offline meta-RL RL algorithm that combines BCQ (Fujimoto et al.) 2018) with a task inference module. Interestingly, Li al. (020)) also describe a problem similar to MDP ambiguity, and resolve it using technique similar to reward relahelling. However, their approach does not take into account the task uncertainty, and cannot plan actions that actively explore reduce this uncertainty this is form of Thompson sampling, where a task-conditional policy reacts to the task inference (see Figure I in Li et al. (020). Our work is the first to tackle offfine meta-learning of Bayes-optimal exploration. In addition, we demonstrate the first offline results on sparse reward tasks, which, compared to the dense reward tasks in Li et al. (202)). require significantly more complicated solution than Thompson sampling (see experiments section). We achieve this by building on BRL theory, which both optimizes for Bayes-optimality and results in a much simpler algorithm. Recent work on meta Q-learning (Fakoor et al. 2019) also does not incorporate task uncertainty, and thus cannot be Bayes-optimal. The very recent work of Mitchell et al. 2020) considers a different offline meta-RL. setting, where an offline dataset from the test environment is available.
Classical works on BRL are comprehensively surveyed in Ghavamzadeh et al. (016)). Our work, in comparison, allows training scalable deep BRL policies. Finally, there is growing interest in offline deep RL (Sarafian et al. [0118] Levine et al.. 2020). Most recent work focus on how to avoid actions that were not sampled enough in the data. In our experiments, state-of-the-art method of this flavor led to minor improvements, though future offline RL developments may possibly benefit OMRL too.
EXPERIMENTS
In our experiments, we aim to answer the following questions: (1) can we learn approximately Bayes-optimal policies in the offline setting? (2) what are the main practical challenges in OMRL? and (3) does our off- off-policy method improve meta-RI. performance in the online setting as well?
Answering (1) is difficult because the Bayes-optimal policy is generally intractable, and because our results crucially depend on the available data. However, in deterministic domains with a single sparse reward, the optimal solution amounts to 'search all possible goal locations as efficiently as possible, and stay goal once found; in subsequent episodes, move directly to goal.'. We therefore chose domains where this behavior can be identified qualitatively. Quantitatively, we compare our offline results with online methods based on Thompson sampling, which are not Bayes-optimal, and aim to show that the performance improvement due to being approximately Bayes-optimal gives an advantage even with the offline data restriction.
Domains and evaluation metric: Our emphasis is on learning to explore efficiently. We thus focus on sparse reward problems where non-trivial exploration is required to identify the task: (1) A discrete 5 x 5 Gridworld (Zintrraf et al.] 2020); (2) A continuous point robot where a sparse reward is located somewhere on semi-circle (see Figure : (3) Ant-Semi-cirele a challenging modification
of the popular Ant- task (Fakoor et al. 2019) to a sparse reward setting similar to the semi-cirele task above; (4) Half-Cheetale Vel (Finn al. 2017), a popular dense reward domain that serves as a comparison to the sparse domains (full details in Appendix C). Note that for all these problems, the tasks differ by their reward function. We also provide an experiment with varying transition functions in Appendix G.3 To evaluate performance, we measure average reward in the first 2 episodes on unseen tasks- this is where efficient exploration makes a critical difference In the supplementary, we report results for more evaluation episodes.
Data collection and organization: For data collection, we used off-the-shelf DQN (Gridworld) and SAC (continuous domains) implementations. To study the effect of data diversity, we diversified the offline dataset by modifying the initial state distribution Ponse to either (1) uniform over large region, (2) uniform over a restricted region, or (3) fixed to single position. At meta-test time, only the single fixed position is used. The tasks are episodic, but we want the agent to maintain its belief between episodes, so that it can continually improve performance (see Figure D.. We follow ZZntgraf ctt al. (2020), and aggregate & consecutive episodes of length H to a long trajectory of length k x H, and we do not reset the hidden state in the VAE recurrent neural network after episode termination. For reward relabelling, we replace the first k/2 trajectories with trajectories from a randomly chosen
For Gridwarld, we measure average reward in the first 4 episades.
7

---

Under review as a conference paper at ICLR 2021
Gridworld
Semi-circle
Ant-Semi-circl
Half-Cheetah-Vel
Durs, web Chrine PEAIL
TEATL
Average return
90
20
who Thompson Sampling
Online PEARL
Train Steps
23
16
Train Steps
Train Steps
Train Steps
Figure 3: Offline performance various domains. Blue: method. Red: our method with reward relabelling ablated. Black: Thompson sampling baselines calculated exactly in Gridworld, and using online PEARL for the other continuous domains.
MDP, and relabel their rewards. Technically, network architectures and hyperparameters were chosen similarly to Zintgraf et al. 22020, as detailed in the supplementary.
Offline Results: In Figure 3 we compare offline algorithm with Thompson sampling based methods, and also with an ablation of the reward relabelling method. For Gridworld, the Thomp- son sampling method is computed exactly, while for the continuous environments, We online PEARL (Rakelly et al. 2019) strong baseline that is not affected by our offline data limitation. In particular, this baseline stronger than the offline meta-RL algorithm of et (2020), For these results the uniform initial state distribution was used to collect data. Note that we significantly out- perform Thompson sampling based methods, demonstrating our claim learning non-trivial exploration from offline data. Results on sparse reward domains signify the severity of MDP ambiguity without reward relabelling performance drops significantly. In the Half-Cheetal Vel environment, on the other hand, the rewards are dense and thus MDP ambiguity is not problem.
In Figure and in Figure 4 we visualize the trajectories the trained agents in the Semi-circle and Ant-Semi-cirele domains, respectively Qualitatively, approximately Bayes-optimal behavior evident: in the first episode, the agents search for the goal along the semi-circle, and in the second episode, the agents maximize reward by moving directly towards the already found goal. Similar behavior for Gridworld is reported in Appendix different from the training data, in which the agents learned to reach specific goals. In Figure 7 in the supplementary. we provide further insight into these results, by showing the belief update during the episode rollout: the belief starts uniform the semi-cirele, and narrows in the target the agent explores the semi-circle. With reward relabelling ablated, however, we show that the belief does not update correctly.
We emphasize that this behavior is very
"Video is provided: bttp:///wwwuubbeeeeeeeeeeeeeeeee 65wg552Y0U4
Episode
Episodle 2
Episode 1
Episode 2
Figure 4: Ant-Semi-circle: trajectories from trained policy on new goal. Left: Trajectory of the center of mass. Right: Visualization of the ant at different steps along the trajectories. Note that in the first episode, the ant searches for the goal, and in the second one it directly moves toward the goal it has previously found. This search behavior different from the goal-reaching behaviors that dominate the training data.
8

---

Under review as a conference paper at ICLR 2021
Data Quality Ablative Study; To evaluate the dependency of our method on the offline data quality, we report results for the different data collection strategies described above (see supplementary for more details). The results summarized in Table [ As expected, data diversity is instrumental to offline training. However, as we qualitatively show in Figure TT in the supplementary, even on the low-diversity datasets, agents learned non-trivial exploration strategies that search for the goal. This is especially remarkable for the fixed- ixxeddistribution dataset, where it is unlikely that any training trajectory traveled along the semi-circle.
One may ask whether OMRL presents the same challenge as standard offline RL, and whether recent offline RL advances can mitigate the de-
Ours
w/ CQL
Uniform
171.8 7.0
176.0 10.2
pendency on data diversity. To investigate this,
Excluding s.c.
102.8 + 32.7
116.6 19.9
we also compare our method with variant that uses CQL (Kumar et al. 2020) state-of-the- art offline RL method to train the critic net- work of the meta-RL agent. Interestingly, while CQL slightly improved results (Table TD. the data diversity is much more significant. Taken together with our results on MDP ambiguity, our investigation shows that OMRI presents differ- ent practical challenges than standard offline RL.
Fixed
99.2 + 27.4
112.4 + 31.3
Table Average return Ant-Semi-cirele for dif- ferent initial state distributions during offline data collection: Uniform distribution, uniform distribu- tion excluding states on the semi-cirele (Excluding S.c.), and fixed initial position (Fixed).
Online Setting: Our method can also be applied to the online setting. In this case, it is simply modification of VariBAD, where the policy gradient optimization is replaced with off-policy RL algorithm. Since MDP ambiguity does not concern online meta-RL, we did not use reward relabelling in this setting. As shown in Figure 5 by exploiting the efficiency of off off-policy RL, our method significantly improves sample-efficiency, without sacrificing final performance.
Gridworld
Semi-cirola
Half-Cheetah-Vel
Averaga return
Figure 5: Online perfor- mance comparison. The off- policy optimization signif- icantly improved VariBAD performance.
Outt VrrBAD
Ourt
Durs VerBAD
Wr1BAD
C4 64 Frames
Frames
Frames
When comparing Figure and Figure 3 the reader may notice that the online algorithm's final performance outperforms the final performance in the offline setting. We emphasize that this phenomenon largely depends on the quality of the offline data, and not on the algorithm itself.
CONCLUSION
We presented the first offline meta-RI. algorithm that is approximately Bayes-optimal, Key to our approach the connection between Bayesian and meta learning, which in principle allows to reduce the problem to standard offline RL. In practice, however, showed that the MDP ambiguity problem prohibits learning, and proposed simple solution based on mixing trajectories in the data and relabelling their rewards. Our results show that this solution is effective several domains.
Offline learning is appealing for domains where data collection costly, such as robotics and healthcare, and there growing interest in applying deep RL this setting (Levine et al. 2020). Key advances in this field will likely play a role in improving OMRL as well an exciting direction for future research. How handle the MDP ambiguity problem when it is caused by changes in the MDP transitions another important research direction.
REFERENCES
Dimitri P Bertsekas. Dynamic programming and optimal control, volume 1. Athena Scientifie, 1995.

---

Under review as a conference paper at ICLR 2021
Anthony R Cassandra, Leslie Pack Kaelbling, and Michael L Littman. Acting optimally in partially observable stochastic domains. In AAAI, volume 94, pp. 1023-1028, 1994.
Ignasi Clavera, Jonas Rothfuss, John Schulman, Yasuhiro Fujita, Tamim Asfour, and Pieter Abbeel. Model-based reinforcement learning via meta-policy optimization. In Conference Robot Learning, pp. 617-629, 2018.
Yan Duan, John Schulman, Xi Chen, Peter L Bartlett, llya Sutskever, and Pieter Abbeel. RL Fast reinforcement learning via slow reinforcement learning. arXiv preprint arXiv: 1.02779, 2016.
Michael O' O'Gordon Duff and Andrew Barto. Optimal Learning: Computational procedures for Bayes-adaptive Markov decision processes. PhD thesis, University of Massachusetts at Amherst, 2002.
Damien Ernst, Pierre Geurts, and Louis Wehenkel. Tree-based batch mode reinforcement learning. Journal of Machine Learning Research, 6(Apr):503-556, 2005.
Rasool Fakoor, Pratik Chaudhari, Stefano Soatto, and Alexander J Smola. Meta-q-learning, arXiv preprint arXiv: 1rXiv:1910000122 2019.
Chelsea Finn, Pieter Abbecl, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In Proceedings of the 34th International Conference on Machine Learning-Voii Volume 70, pp. 1 1126-112 135. JMLR. org, 2017.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. arXiv preprint 1rXiv:1812.029000 2018.
Mohammad Ghavamzadeh, Shie Mannor, Joelle Pineau, and Aviv Tamar. Bayesian reinforcement learning: A survey. arXiv preprint arXiv: wewp::669044445 2016.
Erin Grant, Chelsea Finn, Sergey Levine, Trevor Darrell, and Thomas Griffiths. Recasting gradient- based meta-learning as hicrarchical bayes. arXiv preprint arXiv: 08930, 2018.
Abhishek Gupta, Russell Mendonca, YuXuan Liu, Pieter Abbeel, and Sergey Levine. Meta- reinforcement learning of structured exploration strategies. In Advances in Neural Information Processing Systems, pp. 5302-5311, 2018.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maxi- mum entropy deep reinforcement learning with stochastic actor. arXiv preprint arXiv: 101.01000 2018.
Sepp Hochreiter, A Steven Younger, and Peter R Conwell. Learning to learn using gradient descent. In International Conference on Artificial Neural Networks, pp. 87-94. Springer, 2001.
Jan Humplik, Alexandre Galashov, Leonard Hasenelever, Pedro A Ortega, Yee Whye Teh, and Nicolas Heess. Meta reinforcement learning as task inference. arXiv preprint urXiv:1905.06422 2019.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. arXiv preprint 0et.2000604779, 2020.
Gilwoo Lee, Brian Hou, Aditya Mandalika, Jeongseok Lee, Sanjiban Choudhury, and Siddhartha Srinivasa. Bayesian policy optimization for model uncertainty. arXiv preprint 1101001810.01014, 2018.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives open problems. arXiv preprint arXiv:2005.0 00..00000.004,, 2020.
Jiachen Li, Quan Vuong, Shuang Liu, Minghua Liu, Kamil Ciosek, Henrik Iskov Christensen, and Hao Su. Multi-task batch reinforcement learning with metric learning. arXiv preprint arxiv:1909.1133 2020.
Timothy P Lillicrap, Jonathan J Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint 1509.020 2015.
10

---

Under review as a conference paper at ICLR 2021
Eric Mitchell, Rafael Rafailov, Xue Bin Peng, Sergey Levine, and Chelsea Finn. Offline meta- reinforcement learning with advantage weighting. arXiv preprint arXiv:2008. 06043, 2020.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas Fidjeland, Georg Ostrovski, al. Human-level control through deep reinforcement learning. Nature, (aaeee,,18855441 518(7540):529-533, 2015.
Pedro A Ortega, Jane X Wang, Mark Rowland, Tim Genewein, Zeb Kurth-Nelson, Razvan Pascanu, Nicolas Heess, Joel Veness, Alex Pritzel, Pablo Sprechmann, al. Meta-learning of sequential strategies. arXiv preprint arXiv:1905.0300 2019.
Kate Rakelly, Aurick Zhou, Deirdre Quillen, Chelsea Finn, and Sergey Levine. Efficient off- ofoliy meta-reinforceement learning via probabilistic context variables. arXiv preprint 1K03::9900002255 2019.
Jonas Rothfuss, Dennis Lee, Ignasi Clavera, Tamim Asfour, and Pieter Abbeel. Promp: Proximal meta-policy search. arXiv preprint arXiv:1810.061 0678444, 2018.
Elad Sarafian, Aviv Tamar, and Sarit Kraus. Safe policy learning from observations. arXiv preprint arXiv: 1805:080000000 2018.
Bradly C Stadie, Ge Yang. Rein Houthooft, Xi Chen, Yan Duan, Yuhuai Wu, Pieter Abbeel, and Ilya Sutskever. Some considerations learning to explore via meta-reinforcement learning. arXiv preprint arXiv:1803.01 / 18, 2018.
Sebastian Thrun and Lorien Pratt. Learning to learn: Introduction and overview. In Learning to learn, pp. 3-17. Springer, 1998.
Harm Van Seijen, Hado Van Hasselt, Shimon Whiteson, and Marco Wiering. theoretical and empirical analysis of expected sarsa. In 2009 ieee symposium on adaptive dynamic programming and reinforcement learning, pp. 177-184. IEEE, 2009.
Jane X Wang, Zeb Kurth-Nelson, Dhruva Tirumala, Hubert Soyer, Joel Z Leibo, Remi Munos, Charles Blundell, Dharshan Kumaran, and Matt Botvinick. Learning reinforcement learn. arXiv preprint 1oXiv:611.005766 2016.
Luisa Zintgraf, Kyriacos Shiarlis, Maximilian Igl, Sebastian Schulze, Yarin Gal, Katja Hofmann, and Shimon Whiteson. Varibad: very good method for bayes-adaptive deep via meta- meta-learning. In International Conference on Learning Representation (ICLR), 2020.

---

Under review as a conference paper at ICLR 2021
A PROPOSITION PROOF
For ease of reading, we copy Proposition T from Section 3.2
denote a random trajectory from fixed history dependent policy n, generated according to the following process. First, MDP parameters R, P are drawn from the prior p(R, P). Then, the state trajectory is generated according to so Pinit. a₁ ~ π (- Su, ag. FL. Fi....,,:), St+1 ~ P(-ss a1) and T(+1 ~ R(st, ar). Let by denote the posterior belief at time b, P(R, P|so, : 0o, F,,...,, St). Then
Proposition 1. Let So. dg. F1,81 .1.....H
P(s su. ao. Fi,..., St, at) = b.,, (($1+1|81 at), and, P(re+1180, an, .1...., St, 0e) E R(r++1 at).
Proof. For the transitions, we have that,
/ P P(1141, R. Plsp, (0,000000000 Tor - Tt, 8c, 01 dRRPP
P(8411 8o, do, Cor o to 8c: a+)
/ P( RR P, 80, On, ro rt, St. a4) P(R. Plso. ao, To,-- 81, 0g d'RdP
80: ao; -o.-.-. Pty at) ER P( RR P, 80, ao, ro: ................................ Ft, 8t, (t)
ER.P P [P(*++ 1st, ar) 80, ao; To, T 8t, [z)
ER,P-h, (1t+1//5,, at).
The proof for the rewards proceeds similarly.
B VAE TRAINING OBJECTIVE
For completeness, we follow Zintgraf al. (2020) and outline the full training objective of the VAE. Consider the approximate posterior la (m) h.t) conditioned on the history up to time f. In this case, the ELBO can be derived as follows:
log / P(sp:Ti,3) .... 8,, m an 1)dm
log (S5p, P1, Sm Ag....... an 00.....aa-111
log / P(so, T1, ..... m|ao.........................m..
2o ( hm) go go h
dm
P(sa, F1, 81....... 81, m Ao, 011 -1 go (m
log (mm
2 Emagat (h..) [log (og,p(so,
81, sH m, ao aH-1) + log - log go (ml)
Em-golthil h.)) [log (o r1, 81 s1... sH m, ao . ( aH-1)
-Dilga (m))))
= ELBO, 6).
The prior Po (m)) set to be the previous posterior qo (m)(ht1))) with initial prior chosen to be standard normal po(m) (m) = (0, I). The decoder pe Pe r1, ..... SH mm, ap...., (1-1) factorizes to reward and next state models po(6'|s (ss a, m) and po a, m), according to: log (3(30, 81.81..., SH [m, ao, ... -H)1) log p(solm)
[log Pe (S ar am m) + log (o m)].
As in Zintgraf et al. 22020). we trained only a reward decoder.
The overall training objective of the VAE is to maximize the sum of ELBO terms for different time steps,
max
(L
(4)
0.0
12

---

Under review as a conference paper at ICLR 2021
C ENVIRONMENTS DESCRIPTION
In this section we describe the details of the domains we experimented with.
5 gridworld environment as in Zintgraf al. (2020). The task distribution is defined by the location of a goal, which unobserved and can be anywhere but around the starting state at the bottom-left cell. For each task, the agent receives a reward of -0.1 on non-goal cells and +1 at the goal, i.c.,
Gridworld:
1, 0.11, else,
St
where is the current cell and is the goal cell. Similarly Zintgraf al. (2020), the horizon for this domain set to 15 and we aggregate consecutive episodes to form trajectory of length 60.
Semi-circle: A continuous 2D environment as in Figure where the agent must navigate to an unknown goal, randomly chosen a semi-circle of radius 1 (Rakelly al. 2019). For each task, the agent receives a reward of +1 if is within a small radius 0.2 of the goal, and 0 otherwise,
1, | Zgoal S
T =
0, else,
where x; is the current 2D location. Action space is 2-dimensional and bounded: [-0.1, 0.1]². We set the horizon to 60 and aggregate & = consecutive episodes to form trajectory of length 120.
MuJoCo: For both MuJoCo domains, the horizon is set to H 200 and we consider episodes.
Half-Cheetah-Ve Vel: In this environment, a half-cheetah agent must run at fixed target velocity. Following recent works in meta-RL (Finn et al. [017 Rakelly et al., 2019 Zintgraf jet al., (220).. we consider velocities drawn uniformly between 0.0 and 3.0. The reward in this environment given by
Ti 12 Ugoal 0.05 ||a.12
where b, is the current velocity, and ar is the current action.
2. Ant-Semi-circle: In this environment, an ant needs to navigate to an unknown goal, ran- domly chosen on semi-circle, similarly the Semi-circle task above.
When collecting data for this domain, we found that the standard SAC algorithm (Haarnoja et al. 2018) was not able to solve the task effectively due to the sparse reward (which is deseribed later), and did not produce trajectories that reached the goal. We thus modified the reward only during data collection to be dense, and inversely proportional to the distance from the goal,
dense Y1
1|x, - #ggaa| || 0.1 ||a.|]
where 2, is the current 2D location and a, is the current action. After collecting the data trajectories, we replaced all the dense rewards in the data with the sparse rewards that are given by
{ {0.
||x, Zgoal 2l2 0.2
sparse -0.1 ||a
else.
We note that Rakelly al. (2019) similar approach cope with sparse rewards in the online setting.
D EXPERIMENTAL DETAILS
In this section we outline our training process and hyperparameters.
For the discrete Gridworld domain we used DQN Mmih et al. 2015) with soft target net- work updates, as proposed by |illicrap et al. 2015), which has shown improve the stability of
13

---

Under review as a conference paper at ICLR 2021
learning. For the rest of the continuous domains, used SAC (Haarnoja al. 2018 with the architectures of the actor and critic chosen similarly, and with a fixed entropy coefficient. For both DQN and SAC, set the soft target update parameter to 0.005.
In all our experiments average the performance over random seeds and present the mean and standard deviation.
D.1 OFFLINE SETTING
Our offline training procedure is comprised of separate training steps. First is the training of the data collection RL agents. Each agent is trained on different task from the task distribution.
For the Gridworld domain, we train 21 agents. We note that this covers the entire task distribution, as goals can be anywhere but around the starting state at the bottom-left cell. For the Semi-circle and Ant-Semi-circle domains, we train data collection agents, and for the Half-Cheetah- Vel environment we used 100 agents.
For all tasks, we used a similar architecture of 2 fully-connected (FC) hidden layers of size that depends on the domain with ReLU activations, and set the batch size to 256. The rest of the hyperparameters used for training the data collection RL agents are summarized the following table:
Gridworld (DQN)
Semi-cirele (SAC)
MuJoCo (SAC)
Hidden layers size Num. iterations RL updates per iter.
16
32
128
200
300
1000
500
500
2000
e-greedy, linear annealing from to 0.1
Exploration/entropy coeff.
0.01
0.2
over 100 iterations
Collected episodes per iter. Learning rate/s
5
2
3-10 10
3 10
3 10
Discount factor (7)
0.99
0.9
0.99
The second training step is the VAE training after applying reward relabelling the collected data.
The VAE consists of a recurrent encoder, which at time step t takes as input the tuple (a,,Ti+z T1+1+ 8t+1) The state and reward passed each through different fully-connedd colleceedd (FC) layer. The state FC layer is of size 32 and the reward FC layer is size for the Gridworld and 16 for the rest of the domains, all with ReLU activations. For the MuloCo environments,
also pass the action through FC layer of size 16 with ReLU. Then, the state and reward layers' outputs are concatenated along with the action (or with the output of the action layer in the case of MuJoCo) and passed a GRU of size 64/128 (Gridworld/other domains). The GRU outputs the Gaussian parameters p(h.), (haa) of the latent vector m, whose dimensionality is in all our experiments.
The VAE reward decoder takes as input latent sample m (p(h))) E(h:)) and the states along the trajectory $1.... SH. each state at a time, and outputs (for every timestep
I,..., H) the entire reconstructed/prediieed rewards .I. along the trajectory. In the MuloCo domains the reward decoder also takes as input the actions along the trajectory a1...., ag az and the previous states so...., -1 the reward in these environments generally depends on 8t-1, Qr St. The reward decoder comprised of FC layers, each size 32.
The VAE is trained to optimize Equation 4 but similarly to Zintgraf et al. 2020), we weight the KL term in each of the ELBO terms with some parameter 3, which is not necessarily In our experiments used 8 0.05.
After the VAE is trained, we apply state relabelling to the data collected by the RL agents, to create large offline dataset that effectively comes from the BAMDP. Then, we train off- off-policy RL algorithm, which is our meta-RL agent, using the offline data.
For the offline meta-RL agents training, used similar hyperparameters to those used for the data collection RL agents training. We only enlarge the size of the hidden layers in our models from 16, 32 and 128 64, 128 and 256 for the Gridworld, Semi-cirele and MuJoCo domains, respectively. In
14

---

Under review as a conference paper at ICLR 2021
every iteration perform 1000 parameter updates for all environments expect the in which case we perform 2000 updates per iteration.
D.2 ONLINE SETTING
In the online setting didn't 't apply reward relabelling the data, since, we explained, MDP ambiguity doesn't concern online meta-RL. The hyperparameters used in the online setting are as follows:
Gridworld (DQN)
Semi-circle (SAC)
Cheetah- Vel (SAC)
RL parameters
2 FC layers of size 100.
2 FC layers size 128. 1000
3 FC layers of size 128. 2000
Architecture/s
Num. updates per iter.
250
<-greedy, linear annealing from 1 to 0.1 over 1000 iterations.
Exploration/entropy coeff.
0.01
0.2
Collected episodes per iter. Learning rate/s
25
25 7 10-5
7 10-5
3 10
Discount factor (7)
0.99
0.9
0.99
VAE parameters
state/reward FC layer of size 32/8. GRU of size 64.
state/reward FC layer of size 32/8. GRU of size 128.
state/action/rewaar FC layer of size 32/ 1///16 16. GRU of size 128.
Encoder architecture
2 FC layers of size 32.
2 FC layers of sizes 64 and 32.
2 FC layers of sizes 64 and 32.
Reward decoder architecture
Num. updates per iter.
20
25
20
Learning rate
3 10 1.0
10 10-3
3 10 1.0
Weight of KL term (13)
0.1
E LEARNED BELIEF VISUALIZATION
In this section visualize the learned belief states, in order get more insight into the decision making process of the agent during interaction.
In Figure 6 visualize the interaction of trained agent with the Gridworld environment, exactly as visualized in Figure 3 at Zintgraf et al. (202)). The agent reduces its uncertainty by effectively searching the goal. After the goal is found, the agent stops and in subsequent episodes it directly moves toward it.
5
E10
F15
lalent mean
Episode 1
SPP 11) for eooch cll
Episode 2
20
env steps
latent variance
Episode
20
20
env steps
env steps
Figure 6: Interaction of trained agent with Gridworld. For more details, see Zintgraf et al. (2020).
In Figure 7 we plot the reward belief (obtained from the VAE decoder) at different steps during the agent's 's interaction in the Semi-circle domain. Note how the belief starts uniform the
15

---

Under review as a conference paper at ICLR 2021
semi-circle, and narrows in the target more evidence is collected. Also note that without reward relabelling, the agent fails to find the goal. In this instance of the MDP ambiguity problem. the training data for the meta-RL agent consists of trajectories that mostly reach the goal, and result, the agent believes that the reward is located at the first point it reaches on the semi-circle.
Belief update: when reaching the goal
1.0
0.8
Belief update: when searching for the goal
24
0.6
-0.
Belief update: reward relabelling ablation
12
24
-.2
0.0
Figure 7: Semi-circle belief visualization. The plots show the reward belief over the 2-dimensional state space (obtained from the VAE) different stages of interacting with the system. The red line marks the agent trajectory, and the light blue circle marks the true reward location. Top: Once the agent finds the true goal, it reduces the belief over other possible goals from the task distribution. Middle: As long as the agent doesn't find the goal, it explores efficiently, reducing the uncertainty until the goal found. Bottom: Without reward relabelling, the agent doesn't learn differentiate between different MDPs, and therefore fails to identify the goal.
DATA QUALITY ABLATION
our data quality ablative study, we consider the Ant-Semi-circle domains, for which we modify initial state distribution during the data collection phase. The initial state distributions consider are visualized in Figure 8 Uniform distribution, uniform excluding states on the semi-circle, and fixed initial position.
(a) Uniform
(b) Excluding s.c.
(c) Fixed
Figure Initial state distributions. Red locations indicate non-zero sampling probability.
Figure D shows the learning curves for the results presented in Table For completeness, we add the learning curve for the uniform distribution which also presented in Figure
16

---

Under review as a conference paper at ICLR 2021
Ant-Semi-circle: Uniform
Ant-Semi-circle: Excluding s.c.
Ant-Semi-circle: Fixed
Average retur
Ours Ours, OQL
Ours, w/ OGL
Ours, CGL
Train Steps
Train Steps
Train Steps
Figure 9: Learning curves for the results presented in Table 1 In blue is our method and in red is our method, with critic network trained according to the CQL objective (Kumar et al. 2020). Left: Uniform initial state distribution. Middle: Uniform distribution, excluding states over the semi-cirele. Right: Initial state is fixed.
We also visualize trajectories of trained agents for the 3 different cases as well as for PEARL (Rakelly al. 2019), Figure TO Note that even for the fixed-distribution dataset, our agent learns search for the goal.
Uniform
Excluding s.c.
Fixed
PEARL
pisode
Episade
Episade
Figure 10: Ant-Semi-ccrele trajectories trained agents for diffcrent offline datasets and for PEARL.
G ADDITIONAL RESULTS
G.1 PERFORMANCE vs. ADAPTATION EPISODES
In this part, we present the average reward per-episode as a function of the number of adaptation episodes at the environment. Figure TT shows the performance for the Ant-Semi-circle and Half- Cheetah- Vel domains. Note that within the first few episodes, PEARL does not collect high rewards due to the Thompson sampling-bes based nature of the algorithm. Our method, on the other hand, efficiently explores new tasks and is able to collect rewards within the first episodes of interaction.
Ant-Semi-circle
Half-Cheetah-Vel
Average return
60
-90
-19
Figure 11: Adaptation performance.
Our method outperforms PEARL, collecting high rewards within first adaptation episodes.
60
129
Curs
Ours
PEARL (00
PEARL
Adaptation episodes
Adaptation episodes
G.2 PEARL LEARNING CURVES
We present the training curves PEARL Figure 12
Note that since PEARL an online algorithm, the r-axis represents the number of environment interactions.
17

---

Under review as a conference paper at ICLR 2021
Semi-circle
Half-Cheetah-Vel
Ant-Semi-circle
17.
-100
Avg. Return
15.0
-425
-10
-10
13.0
475
460
25
112
Frames
Frames
Frames
Figure 12:
Learning curves for online PEARL training. Performance is measured by average reward in the first 2 episodes on unseen tasks from the task distribution.
G.3 EXPERIMENT ON MDPs WITH VARYING TRANSITIONS
In the following section we describe experiments on domain in which MDPs differ from each other by their transition function, when MDP ambiguity is not problem. Our goal is show that our method can also be applied MDPs with varying transitions.
We consider a continuous 2-dimensional environment we term Point-Robot- Windd In Point-Robot- Wind, the agent must navigate to fixed (unknown) goal within distance of D 1 from the initial state. The goal position is the same for all tasks. The agent receives reward of +1 is within radius r 0.2 of the goal, and 0 otherwise. For each task in this domain, the agent is experiencing a different 'wind", which results a shift in the transitions, such that when taking action a, E [,,0,1,0.0 0.1]² from state in MDP M, the agent transitions to a new state St41, which is given by
8 + WM
where uxy is a task spec-seeiiic wind, which is randomly drawn for each task from the uniform distribu- tion on [-0.05, 0.05]". To navigate correctly to the goal and stay there, the agent must take actions that cancel the wind effect.
We set the horizon to 25 and evaluate the performance in terms of average return within the first test tasks, sampled from the task distribution. We do not apply reward-
episode of interaction relabelling, as the MDPs differ their transition functions rather than rewards. Figure (3Ja) a) we compare our method with online PEARL. Note that our learned policy performs much better than the online PEARL baseline.
Point-Robot-Wiig
Point-Robot-Wind
17.5
Average return
125
5
10.0
7.8
Ours
26
Online PEARL
02
1.0
Train Steps
Frames
(a) Offline Performance
(b) PEARL. Training Curve
Figure 13: Evaluation Point-Robot- Wind - domain with varying transition function. In (a) is offline performance of our method, compared to the best performance of online PEARL (similar to Figure 3) and in (b) is the learning curve for online PEARL training (similar Figure 12
18

---

Under review as a conference paper at ICLR 2021
In Figure 14 visualize trajectories of the trained agent on different test tasks. As can be seen, after several steps in the environment, our agent learns to adapt to the varying wind, and travels to the goal in straight line. PEARL, the other hand, only adapts after the first episode, and therefore obtains worse results. We believe it is possible to improve PEARL. to update its posterior after every step, and in this case the improved PEARL will obtain similar performance our method in Point-Robot-Wi Wind. However, this will not work in the sparse reward domains described in the main text, where the Bayes adaptive exploration has an inherent advantage over Thompson sampling. We emphasize that in Point-Robot- Robot-Wind, MDP ambiguity is not concern, since the data from all agents is largely centered on the line between the agent's 's initial position and the goal. Thus, the effect of the wind on these states can uniquely be identified in each domain. These results confirm that when MDP ambiguity is not concern, our method learns an effective policy in domains with varying transitions.
Trajectories in Point-Robot-Wind
Test task 1
Test task 2
Test task 3
Figure 14: Point-Robot- Wind: trajectories of trained agent on different tasks from the task distribution.
19


Reviews:
[{"review_id": "jz-Ikw-TJxT", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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. \n\nThe VariBAD algorithm formulates online Meta-RL as a Bayesian Adaptive MDP (BAMDP) and solves it approximately using a variational approach. The authors, due to lack of online data, use partial offline trajectories to create belief over RL environment. Next, this belief is used to map offline Meta-RL to the BAMDP setting -- a.k.a. state relabelling. Once mapped to BAMDP off-the-shelf offline RL techniques are used to solve the problem approximately. The authors further introduce reward relabelling that adds new trajectories by mixing rewards among the RL environments conditioned on the state and action. This is claimed to be useful if the reward distribution is independent of the transition matrix distribution of the RL. \n\nThe authors provide multiple experiments with sparse rewards to show the VariBAD with relabelling works in some established benchmark RL tasks (albeit in the new offline Meta-RL setting). \n\nPros:\n* The offline Meta-RL is a very relevant problem with a lot of applications in the coming years. Designing a methodology to solve this problem using off-the-shelf techniques is a novel pursuit, which is carried out in the paper. \n\n* The idea of state and reward relabeling is novel in the meta-RL setup, as per my understanding (as an outsider to the Meta-RL community but with a background in RL).\n\nCons:\n1. Methodology: \n\n* The paper does not explore properly the (in)-famous \"deadly triad\" which is known to make offline RL difficult. Given the state-space of BAMDP captures beliefs over (possibly) complicated RL environments, the effect of  \"deadly triad\" is conceivably even more prominent in the current formulation.  \n\n* I am not sure why the reward relabeling is able to solve the MDP (a.k.a. aliasing in the offline RL community). In particular, \"Note that our relabelling effectively samples data from an MDP with transitions Pi and reward Ri, which has non-zero prior probability mass under the decoupled prior assumption\" -- the above statement is very vague given finite trajectories. The MDP aliasing (a part of the deadly triad) is closely related to this issue. \n\n* The state relabelling is heavily reliant on the existing VariBAD algorithm, and off-the-shelf RL techniques. Therefore, the algorithmic contribution seems somewhat lacking in my opinion. \n\n2. Experiments:\n* The provided experiments lack any conclusive evidence of the effectiveness of reward-labeling. Why in the Half-Cheetah-Vel experiment reward relabeling performs worse -- is the reward independence assumption not true here?  Will reward relabeling always perform worse? \n\n* The effectiveness of the Thompson sampling baseline used is not clear to me. Does no other simple baseline exist? I am alright with not comparing it with the concurrent works. However, such a comparison would have been much more convincing. \n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "Zg7Cip3ImEm", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\nThis submission studies the meta-learning problem in RL under offline settings. A new algorithm is proposed to address this problem by extending the recent VariBAD algorithm designed for online meta-RL. The key modifications to adapt the original VariBAD to offline settings are the state re-labelling and reward re-labelling tricks, which aim at addressing the MDP ambiguity specifically showing up in offline meta-RL. Experiments are conducted on four sparse reward tasks under both offline and online settings, comparing with PEARL (a recent online approach) and the original VariBAD respectively.\n\nPros:\n1. The problem studied is interesting and important in RL, and largely open.\n\n2. The proposed tricks for adapting VariBAD to offline settings are simple and easy to implement.\n\n3. The ablation study on the Ant-Semi-circle task is helpful.\n\nCons:\n1. Regarding the reward re-labelling trick, the assumption that reward functions for all training environments are known limits the applicability of the proposed algorithm for real-world problems.\n\n2. Regarding the experiments comparing with online PEARL, the training curves of PEARL in Fig. 3 do not seem correct to me, it is strange that PEARL does not make any progress with training, which is very different from the behavior reported in the original PEARL paper.\n\n3. Regarding the soundness of comparison experiments, it would be much more convincing if the following things are taken care of:\n\n (a) All four tasks reported in the original VariBAD paper are included, rather than just picking one from them. \n\n (b) Averaging performance is reported over 5 runs rather than just 3, especially considering that SAC could perform quite differently between different random seeds.\n\n (c) Hyper parameters are set to be more consistent for the same RL backbone (say SAC) across different tasks. It is understandable that two sets of parameters might be needed for online and offline settings.\n\nBased on my main concerns regarding the assumptions on reward function and soundness/correctness of experiments, I lean towards a rejection.\n\nOther questions:\n\n1. Regarding the state re-labelling trick, I have a question regarding the proof of Proposition 1: the first equality of the proof in Appendix A does not seem straightforward to me. It would be helpful if more derivation details are provided.\n\n2. I feel that the claim of approximating Bayes-optimal policy in the offline setting is overly stated. Is there any explicit arguments showing that this is \"approximately\" true?\n\n====================================================\nPost rebuttal:\n\nMy concerns regarding the experiments are mostly addressed, though as pointed out by other reviewers, more convincing experiments under changing transition dynamics would be very helpful. I also stand by the authors explanation regarding limited resource in running RL experiments, especially for novel research directions.\n\nThat said, for the same reason (pioneering research vs large-scale application), I do not agree with the authors explanation regarding the limitation of the proposed approach in assuming a known ground truth reward function. The main contribution of this submission is not in solving a specific real-world RL application problem as the cited references. As one of the first efforts in addressing the meta-RL under offline setting, I feel that setting this constraint is a limitation and should be relaxed by means of estimating the reward function. This should be addressed in future work.\n\nI raise my rating to a weak acceptance conditioned on the wording regarding \"Bayes optimal\" being more precisely presented, I think it is currently over-stated throughout the paper, which could be misleading to the community if published as it is.\n\nIt is important to carefully reword in which sense the proposed algorithm \"approximate Bayes-optimal policy\", as explained in the authors' response, the algorithm is shown to qualitatively behave in a way that a Bayes optimal policy would do under this particular setup, this is far from sufficient to claim any approximation to the Bayes-optimality in a principled sense. I would like to also point out that while Bayes-optimality is generally intractable, it is possible and not uncommon for a method to start from an explicit pursue of Bayes optimal solution and specify where and how approximations are performed to overcome specific intractibilities, and further show quantitatively that such Bayes optimality is indeed achievable under well-controlled toy examples where the true Bayes optimal solution is known. The concept of Bayes optimality is in essence quantitative, rather than qualitative.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "ClbklHjRi0L", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary after discussion period:\n----------------------------------------------\nThe authors have done a good job in toning down their claims to match what the evidence supports. After reading the other review's comments and the updated version of the paper, I feel both my comments and most of the reivewer's comments have been addressed, allowing me to recommend this paper for acceptance.\n\nSummary\n-------------\n\nIn this paper, the authors tackle the problem of offline meta RL, where one observes trajectories from environments drawn from a distribution of environments, then using these past observations finds a policy which will work well on new environments drawn from the same distribution, optimizing both in the face of uncertainty about the environment and about the random transitions. The authors propose a method based on VariBAD which achieves good experimental results.\n\nPros:\n-------\n\nI find the paper to be novel, relevant and well written, a good contribution to the scientific community. I liked it’s introduction to the idea of offline Meta-RL, and thought the experimental results where an agent learns to walk the semi-circle were impressive. Further, I greatly enjoyed the discussion of MDP ambiguity, it’s a genuine problem that I’d never seen before and I think addressing exactly this issue will be important in advancing offline meta-RL.\n\nMajor Concern:\n--------------------\n\nMy major concern is that the evidence does not fully support the author’s central claim: “quickly maximize reward in a new, unseen task from the same distribution” where this distribution is described as varying over both the reward and transition functions. Although I do believe that the author’s method addresses varying reward structures, they have provided no evidence that the method addresses varying transition dynamics. In their experiments, they sample different MDPs, but each has the same transition dynamics, only the reward dynamics vary.\n\nPlease do not misunderstand me, I believe that science is incremental, small steps are valuable, and having a method which addresses only the smaller problem of fixed-transition dynamics with sampled reward dynamics is already an interesting and valuable contribution. Yet I don’t feel anyone is served by a paper which claims it solves the larger problem of offline Meta-RL, but doesn’t provide evidence for this larger claim.\n\nTherefore I would request the authors either provide evidence for the larger claim, or (it’s probably the easier route) simply make a smaller claim that is supported by the evidence, and save the real offline Meta-RL problem for a future submission. This would be that one is interested in the setting where MDPs with varying reward structure are sampled.\n\nI believe that the problem of MDP ambiguity can’t be so easily solved for ambiguous state transitions. In section 3.3, the authors nicely address the issue of MDP ambiguity, (it was a very insightful section, a valuable contribution) and that it’s impossible to know if two trajectories come from different MDPs with different rewards, or one MDP with rewards at both locations. It’s an inherent counterfactual problem, one doesn’t know what would have happened if one had done something else.\n\nThe authors fill in the gaps in this counterfactual knowledge by assuming knowledge of the true reward structure from each past MDP. In this way, for all past trajectories the authors can say what would have happened if the same actions had been taken in a setting with a different rewards structure.\n\nYet the problem of MDP ambiguity isn’t restricted to unknown reward structures. Assume a meta MDP setting where each MDP is a gridworld with a key and a door and the agent collects reward after getting the key then going to the door, thus the state is the agent’s position and a binary variable indicating whether the agent has collected the key. For different MDPs in the distribution, the key is placed in different places. In this way, the reward structure is the same for different MDPs in the distribution, but the transition structure (where the agent collects the key) changes. I don’t see how the author’s proposed MDP ambiguity fix would address ambiguity in the transition dynamics, or even provide a roadmap for addressing this issue.\n\nIndeed, it is telling that the experiments only deal with Meta-RL where the reward distribution changes, and never address changing transition dynamics.\n\nTherefore I would ask the authors to clarify the scope of their contribution, and perhaps address what challenges stand in the way of offline Meta-RL where the state transitions vary too. After that, I'd be happy to update my review.\n\nMinor Concerns:\n----------------------\nIn Prop. 1, why do we have $E_{\\mathcal R, \\mathcal P\\sim b_t}$, it seems one must only sample $\\mathcal P$ in the first line, and $\\mathcal R$ in the second.\n\n“Recall that the VariBAD VAE encoder….” perhaps drop the word ‘encoder’, encoder is already in the ‘E’ of VAE.\n\nIn section 3.3, when describing figure 2 (also in the caption of figure two) you say the reward is in the yellow circle, but the circle looks reddish on my PDF viewer.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "Btrex-_ReSf", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: \n\nThe paper proposes an extension of VariBAD to the offline setting. The main difficulty of such an application in an offline setting is what the paper termed \"MDP ambiguity\" while proposing reward relabeling as a solution. \n\n#########################################################\n\n Pros: \n\n1. The paper is clear and well-written. \n2. It proposes a simple but effective solution to adapt VariBAD to the offline setting for a type of environments  where transition probabilities and reward probabilities are independent. \n3. The paper demonstrates very interesting learned exploration behavior and shows its superior empirical performance compared to baselines (VariBAD and PEARL). \n\n#########################################################\n\nCons: \n\n1. The test environments are not extensive: all 4 test environments used in the paper are quite simple and 2 out of 4 are even similar. What about other openai gym control tasks (or Atari games)? Given the strong independence assumption between transition and reward probabilities, it would be nice to demonstrate how resilient this method is on environments that violate this assumption. \n\n2. There are few baselines being compared against in this paper so it's unclear given some offline data from an unknown environment this method is the sota. For example, one potential area of comparison is model-based RL agents where the prediction models (or even transition models if the environment is simple enough) are trained with supervised learning using offline data.\n\n#########################################################\n\n Questions during rebuttal period: \n\n1. Why is the top right slot of Table 1 blank? ", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper studies offline meta reinforcement learning. Overall the scope of this contribution seems limited. Reviewers have raised concerns about the significance of the presented results given the assumptions, and that the experimental environments are not extensive and do not fully support the claimed advances. 
Decision: Reject