Original: O
Revised: O

Original: FFLINE M ETA L EARNING OF E XPLORATION
Revised: FFLINE M ETA L EARNING OF E XPLORATION

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: This meta-learning problem is an instance of a setting we term Ofﬂine Meta Reinforcement Learning (OMRL).
Revised: This meta-learning problem is an instance of a setting we term Ofﬂine Meta Reinforcement Learning (OMRL).

Original: To solve our challenge, we take a Bayesian RL (BRL) view, and seek to learn a Bayesoptimal policy from the ofﬂine data.
Revised: To solve our challenge, we take a Bayesian RL (BRL) view, and seek to learn a Bayesoptimal policy from the ofﬂine data.

Original: We extend the recently proposed VariBAD
Revised: We extend the recently proposed VariBAD

Original: BRL algorithm to the off-policy setting, and demonstrate learning of approximately
Revised: BRL algorithm to the off-policy setting, and demonstrate learning of approximately

Original: Bayes-optimal exploration strategies from ofﬂine data using deep neural networks.
Revised: Bayes-optimal exploration strategies from ofﬂine data using deep neural networks.

Original: 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.
Revised: 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.

Original: Furthermore, we ﬁnd that when applied to the online meta-RL setting (agent simultaneously collects data and improves its meta-RL policy), our method is signiﬁcantly more sample efﬁcient than the state-of-the-art VariBAD.
Revised: Furthermore, we ﬁnd that when applied to the online meta-RL setting (agent simultaneously collects data and improves its meta-RL policy), our method is signiﬁcantly more sample efﬁcient than the state-of-the-art VariBAD.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: A central question in reinforcement learning (RL) is how to learn quickly (i.e., with few samples) in a new environment.
Revised: A central question in reinforcement learning (RL) is how to learn quickly (i.e., with few samples) in a new environment.

Original: Meta-RL addresses this issue by assuming a distribution over possible environments, and having access to a large set of environments from this distribution during training (Duan et al., 2016;
Revised: Meta-RL addresses this issue by assuming a distribution over possible environments, and having access to a large set of environments from this distribution during training (Duan et al., 2016;

Original: Finn et al., 2017).
Revised: Finn et al., 2017).

Original: Intuitively, the meta-RL agent can learn regularities in the environments, which allow quick learning in any environment that shares a similar structure.
Revised: Intuitively, the meta-RL agent can learn regularities in the environments, which allow quick learning in any environment that shares a similar structure.

Original: Indeed, recent work demonstrated this by training memory-based controllers that can ‘identify’ the domain (Duan et al., 2016;
Revised: Indeed, recent work demonstrated this by training memory-based controllers that can ‘identify’ the domain (Duan et al., 2016;

Original: Rakelly et al., 2019; Humplik et al., 2019), or by learning a parameter initialization that can lead to good performance with only a few gradient updates (Finn et al., 2017).
Revised: Rakelly et al., 2019; Humplik et al., 2019), or by learning a parameter initialization that can lead to good performance with only a few gradient updates (Finn et al., 2017).

Original: Another approach to quick RL is Bayesian RL (Ghavamzadeh et al., 2016, BRL).
Revised: Another approach to quick RL is Bayesian RL (Ghavamzadeh et al., 2016, BRL).

Original: In BRL, the environment parameters are treated as unobserved variables, with a known prior distribution.
Revised: In BRL, the environment parameters are treated as unobserved variables, with a known prior distribution.

Original: 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.
Revised: 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.

Original: Recently, Zintgraf et al.
Revised: Recently, Zintgraf et al.

Original: (2020) showed that meta-RL is in fact an instance of BRL, where the meta-RL environment distribution is simply the BRL prior.
Revised: (2020) showed that meta-RL is in fact an instance of BRL, where the meta-RL environment distribution is simply the BRL prior.

Original: Furthermore, a Bayes-optimal policy can be trained using standard policy gradient methods, simply by adding to the state the posterior belief over the environment parameters.
Revised: Furthermore, a Bayes-optimal policy can be trained using standard policy gradient methods, simply by adding to the state the posterior belief over the environment parameters.

Original: The VariBAD algorithm (Zintgraf et al., 2020) is an implementation of this approach that uses a variational autoencoder (VAE) for parameter estimation and deep neural network policies.
Revised: The VariBAD algorithm (Zintgraf et al., 2020) is an implementation of this approach that uses a variational autoencoder (VAE) for parameter estimation and deep neural network policies.

Original: 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.
Revised: 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.

Original: In domains where data collection is expensive, such as robotics and healthcare to name a few, online training is a limiting factor.
Revised: In domains where data collection is expensive, such as robotics and healthcare to name a few, online training is a limiting factor.

Original: For standard RL, ofﬂine (a.k.a. batch) RL mitigates this problem by learning from data collected beforehand by an arbitrary policy (Ernst et al., 2005; Levine et al., 2020).
Revised: For standard RL, ofﬂine (a.k.a. batch) RL mitigates this problem by learning from data collected beforehand by an arbitrary policy (Ernst et al., 2005; Levine et al., 2020).

Original: In this work we investigate the ofﬂine approach to meta-RL (OMRL).
Revised: In this work we investigate the ofﬂine approach to meta-RL (OMRL).

Original: It is well known that any ofﬂine RL approach is heavily inﬂuenced by the data collection policy.
Revised: It is well known that any ofﬂine RL approach is heavily inﬂuenced by the data collection policy.

Original: To ground our investigation, we focus on the following practical setting: we assume that data has been collected by running standard RL agents on a set of environments from the environment distribution.
Revised: To ground our investigation, we focus on the following practical setting: we assume that data has been collected by running standard RL agents on a set of environments from the environment distribution.

Original: Importantly - we do not allow any modiﬁcation to the RL algorithms used for data collection, and the meta-RL learner must make use of data that was not speciﬁcally collected for the meta-RL task.
Revised: Importantly - we do not allow any modiﬁcation to the RL algorithms used for data collection, and the meta-RL learner must make use of data that was not speciﬁcally collected for the meta-RL task.

Original: Nevertheless, we hypothesize that regularities between the training domains can still be learned, to provide faster learning in new environments.
Revised: Nevertheless, we hypothesize that regularities between the training domains can still be learned, to provide faster learning in new environments.

Original: Figure 1 illustrates our problem: in this navigation task, each RL agent in the data learned to ﬁnd its own goal, and converged to a behavior that quickly navigates toward it.
Revised: Figure 1 illustrates our problem: in this navigation task, each RL agent in the data learned to ﬁnd its own goal, and converged to a behavior that quickly navigates toward it.

Original: The meta-RL agent, on the other hand, needs to learn a completely different behavior that effectively searches for the unknown goal position.
Revised: The meta-RL agent, on the other hand, needs to learn a completely different behavior that effectively searches for the unknown goal position.

Original: 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.
Revised: 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.

Original: This, however, requires some care, as Q-learning applies to states of fully observed systems.
Revised: This, however, requires some care, as Q-learning applies to states of fully observed systems.

Original: We show that the VariBAD approach of augmenting states with the belief in the data applies to the off-policy setting as well, leading to an effective algorithm we term Off-Policy VariBAD.
Revised: We show that the VariBAD approach of augmenting states with the belief in the data applies to the off-policy setting as well, leading to an effective algorithm we term Off-Policy VariBAD.

Original: The ofﬂine setting, however, brings about another challenge – when the agent visits different parts of the state space in different environments, it becomes challenging to obtain an accurate belief estimate, a problem we term MDP ambiguity.
Revised: The ofﬂine setting, however, brings about another challenge – when the agent visits different parts of the state space in different environments, it becomes challenging to obtain an accurate belief estimate, a problem we term MDP ambiguity.

Original: We propose a simple solution, based on a reward relabelling trick that signiﬁcantly improves the performance of the VariBAD VAE trained on ofﬂine data.
Revised: When the ambiguity is due to differences in the reward between the environments, we propose a simple solution, based on a reward relabelling trick that signiﬁcantly improves the performance of the VariBAD VAE trained on ofﬂine data.

Original: Our experimental results show that our method can learn effective exploration policies from ofﬂine data on both discrete and continuous control problems.
Revised: Our experimental results show that our method can learn effective exploration policies from ofﬂine data on both discrete and continuous control problems.

Original: Our main contributions are as follows.
Revised: Our main contributions are as follows.

Original: To our knowledge, this is the ﬁrst study of meta learning exploration in the ofﬂine setting.
Revised: To our knowledge, this is the ﬁrst study of meta learning exploration in the ofﬂine setting.

Original: We provide the necessary theory to extend VariBAD to off-policy RL.
Revised: We provide the necessary theory to extend VariBAD to offpolicy RL.

Original: We show that a key difﬁculty in OMRL is MDP ambiguity, and propose an effective solution.
Revised: We show that a key difﬁculty in OMRL is MDP ambiguity, and propose an effective solution for the case where tasks differ in their rewards.

Original: We show non-trivial empirical results that demonstrate signiﬁcantly 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.
Revised: We show non-trivial empirical results that demonstrate signiﬁcantly 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.

Original: Finally, our method can also be applied in the online setting, and demonstrates signiﬁcantly improved sample efﬁciency compared to conventional VariBAD.
Revised: Finally, our method can also be applied in the online setting, and demonstrates signiﬁcantly improved sample efﬁciency compared to conventional VariBAD.

Original:  B ACKGROUND
Revised:  B ACKGROUND

Original: Our work leverages ideas from meta-RL, BRL and the VariBAD algorithm, as we now recapitulate.
Revised: Our work leverages ideas from meta-RL, BRL and the VariBAD algorithm, as we now recapitulate.

Original: Meta-RL: In meta-RL, a distribution over tasks is assumed.
Revised: Meta-RL: In meta-RL, a distribution over tasks is assumed.

Original: A task T i is described by a Markov
Revised: A task T i is described by a Markov

Original: Decision Process (MDP, Bertsekas, 1995)
Revised: Decision Process (MDP, Bertsekas, 1995)

Original: M i =
Revised: M i =

Original: ( S , A , R i , P i ) , where the state space S and the action space
Revised: ( S , A , R i , P i ) , where the state space S and the action space

Original: A are shared across tasks, and R i and P i are task speciﬁc reward and transition functions.
Revised: A are shared across tasks, and R i and P i are task speciﬁc reward and transition functions.

Original: Thus, we write the task distribution as p ( R , P ) .
Revised: Thus, we write the task distribution as p ( R , P ) .

Original: For simplicity, we assume throughout that the initial state distribution P init ( s 0 ) is the same for all MDPs.
Revised: For simplicity, we assume throughout that the initial state distribution P init ( s 0 ) is the same for all MDPs.

Original: The goal in meta-RL is to train an agent that can quickly maximize reward in new, unseen tasks, drawn from p ( R , P ) .
Revised: The goal in meta-RL is to train an agent that can quickly maximize reward in new, unseen tasks, drawn from p ( R , P ) .

Original: To do so, the agent must leverage any shared structure among tasks, which can typically be learned from a set of training tasks.
Revised: To do so, the agent must leverage any shared structure among tasks, which can typically be learned from a set of training tasks.

Original: Bayesian Reinforcement Learning:
Revised: Bayesian Reinforcement Learning:

