HIDDEN-PARAMETER BLOCK MDPS
Year: 2021

Abstract:
Many control tasks exhibit similar dynamics that can be modeled as having common latent structure. Hidden-Parameter Markov Decision Processes (HiP-MDPs) explicitly model this structure to improve sample efficiency in multi-task settings. However, this setting makes strong assumptions on the observability of the state that limit its application in real-world scenarios with rich observation spaces. 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. We derive instantiations of this new framework for both multi-task reinforcement learning (MTRL) and meta-reinforcement learning (Meta-RL) settings. 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 significant improvement over prior work that use the same environment assumptions. To further demonstrate the efficacy of the proposed method, we empirically compare and show improvement over multi-task and meta-reinforcement learning baselines.

text: 1 INTRODUCTION
A key open challenge in AI research that remains is how to train agents that can learn behaviors that generalize across tasks and environments. 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 definite advantages (in terms of robustness and sample efficiency) over the singletask setting, where the agent independently learns each task. 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. Prior work in MTRL has also
leveraged the idea by sharing representations across tasks (D’Eramo et al., 2020) or providing pertask sample complexity results that show improved sample efficiency from transfer (Brunskill & Li, 2013). However, explicit exploitation of the shared structure across tasks via a unified dynamics has been lacking. Prior works that make use of shared representations use a naive unification approach that posits all tasks lie in a shared domain (Figure 1, left). 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).
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. We frame the problem as a structured super-MDP with a shared state space and universal dynamics model conditioned on a task-specific hidden parameter (Figure 1, right). This additional structure gives us better sample efficiency, both
∗Corresponding author: amy.x.zhang@mail.mcgill.ca
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). We learn a latent representation with smoothness properties for better few-shot generalization to other unseen tasks within this family. 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.
We focus on multi-task settings where dynamics can vary across tasks, but the reward function is shared. We show that this setting can be formalized as a hidden-parameter MDP (HiP-MDP) (DoshiVelez & Konidaris, 2013), where the changes in dynamics can be defined by a latent variable, unifying dynamics across tasks as a single global function. This setting assumes a global latent structure over all tasks (or MDPs). 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. Another example is warehouse robots, where the same tasks are performed in different conditions and warehouse layouts. 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. With this assumed structure, we can provide concrete zero-shot generalization bounds to unseen tasks within this family. 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. This is, again, a highly realistic setting in robotics when we do not always have an amenable, Lipschitz low-dimensional state space. Cameras are a convenient and inexpensive way to acquire state information, and handling pixel observations is key to approaching these problems. A block MDP (Du et al., 2019) provides a concrete way to formalize this observation-based setting. Leveraging this property of the block MDP framework, in combination with the assumption of a unified 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.
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. Further, this framework allows us to compute theoretical generalization results with the incorporation of a learned state representation. Finally, empirical results show that our method outperforms other multi-task and meta-learning baselines in both fast adaptation and zero-shot transfer settings.

heading: 2 BACKGROUND
text: 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.
A finite1, discrete-time Markov Decision Process (MDP) (Bellman, 1957; Puterman, 1995) is a tuple 〈S,A, R, T, γ〉, where S is the set of states, A is the set of actions, R : S × A → R is the reward function, T : S × A → Dist(S) is the environment transition probability function, and γ ∈ [0, 1) is the discount factor. At each time step, the learning agent perceives a state st ∈ S , takes an action at ∈ A drawn from a policy π : S ×A → [0, 1], and with probability T (st+1|st, at) enters next state st+1, receiving a numerical reward Rt+1 from the environment. The value function of policy π is defined as: Vπ(s) = Eπ[ ∑∞ t=0 γ
tRt+1|S0 = s]. The optimal value function V ∗ is the maximum value function over the class of stationary policies.
Hidden-Parameter MDPs (HiP-MDPs) (Doshi-Velez & Konidaris, 2013) can be defined by a tuple M: 〈S,A,Θ, Tθ, R, γ, PΘ〉 where S is a finite state space, A a finite action space, Tθ describes the transition distribution for a specific task described by task parameter θ ∼ PΘ, R is the reward function, γ is the discount factor, and PΘ the distribution over task parameters. This defines a family of MDPs, where each MDP is described by the parameter θ ∼ PΘ. We assume that this parameter θ is fixed for an episode and indicated by an environment id given at the start of the episode.
Block MDPs (Du et al., 2019) are described by a tuple 〈S,A,X , p, q, R〉 with an unobservable state space S , action spaceA, and observable space X . p denotes the latent transition distribution p(s′|s, a) for s, s′ ∈ 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. We are interested in the setting where this
1We use this assumption only for theoretical results, but our method can be applied to continuous domains.
mapping q is one-to-many. 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.
Assumption 1 (Block structure (Du et al., 2019)). Each observation x uniquely determines its generating state s. That is, the observation space X can be partitioned into disjoint blocks Xs, each containing the support of the conditional distribution q(·|s).
Assumption 1 gives the Markov property inX , a key difference from partially observable MDPs (Kaelbling et al., 1998; Zhang et al., 2019), which has no guarantee of determining the generating state from the history of observations. This assumption allows us to compute reasonable bounds for our algorithm in k-order MDPs2 (which describes many real world problems) and avoiding the intractability of true POMDPs, which have no guarantees on providing enough information to sufficiently predict future rewards. A relaxation of this assumption entails providing less information in the observation for predicting future reward, which will degrade performance. We show empirically that our method is still more robust to a relaxation of this assumption compared to other MTRL methods.
Bisimulation is a strict form of state abstraction, where two states are bisimilar if they are behaviorally equivalent. Bisimulation metrics (Ferns et al., 2011) define a distance between states as follows: Definition 1 (Bisimulation Metric (Theorem 2.6 in Ferns et al. (2011))). Let (S,A, P, r) be a finite MDP and met the space of bounded pseudometrics on S equipped with the metric induced by the uniform norm. Define F : met 7→ met by
F (d)(s, s′) = max a∈A (|ras − ras′ |+ γW (d)(P as , P as′)),
where W (d) is the Wasserstein distance between transition probability distributions. Then F has a unique fixed point d̃ which is the bisimulation metric.
A nice property of this metric d̃ is that difference in optimal value between two states is bounded by their distance as defined by this metric.
Theorem 1 (V ∗ is Lipschitz with respect to d̃ (Ferns et al., 2004)). Let V ∗ be the optimal value function for a given discount factor γ. Then V ∗ is Lipschitz continuous with respect to d̃ with Lipschitz constant 11−γ ,
|V ∗(s)− V ∗(s′)| ≤ 1 1− γ d̃(s, s′).
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 fits a parametric soft Q-function Q(st, at) using transitions sampled from the replay buffer D by minimizing the soft Bellman residual,
J(Q) = E(st,st,rt,st+1)∼D [( Q(st, at)− rt − γV̄ (xt+1) )2] .
The target value function V̄ is approximated via a Monte-Carlo estimate of the following expectation, V̄ (xt+1) = Eat+1∼π [ Q̄(xt+1, at+1)− α log π(at+1|st+1) ] ,
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 π(at|st) by minimizing KL divergence between the policy and a Boltzmann distribution induced by the Q-function, producing the following objective,
J(π) = Est∼D [ Eat∼π[α log(π(at|st))−Q(st, at)] ] .
2Any k-order MDP can be made Markov by stacking the previous k observations and actions together.

