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

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Consider the following instance of the Ofﬂine Meta Reinforcement Learning (OMRL) problem: given the complete training logs 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 instance of the Ofﬂine Meta Reinforcement Learning (OMRL) problem: given the complete training logs 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: Here, we take a Bayesian RL (BRL) view, and seek to learn a
Revised: Here, we take a Bayesian RL (BRL) view, and seek to learn a

Original: Bayes-optimal policy from the ofﬂine data.
Revised: Bayes-optimal policy from the ofﬂine data.

Original: Building on the recent VariBAD BRL approach, we develop an off-policy BRL method that learns to plan an exploration strategy based on an adaptive neural belief estimate.
Revised: Building on the recent VariBAD BRL approach, we develop an off-policy BRL method that learns to plan an exploration strategy based on an adaptive neural belief estimate.

Original: However, learning to infer such a belief from ofﬂine data brings a new identiﬁability issue we term MDP ambiguity.
Revised: However, learning to infer such a belief from ofﬂine data brings a new identiﬁability issue we term MDP ambiguity.

Original: We characterize the problem, and suggest resolutions via data collection and modiﬁcation procedures.
Revised: We characterize the problem, and suggest resolutions via data collection and modiﬁcation procedures.

Original: Finally, we evaluate our framework on a diverse set of domains, including difﬁcult sparse reward tasks, and demonstrate learning of effective exploration behavior that is qualitatively different from the exploration used by any RL agent in the data.
Revised: Finally, we evaluate our framework on a diverse set of domains, including difﬁcult sparse reward tasks, and demonstrate learning of effective exploration behavior that is qualitatively different from the exploration used by any RL agent in the data 1 .

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: One possible formulation of quick RL is Bayesian RL (BRL, Ghavamzadeh et al., 2016).
Revised: One possible formulation of quick RL is Bayesian RL (BRL, Ghavamzadeh et al., 2016).

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 Bayesoptimal 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 Bayesoptimal 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: They proposed the VariBAD algorithm – an implementation of this approach that uses a variational autoencoder (VAE) for adaptive belief estimation and deep neural network policies that take as input both state and posterior belief over the environment parameters.
Revised: They proposed the VariBAD algorithm – an implementation of this approach that uses a variational autoencoder (VAE) for adaptive belief estimation and deep neural network policies that take as input both state and posterior belief over the environment parameters.

Original: Most meta-RL studies, including VariBAD, have focused on the online setting, where, during training, the meta-RL policy is 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 updated using data collected from running it in the training environments.

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

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: Figureillustrates 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: Figureillustrates 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 method for solving OMRL is an off-policy variant of the VariBAD algorithm, based on replacing the on-policy policy gradient optimization in VariBAD with an off-policy Q-learning based approach.
Revised: Our method for solving OMRL is an off-policy variant of the VariBAD algorithm, based on replacing the on-policy policy gradient optimization in VariBAD with an off-policy Q-learning based approach.

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

Original: We formalize this problem and propose a general data collection strategy that can mitigate it.
Revised: We formalize this problem and propose a general data collection strategy that can mitigate it.

Original: Further, when ambiguity is only due to reward differences, we show that a simple reward relabelling trick sufﬁces, without changing data collection.
Revised: Further, when ambiguity is only due to reward differences, we show that a simple reward relabelling trick sufﬁces, without changing data collection.

Original: We collectively term our data collection/relabelling and off-policy algorithm as Bayesian Ofﬂine Reinforcement Learning (BOReL) .
Revised: We collectively term our data collection/relabelling and off-policy algorithm as Bayesian Ofﬂine Reinforcement Learning (BOReL) .

Original: The meta-RL agent (right) needs to ﬁnd a policy that quickly ﬁnds the unknown goal, here, by searching across the semi-circle in the ﬁrst episode, and directly reaching it the second – a completely different strategy from the dominant behaviors in the data.

Original:  P ROBLEM D EFINITION
Revised:  P ROBLEM D EFINITION

Original: We follow the standard meta-RL formulation with distribution over MDP rewards and transitions p ( R , P ) .
Revised: We follow the standard meta-RL formulation with distribution over MDP rewards and transitions p ( R , P ) .

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

Original: We assume that each interaction is organized as
Revised: We assume that each interaction is organized as

Original: M trajectories of length H , τ i,j = s i,j 0 , a i,j 0 , r i,j 1 , s i,j 1 . . .
Revised: 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, π iβ .
Revised: i ( ·| s i,jt , a i,jt ) , and the actions are chosen from an arbitrary data collection policy, π iβ .

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

Original: In Appendix F we investigate implications of this assumption, but we emphasize that this is merely an illustration, and our approach does not place any such constraint – the trajectories can also be collected differently.
Revised: In Appendix F we investigate implications of this assumption, but we emphasize that this is merely an illustration, and our approach does not place any such constraint – the trajectories can also be collected differently.

Original: Our goal is to use the data for learning a Bayes-optimal policy (Ghavamzadeh et al., 2016), i.e., a policy π that maximizes
Revised: Our goal is to use the data for learning a Bayes-optimal policy (Ghavamzadeh et al., 2016), i.e., a policy π that maximizes

Original: E π
Revised: E π

Original: [ (cid:80) ∞ t =

Revised: [ (cid:80) ∞ t =

Original: γ t r ( s t , a t )] , 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: γ t r ( s t , a t )] , 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:  I DENTIFIABILITY P ROBLEMS IN OMRL
Revised:  I DENTIFIABILITY P ROBLEMS IN OMRL

Original: We developed an off-policy variant of the VariBAD algorithm, detailed in the appendix.
Revised: We developed an off-policy variant of the VariBAD algorithm, detailed in the appendix.

Original: The main idea is to replace the on-policy policy gradient algorithm with an off-policy method, and to show that when training data is collected as in Section 2, the resulting data tuples can be seen as coming from the correct belief-MDP in expectation.
Revised: The main idea is to replace the on-policy policy gradient algorithm with an off-policy method, and to show that when training data is collected as in Section 2, the resulting data tuples can be seen as coming from the correct belief-MDP in expectation.

Original: However, while in principle, it is possible to simply run off-policy VariBAD on the ofﬂine data, we claim that in many problems this may not work well.
Revised: However, while in principle, it is possible to simply run off-policy VariBAD on the ofﬂine data, we claim that in many problems this may not work well.

Original: The reason is that the VariBAD belief update should reason about the uncertainty in the MDP parameters, which requires to effectively distinguish between the different possible MDPs.
Revised: The reason is that the VariBAD belief update should reason about the uncertainty in the MDP parameters, which requires to effectively distinguish between the different possible MDPs.

Original: Training the VAE to distinguish between MDPs, however, depends on the ofﬂine data , and might not always  be possible.
Revised: Training the VAE to distinguish between MDPs, however, depends on the ofﬂine data , and might not always if there are two MDPs with different rewards (blue and yellow circles), or one MDP with rewards at both be possible.

Original: This problem, which we term MDP ambiguity , is illustrated in Figure 2: consider two
Revised: This problem, which we term MDP ambiguity , is illustrated in Figure 2: consider two

Original: MDPs, one with rewards in the blue circle, and the other with rewards in the yellow circle.
Revised: 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: Accordingly, we cannot expect to learn a meaningful belief update.
Revised: Accordingly, we cannot expect to learn a meaningful belief update.