Original: The goal in BRL is to ﬁnd the optimal policy π in an MDP, when the transitions and rewards are not known in advance.
Revised: The goal in BRL is to ﬁnd the optimal policy π in an MDP, when the transitions and rewards are not known in advance.

Original: Similar to meta-RL, we assume a prior over the MDP parameters p
Revised: Similar to meta-RL, we assume a prior over the MDP parameters p

Original: ( R , P ) , and seek to maximize the expected discounted return,
Revised: ( R , P ) , and seek to maximize the expected discounted return,

Original: [Equation]
Revised: [Equation]

Original: where the expectation is taken with respect to both the uncertainty in state-action transitions s t +1 ∼ P ( ·| s t , a t ) , a t ∼ π, and the uncertainty in the MDP parameters R , P ∼ p ( R , P ) .
Revised: where the expectation is taken with respect to both the uncertainty in state-action transitions s t +1 ∼ P ( ·| s t , a t ) , a t ∼ π, and the uncertainty in the MDP parameters R , P ∼ p ( R , P ) .

Original: 1 A key observation is that this formulation naturally accounts for the exploration/exploitation tradeoff – an optimal agent must plan its actions to reduce uncertainty in the MDP parameters, if such leads to higher rewards.
Revised: 1 A key observation is that this formulation naturally accounts for the exploration/exploitation tradeoff – an optimal agent must plan its actions to reduce uncertainty in the MDP parameters, if such leads to higher rewards.

Original: One way to approach the BRL problem is to model R , P as unobserved state variables in a partially observed MDP (POMDP, Cassandra et al., 1994), reducing the problem to solving a particular POMDP instance where the unobserved variables cannot change in time ( P and R are assumed to be stationary).
Revised: One way to approach the BRL problem is to model R , P as unobserved state variables in a partially observed MDP (POMDP, Cassandra et al., 1994), reducing the problem to solving a particular POMDP instance where the unobserved variables cannot change in time ( P and R are assumed to be stationary).

Original: The belief at time t , b t , denotes the posterior probability over R , P given the history of state transitions and rewards observed until this time b t
Revised: The belief at time t , b t , denotes the posterior probability over R , P given the history of state transitions and rewards observed until this time b t

Original: = P ( R , P| h : t ) , where h :
Revised: = P ( R , P| h : t ) , where h :

Original: t = { s 0 , a 0 , r 1 , s 1 . .
Revised: t = { s 0 , a 0 , r 1 , s 1 . .

Original: , r t , s t } (note that we denote the reward after observing the state and action at time t as r t +1 =
Revised: , r t , s t } (note that we denote the reward after observing the state and action at time t as r t +1 =

Original: r
Revised: r

Original: ( s t , a t ) ).
Revised: ( s t , a t ) ).

Original: The belief can be updated iteratively according to Bayes rule, where b 0 ( R , P ) =
Revised: The belief can be updated iteratively according to Bayes rule, where b 0 ( R , P ) =

Original: p ( R , P ) , and: [Equation] t +
Revised: p ( R , P ) , and: b t +1 ( R , P ) =

Original: [Equation] [Equation] h 
Revised: P ( R , P| h : t +1 )

Original: [Equation] , P [Equation] R , P ) .
Revised: ∝ P ( s t +1 , r t +1 | h : t , R , P ) b t ( R , P ) .

Original: Similar to the idea of solving a POMDP by representing it as an MDP over belief states (Cassandra et al., 1994), the state in BRL can be augmented with the belief to result in the Bayes-Adaptive
Revised: Similar to the idea of solving a POMDP by representing it as an MDP over belief states (Cassandra et al., 1994), the state in BRL can be augmented with the belief to result in the Bayes-Adaptive

Original: MDP (BAMDP) model (Duff & Barto, 2002).
Revised: MDP (BAMDP) model (Duff & Barto, 2002).

Original: Denote the augmented state s + t
Revised: Denote the augmented state s + t

Original: =
Revised: =

Original: ( s t , b t ) and the augmented state space S + = S ×B , where B denotes the belief space.
Revised: ( s t , b t ) and the augmented state space S + = S ×B , where B denotes the belief space.

Original: The transitions in the BAMDP are given by:
Revised: The transitions in the BAMDP are given by:

Original: P + ( s + t +1 | s + t , a t ) =
Revised: P + ( s + t +1 | s + t , a t ) =

Original: E b t [ P ( s t +1 | s t , a t )]
Revised: E b t [ P ( s t +1 | s t , a t )]

Original: δ ( b t +1 =
Revised: δ ( b t +1 =

Original: P
Revised: P

Original: ( R , P| h : t +1 )) , and the reward in the BAMDP is the expected reward with respect to the belief: R + ( s + t , a t ) =
Revised: ( R , P| h : t +1 )) , and the reward in the BAMDP is the expected reward with respect to the belief: R + ( s + t , a t ) =

Original: E b t [ R ( s t , a t )] .
Revised: E b t [ R ( s t , a t )] .

Original: 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.
Revised: 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.

Original: As in standard MDPs, the optimal action-value function in the BAMDP satisﬁes the Bellman equation:
Revised: As in standard MDPs, the optimal action-value function in the BAMDP satisﬁes the Bellman equation:

Original: [Equation]
Revised: [Equation]

Original: Computing a Bayes-optimal agent amounts to solving the BAMDP, where the optimal policy is a function of the augmented state.
Revised: Computing a Bayes-optimal agent amounts to solving the BAMDP, where the optimal policy is a function of the augmented state.

Original: 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.
Revised: 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.

Original: The VariBAD Algorithm: VariBAD (Zintgraf et al., 2020) approximates the Bayes-optimal solution by combining a model for the MDP parameter uncertainty, and an optimization method for the corresponding BAMDP.
Revised: The VariBAD Algorithm: VariBAD (Zintgraf et al., 2020) approximates the Bayes-optimal solution by combining a model for the MDP parameter uncertainty, and an optimization method for the corresponding BAMDP.

Original: The MDP parameters are represented by a vector m ∈ R d , corresponding to the latent variables in a parametric generative model for the statereward trajectory distribution conditioned on the actions P ( s 0 , r 1 , s 1 . . .
Revised: The MDP parameters are represented by a vector m ∈ R d , corresponding to the latent variables in a parametric generative model for the statereward trajectory distribution conditioned on the actions P ( s 0 , r 1 , s 1 . . .

Original: , r H , s H | a 0 , . .
Revised: , r H , s H | a 0 , . .

Original: . , a H − 1 ) =
Revised: . , a H − 1 ) =

Original: (cid:82) p θ ( m ) p
Revised: (cid:82) p θ ( m ) p

Original: θ ( s 0 , r 1 , s 1 . .
Revised: θ ( s 0 , r 1 , s 1 . .

Original: , r H , s H | m, a 0 , . .
Revised: , r H , s H | m, a 0 , . .

Original: . , a H − 1 ) dm .
Revised: . , a H − 1 ) dm .

Original: The model parameters θ are learned by a variational approximation to the maximum likelihood objective, where the variational approximation to the posterior P ( m | s 0 , r 1 , s 1 . .
Revised: The model parameters θ are learned by a variational approximation to the maximum likelihood objective, where the variational approximation to the posterior P ( m | s 0 , r 1 , s 1 . .

Original: , r H , s H , a 0 , . .
Revised: , r H , s H , a 0 , . .

Original: . , a H − 1 ) is chosen to have the structure q φ ( m | s 0 , a 0 , r 1 , s 1 . . .
Revised: . , a H − 1 ) is chosen to have the structure q φ ( m | s 0 , a 0 , r 1 , s 1 . . .

Original: , r t , s t )
Revised: , r t , s t )

Original: =
Revised: =

Original: q φ ( m | h : t ) .
Revised: q φ ( m | h : t ) .

Original: That is, the approximate posterior is conditioned on the history up to time t .
Revised: That is, the approximate posterior is conditioned on the history up to time t .

Original: The evidence lower bound (ELBO) in this case is ELBO t =
Revised: The evidence lower bound (ELBO) in this case is ELBO t =

Original: E m ∼ q φ ( ·| h : t ) [log p θ ( s 0 , r 1 , s 1 . . .
Revised: E m ∼ q φ ( ·| h : t ) [log p θ ( s 0 , r 1 , s 1 . . .

Original: , r H , s H | m, a 0 , . .
Revised: , r H , s H | m, a 0 , . .

Original: . , a H − 1 )]
Revised: . , a H − 1 )]

Original: − D KL ( q φ ( m | h : t ) ||
Revised: − D KL ( q φ ( m | h : t ) ||

Original: p θ ( m )) .
Revised: p θ ( m )) .

Original: The main claim of Zintgraf et al.
Revised: The main claim of Zintgraf et al.

Original: (2020) is that q φ ( m | h : t ) can be taken as an approximation of the belief b t .
Revised: (2020) is that q φ ( m | h : t ) can be taken as an approximation of the belief b t .

Original: In practice, q φ ( m | h : t ) is represented as a Gaussian distribution q ( m | h : t ) =
Revised: In practice, q φ ( m | h : t ) is represented as a Gaussian distribution q ( m | h : t ) =

Original: N ( µ ( h : t ) , Σ( h : t )) , where µ and Σ are learned recurrent neural networks.
Revised: N ( µ ( h : t ) , Σ( h : t )) , where µ and Σ are learned recurrent neural networks.

Original: To approximately solve the BAMDP, Zintgraf et al.
Revised: To approximately solve the BAMDP, Zintgraf et al.

Original: (2020) exploit the fact that an optimal BAMDP policy is a function of the state and belief, and therefore consider neural network policies that take the augmented BAMDP state as input π ( a t | s t , q φ ( m | h : t )) , where the posterior is practically represented by the distribution parameters µ ( h : t ) ,
Revised: (2020) exploit the fact that an optimal BAMDP policy is a function of the state and belief, and therefore consider neural network policies that take the augmented BAMDP state as input π ( a t | s t , q φ ( m | h : t )) , where the posterior is practically represented by the distribution parameters µ ( h : t ) ,

Original: Σ( h : t ) .
Revised: Σ( h : t ) .

Original: To train such policies, Zintgraf et al.
Revised: To train such policies, Zintgraf et al.

Original: (2020) maximize the BRL objective,
Revised: (2020) maximize the BRL objective,

Original: [Equation]
Revised: [Equation]

Original: using policy gradient based methods.
Revised: using policy gradient based methods.

Original: The expectation over MDP parameters in (3) is approximated by averaging over training environments, and the RL agent is trained online, alongside the VAE.
Revised: The expectation over MDP parameters in (3) is approximated by averaging over training environments, and the RL agent is trained online, alongside the VAE.

Original:  OMRL WITH O FF -P OLICY V ARI BAD
Revised:  OMRL WITH O FF -P OLICY V ARI BAD

Original: In this section, we derive an off-policy variant of the VariBAD algorithm, and apply it to the OMRL problem.
Revised: In this section, we derive an off-policy variant of the VariBAD algorithm, and apply it to the OMRL problem.

Original: We begin by describing our OMRL problem setting, and then present our algorithm.
Revised: We begin by describing our OMRL problem setting, and then present our algorithm.

