Original: L EARNING
Revised: L EARNING

Original: R OBUST
Revised: R OBUST

Original: S TATE
Revised: S TATE

Original: A BSTRACTIONS
Revised: A BSTRACTIONS

Original: FOR
Revised: FOR

Original: H IDDEN -P ARAMETER B LOCK MDP S
Revised: H IDDEN -P ARAMETER B LOCK MDP S

Original: Anonymous authors
Revised: Anonymous authors

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

Original: A BSTRACT
Revised: A BSTRACT

Original: Many control tasks exhibit similar dynamics that can be modeled as having common latent structure.
Revised: Many control tasks exhibit similar dynamics that can be modeled as having common latent structure.

Original: Hidden-Parameter Markov Decision Processes (HiP-MDPs) explicitly model this structure to improve sample efﬁciency in multi-task settings.
Revised: Hidden-Parameter Markov Decision Processes (HiP-MDPs) explicitly model this structure to improve sample efﬁciency in multi-task settings.

Original: However, this setting makes strong assumptions on the observability of the state that limit its application in real-world scenarios with rich observation spaces.
Revised: However, this setting makes strong assumptions on the observability of the state that limit its application in real-world scenarios with rich observation spaces.

Original: In this work, we leverage ideas of common structure from the HiP-MDP setting, and extend it to enable robust state abstractions inspired by Block MDPs.
Revised: In this work, we leverage ideas of common structure from the HiP-MDP setting, and extend it to enable robust state abstractions inspired by Block MDPs.

Original: We derive instantiations of this new framework for both multi-task reinforcement learning (MTRL) and meta-reinforcement learning (Meta-RL) settings.
Revised: We derive instantiations of this new framework for both multi-task reinforcement learning (MTRL) and meta-reinforcement learning (Meta-RL) settings.

Original: Further, we provide transfer and generalization bounds based on task and state similarity, along with sample complexity bounds that depend on the aggregate number of samples across tasks, rather than the number of tasks, a signiﬁcant improvement over prior work.
Revised: Further, we provide transfer and generalization bounds based on task and state similarity, along with sample complexity bounds that depend on the aggregate number of samples across tasks, rather than the number of tasks, a signiﬁcant improvement over prior work that use the same environment assumptions.

Original: To further demonstrate efﬁcacy of the proposed method, we empirically compare and show improvement over multi-task and meta-reinforcement learning baselines.
Revised: To further demonstrate efﬁcacy of the proposed method, we empirically compare and show improvement over multi-task and meta-reinforcement learning baselines.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: A key open challenge in AI research that remains is how to train agents that can learn behaviors that generalize across tasks and environments.
Revised: A key open challenge in AI research that remains is how to train agents that can learn behaviors that generalize across tasks and environments.

Original: When there is common structure underlying the tasks, we have seen that multi-task reinforcement learning (MTRL), where the agent learns a set of tasks simultaneously, has deﬁnite advantages (in terms of robustness and sample efﬁciency) over the singletask setting, where the agent independently learns each task.
Revised: When there is common structure underlying the tasks, we have seen that multi-task reinforcement learning (MTRL), where the agent learns a set of tasks simultaneously, has deﬁnite advantages (in terms of robustness and sample efﬁciency) over the singletask setting, where the agent independently learns each task.

Original: There are two ways in which learning multiple tasks can accelerate learning: the agent can learn a common representation of observations, and the agent can learn a common way to behave.
Revised: There are two ways in which learning multiple tasks can accelerate learning: the agent can learn a common representation of observations, and the agent can learn a common way to behave.

Original: Prior work in MTRL has also leveraged the idea by sharing representations across tasks (D’Eramo et al., 2020) or providing per-task sample complexity results that show improved sample efﬁciency from transfer (Brunskill &
Revised: Prior work in MTRL has also leveraged the idea by sharing representations across tasks (D’Eramo et al., 2020) or providing per-task sample complexity results that show improved sample efﬁciency from transfer (Brunskill &

Original: Li, 2013).
Revised: Li, 2013).

Original: However, explicit exploitation of the shared structure across tasks via a uniﬁed dynamics has been lacking.
Revised: However, explicit exploitation of the shared structure across tasks via a uniﬁed dynamics has been lacking.

Original: Prior works that make use of shared representations use a naive uniﬁcation approach that posits all tasks lie in a shared domain (Figure 1, left).
Revised: Prior works that make use of shared representations use a naive uniﬁcation approach that posits all tasks lie in a shared domain (Figure 1, left).

Original: On the other hand, in the single-task setting, research on state abstractions has a much richer history, with several works on improved generalization through the aggregation of behaviorally similar states (Ferns et al., 2004; Li et al., 2006; Luo et al., 2019; Zhang et al., 2020b).
Revised: On the other hand, in the single-task setting, research on state abstractions has a much richer history, with several works on improved generalization through the aggregation of behaviorally similar states (Ferns et al., 2004; Li et al., 2006; Luo et al., 2019; Zhang et al., 2020b).

Original: In this work, we propose to leverage rich state abstraction models from the single-task setting, and explore their potential for the more general multi-task setting.
Revised: In this work, we propose to leverage rich state abstraction models from the single-task setting, and explore their potential for the more general multi-task setting.

Original: We frame the problem as a structured super-MDP with a shared state space and universal dynamics model conditioned on a task-speciﬁc hidden parameter (Figure 1, right).
Revised: We frame the problem as a structured super-MDP with a shared state space and universal dynamics model conditioned on a task-speciﬁc hidden parameter (Figure 1, right).

Original: This additional structure gives us better sample efﬁciency, both theoretically and empirically.
Revised: This additional structure gives us better sample efﬁciency, both theoretically, compared to related bounds (Brunskill & Li, 2013; Tirinzoni et al., 2020) and empirically against relevant baselines (Yu et al., 2020; Rakelly et al., 2019; Chen et al., 2018; Teh et al., 2017).

Original: We learn a latent representation with smoothness properties for better few-shot generalization to other unseen tasks within this family.
Revised: We learn a latent representation with smoothness properties for better few-shot generalization to other unseen tasks within this family.

Original: This allows us to derive new value loss bounds and sample complexity bounds that depend on how far away a new task is from the ones already seen.
Revised: This allows us to derive new value loss bounds and sample complexity bounds that depend on how far away a new task is from the ones already seen.

Original: We focus on multi-task settings where dynamics can vary across tasks, but the reward function is shared.
Revised: We focus on multi-task settings where dynamics can vary across tasks, but the reward function is shared.

Original: We show that this setting can be formalized as a hidden-parameter MDP (HiP-MDP)
Revised: We show that this setting can be formalized as a hidden-parameter MDP (HiP-MDP)

Original: (DoshiVelez & Konidaris, 2013), where the changes in dynamics can be deﬁned by a latent variable, unifying dynamics across tasks as a single global function.
Revised: (DoshiVelez & Konidaris, 2013), where the changes in dynamics can be deﬁned by a latent variable, unifying dynamics across tasks as a single global function.

Original: This setting assumes a global latent structure over all tasks (or MDPs).
Revised: This setting assumes a global latent structure over all tasks (or MDPs).

Original: Many real-world scenarios fall under this framework, such as autonomous driving under different weather and road conditions, or even different vehicles, which change the dynamics of driving.
Revised: Many real-world scenarios fall under this framework, such as autonomous driving under different weather and road conditions, or even different vehicles, which change the dynamics of driving.

Original: Another example is warehouse robots, where the same tasks are performed in different conditions and warehouse layouts.
Revised: Another example is warehouse robots, where the same tasks are performed in different conditions and warehouse layouts.

Original: The setting is also applicable to some cases of RL for medical treatment optimization, where different patient groups have different responses to treatment, yet the desired outcome is the same.
Revised: The setting is also applicable to some cases of RL for medical treatment optimization, where different patient groups have different responses to treatment, yet the desired outcome is the same.

Original: With this assumed structure, we can provide concrete zero-shot generalization bounds to unseen tasks within this family.
Revised: With this assumed structure, we can provide concrete zero-shot generalization bounds to unseen tasks within this family.

Original: Further, we explore the setting where the state space is latent and we have access to only high-dimensional observations, and we show how to recover robust state abstractions in this setting.
Revised: Further, we explore the setting where the state space is latent and we have access to only high-dimensional observations, and we show how to recover robust state abstractions in this setting.

Original: This is, again, a highly realistic setting in robotics when we don’t always have an amenable, Lipschitz low-dimensional state space.
Revised: This is, again, a highly realistic setting in robotics when we do not always have an amenable, Lipschitz low-dimensional state space.

Original: Cameras are a convenient and inexpensive way to acquire state information, and handling pixel observations is key to approaching these problems.
Revised: Cameras are a convenient and inexpensive way to acquire state information, and handling pixel observations is key to approaching these problems.

Original: A block MDP (Du et al., 2019) provides a concrete way to formalize this observation-based setting.
Revised: A block MDP (Du et al., 2019) provides a concrete way to formalize this observation-based setting.

Original: Leveraging this property of the block MDP framework, in combination with the assumption of a uniﬁed dynamical structure of HiP-MDPs, we introduce the hidden-parameter block MDP (HiP-BMDP) to handle settings with high-dimensional observations and structured, changing dynamics.
Revised: Leveraging this property of the block MDP framework, in combination with the assumption of a uniﬁed dynamical structure of HiP-MDPs, we introduce the hidden-parameter block MDP (HiP-BMDP) to handle settings with high-dimensional observations and structured, changing dynamics.

Original: Key contributions of this work are a new viewpoint of the multi-task setting with same reward function as a universal MDP under the HiP-BMDP setting, which naturally leads to a gradientbased representation learning algorithm.
Revised: Key contributions of this work are a new viewpoint of the multi-task setting with same reward function as a universal MDP under the HiP-BMDP setting, which naturally leads to a gradientbased representation learning algorithm.

Original: Further, this framework allows us to compute theoretical generalization results with the incorporation of a learned state representation.
Revised: Further, this framework allows us to compute theoretical generalization results with the incorporation of a learned state representation.

Original: Finally, empirical results show that our method outperforms other multi-task and meta-learning baselines in both fast adaptation and zero-shot transfer settings.
Revised: Finally, empirical results show that our method outperforms other multi-task and meta-learning baselines in both fast adaptation and zero-shot transfer settings.

Original:  B ACKGROUND
Revised:  B ACKGROUND

Original: In this section, we introduce the base environment as well as notation and additional assumptions about the latent structure of the environments and multi-task setup considered in this work.
Revised: In this section, we introduce the base environment as well as notation and additional assumptions about the latent structure of the environments and multi-task setup considered in this work.

Original: A ﬁnite, discrete-time Markov Decision Process (MDP) (Bellman, 1957; Puterman, 1995) is a tuple (cid:104)S , A , R, T, γ (cid:105) , where S is the set of states, A is the set of actions, R : S ×
Revised: A ﬁnite 1 , discrete-time Markov Decision Process (MDP) (Bellman, 1957; Puterman, 1995) is a tuple (cid:104)S , A , R, T, γ (cid:105) , where S is the set of states, A is the set of actions, R : S ×

Original: A → R is the reward function, T : S ×
Revised: A → R is the reward function, T : S ×

Original: A → Dist ( S ) is the environment transition probability function, and γ ∈ [0 , 1) is the discount factor.
Revised: A → Dist ( S ) is the environment transition probability function, and γ ∈ [0 , 1) is the discount factor.

Original: At each time step, the learning agent perceives a state s t ∈ S , takes an action a t ∈ A drawn from a policy π : S ×
Revised: At each time step, the learning agent perceives a state s t ∈ S , takes an action a t ∈ A drawn from a policy π : S ×

Original: A → [0 , 1] , and with probability T ( s t +1 | s t , a t ) enters next state s t +1 , receiving a numerical reward R t +1 from the environment.
Revised: A → [0 , 1] , and with probability T ( s t +1 | s t , a t ) enters next state s t +1 , receiving a numerical reward R t +1 from the environment.

Original: The value function of policy π is deﬁned as: V π ( s ) =
Revised: The value function of policy π is deﬁned as: V π ( s ) =

Original: E π [ (cid:80) ∞ t =

Revised: E π [ (cid:80) ∞ t =

Original: γ t R t +1 | S 

Revised: γ t R t +1 | S

Original: = s ] .
Revised: = s ] .