heading: 3 THE HIP-BMDP SETTING
text: The HiP-MDP setting (as defined in Section 2) assumes full observability of the state space. However, in most real-world scenarios, we only have access to high-dimensional, noisy observations, which often contain irrelevant information to the reward. 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 dimension of complexity to our problem – we first want to learn an amenable state space S, and a universal dynamics model in that representation3. In this section, we formally define the HiP-BMDP family in Section 3.1, propose an algorithm for learning HiP-BMDPs in Section 3.2, and finally provide theoretical analysis for the setting in Section 3.3.

heading: 3.1 THE MODEL
text: A HiP-BMDP family can be described by tuple 〈S,A,Θ, Tθ, R, γ, PΘ,X , q〉, with a graphical model of the framework found in Figure 2. We are given a label k ∈ {1, ..., N} for each of N environments. We plan to learn a candidate Θ that unifies the transition dynamics across all environments, effectively finding T (·, ·, θ). For two environment settings θi, θj ∈ Θ, we define a distance metric:
d(θi, θj) := max s,a∈{S,A}
[ W ( Tθi(s, a), Tθj (s, a) )] . (1)
The Wasserstein-1 metric can be written as Wd(P,Q) = supf∈Fd ∥∥Ex∼P f(x) − Ey∼Qf(y)∥∥1, where Fd is the set of 1-Lipschitz functions under metric d (Müller, 1997). We omit d but use d(x, y) = ‖x − y‖1 in our setting. This ties distance between θ to the maximum difference in the next state distribution of all state-action pairs in the MDP.
Given a HiP-BMDP familyMΘ, we assume a multi-task setting where environments with specific θ ∈ Θ are sampled from this family. We do not have access to θ, and instead get environment labels I1, I2, ..., IN . The goal is to learn a latent space for the hyperparameters θ4. We want θ to be smooth with respect to changes in dynamics from environment to environment, which we can set explicitly through the following objective:
||ψ(I1)− ψ(I2)||1 = max s∈S a∈A
[ W2 ( p(st+1|st, at, ψ(I1)), p(st+1|st, at, ψ(I2)) )] , (2)
given environment labels I1, I2 and ψ : Z+ 7→ Rd, the encoder that maps from environment label, the set of positive integers, to θ.

heading: 3.2 LEARNING HIP-BMDPS
text: The premise of our work is that the HiP-BMDP formulation will improve sample efficiency and generalization performance on downstream tasks. We examine two settings, multi-task reinforcement
3We overload notation here since the true state space is latent. 4We again overload notation here to refer to the learned hyperparameters as θ, as the true ones are latent.
learning (MTRL) and meta-reinforcement learning (meta-RL). In both settings, we have access to N training environments and a held-out set of M evaluation environments, both drawn from a defined family. In the MTRL setting, we evaluate model performance across all N training environments and ability to adapt to new environments. 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. We evaluate against ablations and other MTRL methods. 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. 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).
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 7→ Z , ii) an environment encoder ψ that maps an environment identifier to a hidden parameter θ, iii) a universal dynamics model T conditioned on task parameter θ. Figure 2 shows how the components interact during training. In practice, computing the maximum Wasserstein distance over the entire state-action space is computationally infeasible. Therefore, we relax this requirement by taking the expectation over Wasserstein distance with respect to the marginal state distribution of the behavior policy. We train a probabilistic universal dynamics model T to output the desired next state distributions as Gaussians5, for which the 2-Wasserstein distance has a closed form:
W2(N (m1,Σ1), N (m2,Σ2))2 = ||m1 −m2||22 + ||Σ 1/2 1 − Σ 1/2 2 ||2F ,
where || · ||F is the Frobenius norm. Given that we do not have access to the true universal dynamics function across all environments, it must be learned. The objective in Equation (2) is accompanied by an additional objective to learn T , giving a final loss function:
L(ψ, T ) = MSE (∣∣∣∣∣∣ψ(I1)− ψ(I2)∣∣∣∣∣∣
2 ,W2
( T (sI1t , π(s I1 t ), ψ(I1)), T (s I2 t , π(s I2 t ), ψ(I2)) )) ︸ ︷︷ ︸
Θ learning error
+MSE ( T (sI1t , a I1 t , ψ(I1)), s I1 t+1 ) +MSE ( T (sI2t , a I2 t , ψ(I2)), s I2 t+1 ) ︸ ︷︷ ︸
Model learning error
.
(3)
where red indicates gradients are stopped. Transitions {sI1t , a I1 t , s I1 t+1, I1} and {s I2 t , a I2 t , s I2 t+1, I2} from two different environments (I1 6= I2) are sampled randomly from a replay buffer. In practice, we scale the Θ learning error, our task bisimulation metric loss, using a scalar value denoted as αψ .