Original: 3. OMRL P ROBLEM D EFINITION
Revised: 3. OMRL P ROBLEM D EFINITION

Original: We follow the Meta-RL and BRL formulation described above, with a prior distribution over MDP parameters p
Revised: We follow the Meta-RL and BRL formulation described above, with a prior distribution over MDP parameters p

Original: ( R , P ) .
Revised: ( R , P ) .

Original: We are provided training data of an agent interacting with N different MDPs, {R i , P i } Ni =1 , sampled from the prior.
Revised: We are provided training data of an agent interacting with N different MDPs, {R i , P i } Ni =1 , sampled from the prior.

Original: We assume that each interaction is organized as M trajectories of length H , τ i,j = s i,j 0 , a i,j 0 , r i,j 1 , s i,j 1 . . .
Revised: We assume that each interaction is organized as M trajectories of length H , τ i,j = s i,j 0 , a i,j 0 , r i,j 1 , s i,j 1 . . .

Original: , r i,jH , s i,jH , i ∈ 1 , . . .
Revised: , r i,jH , s i,jH , i ∈ 1 , . . .

Original: , N, j ∈ 1 , . .
Revised: , N, j ∈ 1 , . .

Original: , M , where the rewards satisfy r i,jt +

Revised: , M , where the rewards satisfy r i,jt +

Original: = R i
Revised: = R i

Original: ( s i,jt , a i,jt ) , the transitions satisfy
Revised: ( s i,jt , a i,jt ) , the transitions satisfy

Original: s i,jt +1 ∼ P
Revised: s i,jt +1 ∼ P

Original: i ( ·| s i,jt , a i,jt ) , and the actions are chosen from an arbitrary data collection policy.
Revised: i ( ·| s i,jt , a i,jt ) , and the actions are chosen from an arbitrary data collection policy.

Original: To ground our work in a speciﬁc context, we further assume that the trajectories are obtained from running a conventional RL agent (i.e., the complete RL training history) in each one of the MDPs, which implicitly speciﬁes the data collection policy.
Revised: To ground our work in a speciﬁc context, we further assume that the trajectories are obtained from running a conventional RL agent (i.e., the complete RL training history) in each one of the MDPs, which implicitly speciﬁes the data collection policy.

Original: In our experiments, we will later investigate the implications of this assumption.
Revised: In our experiments, we will later investigate the implications of this assumption.

Original: Our goal is to use the data for learning a Bayes-optimal policy, i.e., a policy π that maximizes Eq.
Revised: Our goal is to use the data for learning a Bayes-optimal policy, i.e., a policy π that maximizes Eq.

Original: 3. O FF -P OLICY V ARI BAD
Revised: 3. O FF -P OLICY V ARI BAD

Original: The online VariBAD algorithm updates the policy using trajectories sampled from the current policy , and thus cannot be applied to our ofﬂine setting.
Revised: The online VariBAD algorithm updates the policy using trajectories sampled from the current policy , and thus cannot be applied to our ofﬂine setting.

Original: Our ﬁrst step is to modify VariBAD to work off-policy.
Revised: Our ﬁrst step is to modify VariBAD to work off-policy.

Original: We start with an observation about the use of the BAMDP formulation in VariBAD, which will motivate our subsequent development.
Revised: We start with an observation about the use of the BAMDP formulation in VariBAD, which will motivate our subsequent development.

Original: Does VariBAD really optimize the BAMDP?
Revised: Does VariBAD really optimize the BAMDP?

Original: Recall that a BAMDP is in fact a reduction of a
Revised: Recall that a BAMDP is in fact a reduction of a

Original: POMDP to an MDP over augmented states s + = ( s, b ) , and with the rewards and transitions given by R + and P + .
Revised: POMDP to an MDP over augmented states s + = ( s, b ) , and with the rewards and transitions given by R + and P + .

Original: Thus, an optimal Markov policy for the BAMDP exists in the form of π ( s + ) .
Revised: Thus, an optimal Markov policy for the BAMDP exists in the form of π ( s + ) .

Original: The
Revised: The

Original: VariBAD policy, as described above, similarly takes as input the augmented state, and is thus capable of representing an optimal BAMDP policy.
Revised: VariBAD policy, as described above, similarly takes as input the augmented state, and is thus capable of representing an optimal BAMDP policy.

Original: However, VariBAD’s policy optimization in Eq.
Revised: However, VariBAD’s policy optimization in Eq.

Original: (3) does not make use of the BAMDP parameters R + and P + !
Revised: (3) does not make use of the BAMDP parameters R + and P + !

Original: While at ﬁrst this may seem counterintuitive,
Revised: While at ﬁrst this may seem counterintuitive,

Original: Eq.
Revised: Eq.

Original: (3) is in fact a sound objective for the BAMDP, as we now show 2 .
Revised: (3) is in fact a sound objective for the BAMDP, as we now show 2 .

Original: Proposition 1.
Revised: Proposition 1.

Original: Let τ = s 0 , a 0 , r 1 , s 1 . .
Revised: Let τ = s 0 , a 0 , r 1 , s 1 . .

Original: , r H , s H denote a random trajectory from a ﬁxed history dependent policy π , generated according to the following process.
Revised: , r H , s H denote a random trajectory from a ﬁxed history dependent policy π , generated according to the following process.

Original: First, MDP parameters R , P are drawn from the prior p ( R , P ) .
Revised: First, MDP parameters R , P are drawn from the prior p ( R , P ) .

Original: Then, the state trajectory is generated according to s 0 ∼ P init , a t ∼ π ( ·| s 0 , a 0 , r 1 , . . .
Revised: Then, the state trajectory is generated according to s 0 ∼ P init , a t ∼ π ( ·| s 0 , a 0 , r 1 , . . .

Original: , s t ) , s t +1 ∼ P ( ·| s t , a t ) and r t +1 ∼ R ( s t , a t ) .
Revised: , s t ) , s t +1 ∼ P ( ·| s t , a t ) and r t +1 ∼ R ( s t , a t ) .

Original: Let b t denote the posterior belief at time t, b t =
Revised: Let b t denote the posterior belief at time t, b t =

Original: P ( R , P| s 0 , a 0 , r 1 , . . .
Revised: P ( R , P| s 0 , a 0 , r 1 , . . .

Original: , s t ) .
Revised: , s t ) .

Original: Then
Revised: Then

Original: P ( s t +1 | s 0 , a 0 , r 1 , . . .
Revised: P ( s t +1 | s 0 , a 0 , r 1 , . . .

Original: , r t , s t , a t ) =
Revised: , r t , s t , a t ) =

Original: E R , P∼ b t P ( s t +1 | s t , a t ) , and,
Revised: E R , P∼ b t P ( s t +1 | s t , a t ) , and,

Original: P ( r t +1 | s 0 , a 0 , r 1 , . . .
Revised: P ( r t +1 | s 0 , a 0 , r 1 , . . .

Original: , s t , a t ) =
Revised: , s t , a t ) =

Original: E R , P∼ b t R ( r t +1 | s t , a t ) .
Revised: E R , P∼ b t R ( r t +1 | s t , a t ) .

Original: 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.
Revised: 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.

Original: 3 is valid.
Revised: 3 is valid.

Original: However, off-policy RL does not take as input trajectories, but tuples of the form (state, action, reward, next state), where states and actions can be sampled from any distribution .
Revised: However, off-policy RL does not take as input trajectories, but tuples of the form (state, action, reward, next state), where states and actions can be sampled from any distribution .

Original: For an arbitrary distribution of augmented states, we must replace the rewards and transitions in our data with R + and P + , which can be difﬁcult to compute.
Revised: For an arbitrary distribution of augmented states, we must replace the rewards and transitions in our data with R + and P + , which can be difﬁcult to compute.

Original: Fortunately, Propositionshows 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.
Revised: Fortunately, Propositionshows 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.

Original: Based on Proposition 1, we can use a state augmentation method similar to VariBAD, which we refer to as state relabelling .
Revised: Based on Proposition 1, we can use a state augmentation method similar to VariBAD, which we refer to as state relabelling .

Original: Consider each trajectory in our data τ i,j = s i,j 0 , a i,j 0 , r i,j 1 , . . .
Revised: Consider each trajectory in our data τ i,j = s i,j 0 , a i,j 0 , r i,j 1 , . . .

Original: , s i,jH , as deﬁned above.
Revised: , s i,jH , as deﬁned above.

Original: Recall that the VariBAD VAE encoder provides an estimate of the belief given the state history q ( m | h : t ) =
Revised: Recall that the VariBAD VAE encoder provides an estimate of the belief given the state history q ( m | h : t ) =

Original: N ( µ ( h : t ) , Σ( h : t )) .
Revised: N ( µ ( h : t ) , Σ( h : t )) .

Original: Thus, we can run the encoder on every partial t -length history τ i,j : t to obtain the belief at each time step.
Revised: Thus, we can run the encoder on every partial t -length history τ i,j : t to obtain the belief at each time step.

Original: Following the BAMDP formulation, we deﬁne the augmented state s + ,i,j t =
Revised: Following the BAMDP formulation, we deﬁne the augmented state s + ,i,j t =

Original: ( s i,jt , b i,jt ) , where b i,jt = µ ( τ i,j : t ) , Σ
Revised: ( s i,jt , b i,jt ) , where b i,jt = µ ( τ i,j : t ) , Σ

Original: ( τ i,j : t ) .
Revised: ( τ i,j : t ) .

Original: We next replace each state in our data s i,jt with s + ,i,j t , effectively transforming the data to as coming from a BAMDP.
Revised: We next replace each state in our data s i,jt with s + ,i,j t , effectively transforming the data to as coming from a BAMDP.

Original: After applying state relabelling, any off-policy RL algorithm can be applied to the modiﬁed data, for learning a Bayes-optimal policy.
Revised: After applying state relabelling, any off-policy RL algorithm can be applied to the modiﬁed data, for learning a Bayes-optimal policy.

Original: 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.
Revised: 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.

Original: 3. O FFLINE V ARI BAD WITH R EWARD R ELABELLING
Revised: 3. O FFLINE M ETA -RL AND MDP A MBIGUITY

Original: We next consider the OMRL problem.
Revised: We next consider the OMRL problem.

Original: While in principle, it is possible to simply run off-policy
Revised: While in principle, it is possible to simply run off-policy

Original: VariBAD on the ofﬂine data, we claim that in most problems this may not work well.
Revised: VariBAD on the ofﬂine data, we claim that in many problems this may not work well.

Original: 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.
Revised: 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.

Original: However, since the trajectories from each MDP can visit different parts of the state space, it becomes impossible to identify if two trajectories come from different MDPs, or actually come from different parts of the same MDP.
Revised: However, if trajectories from the different MDPs visit different parts of the state space, it becomes impossible to identify if two trajectories come from different MDPs, or actually come from different parts of the same MDP.

Original: 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.
Revised: 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.

Original: If the data contains trajectories similar to the ones in the ﬁgure, it is impossible to distinguish between having two different MDPs with the indicated rewards, or a single MDP with rewards at both the blue and yellow circles.
Revised: If the data contains trajectories similar to the ones in the ﬁgure, it is impossible to distinguish between having two different MDPs with the indicated rewards, or a single MDP with rewards at both the blue and yellow circles.