Original: The optimal value function V ∗ is the maximum value function over the class of stationary policies.
Revised: The optimal value function V ∗ is the maximum value function over the class of stationary policies.

Original: Hidden-Parameter MDPs (HiP-MDPs) (Doshi-Velez & Konidaris, 2013) can be deﬁned by a tuple
Revised: Hidden-Parameter MDPs (HiP-MDPs) (Doshi-Velez & Konidaris, 2013) can be deﬁned by a tuple

Original: M : (cid:104)S , A , Θ , T θ , R, γ, P Θ (cid:105) where S is a ﬁnite state space, A a ﬁnite action space, T θ describes the transition distribution for a speciﬁc task described by task parameter θ ∼ P Θ , R is the reward function, γ is the discount factor, and P Θ the distribution over task parameters.
Revised: M : (cid:104)S , A , Θ , T θ , R, γ, P Θ (cid:105) where S is a ﬁnite state space, A a ﬁnite action space, T θ describes the transition distribution for a speciﬁc task described by task parameter θ ∼ P Θ , R is the reward function, γ is the discount factor, and P Θ the distribution over task parameters.

Original: This deﬁnes a family of MDPs, where each MDP is described by the parameter θ ∼ P Θ .
Revised: This deﬁnes a family of MDPs, where each MDP is described by the parameter θ ∼ P Θ .

Original: We assume that this parameter θ is ﬁxed for an episode and indicated by an environment id given at the start of the episode.
Revised: We assume that this parameter θ is ﬁxed for an episode and indicated by an environment id given at the start of the episode.

Original: Block MDPs (Du et al., 2019) are described by a tuple (cid:104)S , A , X , p, q, R (cid:105) with a unobservable state space S , action space A , and observable space X .
Revised: Block MDPs (Du et al., 2019) are described by a tuple (cid:104)S , A , X , p, q, R (cid:105) with an unobservable state space S , action space A , and observable space X .

Original: p denotes the latent transition distribution p ( s (cid:48) | s, a ) for s, s (cid:48) ∈ S , a ∈ A , q is the (possibly stochastic) emission mapping that emits the observations q ( x | s ) for x ∈ X , s ∈ S , and R the reward function.
Revised: p denotes the latent transition distribution p ( s (cid:48) | s, a ) for s, s (cid:48) ∈ S , a ∈ A , q is the (possibly stochastic) emission mapping that emits the observations q ( x | s ) for x ∈ X , s ∈ S , and R the reward function.

Original: We are interested in the setting where this mapping q is one-to-many.
Revised: We are interested in the setting where this mapping q is one-to-many.

Original: This is common in many real world problems with many tasks where the underlying states and dynamics are the same, but the observation space that the agent perceives can be quite different, e.g. navigating a house of the same layout but different decorations and furnishings.
Revised: This is common in many real world problems with many tasks where the underlying states and dynamics are the same, but the observation space that the agent perceives can be quite different, e.g. navigating a house of the same layout but different decorations and furnishings.

Original: Assumption 1 (Block structure (Du et al., 2019)) .
Revised: Assumption 1 (Block structure (Du et al., 2019)) .

Original: Each observation x uniquely determines its generating state s .
Revised: Each observation x uniquely determines its generating state s .

Original: That is, the observation space X can be partitioned into disjoint blocks X s , each containing the support of the conditional distribution q ( ·| s ) .
Revised: That is, the observation space X can be partitioned into disjoint blocks X s , each containing the support of the conditional distribution q ( ·| s ) .

Original: Assumption 1 gives the Markov property in X , a key difference from partially observable MDPs (Kaelbling et al., 1998), which has no guarantee of determining the generating state from the history of observations.
Revised: Assumption 1 gives the Markov property in X , a key difference from partially observable MDPs (Kaelbling et al., 1998), which has no guarantee of determining the generating state from the history of observations.

Original: This assumption allows us to compute reasonable bounds for our algorithm in k -order
Revised: This assumption allows us to compute reasonable bounds for our algorithm in k -order

Original: MDPs 1 (which  many real world problems can be described as) and avoiding the intractability of true POMDPs.
Revised: MDPs 2 (which describes many real world problems) and avoiding the intractability of true POMDPs, which have no guarantees on providing enough information to sufﬁciently predict future rewards.

Revised: A relaxation of this assumption entails providing less information in the observation for predicting future reward, which will degrade performance.

Original: We show empirically that our method is  robust to a relaxation of this assumption.
Revised: We show empirically that our method is still more robust to a relaxation of this assumption compared to other MTRL methods.

Original: Bisimulation is a strict form of state abstraction, where two states are bisimilar if they are behaviorally equivalent.
Revised: Bisimulation is a strict form of state abstraction, where two states are bisimilar if they are behaviorally equivalent.

Original: Bisimulation metrics (Ferns et al., 2011) deﬁne a distance between states as follows:
Revised: Bisimulation metrics (Ferns et al., 2011) deﬁne a distance between states as follows:

Original: Deﬁnition 1 (Bisimulation Metric (Theorem 2.6 in Ferns et al.
Revised: Deﬁnition 1 (Bisimulation Metric (Theorem 2.6 in Ferns et al.

Original: Let ( S , A , P, r ) be a ﬁnite
Revised: Let ( S , A , P, r ) be a ﬁnite

Original: MDP and met the space of bounded pseudometrics on S equipped with the metric induced by the uniform norm.
Revised: MDP and met the space of bounded pseudometrics on S equipped with the metric induced by the uniform norm.

Original: Deﬁne F : met (cid:55)→ met by
Revised: Deﬁne F : met (cid:55)→ met by

Original: [Equation]
Revised: [Equation]

Original: where W ( d ) is the Wasserstein distance between transition probability distributions.
Revised: where W ( d ) is the Wasserstein distance between transition probability distributions.

Original: Then F has a unique ﬁxed point ˜
Revised: Then F has a unique ﬁxed point ˜

Original: d which is the bisimulation metric.
Revised: d which is the bisimulation metric.

Original: A nice property of this metric ˜ d is that difference in optimal value between two states is bounded by their distance as deﬁned by this metric.
Revised: A nice property of this metric ˜ d is that difference in optimal value between two states is bounded by their distance as deﬁned by this metric.

Original: Theorem 1 ( V ∗ is Lipschitz with respect to ˜ d
Revised: Theorem 1 ( V ∗ is Lipschitz with respect to ˜ d

Original: (Ferns et al., 2004)) .
Revised: (Ferns et al., 2004)) .

Original: Let V ∗ be the optimal value function for a given discount factor γ .
Revised: Let V ∗ be the optimal value function for a given discount factor γ .

Original: Then V ∗ is Lipschitz continuous with respect to ˜ d with
Revised: Then V ∗ is Lipschitz continuous with respect to ˜ d with

Original: Lipschitz constant 11 − γ ,
Revised: Lipschitz constant 11 − γ ,

Original: [Equation]
Revised: [Equation]

Original: Therefore, we see that bisimulation metrics give us a Lipschitz value function with respect toa state representation where (cid:96) 2 distance corresponds to ˜ d . 
Revised: Therefore, we see that bisimulation metrics give us a Lipschitz value function with respect to ˜ d . For downstream evaluation of the representations we learn, we use Soft Actor Critic (SAC) (Haarnoja et al., 2018), an off-policy actor-critic method that uses the maximum entropy framework for soft policy iteration. At each iteration, SAC performs soft policy evaluation and improvement steps. The policy evaluation step ﬁts a parametric soft Q-function Q ( s t , a t ) using transitions sampled from the replay buffer D by minimizing the soft Bellman residual,

Revised: [Equation]

Original: T HE H I P-BMDP S ETTING
Revised: The target value function ¯ V is approximated via a Monte-Carlo estimate of the following expectation, ¯ [Equation] where ¯ Q is the target soft Q-function parameterized by a weight vector obtained from an exponentially moving average of the Q-function weights to stabilize training. The policy improvement step then attempts to project a parametric policy π ( a t | s t ) by minimizing KL divergence between the policy and a Boltzmann distribution induced by the Q-function, producing the following objective,

Original:  The HiP-MDP setting (as deﬁned in Section 2) assumes full observability of the state space.
Revised: [Equation]  T HE H I P-BMDP S ETTING The HiP-MDP setting (as deﬁned in Section 2) assumes full observability of the state space.

Original: However, in most real-world scenarios, we only have access to high-dimensional, noisy observations, which often contain irrelevant information to the reward.
Revised: However, in most real-world scenarios, we only have access to high-dimensional, noisy observations, which often contain irrelevant information to the reward.

Original: We combine the Block MDP and HiP-MDP settings to introduce the Hidden-Parameter Block MDP setting (HiP-BMDP), where states are latent, and transition distributions change depending on the task parameters θ . 
Revised: We combine the Block MDP and HiP-MDP settings to introduce the Hidden-Parameter Block MDP setting (HiP-BMDP), where states are latent, and transition distributions change depending on the task parameters θ . This adds an additional [Figure] ψ θ𝑠 a action

Revised: Reward ModelR

Revised: Iθ

Original: This adds an additional dimension of complexity to our problem – we ﬁrst want to learn an amenable state space S , and a universal dynamics model in that representation 2 .
Revised: R(s’) 𝑠′ 𝑇(𝑠,𝑎,𝜃 ! ) 𝑇(𝑠,𝑎,𝜃 " ) 𝐸𝑞. dimension of complexity to our problem – we ﬁrst want to learn an amenable state space S , and a universal dynamics model in that representation 3 .

Original: In this section, we formally deﬁne the HiP-BMDP family in Section 3.1, propose an algorithm for learning HiP-BMDPs in Section 3.2, and ﬁnally provide theoretical analysis for the setting in Section 3.3.
Revised: In this section, we formally deﬁne the HiP-BMDP family in Section 3.1, propose an algorithm for learning HiP-BMDPs in Section 3.2, and ﬁnally provide theoretical analysis for the setting in Section 3.3.

Original: 3. T HE M ODEL
Revised: 3. T HE M ODEL

Original: A HiP-BMDP family can be described by tuple (cid:104)S , A , Θ , T θ , R, γ, P Θ , X , q (cid:105) , with a graphical model of the framework found in Figure 2.
Revised: A HiP-BMDP family can be described by tuple (cid:104)S , A , Θ , T θ , R, γ, P Θ , X , q (cid:105) , with a graphical model of the framework found in Figure 2.

Original: We are given a label k ∈ { 1 , ..., N } for each of N environments.
Revised: We are given a label k ∈ { 1 , ..., N } for each of N environments.

Original: We plan to learn a candidate Θ that uniﬁes the transition dynamics across all environments, effectively ﬁnding T ( · , · , θ ) .
Revised: We plan to learn a candidate Θ that uniﬁes the transition dynamics across all environments, effectively ﬁnding T ( · , · , θ ) .

Original: For two environment settings θ i , θ j ∈ Θ , we deﬁne a distance metric:
Revised: For two environment settings θ i , θ j ∈ Θ , we deﬁne a distance metric:

Original: [Equation]
Revised: [Equation]

Original:  The Wasserstein-1 metric can be written asW d ( P, Q ) =
Revised: [Equation] The Wasserstein-1 metric can be written as

Original: sup f ∈F d (cid:13)(cid:13) E x ∼ P f ( x ) − E y ∼ Q f
Revised: [Equation]

Original: ( y ) (cid:13)(cid:13) 1 , where F d is the set of 1-Lipschitz functions under metric d (M¨uller, 1997).
Revised: ) (cid:13)(cid:13) 1 , where F d is the set of 1-Lipschitz functions under metric d (M¨uller, 1997).

Original: We omit d but use d ( x, y )
Revised: We omit d but use d ( x, y )

Original: = (cid:107) x − y (cid:107) 1 in our setting.
Revised: = (cid:107) x − y (cid:107) 1 in our setting.

Original: This ties distance between θ to the maximum difference in the next state distribution of all state-action pairs in the MDP.
Revised: This ties distance between θ to the maximum difference in the next state distribution of all state-action pairs in the MDP.

Original: Given a HiP-BMDP family M Θ , we assume a multi-task setting where environments with speciﬁc θ ∈ Θ are sampled from this family.
Revised: Given a HiP-BMDP family M Θ , we assume a multi-task setting where environments with speciﬁc θ ∈ Θ are sampled from this family.

Original: We do not have access to θ , and instead get environment labels
Revised: We do not have access to θ , and instead get environment labels

Original: I 1 , I 2 , ..., I N .
Revised: I 1 , I 2 , ..., I N .

Original: The goal is to learn a latent space for the hyperparameters θ 3 .
Revised: The goal is to learn a latent space for the hyperparameters θ 4 .

Original: We want θ to be smooth

Original: [Figure]

Original: ψ θ𝑠 a action

Original: Reward ModelR

Original: Iθ R(s’) 𝑠′ 𝑇(𝑠,𝑎,𝜃 ! ) 𝑇(𝑠,𝑎,𝜃 " ) 𝐸𝑞. with respect to changes in dynamics from environment to environment, which we can set explicitly through the following objective:
Revised: We want θ to be smooth with respect to changes in dynamics from environment to environment, which we can set explicitly through the following objective:

Original: [Equation]
Revised: [Equation]

Original: given environment labelsI 1 , I 2 and ψ : Z (cid:55)→ R d , the encoder that maps from environment label to θ . We now detail an approach for learning a state space using bisimulation objectives and analyze the corresponding error and value bounds. We ﬁrst deﬁne three additional error terms associated with learning a (cid:15) R , (cid:15) T , (cid:15) θ -bisimulation abstraction, [Equation]
Revised: given environment labels

Original: 2 , a ) [Equation] 1 , (cid:15) [Equation]
Revised: I 1 , I 2 and ψ : Z

Original: =

Original: [Equation] We can think of (cid:15) R  ,(cid:15) T as describing a new MDP which is close – but not necessarily the same, if (cid:15) R , (cid:15) T > 0 – to the original Block MDP. These two error terms can be computed empirically over all training environments and are therefore not task-speciﬁc. (cid:15) θ , on the other hand, is measured as a per-task error.
Revised: + (cid:55)→ R d , the encoder that maps from environment label, the set of positive integers, to θ .

Original: Similar methods are used in Jiang et al.

Original: (2015) to bound the loss of a single abstraction, which we extend to the HiP-BMDP setting with a family of tasks.

Original: 3. L EARNING H I P-BMDP S
Revised: 3. L EARNING H I P-BMDP S

Original: The premise of our work is that the HiP-BMDP formulation will improve sample efﬁciency and generalization performance on downstream tasks.
Revised: The premise of our work is that the HiP-BMDP formulation will improve sample efﬁciency and generalization performance on downstream tasks.

Original: We examine two settings, multi-task reinforcement learning (MTRL) and meta-reinforcement learning (meta-RL).
Revised: We examine two settings, multi-task reinforcement learning (MTRL) and meta-reinforcement learning (meta-RL).

Original: In both settings, we have access to N training environments and a held-out set of M evaluation environments, both drawn from a deﬁned family.
Revised: In both settings, we have access to N training environments and a held-out set of M evaluation environments, both drawn from a deﬁned family.

Original: In the MTRL setting, we evaluate model performance across all N training environments and ability to adapt to new environments.
Revised: In the MTRL setting, we evaluate model performance across all N training environments and ability to adapt to new environments.

Original: Adaptation performance is evaluated in both the few-shot regime, where we collect a small number of samples from the evaluation environments to learn each hidden parameter θ , and the zero-shot regime, where we average θ over all training tasks.
Revised: Adaptation performance is evaluated in both the few-shot regime, where we collect a small number of samples from the evaluation environments to learn each hidden parameter θ , and the zero-shot regime, where we average θ over all training tasks.

Original: We evaluate against ablations and other MTRL methods.
Revised: We evaluate against ablations and other MTRL methods.

Original: In the meta-RL setting, the goal for the agent is to leverage knowledge acquired from the previous tasks to adapt quickly to a new task.
Revised: In the meta-RL setting, the goal for the agent is to leverage knowledge acquired from the previous tasks to adapt quickly to a new task.

Original: We evaluate performance in terms of how quickly the agent can achieve a minimum threshold score in the unseen evaluation environments (by learning the correct θ for each new environment).
Revised: We evaluate performance in terms of how quickly the agent can achieve a minimum threshold score in the unseen evaluation environments (by learning the correct θ for each new environment).

Original: Learning a HiP-BMDP approximation of a family of MDPs requires the following components: i) an encoder that maps observations from state space to a learned, latent representation, φ : S (cid:55)→ Z , ii) an environment encoder ψ that maps an environment identiﬁer to a hidden parameter θ , iii) a universal dynamics model T conditioned on task parameter θ .
Revised: Learning a HiP-BMDP approximation of a family of MDPs requires the following components: i) an encoder that maps observations from state space to a learned, latent representation, φ : X (cid:55)→ Z , ii) an environment encoder ψ that maps an environment identiﬁer to a hidden parameter θ , iii) a universal dynamics model T conditioned on task parameter θ .