heading: 3.3 THEORETICAL ANALYSIS
text: In this section, we provide value bounds and sample complexity analysis of the HiP-BMDP approach. We have additional new theoretical analysis of the simpler HiP-MDP setting in Appendix B. We first define three additional error terms associated with learning a R, T , θ-bisimulation abstraction,
R := sup a∈A,
x1,x2∈X ,φ(x1)=φ(x2) ∣∣R(x1, a)−R(x2, a)∣∣, T := sup
a∈A, x1,x2∈X ,φ(x1)=φ(x2) ∥∥ΦT (x1, a)− ΦT (x2, a)∥∥1, θ := ‖θ̂ − θ‖1.
Φ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. We can think of R, T as describing a new MDP which is close – but not necessarily the same, if R, T > 0 – to the original Block MDP. These two error terms can be computed empirically over all training environments and are therefore not task-specific.
5This is not a restrictive assumption to make, as any distribution can be mapped to a Gaussian with an encoder of sufficient capacity.
θ, on the other hand, is measured as a per-task error. Similar methods are used in Jiang et al. (2015) to bound the loss of a single abstraction, which we extend to the HiP-BMDP setting with a family of tasks.
Value Bounds. We first evaluate how the error in θ prediction and the learned bisimulation representation affect the optimal Q∗M̄θ̂
of the learned MDP, by first bounding its distance from the optimal Q∗ of the true MDP for a single-task. Theorem 2 (Q error). Given an MDP M̄θ̂ built on a ( R, T , θ)-approximate bisimulation abstraction of an instance of a HiP-BMDPMθ, we denote the evaluation of the optimal Q function of M̄θ̂ onM as [Q∗M̄θ̂ ]Mθ . The value difference with respect to the optimal Q
∗ M is upper bounded by∥∥Q∗Mθ − [Q∗M̄θ̂ ]Mθ∥∥∞ ≤ R + γ( T + θ) Rmax2(1− γ) .
Proof in Appendix C. As in the HiP-MDP setting, we can measure the transferability of a specific policy π learned on one task to another, now taking into account error from the learned representation. Theorem 3 (Transfer bound). Given two MDPsMθi andMθj , we can bound the difference in Qπ between the two MDPs for a given policy π learned under an R, T , θi -approximate abstraction of Mθi and applied to∥∥Q∗Mθj − [Q∗M̄θ̂i ]Mθj ∥∥∞ ≤ R + γ( T + θi + ‖θi − θj‖1) Rmax2(1− γ) . This result clearly follows directly from Theorem 2. Given a policy learned for task i, Theorem 3 gives a bound on how far from optimal that policy is when applied to task j. Intuitively, the more similar in behavior tasks i and j are, as denoted by ‖θi − θj‖1, the better π performs on task j.
Finite Sample Analysis. 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 finite sample regime. Abstractions are useful in this regime for their generalization capabilities. We can instead perform a counting analysis based on the number of samples of any abstract state-action pair.
We compute a loss bound with abstraction φ which depends on the size of the replay buffer D, collected over all tasks. Specifically, we define the minimal number of visits to an abstract state-action pair, nφ(D) = minx∈φ(S),a∈A |Dx,a|. 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. Theorem 4 (Sample Complexity). For any φ which defines an ( R, T , θ)-approximate bisimulation abstraction on a HiP-BMDP familyMΘ, we define the empirical measurement of Q∗M̄θ̂ over D to be Q∗M̄D
θ̂
. Then, with probability ≥ 1− δ,
∥∥Q∗Mθ − [Q∗M̄D θ̂ ]Mθ ∥∥ ∞ ≤ R + γ( T + θ) Rmax 2(1− γ) + Rmax (1− γ)2
√ 1
2nφ(D) log 2|φ(X )||A| δ . (4)
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 )|. 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. 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.