Original: Note that MDP ambiguity is special to the ofﬂine meta-RL setting, as in online meta-RL, the agent will be driven by the online adapting policy  to explore states that reduce its ambiguity.
Revised: Note that MDP ambiguity is special to the ofﬂine meta-RL setting, as in online meta-RL, the agent may be driven by the online adapting policy (or guided explicitly) to explore states that reduce its ambiguity.

Original: We also emphasize that this problem is not encountered in standard (non-meta) ofﬂine
Revised: We also emphasize that this problem is not encountered in standard (non-meta) ofﬂine

Original: RL, as the problem here concerns the identiﬁcation of the MDP , which in standard RL is unique.
Revised: RL, as the problem here concerns the identiﬁcation of the MDP , which in standard RL is unique.

Revised: We next characterize MDP ambiguity in more detail.

Revised: Different MDPs can differ from each other either in the reward function, the transition function, or both.

Revised: The differences can also take place over the whole state space, or only in parts of it – i.e., in some states the MDP rewards or transitions are similar, while on others they are different.

Original: One way to solve MDP ambiguity is to increase the variability in the data, however, this may not be possible if the data has already been collected.
Revised: Let us denote the set of states where there are such differences as the ‘identifying states’ of an MDP – if the agent has data on 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 to identify from which MDP the data came from (cf. Figure 2).

Revised: 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 3.1, each agent mainly visits the rewarding states in its own MDP.

Revised: It is important to note that there are problems where MDP ambiguity is not an issue.

Revised: For example, MDPs with different transitions, where the set of rewarding states are similar.

Original: Instead, we propose a simple and effective alternative.
Revised: In this case, all agents will mainly visit the same states (though by applying different policies), and therefore the MDPs will be identiﬁable. 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 our discussion to ambiguity due to different rewards, and ambiguity due to different transitions.

Revised: 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.

Revised: Instead, we propose a simple and effective alternative in case where the MDPs differ only in their reward function.

Original: We make the following assumptions:(1) The transition and reward uncertainties are decoupled, i.e., the prior can be written as p ( R , P ) =
Revised: We make the following assumptions: The transition and reward uncertainties are decoupled, i.e., the prior can be written as p ( R , P ) =

Original: p R ( R ) × p P ( P ) ; and (2) For each MDP
Revised: p R ( R ) × p P ( P ) .

Original:  i in the training data, we know the reward function R i .
Revised: For each MDP i in the training data, we know the reward function R i .

Original: These assumption are largely satisﬁed in most meta-RL studies to date.
Revised: These assumptions are largely satisﬁed in most meta-RL studies to date (e.g., Zintgraf et al. 2020;

Revised: Li et al. 2020; Finn et al. 2017; Duan et al. 2016).

Original: The decoupled uncertainty assumption can be hard to verify in practice.
Revised: The decoupled uncertainty assumption can be hard to verify in practice.

Original: 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 p R ( R ) > 0 , p P ( P ) > 0 .
Revised: 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 p R ( R ) > 0 , p P ( P ) > 0 .

Original: Thus, this assumption may lead to less effective estimation (as there are more MDP ‘possibilities’), but should not prohibit the agent from estimating the true posterior with enough data.
Revised: Thus, this assumption may lead to less effective estimation (as there are more MDP ‘possibilities’), but should not prohibit the agent from estimating the true posterior with enough data.

Original: Wenext introduce reward relabelling , a simple solution to the MDP ambiguity problem.
Revised: We introduce reward relabelling , a simple solution to the MDP ambiguity problem.

Original: We propose to make the state distribution in the ofﬂine data approximately uniform across all MDPs, by replacing the rewards in a trajectory from some MDP i in the data with rewards from another randomly chosen MDP i (cid:48) (cid:54)
Revised: We propose to make the state distribution in the ofﬂine data approximately uniform across all MDPs, by replacing the rewards in a trajectory from some MDP i in the data with rewards from another randomly chosen MDP i (cid:48) (cid:54)

Original: =
Revised: =

Original: i .
Revised: i .

Original: That is, for each i ∈ 1 , . . .
Revised: That is, for each i ∈ 1 , . . .

Original: , N , we add to the data M trajectories ˆ τ i,j , j ∈ 1 , . . .
Revised: , N , we add to the data M trajectories ˆ τ i,j , j ∈ 1 , . . .

Original: , M , where ˆ τ i,j = ( s i,j 0 , a i,j 0 , ˆ r i,j 1 , s i,j 1 . . .
Revised: , M , where ˆ τ i,j = ( s i,j 0 , a i,j 0 , ˆ r i,j 1 , s i,j 1 . . .

Original: , ˆ r i,jH , s i,jH ) , where the relabelled rewards ˆ r satisfy ˆ r i,jt +

Revised: , ˆ r i,jH , s i,jH ) , where the relabelled rewards ˆ r satisfy ˆ r i,jt +

Original: = R i (cid:48)
Revised: = R i (cid:48)

Original: ( s i,jt , a i,jt ) .
Revised: ( s i,jt , a i,jt ) .

Original: Note that our relabelling effectively samples data from an MDP with transitions P i and reward R i (cid:48) , which has non-zero prior probability mass under the decoupled prior assumption.
Revised: Note that our relabelling effectively samples data from an MDP with transitions P i and reward R i (cid:48) , which has non-zero prior probability mass under the decoupled prior assumption.

Original: We only use the relabelled data for training the VariBAD VAE.
Revised: We only use the relabelled data for training the VariBAD VAE.

Original: While in principle it could also be used for training the off-policy RL algorithm, we did not ﬁnd it useful in practice.
Revised: While in principle it could also be used for training the off-policy RL algorithm, we did not ﬁnd it useful in practice.

Revised: MDP ambiguity due to transitions: when ambiguity is due to different transitions, the problem becomes signiﬁcantly more complicated, as we cannot simply ‘relabel’ trajectories in the data.

Revised: 4 One solution is to add trajectories that run the policy for agent i on task i (cid:48)

Revised: (this is a generalization of the reward relabelling idea).

Original:   R ELATED W ORK
Revised: This, of course, requires collecting more data, and is not suitable for the ofﬂine 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.  R ELATED W ORK

Original: Meta-learning considers training agents that quickly solve a new learning problem, by exploiting structure in the problem distribution (Thrun & Pratt, 1998; Hochreiter et al., 2001).
Revised: Meta-learning considers training agents that quickly solve a new learning problem, by exploiting structure in the problem distribution (Thrun & Pratt, 1998; Hochreiter et al., 2001).

Original: In this work we focus on meta-RL – quickly learning to solve RL problems.
Revised: In this work we focus on meta-RL – quickly learning to solve RL problems.

Original: Gradient based approaches to meta- RL seek policy parameters that can be updated to the current task with a few gradient steps (Finn et al., 2017; Grant et al., 2018;
Revised: Gradient based approaches to metaRL seek policy parameters that can be updated to the current task with a few gradient steps (Finn et al., 2017; Grant et al., 2018;

Original: Rothfuss et al., 2018; Clavera et al., 2018).
Revised: Rothfuss et al., 2018; Clavera et al., 2018).

Original: These are essentially online methods, and several studies investigated learning of structured exploration strategies in this setting (Gupta et al., 2018;
Revised: These are essentially online methods, and several studies investigated learning of structured exploration strategies in this setting (Gupta et al., 2018;

Original: Rothfuss et al., 2018; Stadie et al., 2018).
Revised: Rothfuss et al., 2018; Stadie et al., 2018).

Original: Memory-based meta-RL, on the other hand, map the observed history in a task h : t to an action (Duan et al., 2016; Wang et al., 2016).
Revised: Memory-based meta-RL, on the other hand, map the observed history in a task h : t to an action (Duan et al., 2016; Wang et al., 2016).

Original: These methods effectively treat the problem as a POMDP, and learn a memory based controller for it.
Revised: These methods effectively treat the problem as a POMDP, and learn a memory based controller for it.

Original: The connection between meta-learning and Bayesian methods, and between meta-RL and Bayesian
Revised: The connection between meta-learning and Bayesian methods, and between meta-RL and Bayesian

Original: RL in particular, has been investigated in a series of recent papers (Lee et al., 2018; Humplik et al., 2019; Ortega et al., 2019; Zintgraf et al., 2020), and our work closely follows these ideas.
Revised: RL in particular, has been investigated in a series of recent papers (Lee et al., 2018; Humplik et al., 2019; Ortega et al., 2019; Zintgraf et al., 2020), and our work closely follows these ideas.

Original: In particular, these works elucidate the difference between Thompson-sampling based strategies, such as in PEARL (Rakelly et al., 2019), and Bayes-optimal policies, such as in VariBAD (Zintgraf et al., 2020), and suggest to estimate the BAMDP belief using the latent state of deep generative models.
Revised: In particular, these works elucidate the difference between Thompson-sampling based strategies, such as in PEARL (Rakelly et al., 2019), and Bayes-optimal policies, such as in VariBAD (Zintgraf et al., 2020), and suggest to estimate the BAMDP belief using the latent state of deep generative models.

Original: Our contribution is an extension of these ideas to the ofﬂine RL setting , which to the best of our knowledge is novel.
Revised: Our contribution is an extension of these ideas to the ofﬂine RL setting , which to the best of our knowledge is novel.

Original: Technically, the VariBAD algorithm in Zintgraf et al.
Revised: Technically, the VariBAD algorithm in Zintgraf et al.

Original: (2020) is limited to on-policy RL, and the off-policy method in Humplik et al.
Revised: (2020) is limited to on-policy RL, and the off-policy method in Humplik et al.

Original: (2019) requires speciﬁc task descriptors during learning, while VariBAD, which our work is based on, does not.
Revised: (2019) requires speciﬁc task descriptors during learning, while VariBAD, which our work is based on, does not.

Original: Concurrently and independently with our work , Li et al.
Revised: Concurrently and independently with our work , Li et al.

Original: (2020) proposed an ofﬂine meta-RL algorithm that combines BCQ (Fujimoto et al., 2018) with a task inference module.
Revised: (2020) proposed an ofﬂine meta-RL algorithm that combines BCQ (Fujimoto et al., 2018) with a task inference module.

Original: Interestingly,
Revised: Interestingly,

Original: Li et al.
Revised: Li et al.

Original: (2020) also describe a problem similar to MDP ambiguity, and resolve it using a technique similar to reward relabelling.
Revised: (2020) also describe a problem similar to MDP ambiguity, and resolve it using a technique similar to reward relabelling.

Original: However, their approach does not take into account the task uncertainty, and cannot plan actions that actively explore to reduce this uncertainty – this is a form of Thompson sampling, where a task-conditional policy reacts to the task inference (see Figure 1 in Li et al. 2020).
Revised: However, their approach does not take into account the task uncertainty, and cannot plan actions that actively explore to reduce this uncertainty – this is a form of Thompson sampling, where a task-conditional policy reacts to the task inference (see Figure 1 in Li et al. 2020).

Original: Our work is the ﬁrst to tackle ofﬂine meta-learning of Bayes-optimal exploration .
Revised: Our work is the ﬁrst to tackle ofﬂine meta-learning of Bayes-optimal exploration .