Original: Figure 2 shows how the components interact during training.
Revised: Figure 2 shows how the components interact during training.

Original: In practice, computing the maximum Wasserstein distance over the entire state-action space is computationally infeasible.
Revised: In practice, computing the maximum Wasserstein distance over the entire state-action space is computationally infeasible.

Original: Therefore, we relax this requirement by taking the expectation over Wasserstein distance with respect to the marginal state distribution of the behavior policy.
Revised: Therefore, we relax this requirement by taking the expectation over Wasserstein distance with respect to the marginal state distribution of the behavior policy.

Original: We train a probabilistic universal dynamics model T to output the desired next state distributions as
Revised: We train a probabilistic universal dynamics model T to output the desired next state distributions as

Original: Gaussians 4 , for which the 2-Wasserstein distance has a closed form:
Revised: Gaussians 5 , for which the 2-Wasserstein distance has a closed form:

Original: [Equation]
Revised: [Equation]

Original: where || · || F is the Frobenius norm.
Revised: where || · || F is the Frobenius norm.

Original: Given that we do not have access to the true universal dynamics function across all environments, it must be learned.
Revised: Given that we do not have access to the true universal dynamics function across all environments, it must be learned.

Original: The objective in Equation (2) is accompanied by an additional objective to learn T , giving a ﬁnal loss function: 
Revised: The objective in Equation (2) is accompanied by an additional objective to learn T , giving a ﬁnal loss function: [Equation] [Table]

Original: [Equation]
Revised: [Equation]

Revised: s I 1 t , a I 1 t , ψ

Revised: ( I 1 )) , s I 1 t

Revised: + (cid:19)

Original:  where red indicates gradients are stopped.
Revised: + MSE (cid:18) T ( s I 2 t , a I 2 t , ψ ( I 2 )) , s I 2 t + (cid:19) (cid:124) (cid:123)(cid:122) (cid:125) Model learning error . (3) where red indicates gradients are stopped.

Original: Transitions { s I 1 t , a I 1 t , s I 1 t
Revised: Transitions { s I 1 t , a I 1 t , s I 1 t

Original: +1 , I 1 } and { s I 2 t , a I 2 t , s I 2 t
Revised: +1 , I 1 } and { s I 2 t , a I 2 t , s I 2 t

Original: +1 , I 2 } from two different environments ( I 1 (cid:54) =
Revised: +1 , I 2 } from two different environments ( I 1 (cid:54) =

Original: I 2 ) are sampled randomly from a replay buffer.
Revised: I 2 ) are sampled randomly from a replay buffer.

Original: In practice, we scale the Θ learning error, our  bisimulation  loss, using a scalar value denoted as α ψ .
Revised: In practice, we scale the Θ learning error, our task bisimulation metric loss, using a scalar value denoted as α ψ .

Original: 3. T HEORETICAL A NALYSIS
Revised: 3. T HEORETICAL A NALYSIS

Original: In this section, we provide value bounds and sample complexity analysis of the HiP-BMDP approach.
Revised: In this section, we provide value bounds and sample complexity analysis of the HiP-BMDP approach.

Original: We have additional new theoretical analysis of the simpler HiP-MDP setting in Appendix B.
Revised: We have additional new theoretical analysis of the simpler HiP-MDP setting in Appendix B.

Revised: We ﬁrst deﬁne three additional error terms associated with learning a (cid:15) R , (cid:15) T , (cid:15) θ -bisimulation abstraction,

Revised: [Equation]

Revised: Φ T denotes the lifted version of T , where we take the next-step transition distribution from observation space X and lift it to latent space S .

Original: Value Bounds.
Revised: We can think of (cid:15) R , (cid:15) T as describing a new MDP which is close – but not necessarily the same, if (cid:15) R , (cid:15) T > 0 – to the original Block MDP. These two error terms can be computed empirically over all training environments and are therefore not task-speciﬁc. (cid:15) θ , on the other hand, is measured as a per-task error. Similar methods are used in Jiang et al.

Revised: (2015) to bound the loss of a single abstraction, which we extend to the HiP-BMDP setting with a family of tasks.

Revised: Value Bounds.

Original: We ﬁrst evaluate how the error in θ prediction and the learned bisimulation representation affect the optimal Q ∗ ¯ M ˆ θ of the learned MDP, by ﬁrst bounding its distance from the optimal
Revised: We ﬁrst evaluate how the error in θ prediction and the learned bisimulation representation affect the optimal Q ∗ ¯ M ˆ θ of the learned MDP, by ﬁrst bounding its distance from the optimal

Original: Q ∗ of the true MDP for a single-task.
Revised: Q ∗ of the true MDP for a single-task.

Original: Theorem 2 ( Q error) .
Revised: Theorem 2 ( Q error) .

Original: Given an MDP ¯ M ˆ θ built on a ( (cid:15) R , (cid:15) T , (cid:15) θ ) -approximate bisimulation abstraction of an instance of a HiP-BMDP M θ , we denote the evaluation of the optimal Q function of ¯ M ˆ θ on M as [ Q ∗ ¯ M ˆ θ ] M θ .
Revised: Given an MDP ¯ M ˆ θ built on a ( (cid:15) R , (cid:15) T , (cid:15) θ ) -approximate bisimulation abstraction of an instance of a HiP-BMDP M θ , we denote the evaluation of the optimal Q function of ¯ M ˆ θ on M as [ Q ∗ ¯ M ˆ θ ] M θ .

Original: The value difference with respect to the optimal Q ∗M is upper bounded by
Revised: The value difference with respect to the optimal Q ∗M is upper bounded by

Original: [Equation]
Revised: [Equation]

Original: Proof in Appendix C. As in the HiP-MDP setting, we can measure the transferability of a speciﬁc policy π learned on one task to another, now taking into account error from the learned representation.
Revised: Proof in Appendix C. As in the HiP-MDP setting, we can measure the transferability of a speciﬁc policy π learned on one task to another, now taking into account error from the learned representation.

Original: Theorem 3 (Transfer bound) .
Revised: Theorem 3 (Transfer bound) .

Original: Given two MDPs M θ i and M θ j , we can bound the difference in Q π between the two MDPs for a given policy π learned under an (cid:15) R , (cid:15) T , (cid:15) θ i -approximate abstraction of
Revised: Given two MDPs M θ i and M θ j , we can bound the difference in Q π between the two MDPs for a given policy π learned under an (cid:15) R , (cid:15) T , (cid:15) θ i -approximate abstraction of

Original: M θ i and applied to
Revised: M θ i and applied to

Original: [Equation]
Revised: [Equation]

Original: This result clearly follows directly from Theorem 2.
Revised: This result clearly follows directly from Theorem 2.

Original: Given a policy learned for task i , Theoremgives a bound on how far from optimal that policy is when applied to task j .
Revised: Given a policy learned for task i , Theoremgives a bound on how far from optimal that policy is when applied to task j .

Original: Intuitively, the more similar in behavior tasks i and j are, as denoted by (cid:107) θ i − θ j (cid:107) 1 , the better π performs on task j .
Revised: Intuitively, the more similar in behavior tasks i and j are, as denoted by (cid:107) θ i − θ j (cid:107) 1 , the better π performs on task j .

Original: Finite Sample Analysis.
Revised: Finite Sample Analysis.