heading: 4 EXPERIMENTS & RESULTS
text: 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.
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. This dual-evaluation setup provides a more nuanced understanding of how well the learned model transfers across the environments. Implementation details can be found in Appendix D and sample videos of policies at https://sites.google.com/view/hip-bmdp.
Environments. We create a family of MDPs using the existing environment-task pairs from DMC and change one environment parameter to sample different MDPs. We denote this parameter as the perturbation-parameter. We consider the following HiP-BMDPs: 1. Cartpole-Swingup-V0: the mass of the pole varies, 2. Cheetah-Run-V0: the size of the torso varies, 3.
Walker-Run-V0: the friction coefficient between the ground and the walker’s legs varies, 4. Walker-Run-V1: the size of left-foot of the walker varies, and 5. Finger-Spin-V0: the size of the finger varies. We show an example of the different pixel observations for Cheetah-Run-V0 in Figure 3. Additional environment details are in Appendix D.
We sample 8 MDPs from each MDP family by sampling different values for the perturbationparameter. 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. We denote the ordered MDPs as A − H . 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). 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). We evaluate the learning agents by computing average reward (over 10 episodes) achieved by the policy after training for a fixed number of steps. All experiments are run for 10 seeds, with mean and standard error reported in the plots.
Multi-Task Setting. We first consider a multi-task setup where the agent is trained on four related, but different environments with pixel observations. 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). For all models, the agent sequentially performs one update per environment. For fair comparison, we ensure that baselines have at least as many parameters as HiP-BMDP. Distral has more parameters as it trains one policy per environment. Additional implementation details about baselines are in Appendix D.1.
In Figures 4, and 10 (in Appendix), we observe that for all the models, performance deteriorates when evaluated on interpolation/extrapolation environments. We only report extrapolation results in the main paper because of space constraints, as they were very similar to the interpolation performance. 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.
At training time (Figure 9 in Appendix), we observe that HiP-BMDP consistently outperforms other baselines on all the environments. The success of our proposed method can not be attributed to task embeddings alone as HiP-BMDP-nobisim also uses task embeddings. 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). We also note that the multi-task learning baselines like Distral, PCGrad, and GradNorm sometimes lag behind even the DeepMDP baseline, perhaps because they do not leverage a shared global dynamics model.
Meta-RL Setting. 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. Specifically, 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 (Finn et al., 2017) and ProMP (Rothfuss et al., 2019) on proprioceptive state. We incorporate our proposed approach in PEARL by training the inference network qφ(z|c) with our additional HiP-BMDP loss. The algorithm pseudocode can be found in Appendix D. 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. We provide additional results in Appendix E.
Evaluating the Universal Transition Model. We investigate how well the transition model performs in an unseen environment by only adapting the task parameter θ. We instantiate a new MDP, sampled from the family of MDPs, and use a behavior policy to collect transitions. These transitions are used to update only the θ parameter, and the transition model is evaluated by unrolling the transition model for k-steps. We report the average, perstep model error in latent space, averaged over 10 environments. 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. In Figure 6, we indeed observe that the proposed HiP-BMDP model adapts much faster than the ablation HiP-BMDP-nobisim.
Relaxing the Block MDP Assumption. We incorporate sticky observations into the environment to determine how HiP-BMDP behaves when the Block MDP assumption is relaxed. For some probability p (set to 0.1 in practice), the current observation is dropped, and the agent sees the previous observation again. In Figure 7, we see that even in this setting the proposed HiP-BMDP model outperforms the other baseline models.