Original: In addition, we demonstrate the ﬁrst ofﬂine results on sparse reward tasks, which, compared to the dense reward tasks in Li et al.
Revised: In addition, we demonstrate the ﬁrst ofﬂine results on sparse reward tasks, which, compared to the dense reward tasks in Li et al.

Original: (2020), require a signiﬁcantly more complicated solution than Thompson sampling (see experiments section).
Revised: (2020), require a signiﬁcantly more complicated solution than Thompson sampling (see experiments section).

Original: We achieve this by building on BRL theory, which both optimizes for
Revised: We achieve this by building on BRL theory, which both optimizes for

Original: Bayes-optimality and results in a much simpler algorithm.
Revised: Bayes-optimality and results in a much simpler algorithm.

Original: Recent work on meta Q-learning (Fakoor et al., 2019) also does not incorporate task uncertainty, and thus cannot be Bayes-optimal.
Revised: Recent work on meta Q-learning (Fakoor et al., 2019) also does not incorporate task uncertainty, and thus cannot be Bayes-optimal.

Original: The very recent work of Mitchell et al.
Revised: The very recent work of Mitchell et al.

Original: (2020) considers a different ofﬂine meta-RL setting, where an ofﬂine dataset from the test environment is available.
Revised: (2020) considers a different ofﬂine meta-RL setting, where an ofﬂine dataset from the test environment is available.

Original: Classical works on BRL are comprehensively surveyed in Ghavamzadeh et al.
Revised: Classical works on BRL are comprehensively surveyed in Ghavamzadeh et al.

Original: Our work, in comparison, allows training scalable deep BRL policies.
Revised: Our work, in comparison, allows training scalable deep BRL policies.

Original: Finally, there is growing interest in ofﬂine deep RL (Saraﬁan et al., 2018; Levine et al., 2020).
Revised: Finally, there is growing interest in ofﬂine deep RL (Saraﬁan et al., 2018; Levine et al., 2020).

Original: Most recent work focus on how to avoid actions that were not sampled enough in the data.
Revised: Most recent work focus on how to avoid actions that were not sampled enough in the data.

Original: In our experiments, a state-of-the-art method of this ﬂavor led to minor improvements, though future ofﬂine RL developments may possibly beneﬁt OMRL too.
Revised: In our experiments, a state-of-the-art method of this ﬂavor led to minor improvements, though future ofﬂine RL developments may possibly beneﬁt OMRL too.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: In our experiments, we aim to answer the following questions: (1) can we learn approximately
Revised: In our experiments, we aim to answer the following questions: (1) can we learn approximately

Original: Bayes-optimal policies in the ofﬂine setting?
Revised: Bayes-optimal policies in the ofﬂine setting?

Original: (2) what are the main practical challenges in OMRL?
Revised: (2) what are the main practical challenges in OMRL?

Original: and (3) does our off-policy method improve meta-RL performance in the online setting as well?
Revised: and (3) does our off-policy method improve meta-RL performance in the online setting as well?

Original: Answering (1) is difﬁcult because the Bayes-optimal policy is generally intractable, and because our results crucially depend on the available data.
Revised: Answering (1) is difﬁcult because the Bayes-optimal policy is generally intractable, and because our results crucially depend on the available data.

Original: However, in deterministic domains with a single sparse reward, the optimal solution amounts to ‘search all possible goal locations as efﬁciently as possible, and stay at goal once found; in subsequent episodes, move directly to goal.’.
Revised: However, in deterministic domains with a single sparse reward, the optimal solution amounts to ‘search all possible goal locations as efﬁciently as possible, and stay at goal once found; in subsequent episodes, move directly to goal.’.

Original: We therefore chose domains where this behavior can be identiﬁed qualitatively.
Revised: We therefore chose domains where this behavior can be identiﬁed qualitatively.

Original: Quantitatively, we compare our ofﬂine 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 ofﬂine data restriction.
Revised: Quantitatively, we compare our ofﬂine 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 ofﬂine data restriction.

Original: Domains and evaluation metric: Our emphasis is on learning to explore efﬁciently.
Revised: Domains and evaluation metric: Our emphasis is on learning to explore efﬁciently.