Original: In MDPs (or families of MDPs) with large state spaces, it can be unrealistic to assume that all states are visited at least once, in the ﬁnite sample regime.
Revised: In MDPs (or families of MDPs) with large state spaces, it can be unrealistic to assume that all states are visited at least once, in the ﬁnite sample regime.

Original: Abstractions are useful in this regime for their generalization capabilities.
Revised: Abstractions are useful in this regime for their generalization capabilities.

Original: We can instead perform a counting analysis based on the number of samples of any abstract state-action pair.
Revised: We can instead perform a counting analysis based on the number of samples of any abstract state-action pair.

Original: We compute a loss bound with abstraction φ which depends on the size of the replay buffer D , collected over all tasks.
Revised: We compute a loss bound with abstraction φ which depends on the size of the replay buffer D , collected over all tasks.

Original: Speciﬁcally, we deﬁne the minimal number of visits to an abstract state-action pair, n φ ( D ) =
Revised: Speciﬁcally, we deﬁne the minimal number of visits to an abstract state-action pair, n φ ( D ) =

Original: min x ∈ φ ( S ) ,a ∈A | D x,a | .
Revised: min x ∈ φ ( S ) ,a ∈A | D x,a | .

Original: This sample complexity bound relies on a Hoeffding-style inequality, and therefore requires that the samples in D be independent, which is usually not the case when trajectories are sampled.
Revised: This sample complexity bound relies on a Hoeffding-style inequality, and therefore requires that the samples in D be independent, which is usually not the case when trajectories are sampled.

Original: Theorem 4 (Sample Complexity) .
Revised: Theorem 4 (Sample Complexity) .

Original: For any φ which deﬁnes an ( (cid:15) R , (cid:15) T , (cid:15) θ ) -approximate bisimulation abstraction on a HiP-BMDP family M Θ , we deﬁne the empirical measurement of Q ∗ ¯ M ˆ θ over D to be
Revised: For any φ which deﬁnes an ( (cid:15) R , (cid:15) T , (cid:15) θ ) -approximate bisimulation abstraction on a HiP-BMDP family M Θ , we deﬁne the empirical measurement of Q ∗ ¯ M ˆ θ over D to be

Original: Q ∗ ¯ M D ˆ θ .
Revised: Q ∗ ¯ M D ˆ θ .

Original: Then, with probability ≥ 1 − δ ,
Revised: Then, with probability ≥ 1 − δ ,

Original: [Equation]
Revised: [Equation]

Original: This performance bound applies to all tasks in the family and has two terms that are affected by using a state abstraction: the number of samples n φ ( D ) , and the size of the state space | φ ( X ) | .
Revised: This performance bound applies to all tasks in the family and has two terms that are affected by using a state abstraction: the number of samples n φ ( D ) , and the size of the state space | φ ( X ) | .

Original: We know that | φ ( X ) | ≤ |X| as behaviorally equivalent states are grouped together under bisimulation, and n φ ( D ) is the minimal number of visits to any abstract state-action pair, in aggregate over all training environments.
Revised: We know that | φ ( X ) | ≤ |X| as behaviorally equivalent states are grouped together under bisimulation, and n φ ( D ) is the minimal number of visits to any abstract state-action pair, in aggregate over all training environments.

Original: This is an improvement over the sample complexity of applying single-task learning without transfer over all tasks, and the method proposed in Brunskill & Li (2013), which both would rely on the number of tasks or number of MDPs seen.
Revised: This is an improvement over the sample complexity of applying single-task learning without transfer over all tasks, and the method proposed in Brunskill & Li (2013), which both would rely on the number of tasks or number of MDPs seen.

Original:  E XPERIMENTS & R ESULTS
Revised:  E XPERIMENTS & R ESULTS

Original: We use environments from Deepmind Control Suite (DMC) (Tassa et al., 2018) to evaluate our method for learning HiP-BMDPs for both multi-task RL and meta-reinforcement learning settings.
Revised: We use environments from Deepmind Control Suite (DMC) (Tassa et al., 2018) to evaluate our method for learning HiP-BMDPs for both multi-task RL and meta-reinforcement learning settings.

Original: We consider two setups for evaluation: i) an interpolation setup and ii) an extrapolation setup where the changes in the dynamics function are interpolations and extrapolations between the changes in the dynamics function of the training environment respectively.
Revised: We consider two setups for evaluation: i) an interpolation setup and ii) an extrapolation setup where the changes in the dynamics function are interpolations and extrapolations between the changes in the dynamics function of the training environment respectively.

Original: This dual-evaluation setup provides a more nuanced understanding of how well the learned model transfers across the environments.
Revised: This dual-evaluation setup provides a more nuanced understanding of how well the learned model transfers across the environments.

Original: Implementation details can be found in Appendix D and sample videos of policies at https://sites.google.com/view/hip-bmdp .
Revised: Implementation details can be found in Appendix D and sample videos of policies at https://sites.google.com/view/hip-bmdp .

Original: Environments.
Revised: Environments.

Original: We create a family of MDPs using the existing environment-task pairs from DMC and change one environment parameter to sample different MDPs.
Revised: We create a family of MDPs using the existing environment-task pairs from DMC and change one environment parameter to sample different MDPs.

Original: We denote this parameter as the perturbation- parameter.
Revised: We denote this parameter as the perturbation- parameter.

Original: We consider the following HiP-BMDPs: 1.
Revised: We consider the following HiP-BMDPs: 1.

Original: Cartpole-Swingup-V0 : the mass of the pole varies, 2.
Revised: Cartpole-Swingup-V0 : the mass of the pole varies, 2.

Original: Cheetah-Run-V0 : the size of the torso varies, 3.
Revised: Cheetah-Run-V0 : the size of the torso varies, 3.

Original: Walker-Run-V0 : the friction coefﬁcient between the ground and the walker’s legs varies, 4.
Revised: Walker-Run-V0 : the friction coefﬁcient between the ground and the walker’s legs varies, 4.

Original: Walker-Run-V1 : the size of left-foot of the walker varies, and 5.
Revised: Walker-Run-V1 : the size of left-foot of the walker varies, and 5.

Original: Finger-Spin-V0 : the size of the ﬁnger varies.
Revised: Finger-Spin-V0 : the size of the ﬁnger varies.

Original: We show an example of the different pixel observations for Cheetah-Run-V in Figure 3.
Revised: We show an example of the different pixel observations for Cheetah-Run-V in Figure 3.

Original: Additional environment details are in Appendix D.
Revised: Additional environment details are in Appendix D.

Original:  We sample 8 MDPs from each MDP family by sampling different values for the perturbationparameter.
Revised: [Figure] We sample 8 MDPs from each MDP family by sampling different values for the perturbationparameter.

Original: The MDPs are arranged in order of increasing values of the perturbation- parameter such that we can induce an order over the family of MDPs.
Revised: The MDPs are arranged in order of increasing values of the perturbation- parameter such that we can induce an order over the family of MDPs.

Original: We denote the ordered MDPs as
Revised: We denote the ordered MDPs as

Original: A − H .
Revised: A − H .

Original: MDPs { B, C, F, G } are training environments and { D, E } are used for evaluating the model in the interpolation setup (i.e. the value of the perturbation- parameter can be obtained by interpolation).
Revised: MDPs { B, C, F, G } are training environments and { D, E } are used for evaluating the model in the interpolation setup (i.e. the value of the perturbation- parameter can be obtained by interpolation).

Original: MDPs { A, H } are for evaluating the model in the extrapolation setup (i.e. the value of the perturbation- parameter can be obtained by extrapolation).
Revised: MDPs { A, H } are for evaluating the model in the extrapolation setup (i.e. the value of the perturbation- parameter can be obtained by extrapolation).

Original: We evaluate the learning agents by computing average reward (over 10 episodes) achieved by the policy after training for a ﬁxed number of steps.
Revised: We evaluate the learning agents by computing average reward (over 10 episodes) achieved by the policy after training for a ﬁxed number of steps.

Original: All experiments are run for 10 seeds, with mean and standard error reported in the plots.
Revised: All experiments are run for 10 seeds, with mean and standard error reported in the plots.

Original: Multi-Task Setting.
Revised: Multi-Task Setting.

Original: We ﬁrst consider a multi-task setup where the agent is trained on four related, but different environments with pixel observations.
Revised: We ﬁrst consider a multi-task setup where the agent is trained on four related, but different environments with pixel observations.

Original: We comparethe our method, HiP-BMDP, with the following baselines and ablations: i) DeepMDP (Gelada et al., 2019) where we aggregate data across all training environments, ii) HiP-BMDP-nobisim , HiP-BMDP without the  bisimulation  loss on task embeddings, iii) Distral , an ensemble of policies trained using the Distral algorithm (Teh et al., 2017) with SAC-AE (Yarats et al., 2019) as the underlying policy, iv) PCGrad (Yu et al., 2020), and v) GradNorm (Chen et al., 2018).
Revised: We compare our method, HiP-BMDP, with the following baselines and ablations: i) DeepMDP (Gelada et al., 2019) where we aggregate data across all training environments, ii) HiP-BMDP-nobisim , HiP-BMDP without the task bisimulation metric loss on task embeddings, iii) Distral , an ensemble of policies trained using the Distral algorithm (Teh et al., 2017) with SAC-AE (Yarats et al., 2019) as the underlying policy, iv) PCGrad (Yu et al., 2020), and v) GradNorm (Chen et al., 2018).

Original: For all models, the agent sequentially performs one update per environment.
Revised: For all models, the agent sequentially performs one update per environment.

Original: For fair comparison, we ensure that baselines have at least as many parameters as HiP-BMDP.
Revised: For fair comparison, we ensure that baselines have at least as many parameters as HiP-BMDP.

Original: Distral has more parameters as it trains one policy per environment.
Revised: Distral has more parameters as it trains one policy per environment.

Original: Additional implementation details about baselines are in Appendix D.1.
Revised: Additional implementation details about baselines are in Appendix D.1.

Revised: In Figures 4, and 10 (in Appendix), we observe that for all the models, performance deteriorates when evaluated on interpolation/extrapolation environments.

Revised: We only report extrapolation results in the main paper because of space constraints, as they were very similar to the interpolation performance.

Revised: The gap between the HiP-BMDP model and other baselines also widens, showing that the proposed approach is relatively more robust to changes in environment dynamics.

Original: At training time (Figure 9 in Appendix), we observe that HiP-BMDP consistently outperforms other baselines on all the environments.
Revised: At training time (Figure 9 in Appendix), we observe that HiP-BMDP consistently outperforms other baselines on all the environments.

Original: The success of our proposed method can not be attributed to task[Figure] embeddings alone as HiP-BMDP-nobisim also uses task embeddings.
Revised: The success of our proposed method can not be attributed to task embeddings alone as HiP-BMDP-nobisim also uses task embeddings.

Original: Moreover, only incorporating the task-embeddings is not guaranteed to improve performance in all the environments (as can be seen in the case of Cheetah-Run-V0 ).
Revised: Moreover, only incorporating the task-embeddings is not guaranteed to improve performance in all the environments (as can be seen in the case of Cheetah-Run-V0 ).

Original: We also note that the multi-task learning baselines like
Revised: We also note that the multi-task learning baselines like

Original: Distral, PCGrad, and GradNorm sometimes lag behind even the DeepMDP baseline, perhaps because they do not leverage a shared global dynamics model.
Revised: Distral, PCGrad, and GradNorm sometimes lag behind even the DeepMDP baseline, perhaps because they do not leverage a shared global dynamics model.

Original: In Figures 4, and 10 (in Appendix), we observe that for all the models, performance deteriorates when evaluated on interpolation/extrapolation environments.

Original: We only report extrapolation results in the main paper because of space constraints, as they were very similar to the interpolation performance.

Original: The gap between the HiP-BMDP model and other baselines also widens, showing that the proposed approach is relatively more robust to changes in environment dynamics.

Original: Meta-RL Setting.
Revised: Meta-RL Setting.

Original: We consider the Meta-RL setup for evaluating the few-shot generalization capabilities of our proposed approach.
Revised: We consider the Meta-RL setup for evaluating the few-shot generalization capabilities of our proposed approach on proprioceptive state, as meta-RL techniques are too timeintensive to train on pixel observations directly.