heading: 5 RELATED WORK
text: Multi-task learning has been extensively studied in RL with assumptions around common properties of different tasks, e.g., reward and transition dynamics. A lot of work has focused on considering tasks as MDPs and learning optimal policies for each task while maximizing shared knowledge. However, in most real-world scenarios, the parameters governing the dynamics are not observed. Moreover, it is not explicitly clear how changes in dynamics across tasks are controlled. The HiP-BMDP setting provides a principled way to change dynamics across tasks via a latent variable.
Much existing work in the multi-task reinforcement learning (MTRL) setting focuses on learning shared representations (Ammar et al., 2014; Parisotto et al., 2016; Calandriello et al., 2014; Maurer et al., 2016; Landolfi et al., 2019). D’Eramo et al. (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. (2020); Tirinzoni et al. (2020) also obtain PAC bounds on sample complexity for the MTRL setting, but Sun et al. (2020) relies on a constructed state-action abstraction that assumes
a discrete and tractably small state space. 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. (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.
Meta-learning, or learning to learn, is also a related framework with a different approach. We focus here on context-based approaches, which are more similar to the shared representation approaches of MTRL and our own method. Rakelly et al. (2019) model and learn latent contexts upon which a universal policy is conditioned. However, no explicit assumption of a universal structure is leveraged. Amit & Meir (2018); Yin et al. (2020) give a PAC-Bayes bound for meta-learning generalization that relies on the number of tasks n. 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. Instead, we assume a shared latent structure underlying all tasks, and seek to exploit that structure for generalization. We find that under this setting, our method indeed outperforms policies initialized through meta-learning.
The ability to extract meaningful information through state abstractions provides a means to generalize across tasks with a common structure. Abel et al. (2018) learn transitive and PAC state abstractions for a distribution over tasks, but they concentrate on finite, tabular MDPs. 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. Prior work also leverages bisimulation for transfer (Castro & Precup, 2010), but on the policy level. Our work instead focuses on learning a latent state representation and established theoretical results for the MTRL setting. 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.

heading: 6 DISCUSSION
text: In this work, we advocate for a new framework, HiP-BMDP, to address the multi-task reinforcement learning setting. Like previous methods, HiP-BMDP assumes a shared state and action space across tasks, but additionally assumes latent structure in the dynamics. We exploit this structure through learning a universal dynamics model with latent parameter θ, which captures the behavioral similarity across tasks. We provide error and value bounds for the HiP-MDP (in appendix) and 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. Our work relies on an assumption that we have access to an environment id, or knowledge of when we have switched environments. This assumption could be relaxed by incorporating an environment identification procedure at training time to cluster incoming data into separate environments. Further, our bounds rely L∞ norms for measuring error and the value and transfer bounds. In future work we will investigate tightening these bounds with Lp norms.

heading: A BISIMULATION BOUNDS
text: We first 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. We can compute approximate error bounds in this setting by denoting φ an ( R, T )-approximate bisimulation abstraction, where
R := sup a∈A,
x1,x2∈X ,φ(x1)=φ(x2) ∣∣R(x1, a)−R(x2, a)∣∣, T := sup
a∈A, x1,x2∈X ,φ(x1)=φ(x2) ∥∥ΦT (x1, a)− ΦT (x2, a)∥∥1. Φ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. Theorem 5. Given an MDP M̄ built on a ( R, T )-approximate bisimulation abstraction of Block MDPM, we denote the evaluation of the optimal Q function of M̄ onM as [Q∗M̄]M. The value difference with respect to the optimal Q∗M is upper bounded by∥∥Q∗M − [Q∗M̄]M∥∥∞ ≤ R + γ T Rmax2(1− γ) . Proof. From Theorem 2 in Jiang (2018).

heading: B THEORETICAL RESULTS FOR THE HIP-MDP SETTING
text: 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.
B.1 VALUE BOUNDS
Given a family of environmentsMΘ, we bound the difference in expected value between two sampled MDPs,Mθi ,Mθj ∈ MΘ using d(θi, θj). Additionally, we make the assumption that we have a behavior policy π that is near both optimal policies π∗θi , π ∗ θj
. We use KL divergence to define this neighborhood for π∗θi ,
dKL(π, π∗θi) = Es∼ρπ [ KL(π(·|s), π∗θi(·|s)) 1/2 ] . (5)
We start with a bound for a specific policy π. One way to measure the difference between two tasks Mθi ,Mθj is to measure the difference in value when that policy is applied in both settings. We show the relationship between the learned θ and this difference in value. 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. Theorem 6. Given policy π, the difference in expected value between two MDPs drawn from the family of MDPsMθi ,Mθj ∈MΘ is bounded by
|V πθi − V π θj | ≤
γ
1− γ ‖θi − θj‖1. (6)
Proof. We use a telescoping sum to prove this bound, which is similar to Luo et al. (2019). First, we let Zk denote the discounted sum of rewards if the first k steps are inMθi , and all steps t > k are in Mθj ,
Zk := E ∀t≥0,at∼π(st) ∀j>t≥0,st+1∼Tθi (st,at) ∀t≥j,st+1∼Tθj (st,at)
[ ∞∑ t=0 γtR(st, at) ] .
By definition, we have Z∞ = V πθi and Z0 = V π θj . Now, the value function difference can be written as a telescoping sum,
V πθi − V π θj = ∞∑ k=0 (Zk+1 − Zk). (7)
Each term can be simplified to
Zk+1 − Zk = γk+1Esk,ak∼π,Tθi [ Esk+1∼Tθj (·|sk,ak), s′k+1∼Tθi (·|sk,ak) [ V πθj (sk+1)− V π θj (s ′ k+1 ]] .
Plugging this back into Equation (7),
V πθi − V π θj =
γ
1− γ Es∼ρπθi , a∼π(s)
[ Es′∼Tθi (·|s,a)V π θj (s ′)− Es′∼Tθj (·|s,a)V π θj (s ′) ] .
This expected value difference is bounded by the Wasserstein distance between Tθi , Tθj ,
|V πθi − V π θj | ≤
γ
1− γ W (Tθi , Tθj )
= γ
1− γ ‖θi − θj‖1 using Equation (1).
Another comparison to make is how different the optimal policies in different tasks are with respect to the distance ‖θi − θj‖. Theorem 7. The difference in expected optimal value between two MDPs Mθi ,Mθj ∈ MΘ is bounded by,
|V ∗θi − V ∗ θj | ≤
γ
(1− γ)2 ‖θi − θj‖1. (8)
Proof.
|V ∗θi(s)− V ∗ θj (s)| = |maxa Q ∗ θi(s, a)−maxa′ Q ∗ θj (s, a ′)|
≤ max a |Q∗θi(s, a)−Q ∗ θj (s, a)|
We can bound the RHS with
sup s,a |Q∗θi(s, a)−Q ∗ θj (s, a)| ≤ sup s,a |rθi(s, a)−rθj (s, a)|+γ sup s,a |Es′∼Tθi (·|s,a)V ∗ θi(s ′)−Es′′∼Tθj (·|s,a)V ∗ θj (s ′′)|
All MDPs inMΘ have the same reward function, so the first term is 0. sup s,a ∣∣Q∗θi(s, a)−Q∗θj (s, a)∣∣ ≤ γ sup s,a |Es′∼Tθi (·|s,a)V ∗ θi(s ′)− Es′′∼Tθj (·|s,a)V ∗ θj (s ′′)|
= γ sup s,a ∣∣∣∣Es′∼Tθi (·|s,a)[V ∗θi(s′)− V ∗θj (s′)]+ Es′′∼Tθj (·|s,a), s′∼Tθi (·|s,a) [ V ∗θj (s ′)− V ∗θj (s ′′) ]∣∣∣∣
≤ γ sup s,a ∣∣Es′∼Tθi (·|s,a)[V ∗θi(s′)− V ∗θj (s′)]∣∣+ γ sups,a ∣∣Es′′∼Tθj (·|s,a), s′∼Tθi (·|s,a) [ V ∗θj (s ′)− V ∗θj (s ′′) ]∣∣
≤ γ sup s,a ∣∣Es′∼Tθi (·|s,a)[V ∗θi(s′)− V ∗θj (s′)]∣∣+ γ1− γ ‖θi − θj‖1 ≤ γmax
s ∣∣V ∗θi(s)− V ∗θj (s)∣∣+ γ1− γ ‖θi − θj‖1 = γmax
s ∣∣max a Q∗θi(s, a)−maxa′ Q ∗ θj (s, a ′) ∣∣+ γ 1− γ ‖θi − θj‖1
≤ γ sup s,a ∣∣Q∗θi(s, a)−Q∗θj (s, a)∣∣+ γ1− γ ‖θi − θj‖1 Solving for sups,a
∣∣Q∗θi(s, a)−Q∗θj (s, a)∣∣, sup s,a
∣∣Q∗θi(s, a)−Q∗θj (s, a)∣∣ ≤ γ(1− γ)2 ‖θi − θj‖1. Plugging this back in,
|V ∗θi(s)− V ∗ θj (s)| ≤
γ
(1− γ)2 ‖θi − θj‖1.
Both these results lend more intuition for casting the multi-task setting under the HiP-MDP formalism. The difference in the optimal performance between any two environments is controlled by the distance between the hidden parameters for corresponding environments. One can interpret the hidden parameter as a knob to allow precise changes across the tasks.

heading: B.2 EXPECTED ERROR BOUNDS
text: In MTRL, we are concerned with the performance over a family of tasks. The empirical risk is typically defined as follows for T tasks (Maurer et al., 2016):
avg(θ) = 1
T T∑ t=1 E[`(ft(h(wt(X))), Y ))]. (9)
Consequently, we bound the expected loss over the family of environments E with respect to θ. In particular, we are interested in the average approximation error and define it as the absolute model error averaged across all environments:
avg(θ) = 1
|E| E∑ i=1 ∣∣∣V ∗ θ̂i (s)− V ∗θi(s) ∣∣∣. (10)
Theorem 8. 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:
avg(θ) ≤ γ
(1− γ)2 , (11)
where each environment’s parameter θi is -close to its approximation θ̂i i.e. d(θ̂i, θi) ≤ , where d is the distance metric defined in Eq. 1.
Proof. We here consider the approximation error averaged across all environments as follows:
avg(θ) = 1
E E∑ i=1 ∣∣∣V ∗ θ̂i (s)− V ∗θi(s) ∣∣∣
avg(θ) = 1
E E∑ i=1 |max a Q∗ θ̂i (s, a)−max a′ Q∗θi(s, a ′)|
≤ 1 E E∑ i=1 max a |Q∗ θ̂i (s, a)−Q∗θi(s, a)|
(12)
Let us consider an environment θi ∈ME for which we can bound the RHS with
sup s,a |Q∗ θ̂i (s, a)−Q∗θi(s, a)| ≤ sup s,a |rθ̂i(s, a)−rθi(s, a)|+γ sup s,a |Es′∼Tθ̂i (·|s,a)V ∗ θ̂i) (s′)−Es′′∼Tθi (·|s,a)V ∗ θi(s ′′)|
Considering the family of environmentsME have the same reward function and is known, resulting in first term to be 0.
sup s,a
∣∣Q∗ θ̂i (s, a)−Q∗θi(s, a) ∣∣ ≤ γ sup s,a |Es′∼Tθ̂i (·|s,a)V ∗ θ̂i (s′)− Es′′∼Tθi (·|s,a)V ∗ θi(s ′′)|
= γ sup s,a ∣∣∣∣Es′∼Tθ̂i (·|s,a)[V ∗θ̂i(s′)− V ∗θi(s′)]+ Es′′∼Tθi (·|s,a), s′∼Tθ̂i (·|s,a) [ V ∗ θ̂i (s′)− V ∗θi(s ′′) ]∣∣∣∣
≤ γ sup s,a ∣∣Es′∼Tθ̂i (·|s,a)[V ∗θ̂i(s′)− V ∗θi(s′)]∣∣+ γ sups,a ∣∣Es′′∼Tθi (·|s,a), s′∼Tθ̂i (·|s,a) [ V ∗ θ̂i (s′)− V ∗θi(s ′′) ]∣∣
≤ γ sup s,a ∣∣Es′∼Tθ̂i (·|s,a)[V ∗θ̂i(s′)− V ∗θi(s′)]∣∣+ γ1− γ |θ̂i − θi| ≤ γmax
s
∣∣V ∗ θ̂i (s)− V ∗θi(s) ∣∣+ γ 1− γ |θ̂i − θi|
= γmax s ∣∣max a Q∗ θ̂i (s, a)−max a′ Q∗θi(s, a ′) ∣∣+ γ 1− γ |θ̂i − θi|
≤ γ sup s,a ∣∣Q∗ θ̂i (s, a)−Q∗θi(s, a) ∣∣+ γ 1− γ |θ̂i − θi|
Solving for sups,a ∣∣Q∗
θ̂i (s, a)−Q∗θi(s, a) ∣∣, sup s,a ∣∣Q∗ θ̂i (s, a)−Q∗θi(s, a) ∣∣ ≤ γ (1− γ)2 |θ̂i − θi| (13)
Plugging Eq. 13 back in Eq. 12,
avg(θ) ≤ 1
E E∑ i=1
γ
(1− γ)2 |θ̂i − θi|
= γ
E(1− γ)2
[ |θ̂i=1 − θi=1|+ |θ̂i=2 − θi=2|+ · · ·+ |θ̂i=E − θi=E | ]
We now consider that the distance between the approximated θ̂i and the underlying hidden parameter θi ∈ME is defined as in Eq. 1, such that: d(θ̂i, θi) ≤ θ Plugging this back concludes the proof,
avg(θ) ≤ γ θ
(1− γ)2 .
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.

heading: C ADDITIONAL RESULTS AND PROOFS FOR HIP-BMDP RESULTS
text: We first compute L∞ norm bounds for Q error under approximate abstractions and transfer bounds. Theorem 9 (Q error). Given an MDP M̄θ̂ built on a ( R, T , θ)-approximate bisimulation abstraction of an instance of a HiP-BMDPMθ, we denote the evaluation of the optimal Q function of M̄θ̂ onM as [Q∗M̄θ̂ ]Mθ . The value difference with respect to the optimal Q
∗ M is upper bounded by∥∥Q∗Mθ − [Q∗M̄θ̂ ]Mθ∥∥∞ ≤ R + γ( T + θ) Rmax2(1− γ) .
Proof. In the HiP-BMDP setting, we have a global encoder φ over all tasks, but the difference in transition distribution also includes θ. The reward functions are the same across tasks, so there is
no change to R. However, we now must incorporate difference in dynamics in T . Assuming we have two environments with hidden parameters θi, θj ∈ Θ, we can compute θi,θj T across those two environments by joining them into a super-MDP:
θi,θj T = sup
a∈A, x1,x2∈X ,φ(x1)=φ(x2) ∥∥ΦTθi(x1, a)− ΦTθj (x2, a)∥∥1 ≤ sup
a∈A, x1,x2∈X ,φ(x1)=φ(x2) (∥∥ΦTθi(x1, a)− ΦTθi(x2, a)∥∥1 + ∥∥ΦTθi(x2, a)− ΦTθj (x2, a)∥∥1) ≤ sup
a∈A, x1,x2∈X ,φ(x1)=φ(x2)
∥∥ΦTθi(x1, a)− ΦTθi(x2, a)∥∥1 + sup a∈A,
x1,x2∈X ,φ(x1)=φ(x2) ∥∥ΦTθi(x2, a)− ΦTθj (x2, a)∥∥1 = θiT + ‖θi − θj‖1
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. We can now extend the single-task bisimulation bound (Theorem 5) to the HiP-BMDP setting by denoting approximation error of θ as ‖θ − θ̂‖1 < θ.
Theorem 4. For any φ which defines an ( R, T , θ)-approximate bisimulation abstraction on a HiP-BMDP familyMΘ, we define the empirical measurement of Q∗M̄θ̂ over D to be Q ∗ M̄D
θ̂
. Then,
with probability ≥ 1− δ,∥∥Q∗Mθ − [Q∗M̄D θ̂ ]Mθ ∥∥ ∞ ≤ R+γ( T + θ) Rmax 2(1− γ) + Rmax (1− γ)2 √ 1 2nφ(D) log 2|φ(X )||A| δ . (14)
Proof. ∥∥Q∗Mθ − [Q∗M̄D θ̂ ]Mθ ∥∥ ∞ ≤ ∥∥Q∗Mθ − [Q∗M̄θ̂ ]Mθ∥∥∞ + ∥∥[Q∗M̄θ̂ ]Mθ − [Q∗M̄Dθ̂ ]Mθ∥∥∞ = ∥∥Q∗Mθ − [Q∗M̄θ̂ ]Mθ∥∥∞ + ∥∥Q∗M̄θ̂ −Q∗M̄Dθ̂ ∥∥∞ The first 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, we write this difference to be a deviation from an expectation in order to apply the concentration inequality. ∥∥Q∗M̄θ̂ −Q∗M̄Dθ̂ ∥∥∞ = ∥∥Q∗M̄θ̂ − T φDQ∗M̄θ̂ + T φDQ∗M̄θ̂ − T φDQ∗M̄Dθ̂ ∥∥∞
≤ ‖Q∗M̄θ̂ − T φ DQ ∗ M̄θ̂ ‖∞ + γ‖Q∗M̄θ̂ −Q ∗ M̄D
θ̂
‖∞
≤ 1 1− γ ‖T φDQ ∗ M̄θ̂ − T φQ∗M̄θ̂‖∞
Now we can apply McDiarmid’s inequality, PD [∣∣Q∗M̄θ̂ −Q∗M̄Dθ̂ ∣∣ ≥ t ] ≤ 2 exp ( − 2t2|Dφ(x),a| R2max/(1− γ)2 ) . Solve for the t that makes this inequality hold for all (φ(x), a) ∈ X ×A with a union bound over all |φ(X )||A| abstract states,
t > Rmax 1− γ
√ 1
2nφ(D) log 2|φ(X )||A| δ .
Combine to get∥∥Q∗Mθ − [Q∗M̄D θ̂ ]Mθ ∥∥ ∞ ≤ R + γ( T + θ) Rmax 2(1− γ) + Rmax (1− γ)2 √ 1 2nφ(D) log 2|φ(X )||A| δ .
Algorithm 1 HiP-BMDP training for the Multi-task RL setting. Require: Along with DeepMDP components (Actor, Critic, Dynamics Model (M )), an additional
environment encoder ψ to generated task-specific θ parameters. 1: for each timestep t = 1..T do 2: for each Ti do 3: ait ∼ πi(·|sit) 4: s′it ∼ pi(·|sit, ait) 5: D ← D ∪ (sit, ait, r(sit, ait), s′it ), i 6: UPDATECRITIC(D, i) (uses data only from ith task) 7: UPDATEACTOR(D, i) (uses data only from ith task) 8: UPDATEUSINGHIP-BMDPLOSS(D, i) 9: end for
10: end for
Algorithm 2 UpdateModelUsingHip-BMDPLoss Require: Batches of data for the different tasks {Ti}i=1...T sampled from the Replay Buffer D,
learning rates α1 and α2, index of the current task i, Transition Model M , environment encoder ψ.
1: for each batch of dataset t = 1..T , t 6= i do 2: Compute L(ψ,M) = Li(ψ,M, i, t) using Equation (3) 3: ψ ← ψ − α1∇θ ∑ i L
4: M ←M − α2∇θ ∑ i L 5: end for

heading: D ADDITIONAL IMPLEMENTATION DETAILS
text: In Figure 8, we show the variation in the left foot of the walker.
MTRL Algorithm The Multitask RL algorithm for the HiP-BMDP setting can be found in Algorithm 1. We take the DeepMDP baseline Gelada et al. (2019) and incorporate our HiP-BMDP objective (text shown in red color)
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)

heading: D.1 BASELINES
text: 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. For GradNorm, we observe that the learned weights wi (for weighing per-task loss) can become negative for some tasks (which means the model tries to unlearn those tasks). In practice, we clamp the wi values to not become smaller than a threshold.

heading: D.2 HYPER PARAMETERS

heading: D.2.1 MTRL ALGORITHM
text: All the hyper parameters (for MTRL algorithm) are listed in Table 1.

heading: D.2.2 METARL ALGORITHM
text: For MetaRL, we use the same hyperparameters as used by PEARL Rakelly et al. (2019). We set αφ = 0.01 for all environments, other than Walker-Stand environments where αφ = 0.001.

heading: E ADDITIONAL RESULTS
text: Along with the environments described in 4, we considered the following additional environments:
1. Walker-Stand-V0: Walker-Stand task where the friction coefficient, between the ground and the walker’s leg, varies across different environments.
2. Walker-Walk-V0: Walker-Walk task where the friction coefficient, between the ground and the walker’s leg, varies across different environments.
3. Walker-Stand-V1: Walker-Stand task where the size of left-foot of the walker varies across different environments.
Algorithm 3 HiP-MDP training for the meta-RL setting. Require: Batch of training tasks {Ti}i=1...T from p(T ), learning rates α1, α2, α3, αφ
1: Initialize replay buffers Bi for each training task 2: while not done do 3: for each Ti do 4: Initialize context Ci = {} 5: for k = 1, . . . ,K do 6: Sample z ∼ qφ(z|Ci) 7: Gather data from πθ(a|s, z) and add to Bi 8: Update Ci = {(sj ,aj , s′j , rj)}j:1...N ∼ Bi 9: end for
10: end for 11: for step in training steps do 12: for each Ti do 13: Sample context Ci ∼ Sc(Bi) and RL batch bi ∼ Bi 14: Sample z ∼ qφ(z|Ci) 15: Liactor = Lactor(bi, z) 16: Licritic = Lcritic(bi, z) 17: LiKL = βDKL(q(z|Ci)||r(z)) 18: Sample a RL batch bj from any other task j 19: Compute LiBiSim = Li(q, T, i, j) using the equation 3 20: end for 21: φ← φ− α1∇φ ∑ i ( Licritic + LiKL + αφ × LiBiSim
) 22: θπ ← θπ − α2∇θ ∑ i Liactor
23: θQ ← θQ − α3∇θ ∑ i Licritic 24: end for 25: end while
4. Walker-Walk-V1: Walker-Walk task where the size of left-foot of the walker varies across different environments.

heading: E.1 MULTI-TASK SETTING
text: In Figure 10, we observe that the HiP-BMDP method consistently outperforms other baselines when evaluated on the interpolation environments (zero-shot transfer). As noted previously, the 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. We hypothesise that the Distral-Ensemble baseline behaves poorly because it cannot leverage a shared global dynamics model.

heading: E.2 META-RL SETTING
text: We provide the Meta-RL results for the additional environments. Recall that we extend the PEARL algorithm (Rakelly et al., 2019) by training the inference network qφ(z|c) with our additional HiPBMDP loss. 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. In some environments (eg Walker-Walk-V1), the proposed approach (blue) converges faster to a threshold reward (green) than the baseline. In the other environments, the gains are quite small.

heading: E.3 EVALUATING THE UNIVERSAL TRANSITION MODEL.
text: We investigate how well the transition model performs in an unseen environment by only adapting the task parameter θ. We instantiate a new MDP, sampled from the family of MDPs, and use a behavior policy to collect transitions. These transitions are used to update only the θ parameter, and the transition model is evaluated by unrolling the transition model for k-steps. In Figures 11 and 12, we report the average, per-step model error in latent space, averaged over 10 environments over 5 and 100 steps respectively. 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. We indeed observe that for both 5 step and 100 step unrolls, the proposed HiP-BMDP model adapts much faster than the baseline HiP-BMDP-nobisim (Figures 11 and 12)

Review summaries:
This paper studies a family of Markov Decision Process (MDP) models with a low-dimensional unobserved state, called the block MDP. The authors assume that system dynamics of the underlying MDP is sufficiently summarized by a parameter $\theta$. This learning setting could be seen as a combination of Block MDP and the Hidden Parameter MDP; hence the name HiP-BMDP.