Original: We thus focus on sparse reward problems where non-trivial exploration is required to identify the task: (1) A discrete 5 × 5 Gridworld (Zintgraf et al., 2020); (2) A continuous point robot where a sparse reward is located somewhere on a semi-circle (see Figure 1); (3) Ant-Semi-circle – a challenging modiﬁcation of the popular Ant-Goal task (Fakoor et al., 2019) to a sparse reward setting similar to the semi-circle task above; (4) Half-Cheetah-Vel (Finn et al., 2017), a popular dense reward domain that serves as a comparison to the sparse domains (full details in Appendix C). 
Revised: We thus focus on sparse reward problems where non-trivial exploration is required to identify the task: (1) A discrete 5 × 5 Gridworld (Zintgraf et al., 2020); (2) A continuous point robot where a sparse reward is located somewhere on a semi-circle (see Figure 1); (3) Ant-Semi-circle – a challenging modiﬁcation of the popular Ant-Goal task (Fakoor et al., 2019) to a sparse reward setting similar to the semi-circle task above; (4) Half-Cheetah-Vel (Finn et 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.

Original: To evaluate performance, we measure average reward in the ﬁrst 2 episodes on unseen tasks – this is where efﬁcient exploration makes a critical difference.
Revised: To evaluate performance, we measure average reward in the ﬁrst 2 episodes on unseen tasks – this is where efﬁcient exploration makes a critical difference.

Original: In the supplementary, we report results for more evaluation episodes.
Revised: In the supplementary, we report results for more evaluation episodes.

Original: Data collection and organization:
Revised: Data collection and organization:

Original: For data collection, we used off-the-shelf DQN (Gridworld) and SAC (continuous domains) implementations.
Revised: For data collection, we used off-the-shelf DQN (Gridworld) and SAC (continuous domains) implementations.

Original: To study the effect of data diversity, we diversiﬁed the ofﬂine dataset by modifying the initial state distribution P init to either (1) uniform over a large region, (2) uniform over a restricted region, or (3) ﬁxed to a single position.
Revised: To study the effect of data diversity, we diversiﬁed the ofﬂine dataset by modifying the initial state distribution P init to either (1) uniform over a large region, (2) uniform over a restricted region, or (3) ﬁxed to a single position.

Original: At meta-test time, only the single ﬁxed position is used.
Revised: At meta-test time, only the single ﬁxed position is used.

Original: The tasks are episodic, but we want the agent to maintain its belief between episodes, so that it can continually improve performance (see Figure 1).
Revised: The tasks are episodic, but we want the agent to maintain its belief between episodes, so that it can continually improve performance (see Figure 1).

Original: We follow Zintgraf et al.
Revised: We follow Zintgraf et al.

Original: (2020), and aggregate k consecutive episodes of length H to a long trajectory of length k × H , and we do not reset the hidden state in the VAE recurrent neural network after episode termination.
Revised: (2020), and aggregate k consecutive episodes of length H to a long trajectory of length k × H , and we do not reset the hidden state in the VAE recurrent neural network after episode termination.

Original: For reward relabelling, we replace the ﬁrst k/ 2 trajectories with trajectories from a randomly chosen
Revised: For reward relabelling, we replace the ﬁrst k/ 2 trajectories with trajectories from a randomly chosen

Original: MDP, and relabel their rewards.
Revised: MDP, and relabel their rewards.

Original: Technically, network architectures and hyperparameters were chosen similarly to Zintgraf et al.
Revised: Technically, network architectures and hyperparameters were chosen similarly to Zintgraf et al.

Original: (2020), as detailed in the supplementary.
Revised: (2020), as detailed in the supplementary.

Original: Ofﬂine Results:
Revised: Ofﬂine Results:

Original: In Figure 3 we compare our ofﬂine algorithm with Thompson sampling based methods, and also with an ablation of the reward relabelling method.
Revised: In Figure 3 we compare our ofﬂine algorithm with Thompson sampling based methods, and also with an ablation of the reward relabelling method.

Original: For Gridworld, the Thompson sampling method is computed exactly, while for the continuous environments, we use online
Revised: For Gridworld, the Thompson sampling method is computed exactly, while for the continuous environments, we use online

Original: PEARL (Rakelly et al., 2019) – a strong baseline that is not affected by our ofﬂine data limitation.
Revised: PEARL (Rakelly et al., 2019) – a strong baseline that is not affected by our ofﬂine data limitation.

Original: In particular, this baseline is stronger than the ofﬂine meta-RL algorithm of Li et al.
Revised: In particular, this baseline is stronger than the ofﬂine meta-RL algorithm of Li et al.

Original: For these results the uniform initial state distribution was used to collect data.
Revised: For these results the uniform initial state distribution was used to collect data.

Original: Note that we signiﬁcantly outperform Thompson sampling based methods, demonstrating our claim of learning non-trivial exploration from ofﬂine data .
Revised: Note that we signiﬁcantly outperform Thompson sampling based methods, demonstrating our claim of learning non-trivial exploration from ofﬂine data .

Original: Results on sparse reward domains signify the severity of MDP ambiguity – without reward relabelling performance drops signiﬁcantly.
Revised: Results on sparse reward domains signify the severity of MDP ambiguity – without reward relabelling performance drops signiﬁcantly.

Original: In the Half-Cheetah-Vel environment, on the other hand, the rewards are dense and thus MDP ambiguity is not a problem.
Revised: In the Half-Cheetah-Vel environment, on the other hand, the rewards are dense and thus MDP ambiguity is not a problem.

Original: In Figure 1 and in Figure 4, we visualize the trajectories of the trained agents in the Semi-circle and
Revised: In Figure 1 and in Figure 4, we visualize the trajectories of the trained agents in the Semi-circle and

Original: Ant-Semi-circle domains, respectively.
Revised: Ant-Semi-circle domains, respectively.

Original: Qualitatively, an approximately Bayes-optimal behavior is evident: in the ﬁrst 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.
Revised: Qualitatively, an approximately Bayes-optimal behavior is evident: in the ﬁrst 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.

Original: Similar behavior for Gridworld is reported in
Revised: Similar behavior for Gridworld is reported in

Original: Appendix E. We emphasize that this behavior is very different from the training data, in which the agents learned to reach speciﬁc goals.
Revised: Appendix E. We emphasize that this behavior is very different from the training data, in which the agents learned to reach speciﬁc goals.

Original: 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 as uniform on the semi-circle, and narrows in on the target as the agent explores the semi-circle.
Revised: 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 as uniform on the semi-circle, and narrows in on the target as the agent explores the semi-circle.

Original: With reward relabelling ablated, however, we show that the belief does not update correctly.
Revised: With reward relabelling ablated, however, we show that the belief does not update correctly.

Original: Data Quality Ablative Study: To evaluate the dependency of our method on the ofﬂine data quality, we report results for the 3 different data collection strategies described above (see supplementary for more details).
Revised: Data Quality Ablative Study: To evaluate the dependency of our method on the ofﬂine data quality, we report results for the 3 different data collection strategies described above (see supplementary for more details).

Original: The results are summarized in Table 1.
Revised: The results are summarized in Table 1.

Original: As expected, data diversity is instrumental to ofﬂine training.
Revised: As expected, data diversity is instrumental to ofﬂine training.

Original: However, as we qualitatively show in Figure 10 in the supplementary, even on the low-diversity datasets, our agents learned non-trivial exploration strategies that search for the goal.
Revised: However, as we qualitatively show in Figure 10 in the supplementary, even on the low-diversity datasets, our agents learned non-trivial exploration strategies that search for the goal.

Original: This is especially remarkable for the ﬁxed-distribution dataset, where it is unlikely that any training trajectory traveled along the semi-circle.
Revised: This is especially remarkable for the ﬁxed-distribution dataset, where it is unlikely that any training trajectory traveled along the semi-circle.

Original: [Table]
Revised: [Table]

Original: [Table][Table] [Table]
Revised: [Table]

Original: [Table]
Revised: [Table]

Original: One may ask whether OMRL presents the same challenge as standard ofﬂine RL, and whether recent ofﬂine RL advances can mitigate the dependency on data diversity.
Revised: One may ask whether OMRL presents the same challenge as standard ofﬂine RL, and whether recent ofﬂine RL advances can mitigate the dependency on data diversity.

Original: To investigate this, we also compare our method with a variant that uses CQL (Kumar et al., 2020) – a state-of-theart ofﬂine RL method – to train the critic network of the meta-RL agent.
Revised: To investigate this, we also compare our method with a variant that uses CQL (Kumar et al., 2020) – a state-of-theart ofﬂine RL method – to train the critic network of the meta-RL agent.

Original: Interestingly, while
Revised: Interestingly, while

Original: CQL slightly improved results (Table 1), the data diversity is much more signiﬁcant.
Revised: CQL slightly improved results (Table 1), the data diversity is much more signiﬁcant.

Original: Taken together with our results on MDP ambiguity, our investigation shows that OMRL presents different practical challenges than standard ofﬂine RL.
Revised: Taken together with our results on MDP ambiguity, our investigation shows that OMRL presents different practical challenges than standard ofﬂine RL.

Original: Online Setting: Our method can also be applied to the online setting.
Revised: Online Setting: Our method can also be applied to the online setting.

Original: In this case, it is simply a modiﬁcation of VariBAD, where the policy gradient optimization is replaced with an off-policy RL algorithm.
Revised: In this case, it is simply a modiﬁcation of VariBAD, where the policy gradient optimization is replaced with an off-policy RL algorithm.

Original: Since MDP ambiguity does not concern online meta-RL, we did not use reward relabelling in this setting.
Revised: Since MDP ambiguity does not concern online meta-RL, we did not use reward relabelling in this setting.

Original: As shown in Figure 5, by exploiting the efﬁciency of off-policy RL, our method signiﬁcantly improves sample-efﬁciency, without sacriﬁcing ﬁnal performance.
Revised: As shown in Figure 5, by exploiting the efﬁciency of off-policy RL, our method signiﬁcantly improves sample-efﬁciency, without sacriﬁcing ﬁnal performance.

Original: When comparing Figure 5 and Figure 3, the reader may notice that the online algorithm’s ﬁnal performance outperforms the ﬁnal performance in the ofﬂine setting.
Revised: When comparing Figure 5 and Figure 3, the reader may notice that the online algorithm’s ﬁnal performance outperforms the ﬁnal performance in the ofﬂine setting.

Original: We emphasize that this phenomenon largely depends on the quality of the ofﬂine data, and not on the algorithm itself.
Revised: We emphasize that this phenomenon largely depends on the quality of the ofﬂine data, and not on the algorithm itself.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: We presented the ﬁrst ofﬂine meta-RL algorithm that is approximately Bayes-optimal.
Revised: We presented the ﬁrst ofﬂine meta-RL algorithm that is approximately Bayes-optimal.

Original: Key to our approach is the connection between Bayesian RL and meta learning, which in principle allows to reduce the problem to standard ofﬂine RL.
Revised: Key to our approach is the connection between Bayesian RL and meta learning, which in principle allows to reduce the problem to standard ofﬂine RL.

Original: In practice, however, we showed that the MDP ambiguity problem prohibits learning, and proposed a simple solution based on mixing trajectories in the data and relabelling their rewards.
Revised: In practice, however, we showed that the MDP ambiguity problem prohibits learning, and proposed a simple solution based on mixing trajectories in the data and relabelling their rewards.

Original: Our results show that this solution is effective on several domains.
Revised: Our results show that this solution is effective on several domains.

Original: Ofﬂine learning is appealing for domains where data collection is costly, such as robotics and healthcare, and there is growing interest in applying deep RL to this setting (Levine et al., 2020).
Revised: Ofﬂine learning is appealing for domains where data collection is costly, such as robotics and healthcare, and there is growing interest in applying deep RL to this setting (Levine et al., 2020).

Original: Key advances in this ﬁeld will likely play a role in improving OMRL as well – an exciting direction for future research.
Revised: Key advances in this ﬁeld will likely play a role in improving OMRL as well – an exciting direction for future research.

Revised: How to handle the MDP ambiguity problem when it is caused by changes in the

Revised: MDP transitions is another important research direction.

Original: A P ROPOSITION P ROOF
Revised: A P ROPOSITION P ROOF

Original: For ease of reading, we copy Proposition 1 from Section 3.2:
Revised: For ease of reading, we copy Proposition 1 from Section 3.2:

Original: Proposition 1.
Revised: Proposition 1.

Original: Let τ = s 0 , a 0 , r 1 , s 1 . .
Revised: Let τ = s 0 , a 0 , r 1 , s 1 . .

Original: , r H , s H denote a random trajectory from a ﬁxed history dependent policy π , generated according to the following process.
Revised: , r H , s H denote a random trajectory from a ﬁxed history dependent policy π , generated according to the following process.

Original: First, MDP parameters R , P are drawn from the prior p ( R , P ) .
Revised: First, MDP parameters R , P are drawn from the prior p ( R , P ) .

Original: Then, the state trajectory is generated according to s 0 ∼ P init , a t ∼ π ( ·| s 0 , a 0 , r 1 , . . .
Revised: Then, the state trajectory is generated according to s 0 ∼ P init , a t ∼ π ( ·| s 0 , a 0 , r 1 , . . .

Original: , s t ) , s t +1 ∼ P ( ·| s t , a t ) and r t +1 ∼ R ( s t , a t ) .
Revised: , s t ) , s t +1 ∼ P ( ·| s t , a t ) and r t +1 ∼ R ( s t , a t ) .

Original: Let b t denote the posterior belief at time t, b t =
Revised: Let b t denote the posterior belief at time t, b t =

Original: P ( R , P| s 0 , a 0 , r 1 , . . .
Revised: P ( R , P| s 0 , a 0 , r 1 , . . .

Original: , s t ) .
Revised: , s t ) .

Original: Then
Revised: Then

Original: P ( s t +1 | s 0 , a 0 , r 1 , . . .
Revised: P ( s t +1 | s 0 , a 0 , r 1 , . . .

Original: , r t , s t , a t ) =
Revised: , r t , s t , a t ) =

Original: E R , P∼ b t P ( s t +1 | s t , a t ) , and,
Revised: E R , P∼ b t P ( s t +1 | s t , a t ) , and,

Original: P ( r t +1 | s 0 , a 0 , r 1 , . . .
Revised: P ( r t +1 | s 0 , a 0 , r 1 , . . .

Original: , s t , a t ) =
Revised: , s t , a t ) =

Original: E R , P∼ b t R ( r t +1 | s t , a t ) .
Revised: E R , P∼ b t R ( r t +1 | s t , a t ) .

Original: Proof.
Revised: Proof.

Original: For the transitions, we have that,
Revised: For the transitions, we have that,

Original: P
Revised: P

Original: [Equation]
Revised: [Equation]

Original: +
Revised: +

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: 0 , r 0 ,. . .
Revised: 0 , r 0 ,. . .

Original: [Equation]
Revised: [Equation]

Original: =
Revised: =

Original: [Equation]
Revised: [Equation]

Original: The proof for the rewards proceeds similarly.
Revised: The proof for the rewards proceeds similarly.

Original: B VAE T RAINING O BJECTIVE
Revised: B VAE T RAINING O BJECTIVE

Original: For completeness
Revised: For completeness

Original: [Equation]
Revised: [Equation]

Original: we follow Zintgraf et al.
Revised: we follow Zintgraf et al.

Original: and outline the full training objective of the VAE.
Revised: and outline the full training objective of the VAE.

Original: Consider the approximate posterior q φ ( m | h : t ) conditioned on the history up to time t .
Revised: Consider the approximate posterior q φ ( m | h : t ) conditioned on the history up to time t .

Original: In this case, the ELBO can be derived as follows:
Revised: In this case, the ELBO can be derived as follows:

Original: [Equation]
Revised: [Equation]

Original: The prior p θ ( m ) is set to be the previous posterior q φ ( m | h : t − 1 ) , with initial prior chosen to be standard normal p θ ( m ) =
Revised: The prior p θ ( m ) is set to be the previous posterior q φ ( m | h : t − 1 ) , with initial prior chosen to be standard normal p θ ( m ) =

Original: N (0 , I ) .
Revised: N (0 , I ) .

Original: The decoder p θ ( s 0 , r 1 , s 1 . . .
Revised: The decoder p θ ( s 0 , r 1 , s 1 . . .

Original: , s H | m, a 0 , . .
Revised: , s H | m, a 0 , . .

Original: . , a H − 1 ) factorizes to reward and next state models p
Revised: . , a H − 1 ) factorizes to reward and next state models p

Original: θ ( s (cid:48) | s, a, m ) and p θ ( r | s, a, m ) , according to: log
Revised: θ ( s (cid:48) | s, a, m ) and p θ ( r | s, a, m ) , according to: log

Original: [Equation]
Revised: [Equation]

Original: , r 
Revised: , r 1 ,

Original: [Equation]
Revised: [Equation]

Original: , s
Revised: , s

Original: [Equation]
Revised: [Equation]

Original: As in Zintgraf et al.
Revised: As in Zintgraf et al.

Original: , we trained only a reward decoder.
Revised: , we trained only a reward decoder.

Original: The overall training objective of the VAE is to maximize the sum of ELBO terms for different time steps,
Revised: The overall training objective of the VAE is to maximize the sum of ELBO terms for different time steps,

Original: [Equation]
Revised: [Equation]

Original: C E NVIRONMENTS D ESCRIPTION
Revised: C E NVIRONMENTS D ESCRIPTION

Original: In this section we describe the details of the domains we experimented with.
Revised: In this section we describe the details of the domains we experimented with.

Original: Gridworld: A 5 × 5 gridworld environment as in Zintgraf et al.
Revised: Gridworld: A 5 × 5 gridworld environment as in Zintgraf et al.

Original: The task distribution is deﬁned by the location of a goal, which is unobserved and can be anywhere but around the starting state at the bottom-left cell.
Revised: The task distribution is deﬁned by the location of a goal, which is unobserved and can be anywhere but around the starting state at the bottom-left cell.

Original: For each task, the agent receives a reward of − 0 .
Revised: For each task, the agent receives a reward of − 0 .

Original: 1 on non-goal cells and +1 at the goal, i.e., r t = (cid:26) 1 , s t = g
Revised: 1 on non-goal cells and +1 at the goal, i.e., r t = (cid:26) 1 , s t = g

Original: [Equation]
Revised: [Equation]

Original: 1 , else, where s t is the current cell and g is the goal cell.
Revised: 1 , else, where s t is the current cell and g is the goal cell.

Original: Similarly to Zintgraf et al.
Revised: Similarly to Zintgraf et al.

Original: (2020), the horizon for this domain is set to 15 and we aggregate k =consecutive episodes to form a trajectory of length 60 .
Revised: (2020), the horizon for this domain is set to 15 and we aggregate k =consecutive episodes to form a trajectory of length 60 .

Original: Semi-circle: A continuous 2D environment as in Figure 1, where the agent must navigate to an unknown goal, randomly chosen on a semi-circle of radius 1 (Rakelly et al., 2019).
Revised: Semi-circle: A continuous 2D environment as in Figure 1, where the agent must navigate to an unknown goal, randomly chosen on a semi-circle of radius 1 (Rakelly et al., 2019).

Original: For each task, the agent receives a reward of +1 if it is within a small radius r = 0 .
Revised: For each task, the agent receives a reward of +1 if it is within a small radius r = 0 .

Original: 2 of the goal, and 0 otherwise,
Revised: 2 of the goal, and 0 otherwise,

Original: [Equation]
Revised: [Equation]

Original: else, where x t is the current 2D location.
Revised: else, where x t is the current 2D location.

Original: Action space is 2-dimensional and bounded: [ − 0 .
Revised: Action space is 2-dimensional and bounded: [ − 0 .

Original: We set the horizon to 60 and aggregate k = 2 consecutive episodes to form a trajectory of length 120 .
Revised: We set the horizon to 60 and aggregate k = 2 consecutive episodes to form a trajectory of length 120 .

Original: MuJoCo: For both MuJoCo domains, the horizon is set to H = 200 and we consider k =episodes.
Revised: MuJoCo: For both MuJoCo domains, the horizon is set to H = 200 and we consider k =episodes.

Original: Half-Cheetah-Vel:
Revised: Half-Cheetah-Vel:

Original: In this environment, a half-cheetah agent must run at a ﬁxed target velocity.
Revised: In this environment, a half-cheetah agent must run at a ﬁxed target velocity.

Original: Following recent works in meta-RL (Finn et al., 2017; Rakelly et al., 2019; Zintgraf et al., 2020), we consider velocities drawn uniformly between 0 . 0 and 3 . 0 .
Revised: Following recent works in meta-RL (Finn et al., 2017; Rakelly et al., 2019; Zintgraf et al., 2020), we consider velocities drawn uniformly between 0 . 0 and 3 . 0 .

Original: The reward in this environment is given by
Revised: The reward in this environment is given by

Original: [Equation]
Revised: [Equation]

Original: where v t is the current velocity, and a t is the current action.
Revised: where v t is the current velocity, and a t is the current action.

Original: Ant-Semi-circle: In this environment, an ant
Revised: Ant-Semi-circle: In this environment, an ant

Original: needs
Revised: needs

Original: to

Original:  navigate to an unknown goal, randomly chosen on a semi-circle, similarly to the Semi-circle task above.
Revised: to navigate to an unknown goal, randomly chosen on a semi-circle, similarly to the Semi-circle task above.

Original: 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 described later), and did not produce trajectories that reached the goal.
Revised: 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 described later), and did not produce trajectories that reached the goal.

Original: We thus modiﬁed the reward only during data collection to be dense, and inversely proportional to the distance from the goal,
Revised: We thus modiﬁed the reward only during data collection to be dense, and inversely proportional to the distance from the goal,

Original: [Equation]
Revised: [Equation]

Original: where x t is the current 2 D location and a t is the current action.
Revised: where x t is the current 2 D location and a t is the current action.

Original: After collecting the data trajectories, we replaced all the dense rewards in the data with the sparse rewards that are given by
Revised: After collecting the data trajectories, we replaced all the dense rewards in the data with the sparse rewards that are given by

Original: [Equation]
Revised: [Equation]

Original: We note that Rakelly et al.
Revised: We note that Rakelly et al.

Original: (2019) use a similar approach to cope with sparse rewards in the online setting.
Revised: (2019) use a similar approach to cope with sparse rewards in the online setting.

Original: D E XPERIMENTAL D ETAILS
Revised: D E XPERIMENTAL D ETAILS

Original: In this section we outline our training process and hyperparameters.
Revised: In this section we outline our training process and hyperparameters.

Original: For the discrete Gridworld domain we used DQN (Mnih et al., 2015) with soft target network updates, as proposed by Lillicrap et al.
Revised: For the discrete Gridworld domain we used DQN (Mnih et al., 2015) with soft target network updates, as proposed by Lillicrap et al.

Original: (2015), which has shown to improve the stability of learning.
Revised: (2015), which has shown to improve the stability of learning.

Original: For the rest of the continuous domains, we used SAC (Haarnoja et al., 2018) with the architectures of the actor and critic chosen similarly, and with a ﬁxed entropy coefﬁcient.
Revised: For the rest of the continuous domains, we used SAC (Haarnoja et al., 2018) with the architectures of the actor and critic chosen similarly, and with a ﬁxed entropy coefﬁcient.

Original: For both
Revised: For both

Original: DQN and SAC, we set the soft target update parameter to 0 .
Revised: DQN and SAC, we set the soft target update parameter to 0 .

Original: In all our experiments we average the performance over 3 random seeds and present the mean and standard deviation.
Revised: In all our experiments we average the performance over 3 random seeds and present the mean and standard deviation.

Original: D. O FFLINE S ETTING
Revised: D. O FFLINE S ETTING

Original: Our ofﬂine training procedure is comprised of 3 separate training steps.
Revised: Our ofﬂine training procedure is comprised of 3 separate training steps.

Original: First is the training of the data collection RL agents.
Revised: First is the training of the data collection RL agents.

Original: Each agent is trained on a different task from the task distribution.
Revised: Each agent is trained on a different task from the task distribution.

Original: For the Gridworld domain, we train 21 agents.
Revised: For the Gridworld domain, we train 21 agents.

Original: We note that this covers the entire task distribution, as goals can be anywhere but around the starting state at the bottom-left cell.
Revised: We note that this covers the entire task distribution, as goals can be anywhere but around the starting state at the bottom-left cell.

Original: For the Semi-circle and Ant-Semi-circle domains, we train 80 data collection agents, and for the Half-Cheetah-Vel environment we used 100 agents.
Revised: For the Semi-circle and Ant-Semi-circle domains, we train 80 data collection agents, and for the Half-Cheetah-Vel environment we used 100 agents.

Original: 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 .
Revised: 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 .

Original: The rest of the hyperparameters used for training the data collection RL agents are summarized in the following table:
Revised: The rest of the hyperparameters used for training the data collection RL agents are summarized in the following table:

Original: [Table]
Revised: [Table]

Original: Exploration/
Revised: Exploration/

Original: [Table] [Table]
Revised: [Table] [Table]

Original: Collected episodes
Revised: Collected episodes

Original: [Table]
Revised: [Table]

Original: Discount
Revised: Discount

Original: [Table]
Revised: [Table]

Original: γ )
Revised: γ )