Original: Speciﬁcally, we use PEARL (Rakelly et al., 2019), an offpolicy meta-learning algorithm that uses probabilistic context variables, and is shown to outperform common meta-RL baselines like MAML-TRPO
Revised: Speciﬁcally, we use PEARL (Rakelly et al., 2019), an off-policy meta-learning algorithm that uses probabilistic context variables, and is shown to outperform common meta-RL baselines like MAML-TRPO

Original: (Finn et al., 2017) and ProMP (Rothfuss et al., 2019) on proprioceptive state 5 .
Revised: (Finn et al., 2017) and ProMP (Rothfuss et al., 2019) on proprioceptive state.

Original: We incorporate our proposed approach in PEARL by training  the inference network q φ ( z | c ) with our additional HiP-BMDP loss.
Revised: We incorporate our proposed approach in PEARL by training [Figure] the inference network q φ ( z | c ) with our additional HiP-BMDP loss.

Original: The algorithm pseudocode can be found in Appendix D.
Revised: The algorithm pseudocode can be found in Appendix D.

Original: In Figure 5 we see that the proposed approach (blue) converges faster to a threshold reward (green) than the baseline for Cartpole-Swingup-V0 and Walker-Walk-V1 .
Revised: In Figure 5 we see that the proposed approach (blue) converges faster to a threshold reward (green) than the baseline for Cartpole-Swingup-V0 and Walker-Walk-V1 .

Original: We provide additional results in Appendix E.
Revised: We provide additional results in Appendix E.

Original: [Figure]

Original: shows a threshold reward.

Original: Evaluating the Universal Transition Model.
Revised: Evaluating the Universal Transition Model.

Original: We investigate how well the transition model performs in an unseen environment by only adapting the task parameter θ .
Revised: We investigate how well the transition model performs in an unseen environment by only adapting the task parameter θ .

Original: We instantiate a new MDP, sampled from the family of MDPs, and use a behavior policy to collect transitions.
Revised: We instantiate a new MDP, sampled from the family of MDPs, and use a behavior policy to collect transitions.

Original: These transitions are used to update only the θ parameter, and the transition model is evaluated by unrolling the transition model for k -steps.
Revised: These transitions are used to update only the θ parameter, and the transition model is evaluated by unrolling the transition model for k -steps.

Original: We report the average, perstep model error in latent space, averaged over 10 environments.
Revised: We report the average, perstep model error in latent space, averaged over 10 environments.

Original: While we expect both the proposed setup and baseline setups to adapt to the new environment, we expect the proposed setup to adapt faster because of the exploitation of underlying structure.
Revised: While we expect both the proposed setup and baseline setups to adapt to the new environment, we expect the proposed setup to adapt faster because of the exploitation of underlying structure.

Original: In Figure 6, we indeed observe that the proposed HiP-BMDP model adapts much faster than the ablation HiP-BMDP-nobisim .
Revised: In Figure 6, we indeed observe that the proposed HiP-BMDP model adapts much faster than the ablation HiP-BMDP-nobisim .

Original: Relaxing the Block MDP Assumption.
Revised: Relaxing the Block MDP Assumption.

Original: We incorporate sticky observations into the environment to determine how HiP-BMDP behaves when the Block MDP assumption is relaxed.
Revised: We incorporate sticky observations into the environment to determine how HiP-BMDP behaves when the Block MDP assumption is relaxed.

Original: For some probability p (set to 0.1 in practice), the current observation is dropped, and the agent sees the previous observation again.
Revised: For some probability p (set to 0.1 in practice), the current observation is dropped, and the agent sees the previous observation again.

Original: In Figure 7, we see that even in this setting the proposed HiP-BMDP model outperforms the other baseline models. 
Revised: In Figure 7, we see that even in this setting the proposed HiP-BMDP model outperforms the other baseline models. [Figure]

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: [Figure]

Original: Multi-task learning has been extensively studied in RL with assumptions around common properties of different tasks, e.g., reward and transition dynamics.
Revised: Multi-task learning has been extensively studied in RL with assumptions around common properties of different tasks, e.g., reward and transition dynamics.

Original: A lot of work has focused on considering tasks as MDPs and learning optimal policies for each task while maximizing shared knowledge.
Revised: A lot of work has focused on considering tasks as MDPs and learning optimal policies for each task while maximizing shared knowledge.

Original: However, in most real-world scenarios, the parameters governing the dynamics are not observed.
Revised: However, in most real-world scenarios, the parameters governing the dynamics are not observed.

Original: Moreover, it is not explicitly clear how changes in dynamics across tasks are controlled.
Revised: Moreover, it is not explicitly clear how changes in dynamics across tasks are controlled.

Original: The HiP-BMDP setting provides a principled way to change dynamics across tasks via a latent variable.
Revised: The HiP-BMDP setting provides a principled way to change dynamics across tasks via a latent variable.

Original: Much existing work in the multi-task reinforcement learning (MTRL) setting focuses on learning shared representations (Ammar et al., 2014;
Revised: Much existing work in the multi-task reinforcement learning (MTRL) setting focuses on learning shared representations (Ammar et al., 2014;

Original: Parisotto et al., 2016; Calandriello et al., 2014; Maurer et al., 2016; Landolﬁ et al., 2019).
Revised: Parisotto et al., 2016; Calandriello et al., 2014; Maurer et al., 2016; Landolﬁ et al., 2019).

Original: D’Eramo et al.
Revised: D’Eramo et al.

Original: (2020) extend approximate value iteration bounds in the single-task setting to the multi-task by computing the average loss across tasks and Brunskill & Li (2013) offer sample complexity results, which still depend on the number of tasks, unlike ours. 
Revised: (2020) extend approximate value iteration bounds in the single-task setting to the multi-task by computing the average loss across tasks and Brunskill & Li (2013) offer sample complexity results, which still depend on the number of tasks, unlike ours. Sun et al.

Revised: Tirinzoni et al.

Revised: (2020) also obtain PAC bounds on sample complexity for the MTRL setting, but Sun et al.

Revised: (2020) relies on a constructed state-action abstraction that assumes a discrete and tractably small state space.

Original:  Similar to our work, Perez et al.
Revised: Tirinzoni et al. (2020) assumes access to a generative model for any state-action pair and scales with the minimum of number of tasks or state space. In the rich observation setting, this minimum will almost always be the number of tasks. Similar to our work, Perez et al.

Original: (2020) also treats the multi-task setting as a HiP-MDP by explicitly designing latent variable models to model the latent parameters, but require knowledge of the structure upfront, whereas our approach does not make any such assumptions.
Revised: (2020) also treats the multi-task setting as a HiP-MDP by explicitly designing latent variable models to model the latent parameters, but require knowledge of the structure upfront, whereas our approach does not make any such assumptions.

Original: Meta-learning , or learning to learn, is also a related framework with a different approach.
Revised: Meta-learning , or learning to learn, is also a related framework with a different approach.

Original: We focus here on the context-based approaches, which are more similar to the shared representation approaches of MTRL, and our own method.
Revised: We focus here on the context-based approaches, which are more similar to the shared representation approaches of MTRL, and our own method.

Original: Rakelly et al.
Revised: Rakelly et al.

Original: (2019) model and learn latent contexts upon which a universal policy is conditioned.
Revised: (2019) model and learn latent contexts upon which a universal policy is conditioned.

Original: However, no explicit assumption of a universal structure is leveraged.
Revised: However, no explicit assumption of a universal structure is leveraged.

Original: Amit & Meir (2018); Yin et al.
Revised: Amit & Meir (2018); Yin et al.

Original: (2020) give a PAC-Bayes bound for meta-learning generalization that relies on the number of tasks n .
Revised: (2020) give a PAC-Bayes bound for meta-learning generalization that relies on the number of tasks n .

Original: Our setting is quite different from the typical assumptions of the meta-learning framework, which stresses that the tasks must be mutually exclusive to ensure a single model cannot solve all tasks.
Revised: Our setting is quite different from the typical assumptions of the meta-learning framework, which stresses that the tasks must be mutually exclusive to ensure a single model cannot solve all tasks.

Original: Instead, we assume a shared latent structure underlying all tasks, and seek to exploit that structure for generalization.
Revised: Instead, we assume a shared latent structure underlying all tasks, and seek to exploit that structure for generalization.

Original: We ﬁnd that under this setting, our method indeed outperforms policies initialized through meta-learning.
Revised: We ﬁnd that under this setting, our method indeed outperforms policies initialized through meta-learning.

Original: The ability to extract meaningful information through state abstractions provides a means to generalize across tasks with a common structure.
Revised: The ability to extract meaningful information through state abstractions provides a means to generalize across tasks with a common structure.

Original: Abel et al.
Revised: Abel et al.

Original: (2018) learn transitive and PAC state abstractions for a distribution over tasks, but they concentrate on ﬁnite, tabular MDPs.
Revised: (2018) learn transitive and PAC state abstractions for a distribution over tasks, but they concentrate on ﬁnite, tabular MDPs.

Original: One approach to form such abstractions is via bisimulation metrics (Givan et al., 2003; Ferns et al., 2004) which formalize a concrete way to group behaviorally equivalent states.
Revised: One approach to form such abstractions is via bisimulation metrics (Givan et al., 2003; Ferns et al., 2004) which formalize a concrete way to group behaviorally equivalent states.

Original: Prior work also leverages bisimulation for transfer (Castro & Precup, 2010), but on the policy level.
Revised: Prior work also leverages bisimulation for transfer (Castro & Precup, 2010), but on the policy level.

Original: Our work instead focuses on learning a latent state representation and established theoretical results for the MTRL setting.
Revised: Our work instead focuses on learning a latent state representation and established theoretical results for the MTRL setting.

Original: Recent work (Gelada et al., 2019) also learns a latent dynamics model and demonstrates connections to bisimulation metrics, but does not address multi-task learning.
Revised: Recent work (Gelada et al., 2019) also learns a latent dynamics model and demonstrates connections to bisimulation metrics, but does not address multi-task learning.

Original:  D ISCUSSION
Revised:  D ISCUSSION

Original: In this work, we advocate for a new framework, HiP-BMDP, to address the multi-task reinforcement learning setting.
Revised: In this work, we advocate for a new framework, HiP-BMDP, to address the multi-task reinforcement learning setting.

Original: Like previous methods, HiP-BMDP assumes a shared state and action space across tasks, but additionally assumes latent structure in the dynamics.
Revised: Like previous methods, HiP-BMDP assumes a shared state and action space across tasks, but additionally assumes latent structure in the dynamics.

Original: We exploit this structure through learning a universal dynamics model with latent parameter θ , which captures the behavioral similarity across tasks.
Revised: We exploit this structure through learning a universal dynamics model with latent parameter θ , which captures the behavioral similarity across tasks.

Original: We provide error and value bounds for the HiP-MDP (in appendix) and
Revised: We provide error and value bounds for the HiP-MDP (in appendix) and

Original: HiP-BMDP settings, showing improvements in sample complexity over prior work by producing a bound that depends on the number of samples in aggregate over tasks, rather than number of tasks seen at training time.
Revised: HiP-BMDP settings, showing improvements in sample complexity over prior work by producing a bound that depends on the number of samples in aggregate over tasks, rather than number of tasks seen at training time.

Original: Our work relies on an assumption that we have access to an environment id, or knowledge of when we have switched environments.
Revised: Our work relies on an assumption that we have access to an environment id, or knowledge of when we have switched environments.

Original: This assumption could be relaxed by incorporating an environment identiﬁcation procedure at training time to cluster incoming data into separate environments.
Revised: This assumption could be relaxed by incorporating an environment identiﬁcation procedure at training time to cluster incoming data into separate environments.

Original: Further, our bounds rely L ∞ norms for measuring error and the value and transfer bounds.
Revised: Further, our bounds rely L ∞ norms for measuring error and the value and transfer bounds.

Original: In future work we will investigate tightening these bounds with L p norms.
Revised: In future work we will investigate tightening these bounds with L p norms.

Original: A B ISIMULATION B OUNDS
Revised: A B ISIMULATION B OUNDS

Original: We ﬁrst look at the Block MDP case only (Zhang et al., 2020a), which can be thought of as the single-task setting in a HiP-BMDP.
Revised: We ﬁrst look at the Block MDP case only (Zhang et al., 2020a), which can be thought of as the single-task setting in a HiP-BMDP.