Original: In the following, we formalize MDP ambiguity, and how it can be avoided.
Revised: In the following, we formalize MDP ambiguity, and how it can be avoided.

Original: For an MDP deﬁned by {R , P} , denote by P R , P ,π ( s, a, r, s (cid:48) ) and P P ,π ( s, a ) the distribution over ( s, a, r, s (cid:48) ) and ( s, a ) , respectively, induced by a policy π .yellow circles), or one MDP with rewards at both locations.
Revised: For an MDP deﬁned by {R , P} , denote by P R , P ,π ( s, a, r, s (cid:48) ) and P P ,π ( s, a ) the distribution over ( s, a, r, s (cid:48) ) and ( s, a ) , respectively, induced by a policy π .

Original: Deﬁnition 1 (MDP Ambiguity) .
Revised: Deﬁnition 1 (MDP Ambiguity) .

Original: Consider data coming from a set of N different MDPs M =
Revised: Consider data coming from a set of N different MDPs M =

Original: {R i , P i } Ni =1 ⊂ M , where M is an hypothesis set of possible MDPs, and corresponding data collection policies { π iβ } Ni =1 , resulting in N different data distributions D =
Revised: {R i , P i } Ni =1 ⊂ M , where M is an hypothesis set of possible MDPs, and corresponding data collection policies { π iβ } Ni =1 , resulting in N different data distributions D =

Original: { P R i , P i ,π iβ ( s, a, r, s (cid:48) ) }
Revised: { P R i , P i ,π iβ ( s, a, r, s (cid:48) ) }

Original: Ni =1 .
Revised: Ni =1 .

Original: We say that the data is ambiguous if there is an MDP {R , P} ∈ M and two policies π and π (cid:48) such that P R i , P i ,π iβ ( s, a, r, s (cid:48) )
Revised: We say that the data is ambiguous if there is an MDP {R , P} ∈ M and two policies π and π (cid:48) such that P R i , P i ,π iβ ( s, a, r, s (cid:48) )

Original: =
Revised: =

Original: P R , P ,π ( s, a, r, s (cid:48) ) and P R j , P j ,π jβ ( s, a, r, s (cid:48) )
Revised: P R , P ,π ( s, a, r, s (cid:48) ) and P R j , P j ,π jβ ( s, a, r, s (cid:48) )

Original: =
Revised: =

Original: P R , P ,π (cid:48) ( s, a, r, s (cid:48) ) , for some i (cid:54) =
Revised: P R , P ,π (cid:48) ( s, a, r, s (cid:48) ) , for some i (cid:54) =

Original: j .
Revised: j .

Original: Otherwise, the data is termed identiﬁable.
Revised: Otherwise, the data is termed identiﬁable.

Original: The essence of identiﬁability, as expressed in Deﬁnition 1, is that there is no single MDP in the hypothesis set that can explain data from multiple MDPs in the data, as in this case it will be impossible to learn an inference model that accurately distinguishes MDPs, even with inﬁnite data.
Revised: The essence of identiﬁability, as expressed in Deﬁnition 1, is that there is no single MDP in the hypothesis set that can explain data from multiple MDPs in the data, as in this case it will be impossible to learn an inference model that accurately distinguishes MDPs, even with inﬁnite data.

Original: Note that MDP ambiguity is special to ofﬂine meta-RL; in online meta-RL, the agent may be driven by the online adapting policy (or guided explicitly) to explore states that reduce its ambiguity.
Revised: Note that MDP ambiguity is special to ofﬂine meta-RL; 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 RL, as the problem here concerns the identiﬁcation of the MDP , which in standard RL is unique.
Revised: We also emphasize that this problem is not encountered in standard (non-meta) ofﬂine RL, as the problem here concerns the identiﬁcation of the MDP , which in standard RL is unique.

Original: Proposition 1.
Revised: Proposition 1.

Original: Consider the setting described in Deﬁnition 1.
Revised: Consider the setting described in Deﬁnition 1.

Original: For a pair of MDPs i and j , we deﬁne the identifying state-action pairs as the state-action pairs that satisfy R i ( s, a ) (cid:54) =
Revised: For a pair of MDPs i and j , we deﬁne the identifying state-action pairs as the state-action pairs that satisfy R i ( s, a ) (cid:54) =

Original: R j
Revised: R j

Original: ( s, a ) and/or
Revised: ( s, a ) and/or

Original: P i ( ·| s, a ) (cid:54) =
Revised: P i ( ·| s, a ) (cid:54) =

Original: P j ( ·| s, a ) .
Revised: P j ( ·| s, a ) .

Original: If for every i (cid:54) =
Revised: If for every i (cid:54) =

Original: j there exists an identifying state-action pair that has positive probability under both i and j , i.e., P P i ,π iβ ( s, a ) , P P j ,π jβ ( s, a ) > 0 , then the data is identiﬁable.
Revised: j there exists an identifying state-action pair that has positive probability under both i and j , i.e., P P i ,π iβ ( s, a ) , P P j ,π jβ ( s, a ) > 0 , then the data is identiﬁable.

Original: Thus, if the agent has data on identifying state-actions obtained from different MDPs , it has the capability to identify which data samples belong to which MDP, regardless of the hypothesis set M .
Revised: Thus, if the agent has data on identifying state-actions obtained from different MDPs , it has the capability to identify which data samples belong to which MDP, regardless of the hypothesis set M .

Original: How can one collect data to mitigate MDP ambiguity?
Revised: How can one collect data to mitigate MDP ambiguity?

Original: We present a simple, general modiﬁcation to the data collection scheme we term policy replaying , which, under mild conditions on the original data collection policies, guarantees that the resulting data will be identiﬁable.
Revised: We present a simple, general modiﬁcation to the data collection scheme we term policy replaying , which, under mild conditions on the original data collection policies, guarantees that the resulting data will be identiﬁable.

Original: We importantly note that changing the data collection method in-hindsight is not suitable for the ofﬂine setting.
Revised: We importantly note that changing the data collection method in-hindsight is not suitable for the ofﬂine setting.

Original: Therefore, the proposed scheme should be viewed as a guideline for effective OMRL data collection .
Revised: Therefore, the proposed scheme should be viewed as a guideline for effective OMRL data collection .

Original: For each MDP, we propose collecting data in the following manner: randomly draw a data collection policy from { π iβ } Ni =1 , collect a trajectory following that policy, and repeat.
Revised: For each MDP, we propose collecting data in the following manner: randomly draw a data collection policy from { π iβ } Ni =1 , collect a trajectory following that policy, and repeat.

Original: After this procedure, the new data distributions are all associated with the same data collection policy, which we denote π r .
Revised: After this procedure, the new data distributions are all associated with the same data collection policy, which we denote π r .

Original: Proposition 2.
Revised: Proposition 2.

Original: For every i (cid:54) =
Revised: For every i (cid:54) =

Original: j , denote the set of identifying state-action pairs by I i,j .
Revised: j , denote the set of identifying state-action pairs by I i,j .

Original: If for every i and every j exists ( s i,j , a i,j ) ∈
Revised: If for every i and every j exists ( s i,j , a i,j ) ∈