Original: [Table]
Revised: [Table]

Original: The second training step is the VAE training after applying reward relabelling to the collected data.
Revised: The second training step is the VAE training after applying reward relabelling to the collected data.

Original: The VAE consists of a recurrent encoder, which at time step t takes as input the tuple ( a t , r t +1 , s t +1 ) .
Revised: The VAE consists of a recurrent encoder, which at time step t takes as input the tuple ( a t , r t +1 , s t +1 ) .

Original: The state and reward are passed each through a different fully-connected (FC) layer.
Revised: The state and reward are passed each through a different fully-connected (FC) layer.

Original: The state FC layer is of size 32 and the reward FC layer is of size 8 for the Gridworld and 16 for the rest of the domains, all with ReLU activations.
Revised: The state FC layer is of size 32 and the reward FC layer is of size 8 for the Gridworld and 16 for the rest of the domains, all with ReLU activations.

Original: For the MuJoCo environments, we also pass the action through a
Revised: For the MuJoCo environments, we also pass the action through a

Original: FC layer of size 16 with ReLU.
Revised: FC layer of size 16 with ReLU.

Original: 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 to a GRU of size 64 / 128 (Gridworld/other domains).
Revised: 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 to a GRU of size 64 / 128 (Gridworld/other domains).

Original: The GRU outputs the Gaussian parameters µ ( h : t ) , Σ( h : t ) of the latent vector m , whose dimensionality is 5 in all our experiments.
Revised: The GRU outputs the Gaussian parameters µ ( h : t ) , Σ( h : t ) of the latent vector m , whose dimensionality is 5 in all our experiments.

Original: The VAE reward decoder takes as input a latent sample m ∼ N ( µ ( h : t ) , Σ( h : t )) and the states along the trajectory s 1 , . . .
Revised: The VAE reward decoder takes as input a latent sample m ∼ N ( µ ( h : t ) , Σ( h : t )) and the states along the trajectory s 1 , . . .

Original: , s H , each state at a time, and outputs (for every timestep t = 1 , . . .
Revised: , s H , each state at a time, and outputs (for every timestep t = 1 , . . .

Original: , H ) the entire reconstructed/predicted rewards r 1 , . . .
Revised: , H ) the entire reconstructed/predicted rewards r 1 , . . .

Original: , r H along the trajectory.
Revised: , r H along the trajectory.

Original: In the MuJoCo domains the reward decoder also takes as input the actions along the trajectory a 1 , . .
Revised: In the MuJoCo domains the reward decoder also takes as input the actions along the trajectory a 1 , . .

Original: . , a H and the previous states s 0 , . .
Revised: . , a H and the previous states s 0 , . .

Original: , s H − 1 , as the reward r t in these environments generally depends on s t − 1 , a t , s t .
Revised: , s H − 1 , as the reward r t in these environments generally depends on s t − 1 , a t , s t .

Original: The reward decoder is comprised of 2 FC layers, each of size 32 .
Revised: The reward decoder is comprised of 2 FC layers, each of size 32 .

Original: The VAE is trained to optimize Equation 4 , but similarly to Zintgraf et al.
Revised: The VAE is trained to optimize Equation 4 , but similarly to Zintgraf et al.

Original: (2020), we weight the
Revised: (2020), we weight the

Original: KL term in each of the ELBO terms with some parameter β , which is not necessarily 1 .
Revised: KL term in each of the ELBO terms with some parameter β , which is not necessarily 1 .

Original: In our experiments we used β = 0 .
Revised: In our experiments we used β = 0 .

Original: After the VAE is trained, we apply state relabelling to the data collected by the RL agents, to create a large ofﬂine dataset that effectively comes from the BAMDP.
Revised: After the VAE is trained, we apply state relabelling to the data collected by the RL agents, to create a large ofﬂine dataset that effectively comes from the BAMDP.

Original: Then, we train an off-policy RL algorithm, which is our meta-RL agent, using the ofﬂine data.
Revised: Then, we train an off-policy RL algorithm, which is our meta-RL agent, using the ofﬂine data.