Original: We can compute approximate error bounds in this setting by denoting φ an ( (cid:15) R , (cid:15) T ) -approximate bisimulation abstraction, where
Revised: We can compute approximate error bounds in this setting by denoting φ an ( (cid:15) R , (cid:15) T ) -approximate bisimulation abstraction, where

Original: [Equation]
Revised: [Equation]

Original: Φ T denotes the lifted version of T , where we take the next-step transition distribution from observation space X and lift it to latent space S .
Revised: Φ T denotes the lifted version of T , where we take the next-step transition distribution from observation space X and lift it to latent space S .

Original: Theorem 5.
Revised: Theorem 5.

Original: Given an MDP ¯ M built on a ( (cid:15) R , (cid:15) T ) -approximate bisimulation abstraction of Block
Revised: Given an MDP ¯ M built on a ( (cid:15) R , (cid:15) T ) -approximate bisimulation abstraction of Block

Original: MDP M , we denote the evaluation of the optimal Q function of ¯ M on M as [ Q ∗ ¯ M ] M .
Revised: MDP M , we denote the evaluation of the optimal Q function of ¯ M on M as [ Q ∗ ¯ M ] M .

Original: The value difference with respect to the optimal Q ∗M is upper bounded by
Revised: The value difference with respect to the optimal Q ∗M is upper bounded by

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: From Theorem 2 in Jiang (2018).
Revised: From Theorem 2 in Jiang (2018).

Original: B T HEORETICAL R ESULTS FOR THE H I P-MDP S ETTING
Revised: B T HEORETICAL R ESULTS FOR THE H I P-MDP S ETTING

Original: We explore the HiP-MDP setting, where a low-dimensional state space is given, to highlight the results that can be obtained just from assuming this hierarchical structure of the dynamics.
Revised: We explore the HiP-MDP setting, where a low-dimensional state space is given, to highlight the results that can be obtained just from assuming this hierarchical structure of the dynamics.

Original: B. V ALUE B OUNDS
Revised: B. V ALUE B OUNDS

Original: Given a family of environments M Θ , we bound the difference in expected value between two sampled
Revised: Given a family of environments M Θ , we bound the difference in expected value between two sampled

Original: MDPs, M θ i , M θ j ∈ M Θ using d ( θ i , θ j ) .
Revised: MDPs, M θ i , M θ j ∈ M Θ using d ( θ i , θ j ) .

Original: Additionally, we make the assumption that we have a behavior policy π that is near both optimal policies π ∗
Revised: Additionally, we make the assumption that we have a behavior policy π that is near both optimal policies π ∗

Original: θ i , π ∗ θ j .
Revised: θ i , π ∗ θ j .

Original: We use KL divergence to deﬁne this neighborhood for π ∗
Revised: We use KL divergence to deﬁne this neighborhood for π ∗

Original: θ i ,
Revised: θ i ,

Original: [Equation]
Revised: [Equation]

Original: We start with a bound for a speciﬁc policy π .
Revised: We start with a bound for a speciﬁc policy π .

Original: One way to measure the difference between two tasks
Revised: One way to measure the difference between two tasks

Original: M θ i , M θ j is to measure the difference in value when that policy is applied in both settings.
Revised: M θ i , M θ j is to measure the difference in value when that policy is applied in both settings.

Original: We show the relationship between the learned θ and this difference in value.
Revised: We show the relationship between the learned θ and this difference in value.

Original: The following results are similar to error bounds in approximate value iteration (Munos, 2005; Bertsekas & Tsitsiklis, 1996), but instead of tracking model error, we apply these methods to compare tasks with differences in dynamics.
Revised: The following results are similar to error bounds in approximate value iteration (Munos, 2005; Bertsekas & Tsitsiklis, 1996), but instead of tracking model error, we apply these methods to compare tasks with differences in dynamics.

Original: Theorem 6.
Revised: Theorem 6.

Original: Given policy π , the difference in expected value between two MDPs drawn from the family of MDPs M θ i , M θ j ∈ M Θ is bounded by
Revised: Given policy π , the difference in expected value between two MDPs drawn from the family of MDPs M θ i , M θ j ∈ M Θ is bounded by

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: We use a telescoping sum to prove this bound, which is similar to Luo et al.
Revised: We use a telescoping sum to prove this bound, which is similar to Luo et al.

Original: First, we let Z k denote the discounted sum of rewards if the ﬁrst k steps are in M θ i , and all steps t > k are in
Revised: First, we let Z k denote the discounted sum of rewards if the ﬁrst k steps are in M θ i , and all steps t > k are in

Original: M θ j ,
Revised: M θ j ,

Original: [Equation]
Revised: [Equation]

Original: By deﬁnition, we have Z ∞ =
Revised: By deﬁnition, we have Z ∞ =

Original: V πθ i and Z 0 =
Revised: V πθ i and Z 0 =

Original: V πθ j .
Revised: V πθ j .

Original: Now, the value function difference can be written as a telescoping sum,
Revised: Now, the value function difference can be written as a telescoping sum,

Original: [Equation]
Revised: [Equation]

Original: Each term can be simpliﬁed to
Revised: Each term can be simpliﬁed to

Original: [Equation]
Revised: [Equation]

Original: Plugging this back into Equation (7),
Revised: Plugging this back into Equation (7),

Original: [Equation]
Revised: [Equation]

Original: This expected value difference is bounded by the Wasserstein distance between T θ i
Revised: This expected value difference is bounded by the Wasserstein distance between T θ i

Original: [Equation]
Revised: [Equation]

Original: T θ
Revised: T θ

Original: [Equation]
Revised: [Equation]

Original: using Equation (1) .
Revised: using Equation (1) .

Original: Another comparison to make is how different the optimal policies in different tasks are with respect to the distance (cid:107) θ i − θ j (cid:107) .
Revised: Another comparison to make is how different the optimal policies in different tasks are with respect to the distance (cid:107) θ i − θ j (cid:107) .

Original: Theorem 7.
Revised: Theorem 7.

Original: The difference in expected optimal value between two MDPs M θ i , M θ j ∈ M Θ is bounded by,
Revised: The difference in expected optimal value between two MDPs M θ i , M θ j ∈ M Θ is bounded by,

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: [Equation]
Revised: [Equation]

Original: We can bound the RHS with
Revised: We can bound the RHS with

Original: [Equation]
Revised: [Equation]

Original: All MDPs in M Θ have the same reward function, so the ﬁrst term is
Revised: All MDPs in M Θ have the same reward function, so the ﬁrst term is

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: Plugging this back in,
Revised: Plugging this back in,

Original: [Equation]
Revised: [Equation]

Original: Both these results lend more intuition for casting the multi-task setting under the HiP-MDP formalism.
Revised: Both these results lend more intuition for casting the multi-task setting under the HiP-MDP formalism.

Original: The difference in the optimal performance between any two environments is controlled by the distance between the hidden parameters for corresponding environments.
Revised: The difference in the optimal performance between any two environments is controlled by the distance between the hidden parameters for corresponding environments.

Original: One can interpret the hidden parameter as a knob to allow precise changes across the tasks.
Revised: One can interpret the hidden parameter as a knob to allow precise changes across the tasks.

Original: B. E XPECTED E RROR B OUNDS
Revised: B. E XPECTED E RROR B OUNDS

Original: In MTRL, we are concerned with the performance over a family of tasks.
Revised: In MTRL, we are concerned with the performance over a family of tasks.

Original: The empirical risk is typically deﬁned as follows for T tasks (Maurer et al., 2016):
Revised: The empirical risk is typically deﬁned as follows for T tasks (Maurer et al., 2016):

Original: [Equation]
Revised: [Equation]

Original: Consequently, we bound the expected loss over the family of environments E with respect to θ .
Revised: Consequently, we bound the expected loss over the family of environments E with respect to θ .

Original: In particular, we are interested in the average approximation error and deﬁne it as the absolute model error averaged across all environments:
Revised: In particular, we are interested in the average approximation error and deﬁne it as the absolute model error averaged across all environments:

Original: [Equation]
Revised: [Equation]

Original: Theorem 8.
Revised: Theorem 8.

Original: Given a family of environments M Θ , each parameterized with an underlying true hidden parameter θ 1 , θ 2 , · · · , θ E , and let ˆ θ 1 , ˆ θ 2 , · · · , ˆ θ E be their respective approximations such that the average approximation error across all environments is bounded as follows:
Revised: Given a family of environments M Θ , each parameterized with an underlying true hidden parameter θ 1 , θ 2 , · · · , θ E , and let ˆ θ 1 , ˆ θ 2 , · · · , ˆ θ E be their respective approximations such that the average approximation error across all environments is bounded as follows:

Original: [Equation]
Revised: [Equation]

Original: where each environment’s parameter
Revised: where each environment’s parameter

Original: θ
Revised: θ

Original: i is (cid:15) -close to its approximation ˆ θ i i.e. d ( ˆ θ i , θ i ) ≤ (cid:15) , where d is the distance metric deﬁned in Eq.
Revised: i is (cid:15) -close to its approximation ˆ θ i i.e. d ( ˆ θ i , θ i ) ≤ (cid:15) , where d is the distance metric deﬁned in Eq.

Original: Proof.
Revised: Proof.

Original: We here consider the approximation error averaged across all environments as follows:
Revised: We here consider the approximation error averaged across all environments as follows:

Original: [Equation]
Revised: [Equation]

Original: Let us consider an environment θ i ∈ M E for which we can bound the RHS with
Revised: Let us consider an environment θ i ∈ M E for which we can bound the RHS with

Original: [Equation]
Revised: [Equation]

Original: Considering the family of environments M E have the same reward function and is known
Revised: Considering the family of environments M E have the same reward function and is known

Original: [Equation]
Revised: [Equation]

Original: resulting in ﬁrst term to be
Revised: resulting in ﬁrst term to be

Original: [Equation]
Revised: [Equation]

Original: ·| s,a ) , s (cid:48) ∼ T ˆ θi ( ·| s,a ) (cid:2) V ∗ ˆ θ
Revised: ·| s,a ) , s (cid:48) ∼ T ˆ θi ( ·| s,a ) (cid:2) V ∗ ˆ θ

Original: i ( s (cid:48) )
Revised: i ( s (cid:48) )

Original: − V ∗
Revised: − V ∗

Original: θ i ( s (cid:48)(cid:48) ) (cid:3) (cid:12) (cid:12)(cid:12)(cid:12) ≤ γ sup s,a (cid:12)(cid:12) E s (cid:48) ∼ T ˆ θi ( ·| s,a ) (cid:2) V ∗ ˆ θ i ( s (cid:48) )
Revised: θ i ( s (cid:48)(cid:48) ) (cid:3) (cid:12) (cid:12)(cid:12)(cid:12) ≤ γ sup s,a (cid:12)(cid:12) E s (cid:48) ∼ T ˆ θi ( ·| s,a ) (cid:2) V ∗ ˆ θ i ( s (cid:48) )

Original: − V ∗
Revised: − V ∗

Original: θ i ( s (cid:48) )
Revised: θ i ( s (cid:48) )

Original: (cid:3)(cid:12)(cid:12)
Revised: (cid:3)(cid:12)(cid:12)

Original: + γ sup s,a (cid:12)(cid:12)
Revised: + γ sup s,a (cid:12)(cid:12)

Original: E s (cid:48)(cid:48) ∼ T θi ( ·| s,a ) , s (cid:48) ∼ T ˆ θi ( ·| s,a ) (cid:2)
Revised: E s (cid:48)(cid:48) ∼ T θi ( ·| s,a ) , s (cid:48) ∼ T ˆ θi ( ·| s,a ) (cid:2)

Original: V ∗ ˆ θ i ( s (cid:48) )
Revised: V ∗ ˆ θ i ( s (cid:48) )

Original: − V ∗
Revised: − V ∗

Original: θ i ( s (cid:48)(cid:48) )
Revised: θ i ( s (cid:48)(cid:48) )

Original: (cid:3)(cid:12)(cid:12) ≤ γ sup s,a (cid:12)(cid:12)
Revised: (cid:3)(cid:12)(cid:12) ≤ γ sup s,a (cid:12)(cid:12)