Original: I i,j such that P P i ,π iβ ( s i,j , a i,j ) > 0 , then replacing π iβ with π r for all i results in identiﬁable data.
Revised: I i,j such that P P i ,π iβ ( s i,j , a i,j ) > 0 , then replacing π iβ with π r for all i results in identiﬁable data.

Original: Note that the requirement on identifying state-actions in Proposition 2 is minimal – without it, the original data collecting policies π iβ are useless, as they do not visit any identifying state-action pair.
Revised: Note that the requirement on identifying state-actions in Proposition 2 is minimal – without it, the original data collecting policies π iβ are useless, as they do not visit any identifying state-action pair.

Original: When the tasks only differ in their reward function, and the reward functions for the training environments are known, policy replaying can be implemented in hindsight, without changing the data collection process .
Revised: When the tasks only differ in their reward function, and the reward functions for the training environments are known, policy replaying can be implemented in hindsight, without changing the data collection process .

Original: This technique, which we term Reward Relabelling (RR) , is applicable under the ofﬂine setting, and described next.
Revised: This technique, which we term Reward Relabelling (RR) , is applicable under the ofﬂine setting, and described next.

Original: In RR, we replace the rewards in a trajectory from some MDP i in the data with rewards from another randomly chosen MDP j (cid:54)
Revised: In RR, we replace the rewards in a trajectory from some MDP i in the data with rewards from another randomly chosen MDP j (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 K trajectories ˆ τ i,k
Revised: , N , we add to the data K trajectories ˆ τ i,k

Original: = ( s i,k 0 , a i,k 0 , ˆ r i,k 1 , s i,k 1 . .
Revised: = ( s i,k 0 , a i,k 0 , ˆ r i,k 1 , s i,k 1 . .

Original: , ˆ r i,kH , s i,kH ) , k ∈ 1 , . .
Revised: , ˆ r i,kH , s i,kH ) , k ∈ 1 , . .

Original: , K , where the relabelled rewards ˆ r satisfy ˆ r i,kt +

Revised: , K , where the relabelled rewards ˆ r satisfy ˆ r i,kt +

Original: =
Revised: =

Original: R j
Revised: R j

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

Original: Thus, our relabelling effectively runs π iβ on
Revised: Thus, our relabelling effectively runs π iβ on

Original: MDP j , which is equivalent to performing policy replaying (in hindsight).
Revised: MDP j , which is equivalent to performing policy replaying (in hindsight).

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: In our experiments, we aim to demonstrate: (1) Learning approximately Bayes-optimal policies in the ofﬂine setting; and (2) The severity of MDP ambiguity, and the effectiveness of our proposed resolutions.
Revised: In our experiments, we aim to demonstrate: (1) Learning approximately Bayes-optimal policies in the ofﬂine setting; and (2) The severity of MDP ambiguity, and the effectiveness of our proposed resolutions.

Original: 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 under the ofﬂine data restriction.
Revised: 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 under the ofﬂine data restriction.

Original: We further describe our method, domains, evaluation metric, and data collection process in the supplementary.
Revised: We further describe our method, domains, evaluation metric, and data collection process in the supplementary.

Original: Main Results:
Revised: Main 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: 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: These results are further explained qualitatively by observing the exploration behavior of our learned agents.
Revised: These results are further explained qualitatively by observing the exploration behavior of our learned agents.

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

Original: 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: 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: In contrast, a Thompson sampling based agent will never display such search behavior, as it does not plan to proactively reduce uncertainty.
Revised: In contrast, a Thompson sampling based agent will never display such search behavior, as it does not plan to proactively reduce uncertainty.

Original: Instead, such an agent will randomly choose an un-visited possible goal at each episode and directly navigate towards it (cf.
Revised: Instead, such an agent will randomly choose an un-visited possible goal at each episode and directly navigate towards it (cf.

Original: Figure 1 in Li et al. 2020).
Revised: Figure 1 in Li et al. 2020).

Original: We further emphasize that the approximately Bayes-optimal search behavior is different from the training data, in which the agents learned to reach speciﬁc goals.
Revised: We further emphasize that the approximately Bayes-optimal search behavior is different from the training data, in which the agents learned to reach speciﬁc goals

Original: sparse is somewhere from trained policy on a new goal

Original: Note in the ﬁrst episode, the ant searches for the goal, and in the second one it directly moves toward the goal it has previously found search behavior is different from the goal-reaching behaviors that dominate the training data.

Original: Our results also signify the severity of MDP ambiguity, and align with the theory in Sec. 3.
Revised: Our results also signify the severity of MDP ambiguity, and align with the theory in Sec. 3.

Original: In domains with non-overlapping identifying states (Gridworld, Semi-circle, and Ant-Semi-circle), as expected, performance without reward relabelling is poor, while in domains with overlapping identifying states reward relabelling has little effect.
Revised: In domains with non-overlapping identifying states (Gridworld, Semi-circle, and Ant-Semi-circle), as expected, performance without reward relabelling is poor, while in domains with overlapping identifying states reward relabelling has little effect.

Original: In Figure 6 in the supplementary, we provide further insight into these results, by plotting the belief update during the episode rollout for Semi-circle: 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 6 in the supplementary, we provide further insight into these results, by plotting the belief update during the episode rollout for Semi-circle: 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:  C ONCLUSION AND F UTURE W ORK
Revised:  C ONCLUSION AND F UTURE W ORK

Original: We presented the ﬁrst ofﬂine meta-RL algorithm that is approximately Bayes-optimal, allowing to solve problems where efﬁcient exploration is crucial.
Revised: We presented the ﬁrst ofﬂine meta-RL algorithm that is approximately Bayes-optimal, allowing to solve problems where efﬁcient exploration is crucial.

Original: The connection between Bayesian RL and meta learning allows to reduce the problem to ofﬂine RL on belief-augmented states.
Revised: The connection between Bayesian RL and meta learning allows to reduce the problem to ofﬂine RL on belief-augmented states.

Original: However, learning a neural belief update from ofﬂine data is prone to the MDP ambiguity problem.
Revised: However, learning a neural belief update from ofﬂine data is prone to the MDP ambiguity problem.

Original: We formalized the problem, and proposed a simple data collection protocol that guarantees identiﬁability.
Revised: We formalized the problem, and proposed a simple data collection protocol that guarantees identiﬁability.

Original: In the particular case of tasks that differ in their rewards, our protocol can be implemented in hindsight, for arbitrarily ofﬂine data.
Revised: In the particular case of tasks that differ in their rewards, our protocol can be implemented in hindsight, for arbitrarily ofﬂine data.

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

Original: It is intriguing whether other techniques can mitigate MDP ambiguity.
Revised: It is intriguing whether other techniques can mitigate MDP ambiguity.

Original: For example, designing data collection policies that diversify the data or injecting prior knowledge by controlling the hypothesis set of the neural belief update.
Revised: For example, designing data collection policies that diversify the data or injecting prior knowledge by controlling the hypothesis set of the neural belief update.

Original: A E XTENDED B ACKGROUND
Revised: A E XTENDED B ACKGROUND

Original: We recapitulate BRL and the VariBAD algorithm.
Revised: We recapitulate BRL and the VariBAD algorithm.

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: Key here is that this formulation naturally accounts for the exploration/exploitation tradeoff – an optimal agent must plan its actions to reduce uncertainty in the MDP parameters, if such leads to higher rewards.
Revised: Key here is that this formulation naturally accounts for the exploration/exploitation tradeoff – an optimal agent must plan its actions to reduce uncertainty in the MDP parameters, if such leads to higher rewards.

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
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

Original: POMDP instance where the unobserved variables do not change in time.
Revised: POMDP instance where the unobserved variables do not change in time.

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: b t +1 ( R , P ) =
Revised: p ( R , P ) , and: b t +1 ( R , P ) =

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

Original: ∝ P ( s t +1 , r t +1 | h : t , R , P ) b t ( 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, the state in BRL can be augmented with the belief to result in the Bayes-Adaptive MDP model (BAMDP,
Revised: Similar to the idea of solving a POMDP by representing it as an MDP over belief states, the state in BRL can be augmented with the belief to result in the Bayes-Adaptive MDP model (BAMDP,

Original: Duff & Barto, 2002).
Revised: 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
Revised: ( s t , b t ) and the augmented state space

Original: S + = S × B , where B denotes the belief space.
Revised: 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
Revised: P + ( s + t

Original: +

Revised: +

Original: | s + t , a t ) =
Revised: | 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: ∀ s + ∈ S + , a ∈ A we have that
Revised: As in standard MDPs, the optimal action-value function in the BAMDP satisﬁes the Bellman equation: ∀ s + ∈ S + , a ∈ A we have that

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: 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: 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: The policies are trained using policy gradients, optimizing
Revised: The policies are trained using policy gradients, optimizing

Original: [Equation]
Revised: [Equation]

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: B OMRL AND O FF -P OLICY V ARI BAD
Revised: B OMRL AND 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 .
Revised: (3) is in fact a sound objective for the BAMDP, as we now show .

Original: Proposition 3.
Revised: Proposition 3.

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

Revised: P

Original: [Equation]
Revised: [Equation]

Original:  , a t ) , and, 
Revised: +1 | s t , a t ) , and, P [Equation] r

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

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

Original: [Equation]
Revised: [Equation]

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

Original: For on-policy VariBAD, Proposition 3 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 3 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 ( s, a, r, s (cid:48) )
Revised: However, off-policy RL does not take as input trajectories, but tuples of the form ( s, a, r, s (cid:48) )

Original: ≡ ( state, action, reward, next state ) , where states and actions can be sampled from any distribution .
Revised: ≡ ( 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,
Revised: Fortunately,

Original: Proposition 3 shows that when collecting data by sampling complete trajectories, this is not necessary, as in expectation, the rewards and transitions are correctly sampled from the BAMDP.
Revised: Proposition 3 shows that when collecting data by sampling complete trajectories, this is not necessary, as in expectation, the rewards and transitions are correctly sampled from the BAMDP.

Original: In the following, we therefore focus on settings where data can be collected that way, for example, by collecting logs of RL agents trained on the different training tasks.
Revised: In the following, we therefore focus on settings where data can be collected that way, for example, by collecting logs of RL agents trained on the different training tasks.

Original: Based on Proposition 3, we can use a state augmentation method similar to VariBAD, which we refer to as state relabelling .
Revised: Based on Proposition 3, 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: C E XTENDED D EFINITIONS AND P ROOFS FOR S ECTION
Revised: C E XTENDED D EFINITIONS AND P ROOFS FOR S ECTION

Original: For the proofs of identiﬁability, we start by elaborating the formal deﬁnition of our setting.
Revised: For the proofs of identiﬁability, we start by elaborating the formal deﬁnition of our setting.

Original: For simplicity, we assume that the MDPs {R i , P i } Ni =1 are deﬁned over ﬁnite state-action spaces ( |S| , |A| < ∞ ).
Revised: For simplicity, we assume that the MDPs {R i , P i } Ni =1 are deﬁned over ﬁnite state-action spaces ( |S| , |A| < ∞ ).

Original: For every i = 1 , . . .
Revised: For every i = 1 , . . .

Original: , N , let π iβ be a general stationary, stochastic, history-dependent policy 3 .
Revised: , N , let π iβ be a general stationary, stochastic, history-dependent policy 4 .

Original: The initial state distribution P init is the same across all MDPs.
Revised: The initial state distribution P init is the same across all MDPs.

Original: We assume that data is collected from trajectories of length at most T max .
Revised: We assume that data is collected from trajectories of length at most T max .

Original: This is a convenient assumption that holds in every practical scenario, and allows us to side step issues of deﬁning visitation frequencies when t → ∞ .
Revised: This is a convenient assumption that holds in every practical scenario, and allows us to side step issues of deﬁning visitation frequencies when t → ∞ .

Original: For some 0 ≤ t ≤ T max , denote by P P i ,π iβ ,t ( s, a ) =
Revised: For some 0 ≤ t ≤ T max , denote by P P i ,π iβ ,t ( s, a ) =

Original: P P i ,π iβ ( s t = s, a t = a ) the probability of visiting the state-action pair ( s, a ) at time t by running the policy π iβ on MDP with transition function
Revised: P P i ,π iβ ( s t = s, a t = a ) the probability of visiting the state-action pair ( s, a ) at time t by running the policy π iβ on MDP with transition function

Original: P i and initial state distribution P init .
Revised: P i and initial state distribution P init .

Original: Now, we deﬁne:
Revised: Now, we deﬁne:

Original: [Equation]
Revised: [Equation]

Original: that is, P P i ,π β ( s, a ) is the probability of observing state-action ( s, a ) in the data from MDP i .
Revised: that is, P P i ,π β ( s, a ) is the probability of observing state-action ( s, a ) in the data from MDP i .

Original: Similarly, we deﬁne
Revised: Similarly, we deﬁne

Original: [Equation]
Revised: [Equation]

Original: the probability of observing the tuple ( s, a, r, s [Equation] ) in the data from MDP i .
Revised: the probability of observing the tuple ( s, a, r, s (cid:48) ) in the data from MDP i .

Original: A trajectory from the replay policy π r in MDP i is generated as follows.
Revised: A trajectory from the replay policy π r in MDP i is generated as follows.

Original: Let x be a discrete random variable deﬁned on 1 , . . .
Revised: Let x be a discrete random variable deﬁned on 1 , . . .

Original: , N with probability P x ( · ) that satisﬁes P x ( x = k ) >
Revised: , N with probability P x ( · ) that satisﬁes P x ( x = k ) >

Original: for every k = 1 , . . .
Revised: for every k = 1 , . . .

Original: , N .
Revised: , N .

Original: First, we draw x .
Revised: First, we draw x .

Original: Then, we sample a trajectory from MDP i using policy π xβ .
Revised: Then, we sample a trajectory from MDP i using policy π xβ .

Original: For ease of reading, we copy here the propositions from the main text.
Revised: For ease of reading, we copy here the propositions from the main text.

Original: Proposition 1.
Revised: Proposition 1.

Original: Consider the setting described in Deﬁnition 1.
Revised: Consider the setting described in Deﬁnition 1.

Original: For a pair of MDPs i and j , we deﬁne the identifying state-action pairs as the state-action pairs that satisfy R i ( s, a ) (cid:54) =
Revised: For a pair of MDPs i and j , we deﬁne the identifying state-action pairs as the state-action pairs that satisfy R i ( s, a ) (cid:54) =

Original: R j
Revised: R j

Original: ( s, a ) and/or P i ( s (cid:48) | s, a ) (cid:54) =
Revised: ( s, a ) and/or P i ( s (cid:48) | s, a ) (cid:54) =

Original: P j ( s (cid:48) | s, a ) .
Revised: P j ( s (cid:48) | s, a ) .

Original: If for every i (cid:54) =
Revised: If for every i (cid:54) =

Original: j there exists an identifying state-action pair that has positive probability under both i and j , i.e., P P i ,π iβ ( s, a ) , P P j ,π jβ ( s, a ) > 0 , then the data is identiﬁable.
Revised: j there exists an identifying state-action pair that has positive probability under both i and j , i.e., P P i ,π iβ ( s, a ) , P P j ,π jβ ( s, a ) > 0 , then the data is identiﬁable.

Original: Before we prove Proposition 2, we present the following lemma, which will be used later in the proof.
Revised: Before we prove Proposition 2, we present the following lemma, which will be used later in the proof.

Original: Lemma 1. Consider a pair of MDPs ( R , P ) and ( R (cid:48) , P (cid:48) ) , and two policies π and π (cid:48) .
Revised: Lemma 1. Consider a pair of MDPs ( R , P ) and ( R (cid:48) , P (cid:48) ) , and two policies π and π (cid:48) .

Original: If there exists an identifying state-action pair of the MDPs (¯ s, ¯ a ) that has positive probability under both ( P , π ) and ( P (cid:48) , π (cid:48) ) , i.e., P P ,π (¯ s, ¯ a ) , P P (cid:48) ,π (cid:48) (¯ s, ¯ a ) > 0 , then P R , P ,π ( s, a, r, s (cid:48) ) (cid:54) =
Revised: If there exists an identifying state-action pair of the MDPs (¯ s, ¯ a ) that has positive probability under both ( P , π ) and ( P (cid:48) , π (cid:48) ) , i.e., P P ,π (¯ s, ¯ a ) , P P (cid:48) ,π (cid:48) (¯ s, ¯ a ) > 0 , then P R , P ,π ( s, a, r, s (cid:48) ) (cid:54) =

Original: P R (cid:48) , P (cid:48) ,π (cid:48)
Revised: P R (cid:48) , P (cid:48) ,π (cid:48)

Original: ( s, a, r, s (cid:48) ) .
Revised: ( s, a, r, s (cid:48) ) .

Original: Proof.
Revised: Proof.

Original: Assume to the contrary that P R , P ,π
Revised: Assume to the contrary that P R [Equation] π

Original: [Equation]
Revised: [Equation]

Original: s, a, r, s
Revised: s, a, r, s

Original: [Equation]
Revised: [Equation]

Original: ) =
Revised: =

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: ( s, a, r, s (cid:48) ) .
Revised: ( s, a, r, s (cid:48) ) .

Original: Marginalizing over r and s (cid:48) , we obtain:
Revised: Marginalizing over r and s (cid:48) , we obtain:

Original: [Equation]
Revised: [Equation]

Original: Speciﬁcally, we have P P ,
Revised: Speciﬁcally, we have P P ,π (

Revised: [Equation]

Revised: s,

Revised: [Equation]

Revised: a

Original: [Equation]
Revised: [Equation]

Original: Since P R , P ,π (¯ s, ¯ a, r, s (cid:48) )
Revised: Since P R , P ,π (¯ s, ¯ a, r, s (cid:48) )

Original: =
Revised: =

Original: P R , P ( r, s (cid:48) | ¯ s, ¯ a ) P P ,π (¯ s, ¯ a ) for every r and s (cid:48) , and P P ,π (¯ s, ¯ a ) =
Revised: P R , P ( r, s (cid:48) | ¯ s, ¯ a ) P P ,π (¯ s, ¯ a ) for every r and s (cid:48) , and P P ,π (¯ s, ¯ a ) =

Original: P P (cid:48) ,π (cid:48) (¯ s, ¯ a ) > 0 , it holds that P R , P ( r, s (cid:48)
Revised: P P (cid:48) ,π (cid:48) (¯ s, ¯ a ) > 0 , it holds that P R , P ( r, s (cid:48)

Original: | ¯ s, ¯ a ) =
Revised: | ¯ s, ¯ a ) =

Original: P R (cid:48) , P (cid:48)
Revised: P R (cid:48) , P (cid:48)

Original: ( r, s (cid:48) | ¯ s, ¯ a ) for every r and s (cid:48) .
Revised: ( r, s (cid:48) | ¯ s, ¯ a ) for every r and s (cid:48) .

Original: By marginalizing over s (cid:48)
Revised: By marginalizing over s (cid:48)

Original: we get that (cid:88)
Revised: we get that (cid:88)

Original: [Equation]
Revised: [Equation]

Original: Similarly, by marginalizing over r , we get P
Revised: Similarly, by marginalizing over r , we get P

Original: i ( s (cid:48) | ¯ s, ¯ a ) =
Revised: i ( s (cid:48) | ¯ s, ¯ a ) =

Original: P j ( s (cid:48) | ¯ s, ¯ a ) .
Revised: P j ( s (cid:48) | ¯ s, ¯ a ) .

Original: Overall, both reward and transition function do not differ in (¯ s, ¯ a ) , which contradicts the fact that (¯ s, ¯ a ) is an identifying state-action pair.
Revised: Overall, both reward and transition function do not differ in (¯ s, ¯ a ) , which contradicts the fact that (¯ s, ¯ a ) is an identifying state-action pair.

Original: We now prove Proposition 1.
Revised: We now prove Proposition 1.

Original: Proof.
Revised: Proof.

Original: Consider some i (cid:54)
Revised: Consider some i (cid:54)

Original: =
Revised: =

Original: j .
Revised: j .

Original: Let ( s i,j , a i,j ) be an identifying state-action pair that has positive probability under both i and j .
Revised: Let ( s i,j , a i,j ) be an identifying state-action pair that has positive probability under both i and j .

Original: Assume to the contrary that there exists an MDP {R , P} ∈ M and two policies π and π (cid:48) such that P R i , P i ,π iβ ( s, a, r, s (cid:48) )
Revised: Assume to the contrary that there exists an MDP {R , P} ∈ M and two policies π and π (cid:48) such that P R i , P i ,π iβ ( s, a, r, s (cid:48) )

Original: =
Revised: =

Original: P R , P ,π ( s, a, r, s (cid:48) ) and P R j , P j ,π jβ ( s, a, r, s (cid:48) )
Revised: P R , P ,π ( s, a, r, s (cid:48) ) and P R j , P j ,π jβ ( s, a, r, s (cid:48) )

Original: =
Revised: =

Original: P R , P ,π (cid:48)
Revised: P R , P ,π (cid:48)

Original: ( s, a, r, s (cid:48) ) .
Revised: ( s, a, r, s (cid:48) ) .

Original: Since ( s i,j , a i,j ) has positive probability under ( P i , π iβ ) and P R i , P i ,π iβ ( s, a, r, s (cid:48) )
Revised: Since ( s i,j , a i,j ) has positive probability under ( P i , π iβ ) and P R i , P i ,π iβ ( s, a, r, s (cid:48) )

Original: =
Revised: =

Original: P R , P ,π ( s, a, r, s (cid:48) ) , then ( s i,j , a i,j ) must also have positive probability under ( P , π ) (otherwise, there are r and s (cid:48) for which
Revised: P R , P ,π ( s, a, r, s (cid:48) ) , then ( s i,j , a i,j ) must also have positive probability under ( P , π ) (otherwise, there are r and s (cid:48) for which

Original: P R i , P i ,π iβ ( s i,j , a i,j , r, s (cid:48) ) > 0 , while P R , P ,π ( s i,j , a i,j , r, s (cid:48) )
Revised: P R i , P i ,π iβ ( s i,j , a i,j , r, s (cid:48) ) > 0 , while P R , P ,π ( s i,j , a i,j , r, s (cid:48) )

Original: = 0 ).
Revised: = 0 ).

Original: Now, since ( s i,j , a i,j ) has positive probability under both ( P i , π iβ ) and ( P , π ) , and P R i , P i ,π iβ ( s, a, r, s (cid:48) )
Revised: Now, since ( s i,j , a i,j ) has positive probability under both ( P i , π iβ ) and ( P , π ) , and P R i , P i ,π iβ ( s, a, r, s (cid:48) )

Original: =
Revised: =

Original: P R , P ,π ( s, a, r, s (cid:48) ) , according to Lemma 1, it cannot be an identifying state-action pair of ( R i , P i ) and ( R , P ) .
Revised: P R , P ,π ( s, a, r, s (cid:48) ) , according to Lemma 1, it cannot be an identifying state-action pair of ( R i , P i ) and ( R , P ) .

Original: Therefore, the MDP {R , P} must satisfy P ( ·| s i,j , a i,j ) =
Revised: Therefore, the MDP {R , P} must satisfy P ( ·| s i,j , a i,j ) =

Original: P i ( ·| s i,j , a i,j ) and R ( s i,j , a i,j )
Revised: P i ( ·| s i,j , a i,j ) and R ( s i,j , a i,j )

Original: =
Revised: =

Original: R i
Revised: R i

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

Original: The same argument can be made for ( R j , P j , π jβ ) and ( R , P , π (cid:48) ) , resulting in P ( ·| s i,j , a i,j ) =
Revised: The same argument can be made for ( R j , P j , π jβ ) and ( R , P , π (cid:48) ) , resulting in P ( ·| s i,j , a i,j ) =

Original: P j ( ·| s i,j , a i,j ) and R ( s i,j , a i,j )
Revised: P j ( ·| s i,j , a i,j ) and R ( s i,j , a i,j )

Original: =
Revised: =

Original: R j
Revised: R j

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

Original: Overall, we get P
Revised: Overall, we get P

Original: i ( ·| s i,j , a i,j ) =
Revised: i ( ·| s i,j , a i,j ) =

Original: P ( ·| s i,j , a i,j ) =
Revised: P ( ·| s i,j , a i,j ) =

Original: P j ( ·| s i,j , a i,j ) and R i ( s i,j , a i,j )
Revised: P j ( ·| s i,j , a i,j ) and R i ( s i,j , a i,j )

Original: =
Revised: =

Original: R ( s i,j , a i,j ) =
Revised: R ( s i,j , a i,j ) =

Original: R j
Revised: R j

Original: ( s i,j , a i,j ) , which is a contradiction, as ( s i,j , a i,j ) is an identifying state-action pair of MDPs i and j .
Revised: ( s i,j , a i,j ) , which is a contradiction, as ( s i,j , a i,j ) is an identifying state-action pair of MDPs i and j .

Original: Proposition 2.
Revised: Proposition 2.

Original: For every i (cid:54) =
Revised: For every i (cid:54) =

Original: j , denote the set of identifying state-action pairs by I i,j .
Revised: j , denote the set of identifying state-action pairs by I i,j .

Original: If for every i and every j exists ( s i,j , a i,j ) ∈
Revised: If for every i and every j exists ( s i,j , a i,j ) ∈

Original: I i,j such that P P i ,π iβ ( s i,j , a i,j ) > 0 , then replacing π iβ with π r for all i results in identiﬁable data.
Revised: I i,j such that P P i ,π iβ ( s i,j , a i,j ) > 0 , then replacing π iβ with π r for all i results in identiﬁable data.

Original: Proof.
Revised: Proof.

Original: Consider some i (cid:54) =
Revised: Consider some i (cid:54) =

Original: j .
Revised: j .

Original: We observe that by the construction of π r , for every ( s, a ) pair that satisﬁes P P i ,π iβ ( s, a ) > 0 , we also have P P
Revised: We observe that by the construction of π r , for every ( s, a ) pair that satisﬁes P P i ,π iβ ( s, a ) > 0 , we also have P P

Original: i ,π r ( s, a ) > 0 .
Revised: i ,π r ( s, a ) > 0 .

Original: In particular, we have P P
Revised: In particular, we have P P

Original: i ,π r ( s i,j , a i,j ) > 0 .
Revised: i ,π r ( s i,j , a i,j ) > 0 .

Original: We will show that either ( s i,j , a i,j ) also has positive probability under ( P j , π r ) or there must exist some other state-action pair that has positive probability under both ( P i , π r ) and ( P j , π r ) .
Revised: We will show that either ( s i,j , a i,j ) also has positive probability under ( P j , π r ) or there must exist some other state-action pair that has positive probability under both ( P i , π r ) and ( P j , π r ) .

Original: This, according to Proposition 1, will result in identiﬁability of the data.
Revised: This, according to Proposition 1, will result in identiﬁability of the data.

Original: We deﬁne the following sets of state-action pairs:
Revised: We deﬁne the following sets of state-action pairs:

Original: [Equation]
Revised: [Equation]

Original: Note that Σ i 

Revised: Note that Σ i

Original: = Σ
Revised: = Σ

Original: i,j 0 , as the initial state distribution P init and π r are ﬁxed across all MDPs.
Revised: i,j 0 , as the initial state distribution P init and π r are ﬁxed across all MDPs.

Original: First, consider the case where for every t = 0 , 1 , . .
Revised: First, consider the case where for every t = 0 , 1 , . .

Original: , T max we have Σ it =
Revised: , T max we have Σ it =

Original: Σ i,jt .
Revised: Σ i,jt .

Original: Given that ( s i,j , a i,j ) has positive probability under ( P i , π r ) , there exists some t for which ( s i,j , a i,j ) ∈
Revised: Given that ( s i,j , a i,j ) has positive probability under ( P i , π r ) , there exists some t for which ( s i,j , a i,j ) ∈

Original: I i,j ∩ Σ it .
Revised: I i,j ∩ Σ it .

Original: Since Σ it =
Revised: Since Σ it =

Original: Σ i,jt , we have ( s i,j , a i,j ) ∈
Revised: Σ i,jt , we have ( s i,j , a i,j ) ∈

Original: I i,j ∩ Σ i,jt , which means ( s i,j , a i,j ) also has positive probability under ( P j , π r ) .
Revised: I i,j ∩ Σ i,jt , which means ( s i,j , a i,j ) also has positive probability under ( P j , π r ) .

Original: Next, consider the case where there exists some t ∈ { 1 , . .
Revised: Next, consider the case where there exists some t ∈ { 1 , . .

Original: , T max } for which Σ
Revised: , T max } for which Σ

Original: it (cid:54) =
Revised: it (cid:54) =

Original: Σ i,jt and let ˆ t = min { t : Σ it (cid:54) =
Revised: Σ i,jt and let ˆ t = min { t : Σ it (cid:54) =

Original: Σ i,jt } .
Revised: Σ i,jt } .

Original: Note that ˆ t > 0 , since we have already shown that Σ i 

Revised: Note that ˆ t > 0 , since we have already shown that Σ i

Original: = Σ
Revised: = Σ

Original: i,j 0 .
Revised: i,j 0 .

Original: Thus, for every t < ˆ t we have Σ it =
Revised: Thus, for every t < ˆ t we have Σ it =

Original: Σ i,jt = Σ jt , and for ˆ t it holds that P P i ,π r , ˆ t ( s, a ) (cid:54) =
Revised: Σ i,jt = Σ jt , and for ˆ t it holds that P P i ,π r , ˆ t ( s, a ) (cid:54) =

Original: P P j ,π r , ˆ t ( s, a ) .
Revised: P P j ,π r , ˆ t ( s, a ) .

Original: If there exits a t (cid:48) <
Revised: If there exits a t (cid:48) <

Original: ˆ t − 1 and ( s, a ) ∈ Σ it (cid:48) such that P i ( ·| s, a ) (cid:54) =
Revised: ˆ t − 1 and ( s, a ) ∈ Σ it (cid:48) such that P i ( ·| s, a ) (cid:54) =

Original: P j ( ·| s, a ) , then we are done as Σ it (cid:48)
Revised: P j ( ·| s, a ) , then we are done as Σ it (cid:48)

Original: =
Revised: =

Original: Σ i,jt (cid:48) , which means that ( s, a ) is an identifying state-action pair that has positive probability under both ( P i , π r ) and ( P j , π r ) .
Revised: Σ i,jt (cid:48) , which means that ( s, a ) is an identifying state-action pair that has positive probability under both ( P i , π r ) and ( P j , π r ) .

Original: Therefore, consider the case where for every t < ˆ t − 1 and every ( s, a ) ∈
Revised: Therefore, consider the case where for every t < ˆ t − 1 and every ( s, a ) ∈

Original: Σ
Revised: Σ

Original: it we have P i ( ·| s, a ) =
Revised: it we have P i ( ·| s, a ) =

Original: P j ( ·| s, a ) .
Revised: P j ( ·| s, a ) .

Original: We will show that there exists ( s, a ) ∈ Σ
Revised: We will show that there exists ( s, a ) ∈ Σ

Original: i ˆ t − 1 such that
Revised: i ˆ t − 1 such that

Original: P i ( ·| s, a ) (cid:54) =
Revised: P i ( ·| s, a ) (cid:54) =

Original: P j ( ·| s, a ) .
Revised: P j ( ·| s, a ) .

Original: Assume to the contrary that for every ( s, a ) ∈ Σ
Revised: Assume to the contrary that for every ( s, a ) ∈ Σ

Original: i ˆ t − 1 we have P i ( ·| s, a ) =
Revised: i ˆ t − 1 we have P i ( ·| s, a ) =

Original: P j ( ·| s, a ) , i.e., the transition function is also equivalent for t
Revised: P j ( ·| s, a ) , i.e., the transition function is also equivalent for t

Original: = ˆ t − 1 .
Revised: = ˆ t − 1 .

Original: Let h ˆ t
Revised: Let h ˆ t

Original: =
Revised: =

Original: ( x, s 0 , a 0 , . .
Revised: ( x, s 0 , a 0 , . .

Original: , s ˆ t , a ˆ t ) be the state-action history up to time ˆ t , including the random variable x that was used to choose a policy.
Revised: , s ˆ t , a ˆ t ) be the state-action history up to time ˆ t , including the random variable x that was used to choose a policy.

Original: We next consider the probability of observing a history under ( P i , π r ) ,
Revised: We next consider the probability of observing a history under ( P i , π r ) ,

Original: [Equation]
Revised: [Equation]

Original: where the last equality holds according to the Markov property, P P i ,π r 
Revised: where the last equality holds according to the Markov property, P P i ,π r ( s

Original: [Equation]
Revised: [Equation]

Original: s
Revised: s

Original: [Equation]
Revised: [Equation]

Original: s

Original: [Equation]
Revised: [Equation]

Original: , s [Equation] [Equation] − 1 ) 
Revised: , s t − 1 , a t − 1 ) =

Original: = P i ( s t | s t − 1 , a t − 1 ) .
Revised: P i ( s t | s t − 1 , a t − 1 ) .

Original: Since π r ( a t | x, s 0 , a 0 , . . .
Revised: Since π r ( a t | x, s 0 , a 0 , . . .

Original: , s t ) is the same replaying policy for all MDPs, and for every t ≤ ˆ t − 1 and ( s, a ) ∈
Revised: , s t ) is the same replaying policy for all MDPs, and for every t ≤ ˆ t − 1 and ( s, a ) ∈

Original: Σ
Revised: Σ

Original: it we have P i ( ·| s, a ) =
Revised: it we have P i ( ·| s, a ) =

Original: P j ( ·| s, a ) , then P P i ,π r ( h ˆ t ) =
Revised: P j ( ·| s, a ) , then P P i ,π r ( h ˆ t ) =

Original: P P j ,π r ( h ˆ t ) .
Revised: P P j ,π r ( h ˆ t ) .

Original: By marginalizing over x, s 0 , a 0 , . . .
Revised: By marginalizing over x, s 0 , a 0 , . . .

Original: , s ˆ t − 1 , a ˆ t − 1 we obtain:
Revised: , s ˆ t − 1 , a ˆ t − 1 we obtain:

Original: [Equation]
Revised: [Equation]

Original:  ˆ t − 1 ,a ˆ t −
Revised: ...,s ˆ t − 1 ,a ˆ t −

Original: P P j ,π r ( x, s 0 , a 0 , . .
Revised: P P j ,π r ( x, s 0 , a 0 , . .

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

Original: P P i ,π r ( s ˆ t , a ˆ t ) =
Revised: P P i ,π r ( s ˆ t , a ˆ t ) =

Original: P P j ,π r ( s ˆ t , a ˆ t ) , which means that Σ
Revised: P P j ,π r ( s ˆ t , a ˆ t ) , which means that Σ

Original: i ˆ t = Σ
Revised: i ˆ t = Σ

Original: i,j ˆ t , which contradicts the deﬁnition of ˆ t .
Revised: i,j ˆ t , which contradicts the deﬁnition of ˆ t .

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

Original: In this section we describe the evaluation metric, data collection process, and the details of the domains we experimented with.
Revised: In this section we describe the evaluation metric, data collection process, and the details of the domains we experimented with.

Original: E VALUATION M ETRIC AND D ATA COLLECTION
Revised: E VALUATION M ETRIC AND D ATA COLLECTION

Original: To evaluate performance, we measure average reward in the ﬁrst few 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 few episodes on unseen tasks – this is where efﬁcient exploration makes a critical difference.

Original: For Gridworld, we measure average reward in the ﬁrst 4 episodes, for Wind only the ﬁrst episode reward is measured, and for the rest of the domains, measure average performance in the ﬁrst 2 episodes.
Revised: For Gridworld, we measure average reward in the ﬁrst 4 episodes, for Wind only the ﬁrst episode reward is measured, and for the rest of the domains, measure average performance in the ﬁrst 2 episodes.

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: 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 either the ﬁrst or last k/ 2 trajectories with trajectories from a randomly chosen MDP, and relabel their rewards.
Revised: For reward relabelling, we replace either the ﬁrst or last k/ 2 trajectories with trajectories from a randomly chosen MDP, and relabel their rewards.

Original: For policy replay, we replace trajectories by sampling a new trajectory using the trained RL policy of another MDP.
Revised: For policy replay, we replace trajectories by sampling a new trajectory using the trained RL policy of another MDP.

Original: D OMAINS WITH V ARYING R EWARD F UNCTION
Revised: D OMAINS WITH V ARYING R EWARD F UNCTION

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.,
Revised: 1 on non-goal cells and +1 at the goal, i.e.,

Original: [Equation]
Revised: [Equation]

Original: where s t is the current cell and g is the goal cell.
Revised: 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

Original: [Equation]

Original: where x t is the current 2D location.
Revised: 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:
Revised: MuJoCo:

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: The horizon is set to 200 and we aggregate k = 2 consecutive episodes.
Revised: The horizon is set to 200 and we aggregate k = 2 consecutive episodes.

Original: Ant-Semi-circle: In this environment, an ant needs to navigate to an unknown goal, randomly chosen on a semi-circle, similarly to the Semi-circle task above.
Revised: Ant-Semi-circle: In this environment, an ant needs 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: The horizon 
Revised: The horizon is set

Original: is set to200 and we aggregate 
Revised: to and we aggregate k

Original: k = 2 consecutive episodes
Revised: = 2 consecutive episodes.

Original: Reacher-Image:
Revised: Reacher-Image:

Original: In this environment, a two-link
Revised: In this environment

Original:  planar robot needs to reach an unknown goal, randomly chosen on a quarter circle.
Revised: , a two-link planar robot needs to reach an unknown goal, randomly chosen on a quarter circle.

Original: The robot receives dense reward which is given by
Revised: The robot receives dense reward which is given by

Original: [Equation]
Revised: [Equation]

Original: where x t is the location of the robot’s end effector.
Revised: where x t is the location of the robot’s end effector.

Original: The agent observes single-channel images of size 64 × 64 of the environment.
Revised: The agent observes single-channel images of size 64 × 64 of the environment.

Original: The horizon is set to 100 and we aggregate k = 2 consecutive episodes.
Revised: The horizon is set to 100 and we aggregate k = 2 consecutive episodes.

Original: D OMAINS WITH V ARYING T RANSITION F UNCTION
Revised: D OMAINS WITH V ARYING T RANSITION F UNCTION

Original: Wind: A continuous 2D domain with varying transitions, where the agent must navigate to a ﬁxed (unknown) goal within a distance of D = 1 from its initial state (the goal location is the same for all tasks).
Revised: Wind: A continuous 2D domain with varying transitions, where the agent must navigate to a ﬁxed (unknown) goal within a distance of D = 1 from its initial state (the goal location is the same for all tasks).

Original: Similarly to Semi-circle, the agent receives a reward of +1 if it is within a radius r = 0 .
Revised: Similarly to Semi-circle, the agent receives a reward of +1 if it is within a radius r = 0 .

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

Original: [Equation]
Revised: [Equation]

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

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

Original: [Equation]
Revised: [Equation]

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

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

Original: We set the horizon to 25 and evaluate the performance in terms of average return within the ﬁrst episode of interaction on test tasks ( k = 1 ).
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 ( k = 1 ).

Original: Escape-Room:
Revised: Escape-Room:

Original: A continuous 2D domain where the agent must navigate outside a circular room of radius R = 1 through an opening, whose location is unknown.
Revised: A continuous 2D domain where the agent must navigate outside a circular room of radius R = 1 through an opening, whose location is unknown.

Original: For all tasks, the central angle of the opening is π/ 8 .
Revised: For all tasks, the central angle of the opening is π/ 8 .

Original: The tasks differ by the location of the opening – the center point of the opening is sampled uniformly from [0 , π ] .
Revised: The tasks differ by the location of the opening – the center point of the opening is sampled uniformly from [0 , π ] .

Original: The reward function is sparse, task-independent, and given by
Revised: The reward function is sparse, task-independent, and given by

Original: [Equation]
Revised: [Equation]

Original: else
Revised: else.

Original: [Equation]

Original: The transition function, however, is task-dependent and given by
Revised: The transition function, however, is task-dependent and given by

Original: [Equation]
Revised: [Equation]

Original: occurs
Revised: occurs

Original: [Equation]
Revised: [Equation]

Original: where intersection occurs means that the line that connects s t and s t + a t and the wall of the circular room intersects.
Revised: where intersection occurs means that the line that connects s t and s t + a t and the wall of the circular room intersects.

Original: To solve a task, the agent must search for the opening by moving along the wall until he ﬁnds it.
Revised: To solve a task, the agent must search for the opening by moving along the wall until he ﬁnds it.

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

Original: In Figure 5 we compare our ofﬂine algorithm with PEARL (Rakelly et al., 2019), and also with an ablation of the policy replaying method.
Revised: In Figure 5 we compare our ofﬂine algorithm with PEARL (Rakelly et al., 2019), and also with an ablation of the policy replaying method.

Original: In the Escape-Room domain, where the identifying states do not overlap, policy replaying indeed improves performance, while in Wind, where the identifying states do overlap, policy replaying has little effect.
Revised: In the Escape-Room domain, where the identifying states do not overlap, policy replaying indeed improves performance, while in Wind, where the identifying states do overlap, policy replaying has little effect.

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 Semi-circle domain, 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 Semi-circle domain, in order to get more insight into the decision making process of the agent during interaction.

Original: In Figure 6, 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 6, 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 domain 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 domain for which we modify the initial state distribution during the data collection phase.

Original: 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 (excluding state on the semi-circle), or (3) ﬁxed to a single position.
Revised: 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 (excluding state on the semi-circle), 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 initial state distributions we consider are visualized in Figure 7.
Revised: The initial state distributions we consider are visualized in Figure 7.

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original:  We report results for the 3 different data collection strategies described above, summarized in Table 1.
Revised: collection Uniform We report results for the 3 different data collection strategies described above, 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 8, even on the low-diversity datasets, our agents learned nontrivial exploration strategies that search for the goal.
Revised: However, as we qualitatively show in Figure 8, even on the low-diversity datasets, our agents learned nontrivial exploration strategies that search for the goal.

Original: This is especially remarkable for the ﬁxeddistribution dataset, where it is unlikely that any training trajectory traveled along the semicircle.
Revised: This is especially remarkable for the ﬁxeddistribution dataset, where it is unlikely that any training trajectory traveled along the semicircle.

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-the-art 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-the-art ofﬂine RL method – to train the critic network of the meta-RL agent.

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

Original: Together with our results on MDP ambiguity, our investigation highlights the particular challenges of the OMRL problem.
Revised: Together with our results on MDP ambiguity, our investigation highlights the particular challenges of the OMRL problem.