Original: For the ofﬂine meta-RL agents training, we used similar hyperparameters to those used for the data collection RL agents training.
Revised: For the ofﬂine meta-RL agents training, we used similar hyperparameters to those used for the data collection RL agents training.

Original: We only enlarge the size of the hidden layers in our models from 16 , 32 and 128 to 64 , 128 and 256 for the Gridworld, Semi-circle and MuJoCo domains, respectively.
Revised: We only enlarge the size of the hidden layers in our models from 16 , 32 and 128 to 64 , 128 and 256 for the Gridworld, Semi-circle and MuJoCo domains, respectively.

Original: In every iteration we perform 1000 parameter updates for all environments expect the Ant-Semi-circle, in which case we perform 2000 updates per iteration.
Revised: In every iteration we perform 1000 parameter updates for all environments expect the Ant-Semi-circle, in which case we perform 2000 updates per iteration.

Original: D. O NLINE S ETTING
Revised: D. O NLINE S ETTING

Original: In the online setting we didn’t apply reward relabelling to the data, since, as we explained, MDP ambiguity doesn’t concern online meta-RL.
Revised: In the online setting we didn’t apply reward relabelling to the data, since, as we explained, MDP ambiguity doesn’t concern online meta-RL.

Original: The hyperparameters used in the online setting are as follows:
Revised: The hyperparameters used in the online setting are as follows:

Original: Gridworld (DQN)
Revised: Gridworld (DQN)

Original: Semi-circle (SAC)
Revised: Semi-circle (SAC)

Original: Cheetah-Vel (SAC)
Revised: Cheetah-Vel (SAC)

Original: RL parameters
Revised: RL parameters

Original: Architecture/s 2 FC layers of size 100 .
Revised: Architecture/s 2 FC layers of size 100 .

Original: 2 FC layers of size 128 .
Revised: 2 FC layers of size 128 .

Original: 3 FC layers of size 128 .
Revised: 3 FC layers of size 128 .

Original: Num. updates per iter.
Revised: Num. updates per iter.

Original: Exploration/entropy coeff.
Revised: Exploration/entropy coeff.

Original: (cid:15) -greedy, linear annealing from 1 to 0 .over 1000 iterations.
Revised: (cid:15) -greedy, linear annealing from 1 to 0 .over 1000 iterations.

Original: Collected episodes per iter.
Revised: Collected episodes per iter.

Original: Learning rate/s 7 · 10 −7 · 10 −3 · 10 −
Revised: Learning rate/s 7 · 10 −7 · 10 −3 · 10 −

Original: Discount factor ( γ ) 0 .0 .0 .
Revised: Discount factor ( γ ) 0 .0 .0 .

Original: VAE parameters
Revised: VAE parameters

Original: Encoder architecture state/reward FC layer of size 32 / 8 .
Revised: Encoder architecture state/reward FC layer of size 32 / 8 .

Original: GRU of size 64 .
Revised: GRU of size 64 .

Original: state/reward FC layer of size 32 / 8 .
Revised: state/reward FC layer of size 32 / 8 .

Original: GRU of size 128 .
Revised: GRU of size 128 .

Original: state/action/reward FC layer of size 32 / 16 / 16 .
Revised: state/action/reward FC layer of size 32 / 16 / 16 .

Original: GRU of size 128 .
Revised: GRU of size 128 .

Original: Reward decoder architecture 2 FC layers of size 32 .
Revised: Reward decoder architecture 2 FC layers of size 32 .

Original: 2 FC layers of sizes 64 and 32 . 2 FC layers of sizes 64 and 32 .
Revised: 2 FC layers of sizes 64 and 32 . 2 FC layers of sizes 64 and 32 .

Original: Num. updates per iter.
Revised: Num. updates per iter.

Original: Learning rate 3 · 10 −10 −3 · 10 −
Revised: Learning rate 3 · 10 −10 −3 · 10 −

Original: Weight of KL term ( β ) 1 .0 .1 .
Revised: Weight of KL term ( β ) 1 .0 .1 .

Original: E L EARNED B ELIEF V ISUALIZATION
Revised: E L EARNED B ELIEF V ISUALIZATION

Original: In this section we visualize the learned belief states, in order to get more insight into the decision making process of the agent during interaction.
Revised: In this section we visualize the learned belief states, in order to get more insight into the decision making process of the agent during interaction.

Original: In Figure 6, we visualize the interaction of a trained agent with the Gridworld environment, exactly as visualized in Figure 3 at Zintgraf et al.
Revised: In Figure 6, we visualize the interaction of a trained agent with the Gridworld environment, exactly as visualized in Figure 3 at Zintgraf et al.

Original: The agent reduces its uncertainty by effectively searching the goal.
Revised: The agent reduces its uncertainty by effectively searching the goal.

Original: After the goal is found, the agent stops and in subsequent episodes it directly moves toward it.
Revised: After the goal is found, the agent stops and in subsequent episodes it directly moves toward it.

Original: In Figure 7, we plot the reward belief (obtained from the VAE decoder) at different steps during the agent’s interaction in the Semi-circle domain.
Revised: In Figure 7, we plot the reward belief (obtained from the VAE decoder) at different steps during the agent’s interaction in the Semi-circle domain.

Original: Note how the belief starts as uniform over the semi-circle, and narrows in on the target as more evidence is collected.
Revised: Note how the belief starts as uniform over the semi-circle, and narrows in on the target as more evidence is collected.

Original: Also note that without reward relabelling, the agent fails to ﬁnd the goal.
Revised: Also note that without reward relabelling, the agent fails to ﬁnd the goal.

Original: 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 as a result, the agent believes that the reward is located at the ﬁrst point it reaches on the semi-circle.
Revised: 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 as a result, the agent believes that the reward is located at the ﬁrst point it reaches on the semi-circle.

Original: F D ATA Q UALITY A BLATION
Revised: F D ATA Q UALITY A BLATION

Original: In our data quality ablative study, we consider the Ant-Semi-circle domains, for which we modify the initial state distribution during the data collection phase.
Revised: In our data quality ablative study, we consider the Ant-Semi-circle domains, for which we modify the initial state distribution during the data collection phase.

Original: The initial state distributions we consider are visualized in Figure 8: Uniform distribution, uniform excluding states on the semi-circle, and ﬁxed initial position.
Revised: The initial state distributions we consider are visualized in Figure 8: Uniform distribution, uniform excluding states on the semi-circle, and ﬁxed initial position.

Original: [Figure]
Revised: [Figure]

Original: Figure 9 shows the learning curves for the results presented in Table 1.
Revised: Figure 9 shows the learning curves for the results presented in Table 1.

Original: For completeness, we add the learning curve for the uniform distribution which is also presented in Figure 3.
Revised: For completeness, we add the learning curve for the uniform distribution which is also presented in Figure 3.

Original: We also visualize trajectories of trained agents for the 3 different cases as well as for PEARL (Rakelly et al., 2019), in Figure 10.
Revised: We also visualize trajectories of trained agents for the 3 different cases as well as for PEARL (Rakelly et al., 2019), in Figure 10.

Original: Note that even for the ﬁxed-distribution dataset, our agent learns to search for the goal.
Revised: Note that even for the ﬁxed-distribution dataset, our agent learns to search for the goal.

Original: G A DDITIONAL R ESULTS G. P ERFORMANCE VS . A DAPTATION E PISODES
Revised: G A DDITIONAL R ESULTS G. P ERFORMANCE VS . A DAPTATION E PISODES

Original: In this part, we present the average reward per-episode as a function of the number of adaptation episodes at the environment.
Revised: In this part, we present the average reward per-episode as a function of the number of adaptation episodes at the environment.

Original: Figure 11 shows the performance for the Ant-Semi-circle and HalfCheetah-Vel domains.
Revised: Figure 11 shows the performance for the Ant-Semi-circle and HalfCheetah-Vel domains.

Original: Note that within the ﬁrst few episodes, PEARL does not collect high rewards due to the Thompson sampling-based nature of the algorithm.
Revised: Note that within the ﬁrst few episodes, PEARL does not collect high rewards due to the Thompson sampling-based nature of the algorithm.

Original: Our method, on the other hand, efﬁciently explores new tasks and is able to collect rewards within the ﬁrst episodes of interaction.
Revised: Our method, on the other hand, efﬁciently explores new tasks and is able to collect rewards within the ﬁrst episodes of interaction.

Original: Our method outperforms PEARL, collecting high rewards within the ﬁrst adaptation episodes.
Revised: Our method outperforms PEARL, collecting high rewards within the ﬁrst adaptation episodes.

Original: G. PEARL L EARNING C URVES
Revised: G. PEARL L EARNING C URVES

Original: We present the training curves of PEARL in Figure 12.
Revised: We present the training curves of PEARL in Figure 12.

Original: Note that since PEARL is an online algorithm, the x -axis represents the number of environment interactions.
Revised: Note that since PEARL is an online algorithm, the x -axis represents the number of environment interactions.

Revised: G. E XPERIMENT ON MDP S WITH V ARYING T RANSITIONS

Revised: In the following section we describe experiments on a domain in which MDPs differ from each other by their transition function, when MDP ambiguity is not a problem.

Revised: Our goal is to show that our method can also be applied to MDPs with varying transitions.

Revised: We consider a continuous 2 -dimensional environment we term Point-Robot-Wind.

Revised: In Point-RobotWind, the agent must navigate to a ﬁxed (unknown) goal within a distance of D = 1 from the initial state.

Revised: The goal position is the same for all tasks.

Revised: The agent receives a reward of +1 if it is within a radius r = 0 .

Revised: 2 of the goal, and 0 otherwise.

Revised: For each task in this domain, the agent is experiencing a different ‘wind’, which results in a shift in the transitions, such that when taking action a t ∈ [ − 0 .

Revised: 1 , 0 . 1] 2 from state s t in MDP M , the agent transitions to a new state s t +1 , which is given by

Revised: [Equation]

Revised: where w M is a task-speciﬁc wind, which is randomly drawn for each task from the uniform distribution on [ − 0 .

Revised: To navigate correctly to the goal and stay there, the agent must take actions that cancel the wind effect.

Revised: We set the horizon to 25 and evaluate the performance in terms of average return within the ﬁrst episode of interaction on test tasks, sampled from the task distribution.

Revised: We do not apply rewardrelabelling, as the MDPs differ in their transition functions rather than rewards.

Revised: In Figure 13(a) we compare our method with online PEARL.

Revised: Note that our learned policy performs much better than the online PEARL baseline.

Revised: [Figure]

Revised: , and in (b) is the learning curve for online PEARL training (similar to Figure 12).

Revised: In Figure 14 we visualize trajectories of the trained agent on different test tasks.

Revised: 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 a straight line.

Revised: PEARL, on the other hand, only adapts after the ﬁrst episode, and therefore obtains worse results.

Revised: 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 as our method in Point-Robot-Wind.

Revised: 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.

Revised: We emphasize that in Point-Robot-Wind, MDP ambiguity is not a concern, since the data from all agents is largely centered on the line between the agent’s initial position and the goal.

Revised: Thus, the effect of the wind on these states can uniquely be identiﬁed in each domain.

Revised: These results conﬁrm that when MDP ambiguity is not a concern, our method learns an effective policy in domains with varying transitions.