Original: E s (cid:48) ∼ T ˆ θi ( ·| s,a ) (cid:2) V ∗ ˆ
Revised: E s (cid:48) ∼ T ˆ θi ( ·| s,a ) (cid:2) V ∗ ˆ

Original: θ i ( s (cid:48) )
Revised: θ i ( s (cid:48) )

Original: − V ∗
Revised: − V ∗

Original: θ i ( s (cid:48) )
Revised: θ i ( s (cid:48) )

Original: (cid:3)(cid:12)(cid:12) + γ 1 − γ | ˆ θ i − θ i | ≤ γ max s (cid:12)(cid:12)
Revised: (cid:3)(cid:12)(cid:12) + γ 1 − γ | ˆ θ i − θ i | ≤ γ max s (cid:12)(cid:12)

Original: V ∗ ˆ θ i ( s ) − V ∗
Revised: V ∗ ˆ θ i ( s ) − V ∗

Original: θ i ( s ) (cid:12)(cid:12)
Revised: θ i ( s ) (cid:12)(cid:12)

Original: + γ 1 − γ | ˆ θ
Revised: + γ 1 − γ | ˆ θ

Original: i − θ i | =
Revised: i − θ i | =

Original: γ max s (cid:12)(cid:12) max a Q ∗ ˆ θ i ( s, a ) − max a (cid:48)
Revised: γ max s (cid:12)(cid:12) max a Q ∗ ˆ θ i ( s, a ) − max a (cid:48)

Original: Q ∗ θ i ( s, a (cid:48) )
Revised: Q ∗ θ i ( s, a (cid:48) )

Original: (cid:12)(cid:12) + γ 1 − γ |
Revised: (cid:12)(cid:12) + γ 1 − γ |

Original: ˆ
Revised: ˆ

Original: θ
Revised: θ

Original: i −
Revised: i −

Original: [Equation]
Revised: [Equation]

Original: Solving for sup s,a (cid:12)(cid:12) Q ∗ ˆ θ
Revised: Solving for sup s,a (cid:12)(cid:12) Q ∗ ˆ θ

Original: i
Revised: i

Original: [Equation]
Revised: [Equation]

Original: − Q ∗ θ
Revised: − Q ∗ θ

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

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

Original: Plugging Eq.
Revised: Plugging Eq.

Original: 13 back in Eq.
Revised: 13 back in Eq.

Original: [Equation]
Revised: [Equation]

Original: We now consider that the distance between the approximated ˆ θ i and the underlying hidden parameter θ i
Revised: We now consider that the distance between the approximated ˆ θ i and the underlying hidden parameter θ i

Original: ∈
Revised: ∈

Original: M E is deﬁned as in Eq.
Revised: M E is deﬁned as in Eq.

Original: 1, such that: d ( ˆ θ i , θ i ) ≤ (cid:15) θ
Revised: 1, such that: d ( ˆ θ i , θ i ) ≤ (cid:15) θ

Original: Plugging this back concludes the proof,
Revised: Plugging this back concludes the proof,

Original: [Equation]
Revised: [Equation]

Original: It is interesting to note that the average approximation error across all environments is independent of the number of environments and primarily governed by the error in approximating the hidden parameter θ for each environment.
Revised: It is interesting to note that the average approximation error across all environments is independent of the number of environments and primarily governed by the error in approximating the hidden parameter θ for each environment.

Original: C A DDITIONAL R ESULTS AND P ROOFS FOR H I P-BMDP R ESULTS
Revised: C A DDITIONAL R ESULTS AND P ROOFS FOR H I P-BMDP R ESULTS

Original: We ﬁrst compute L ∞ norm bounds for Q error under approximate abstractions and transfer bounds.
Revised: We ﬁrst compute L ∞ norm bounds for Q error under approximate abstractions and transfer bounds.

Original: Theorem 9 ( Q error with L ∞ norm) .
Revised: Theorem 9 ( Q error) .

Original: Given an MDP ¯ M ˆ θ built on a ( (cid:15) R , (cid:15) T , (cid:15) θ ) -approximate bisimulation abstraction of an instance of a HiP-BMDP M θ , we denote the evaluation of the optimal Q function of ¯ M ˆ θ on M as [ Q ∗ ¯ M ˆ θ ] M θ .
Revised: Given an MDP ¯ M ˆ θ built on a ( (cid:15) R , (cid:15) T , (cid:15) θ ) -approximate bisimulation abstraction of an instance of a HiP-BMDP M θ , we denote the evaluation of the optimal Q function of ¯ M ˆ θ on M as [ Q ∗ ¯ M ˆ θ ] M θ .

Original: The value difference with respect to the optimal Q ∗M is upper bounded by
Revised: The value difference with respect to the optimal Q ∗M is upper bounded by

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: In the HiP-BMDP setting, we have a global encoder φ over all tasks, but the difference in transition distribution also includes θ .
Revised: In the HiP-BMDP setting, we have a global encoder φ over all tasks, but the difference in transition distribution also includes θ .

Original: The reward functions are the same across tasks, so there is no change to (cid:15) R .
Revised: The reward functions are the same across tasks, so there is no change to (cid:15) R .

Original: However, we now must incorporate difference in dynamics in (cid:15) T .
Revised: However, we now must incorporate difference in dynamics in (cid:15) T .

Original: Assuming we have two environments with hidden parameters θ i , θ j ∈ Θ , we can compute (cid:15) θ i ,θ j T across those two environments by joining them into a super-MDP:
Revised: Assuming we have two environments with hidden parameters θ i , θ j ∈ Θ , we can compute (cid:15) θ i ,θ j T across those two environments by joining them into a super-MDP:

Original: [Equation]
Revised: [Equation]

Original: i 
Revised: i ,

Original: [Equation]
Revised: [Equation]

Original: This result is intuitive in that with a shared encoder learning a per-task bisimulation relation, the distance between bisimilar states from another task depends on the change in transition distribution between those two tasks.
Revised: This result is intuitive in that with a shared encoder learning a per-task bisimulation relation, the distance between bisimilar states from another task depends on the change in transition distribution between those two tasks.

Original: We can now extend the single-task bisimulation bound (Theorem 5) to the
Revised: We can now extend the single-task bisimulation bound (Theorem 5) to the

Original: HiP-BMDP setting by denoting approximation error of θ as (cid:107) θ − ˆ θ (cid:107) 1 < (cid:15) θ .
Revised: HiP-BMDP setting by denoting approximation error of θ as (cid:107) θ − ˆ θ (cid:107) 1 < (cid:15) θ .

Original: Theorem 10 (Transfer bound with L ∞ norm) .

Original: Given two MDPs M θ i and M θ j , we can bound the difference in Q π between the two MDPs for a given policy π learned under an (cid:15) R , (cid:15) T , (cid:15) θ i -approximate abstraction of M θ i and applied to

Original: [Equation]

Original: This result clearly follows directly from Theorem 2.

Original: Given a policy learned for task i , Theoremgives a bound on how far from optimal that policy is when applied to task j .

Original: Intuitively, the more similar in behavior tasks i and j are, as denoted by (cid:107) θ i − θ j (cid:107) 1 , the better π performs on task j .

Original: Theorem 4.
Revised: Theorem 4.

Original: For any φ which deﬁnes an ( (cid:15) R , (cid:15) T , (cid:15) θ ) -approximate bisimulation abstraction on a
Revised: For any φ which deﬁnes an ( (cid:15) R , (cid:15) T , (cid:15) θ ) -approximate bisimulation abstraction on a

Original: HiP-BMDP family M Θ , we deﬁne the empirical measurement of Q ∗ ¯ M ˆ θ over D to be Q ∗ ¯ M D ˆ θ .
Revised: HiP-BMDP family M Θ , we deﬁne the empirical measurement of Q ∗ ¯ M ˆ θ over D to be Q ∗ ¯ M D ˆ θ .

Original: Then, with probability ≥ 1 − δ ,
Revised: Then, with probability ≥ 1 − δ ,

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: [Equation]
Revised: [Equation]

Original: The ﬁrst term is solved by Theorem 2, so we only need to solve the second term using McDiarmid’s inequality and the knowledge that the value function of a bisimulation representation is 11 − γ -Lipschitz from Theorem 1. 
Revised: The ﬁrst term is solved by Theorem 2, so we only need to solve the second term using McDiarmid’s inequality and the knowledge that the value function of a bisimulation representation is 11 − γ -Lipschitz from Theorem 1. First [Equation]

Original: First, we write this difference to be a deviation from an expectation in order to apply the concentration inequality.
Revised: we write this difference to be a deviation from an expectation in order to apply the concentration inequality.

Original: [Equation]
Revised: [Equation]

Original: Now we can apply McDiarmid’s inequality,
Revised: Now we can apply McDiarmid’s inequality,

Original: [Equation]
Revised: [Equation]

Original: Solve for the t that makes this inequality hold for all ( φ ( x ) , a ) ∈
Revised: Solve for the t that makes this inequality hold for all ( φ ( x ) , a ) ∈

Original: X × A with a union bound over all | φ ( X ) ||A| abstract states,
Revised: X × A with a union bound over all | φ ( X ) ||A| abstract states,

Original: [Equation]
Revised: [Equation]

Original: Combine to get
Revised: Combine to get

Original: [Equation]
Revised: [Equation]

Original: (1 − γ )(cid:115) 2 n φ ( D ) log 2 | φ ( X ) ||A| δ .
Revised: [Equation]

Original: D A DDITIONAL I MPLEMENTATION D ETAILS
Revised: D A DDITIONAL I MPLEMENTATION D ETAILS

Original: In Figure 8, we show the variation in the left foot of the walker.
Revised: In Figure 8, we show the variation in the left foot of the walker.

Original: MTRL Algorithm
Revised: MTRL Algorithm

Original: The Multitask RL algorithm for the HiP-BMDP setting can be found in Algorithm 1.
Revised: The Multitask RL algorithm for the HiP-BMDP setting can be found in Algorithm 1.

Original: We take the DeepMDP baseline Gelada et al.
Revised: We take the DeepMDP baseline Gelada et al

Original:  and incorporate our HiP-BMDP objective (text shown in red color)
Revised: ) and incorporate our HiP-BMDP objective (text shown in red color)

Revised: Meta-RL Algorithm

Revised: The meta-RL algorithm for the HiP-MDP setting can be found in Algorithm 3.

Revised: We take the PEARL algorithm (Rakelly et al., 2019) and incorporate our HiP-MDP objective (text shown in red color)

Revised: D. B ASELINES

Original: Require:
Revised: For PCGrad, the authors recommend projecting the gradient with respect to all previous tasks. In practice, that leads to very poor training. Instead, we observe that it is better to project the gradients with respect to any one task (randomly selected per update). We use this scheme in all the experiments.

Revised: For GradNorm, we observe that the learned weights w i (for weighing per-task loss) can become

Original:  Along with DeepMDP components (Actor, Critic, Dynamics Model ( M )), an additional environment encoder ψ to generated task-speciﬁc θ parameters.
Revised: Algorithm Require: Along with DeepMDP components (Actor, Critic, Dynamics Model ( M )), an additional environment encoder ψ to generated task-speciﬁc θ parameters.

Original: 1: for each timestep t = 1 ..T do 2: for each T i do 3: a it ∼ π i ( ·| s it ) 4: s (cid:48) it ∼ p
Revised: 1: for each timestep t = 1 ..T do 2: for each T i do 3: a it ∼ π i ( ·| s it ) 4: s (cid:48) it ∼ p

Original: i ( ·| s it , a it ) 5:
Revised: i ( ·| s it , a it ) 5:

Original: D ← D ∪ ( s it , a it , r
Revised: D ← D ∪ ( s it , a it , r

Original: ( s it , a it ) , s (cid:48) it ) , i 6:
Revised: ( s it , a it ) , s (cid:48) it ) , i 6:

Original: U PDATE C RITIC ( D , i ) (uses data only from i th task) 7:
Revised: U PDATE C RITIC ( D , i ) (uses data only from i th task) 7:

Original: U PDATE A CTOR
Revised: U PDATE A CTOR

Original: ( D , i ) (uses data only from i th task) 8:
Revised: ( D , i ) (uses data only from i th task) 8:

Original: U PDATE U SING H IP -BMDPL OSS ( D , i ) 9: end for 10: end for 
Revised: U PDATE U SING H IP -BMDPL OSS ( D , i ) 9: end for 10: end for Algorithm

Original: Require:
Revised: Require:

Original: Batches of data for the different tasks {T i } i =1 ...
Revised: Batches of data for the different tasks {T i } i =1 ...

Original: T sampled from the Replay Buffer D , learning rates α 

Revised: T sampled from the Replay Buffer D , learning rates α

Original: and α 2 , index of the current task i , Transition Model M ,
Revised: and α 2 , index of the current task i , Transition Model M ,

Original: environment encoder ψ .
Revised: environment encoder ψ .

Original: 1: for each batch of dataset t = 1 ..T , t (cid:54) =
Revised: 1: for each batch of dataset t = 1 ..T , t (cid:54) =

Original: i do 2:
Revised: i do 2:

Original: Compute L ( ψ, M ) =
Revised: Compute L ( ψ, M ) =

Original: L i ( ψ, M, i, t ) using Equation (3) 3: ψ ← ψ − α 1 ∇ θ (cid:80) i
Revised: L i ( ψ, M, i, t ) using Equation (3) 3: ψ ← ψ − α 1 ∇ θ (cid:80) i

Original: L 4:
Revised: L 4:

Original: M ← M − α 2 ∇ θ (cid:80) i
Revised: M ← M − α 2 ∇ θ (cid:80) i

Original: L 5: end for Meta-RL Algorithm The meta-RL algorithm for  the HiP-MDP setting can be found in Algorithm 3. We take the PEARL algorithm (Rakelly et al., 2019) and incorporate our HiP-MDP objective (text shown in red color)
Revised: L 5: end for negative for some tasks (which means the model tries to unlearn those tasks).

Original: D. B ASELINES

Original: For PCGrad, the authors recommend projecting the gradient with respect to all previous tasks.

Original: In practice, that leads to very poor training.

Original: Instead, we observe that it is better to project the gradients with respect to any one task (randomly selected per update).

Original: We use this scheme in all the experiments.

Original: For GradNorm, we observe that the learned weights w i (for weighing per-task loss) can become negative for some tasks (which means the model tries to unlearn those tasks).

Original: In practice, we clamp the w i values to not become smaller than a threshold.
Revised: In practice, we clamp the w i values to not become smaller than a threshold.

Original: D. H YPER P ARAMETERS D.2. MTRL A LGORITHM
Revised: D. H YPER P ARAMETERS D.2. MTRL A LGORITHM

Original: All the hyper parameters (for MTRL algorithm) are listed in Table 1.
Revised: All the hyper parameters (for MTRL algorithm) are listed in Table 1.

Original: D.2. M ETA RL A LGORITHM
Revised: D

Revised: [Equation]

Original: For MetaRL, we use the same hyperparameters as used by PEARL Rakelly et al.
Revised: For MetaRL, we use the same hyperparameters as used by PEARL Rakelly et al.

Original: We set α φ = 0 .
Revised: We set α φ = 0 .

Original: 01 for all environments, other than Walker-Stand environments where α φ = 0 .
Revised: 01 for all environments, other than Walker-Stand environments where α φ = 0 .

Original: E A DDITIONAL R ESULTS
Revised: E A DDITIONAL R ESULTS

Original: Along with the environments described in 4, we considered the following additional environments:
Revised: Along with the environments described in 4, we considered the following additional environments:

Original: Walker-Stand-V0 : Walker-Stand task where the friction coefﬁcient, between the ground and the walker’s leg, varies across different environments.
Revised: Walker-Stand-V0 : Walker-Stand task where the friction coefﬁcient, between the ground and the walker’s leg, varies across different environments.

Original: Walker-Walk-V0 : Walker-Walk task where the friction coefﬁcient, between the ground and the walker’s leg, varies across different environments.
Revised: Walker-Walk-V0 : Walker-Walk task where the friction coefﬁcient, between the ground and the walker’s leg, varies across different environments.

Original: [Table]
Revised: Walker-Stand-V1 :

Original:  the foot of  varies 
Revised: Walker-Stand task where the size of left-foot of the walker varies across different environments.

Original: Walker-Walk-V1 :
Revised: Walker-Walk-V1 :

Original: Walker-Walk task where the size of left-foot of the walker varies across different environments.
Revised: Walker-Walk task where the size of left-foot of the walker varies across different environments.

Original: E. M ULTI -T ASK S ETTING
Revised: E. M ULTI -T ASK S ETTING

Original: In Figure 
Revised: In Figure 10, we observe that the HiP-BMDP method consistently outperforms other baselines when evaluated on the the interpolation environments (zero-shot transfer). As noted previously, the

Original: [Table]
Revised: [Table]

Original: we observe that the HiP-BMDP method consistently outperforms other baselines when evaluated on the the interpolation environments (zero-shot transfer).

Original: As noted previously, the effectiveness of our proposed model can not be attributed to task-embeddings alone as HiPBMDP-nobisim model uses the same architecture as the HiP-BMDP model but does not include the  bisimulation  loss.
Revised: effectiveness of our proposed model can not be attributed to task-embeddings alone as HiP-BMDPnobisim model uses the same architecture as the HiP-BMDP model but does not include the task bisimulation metric loss.

Original: We hypothesise that the Distral-Ensemble baseline behaves poorly because it cannot leverage a shared global dynamics model. 
Revised: We hypothesise that the Distral-Ensemble baseline behaves poorly because it cannot leverage a shared global dynamics model. [Figure]

Original: E. M ETA -RL S ETTING
Revised: E. M ETA -RL S ETTING

Original: We provide the Meta-RL results for the additional environments.
Revised: We provide the Meta-RL results for the additional environments.

Original: Recall that we extend the PEARL algorithm (Rakelly et al., 2019) by training the inference network q φ ( z | c ) with our additional HiPBMDP loss.
Revised: Recall that we extend the PEARL algorithm (Rakelly et al., 2019) by training the inference network q φ ( z | c ) with our additional HiP-

Original: The algorithm pseudocode can be found in Appendix D. In Figure 13, we show the results on the interpolation setup and in Figure 14, we show the results on the extrapolation setup.

Original: Require:
Revised: Require:

Original: Batch of training tasks {T i } i =1 ...T from p ( T ) , learning rates α 1 , α 2 , α 3 , α φ 1: Initialize replay buffers B i for each training task 2: while not done do 3: for each T i do 4:
Revised: Batch of training tasks {T i } i =1 ...T from p ( T ) , learning rates α 1 , α 2 , α 3 , α φ 1: Initialize replay buffers B i for each training task 2: while not done do 3: for each T i do 4:

Original: Initialize context C i = {} 5: for k = 1 , . .
Revised: Initialize context C i = {} 5: for k = 1 , . .

Original: , K do 6:
Revised: , K do 6:

Original: Sample z ∼ q φ ( z | C i ) 7:
Revised: Sample z ∼ q φ ( z | C i ) 7:

Original: Gather data from π θ ( a | s , z ) and add to B i 8:
Revised: Gather data from π θ ( a | s , z ) and add to B i 8:

Original: Update C i
Revised: Update C i

Original: = { ( s j , a j , s (cid:48)
Revised: = { ( s j , a j , s (cid:48)

Original: j , r j ) } j :1 ...
Revised: j , r j ) } j :1 ...

Original: N ∼ B
Revised: N ∼ B

Original: i 9: end for 10: end for 11: for step in training steps do 12: for each T i do 13:
Revised: i 9: end for 10: end for 11: for step in training steps do 12: for each T i do 13:

Original: Sample context C i ∼ S c ( B i ) and RL batch b i ∼ B i 14:
Revised: Sample context C i ∼ S c ( B i ) and RL batch b i ∼ B i 14:

Original: Sample z ∼ q φ ( z | C i ) 15:
Revised: Sample z ∼ q φ ( z | C i ) 15:

Original: L iactor =
Revised: L iactor =

Original: L actor ( b i , z ) 16:
Revised: L actor ( b i , z ) 16:

Original: L i critic =
Revised: L i critic =

Original: L critic ( b i , z ) 17:
Revised: L critic ( b i , z ) 17:

Original: L iKL = βD KL ( q ( z | C i )
Revised: L iKL = βD KL ( q ( z | C i )

Original: ||
Revised: ||

Original: r ( z )) 18:
Revised: r ( z )) 18:

Original: Sample a RL batch b j from any other task j 19:
Revised: Sample a RL batch b j from any other task j 19:

Original: Compute L iBiSim =
Revised: Compute L iBiSim =

Original: L i ( q, T, i, j ) using the equation20: end for 21:
Revised: L i ( q, T, i, j ) using the equation20: end for 21:

Original: φ ← φ −

Original: [Equation]
Revised: [Equation]

Original: (cid:0)

Original: [Equation] icritic [Equation] ×
Revised: ←

Original: [Equation]
Revised: [Equation]

Original: iBiSim (cid:1) 22: θ
Revised: (cid:1) 22: θ

Original: [Equation]
Revised: [Equation]

Original: ←θ
Revised: ←

Original: [Equation]
Revised: [Equation]

Original: − α

Original: [Equation]

Original: ∇
Revised: ∇

Original: [Equation]
Revised: [Equation]

Original: (cid:80)i
Revised: (cid:80)

Revised: [Equation]

Original: L iactor 23: θ Q ← θ Q − α 3 ∇ θ (cid:80)i
Revised: iactor 23: θ Q ← θ Q − α [Equation] ∇ θ (cid:80)

Revised: [Equation]

Original: L icritic 24: end for 25: end while
Revised: L icritic 24: end for 25: end while

Original: [Figure]
Revised: [Figure]

Original: i s o d i c
Revised: BMDP loss.

Original: R e [Figure] a r d [Figure]
Revised: The algorithm pseudocode can be found in Appendix D. In Figure 13, we show the results on the interpolation setup and in Figure 14, we show the results on the extrapolation setup.

Original: In some environments (eg Walker-Walk-V1), the proposed approach (blue) converges faster to a threshold reward (green) than the baseline.
Revised: In some environments (eg Walker-Walk-V1), the proposed approach (blue) converges faster to a threshold reward (green) than the baseline.

Original: In the other environments, the gains are quite small.
Revised: In the other environments, the gains are quite small.

Original: E. E VALUATING THE U NIVERSAL T RANSITION M ODEL .
Revised: E. E VALUATING THE U NIVERSAL T RANSITION M ODEL .

Original: We investigate how well the transition model performs in an unseen environment by only adapting the task parameter θ .
Revised: We investigate how well the transition model performs in an unseen environment by only adapting the task parameter θ .

Original: We instantiate a new MDP, sampled from the family of MDPs, and use a behavior policy to collect transitions.
Revised: We instantiate a new MDP, sampled from the family of MDPs [Figure] and use a behavior policy to collect transitions.

Original: These transitions are used to update only the θ parameter, and the transition model is evaluated by unrolling the transition model for k -steps.
Revised: These transitions are used to update only the θ parameter [Figure] and the transition model is evaluated by unrolling the transition model for k -steps.

Original: In Figures 11 and 12, we report the average, per-step model error in latent space, averaged over 10 environments overand 100 steps respectively.
Revised: In Figures 11 and 12, we report the average, per-step model error in latent space, averaged over 10 environments overand 100 steps respectively.

Original: While we expect both the proposed setup and baseline setups to adapt to the new environment, we expect the proposed setup to adapt faster because of the exploitation of underlying structure.
Revised: While we expect both the proposed setup and baseline setups to adapt to the new environment, we expect the proposed setup to adapt faster because of the exploitation of underlying structure.

Original: We indeed observe that for both 5 step and 100 step unrolls, the proposed
Revised: We indeed observe that for both 5 step and 100 step unrolls, the proposed

Original: HiP-BMDP model adapts much faster than the baseline HiP-BMDP-nobisim (Figures 11 and 12)
Revised: HiP-BMDP model adapts much faster than the baseline HiP-BMDP-nobisim (Figures 11 and 12)

Original: [Figure] [Figure][Figure]
Revised: [Figure] [Figure]

Original: are quite small.