Learning Robust State Abstractions for Hidden-Parameter Block MDPs
Authors: Amy Zhang
Shagun Sodhani
Khimya Khetarpal
Joelle Pineau
Venue: ICLR 2021 Poster (2021)
Keywords: multi-task reinforcement learning
bisimulation
hidden-parameter mdp
block mdp

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. To further demonstrate efficacy of the proposed method, we empirically compare and show improvement over multi-task and meta-reinforcement learning baselines.

prior work Published as a conference paper at ICLR 2021
LEARNING ROBUST STATE ABSTRACTIONS FOR HIDDEN-PARAMETER BLOCK MDPs
Amy Zhang 123 McGiil University 2 Facebboo Al Research
Shagun Sodhani
Khimya Khetarpal
Joelle Pineau 123
Mila
ABSTRACT
Many control tasks exhibit similar dynamics that be modeled as having com- mon latent structure. Hidden-Paramete Markov Decision Processes (HiP-MDPs) explicitly model this structure improve sample efficiency in multi-task settings. However, this setting makes strong assumptions on the observability of the state that limit its application this work, we leverage ideas of common structure from the HiP-MDP setting, and extend to enable robust state abstractions inspired by Block MDPs. We derive instantiations of this new framework for both multi-task reinforcement learning (MTRL. and rentoreeifoor learning (Meta-RI.) 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, significant improvement that use the same environment assumptions. To further demonstrate the efficacy of the proposed method, empirically compare and show improvement over multi-task and meta-reinforcemmnt learning baselines.
real-world scenarios with rich observation spaces. In
INTRODUCTION
key open challenge in Al research that remains is how to train agents that can learn behaviors that generalize across tasks and environments. When there common structure underlying the tasks, we have seen that multi-task reinforcement learning (MTRL), where the agent learns set of tasks simultaneously. has definite advantages (in terms of robustness and sample efficiency) over the single- task setting, where the agent independently learns each task. There are two ways which learning multiple tasks can accelerate learning: the agent can learn 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' ErEmoo et al. 2020) or providing per- task 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 of shared representations use naive unification approach that posits all tasks lie in shared domain (Figure I left). On the other hand, in the single-task setting, research on state abstractions has much richer history, with several works on improved generalization through the aggregation of behaviorally similar states (Ferns et al. 2004 Li et al.
MEP
%
NOP
HIP-MDP
Figure 1:
Visualizations of the typical MTRL setting and the HiP-MDP setting.
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 shared state space and universal dynamics model conditioned hidden parameter (Figure right). This additional structure gives us better sample efficiency, both
task-specific
"Corresponding author: amy x. zhang@mai megiiil

---

Published as a conference paper at ICLR 2021
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 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--prrmmeer MDP (HiP-MDP) (Doshi- Velez & 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 RI 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 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 nigh-dimensional observations and structured, changing dynamics.
Key contributions of this work are new viewpoint of the multi-task setting with same reward function as a universal MDP under the HiP-BMDP setting, which naturally leads to gradient- based representation learning algorithm. Further, this framework allows us to compute theoretical generalization results with the incorporation of 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.
2 BACKGROUND
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 this work.
A finite] [ discrete-tin time Markov Decision Process (MDP) (Bellman 1957 Puterman 1995) is tuple (S,A T, 7). where S is the set of states, AA is the set of actions, R S x A + R is the reward function, T' S A + Dist(S) is the environment transition probability function, and E [0. 1) is the discount factor. At each time step, the learning agent perceives a state 8 E S, takes an action 0 A drawn from a policy π S x A - [0, 1), and with probability T(s++1|s az) enters next state S14 receiving a numerical reward Re+1 from the environment. The value function of policy is defined as: (s() = Ex E2r 7' = s). The optimal value function V is the maximum value function over the class of stationary policies.
Hidden-Parann Parameter MDPs (HiP-MDPs) (Doshi-Velez & Konidaris 2013) can be defined by tuple M: (S, A, θ, To, R. 7, Pe) where S is a finite state space, A a finite action space, Ta describes the transition distribution for a specific task described by task parameter 0 ~ Pe. function, is the discount factor, and Pa the distribution over task parameters. This defines a family of MDPs, where each MDP is described by the parameter 0 ~ Pg. We assume that this parameter is fixed for an episode and indicated by an environment id given at the start of the episode.
the reward
Block MDPs (Du et al. 2019) are described by tuple AS A, q, R) with an unobservable state space S, action space A, and observable space X. denotes the latent transition distribution p(s'|s, a) for s E S,a E A. q is the (possibly stochastic) emission mapping that emits the observations q(z[s) for E X, E S, and R the reward function. We are interested in the setting where this
Wee use this assumption only for theoretical results, hut our method can be applied continuous domains.
2

---

Published as a conference paper at ICLR 2021
mapping q is one-to 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 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(-1s).
Assumption gives the Markov property in X, a key difference from partially observable MDPs (Kael bling 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 compute reasonable bounds for algorithm in k-order MDP (which describes many real world problems) and avoiding the intractabil- ity of true POMDPs, which have no guarantees 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 relaxation of this assumption compared to other MTRL methods.
providing enough information
sufficiently
Bisimulation a strict form of state abstraction, where two states are bisimilar they are behaviorally equivalent. Bisimulation metrics (Ferns et al., 2011) define distance between states as follows:
Definition (Bisimulation Metric (Theorem 2.6 in Ferns et al. (2011)) )), Let (S,.. be finite MDP and met the space of bounded pseudometrics on equipped with the metric induced by the uniform norm. Define F: met + met by
F(d)(s, max((r; + aEA
where (d) is the Wasserstein distance between transition probability distributions. Then F has unique fixed point d which is the bisimulation metric.
nice property of this metric d is that difference in optimal value between two states bounded by their distance as defined by this metric.
Theorem (V* Lipschitz with respect to d (Ferns al. 2004)). Let V be the optimal value function for a given discount facter 7; Then V* is Lipschitz continuous with respect d with Lipschitz constant
VV (s*
1 1-7 ā(s.s').
Therefore, that bisimulation metrics give Lipschitz value function with respect to d.
For downstream evaluation of the representations we learn, we use Soft Actor Critic (SAC) (Haamoja al. 2018), 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 parametric soft Q-function Q(st, az) using transitions sampled from the replay buffer D by minimizing the soft Bellman residual,
J(Q) = (( ( (((5.
- ( (((s:))-nn-(())))
The target value function V is approximated via Monte-Carlo estimate of the following expectation,
V (z++1) Eet,i>s [((=011, 1c41) log π ((c++1/s4+1
where Q is the target soft Q-function parameterized by weight vector obtained from an exponentially moving average of the Q-function weights to stabilize training. The policy improvement step then attempts to project parametric policy (s,]ss) by minimizing KL divergence between the policy and Boltzmann distribution induced by the Q-function, producing the following objective,
J (m) ) = Ea,-D Ea,~ (a
- Q(+e, az)]
Any k-order MDP he made Markov by stacking the previous observations and actions together.
3

---

Published as a conference paper at ICLR 2021
action
Environment Encoder
8,
0
%
environment iss
Dynamics Model
T(),a,0))
Eq.
s
7(s.
State Encoder
S+t
Reward Model
R/=")
R
abservation 3x(3884584)
Figure 2: Graphical model of HiP-BMDP setting (left). Flow diagram of learning HiP-BMDD (right). Two environment ids are selected by permuting randomly sampled batch of data from the replay buffer, and the loss objective requires computing the Wasserstein distance of the predicted next-step distribution for those states.
3 THE HIP-BMDP SETTING
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 digensimesiinnl, noisy observations, which often contain irrelevant information to the reward. We combine the Block MDP and HiP-MDP Block MDP setting (HiP-BMDP), where states
settings to introduce the Hfidden-Paramee latent, and transition distributions change depending on the task parameters O. This adds an additional dimension of complexity to our problem we first want to learn an amenable state space S, and universal dynamics model in that representation 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 Seection 3.3
3.1 THE MODEL
HiP-BMDP family can described by tuple (S, A, R,7: Po, Po.XX'l X.4). with graphical model of the framework found in Figure 2 We are given label & E {1,..., N. for each of N environments. We plan to learn a candidate that unifies the transition dynamics across all environments, effectively finding T).:00). For two environment settings 0,, 0, O we define distance metric:
d(0,, 0,) = max (S,A(S
W (s. To, (s.
The Wasserstein-1 metric can be written as Wa(P.Q) sup E E fxx R where Fa is the set of I-Lipschitz functions under metric (Müller 1997). We omit d but use d(x, y) = || 2/l1 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 family Me, we assume a multi-task setting where environments with specific Ô E θ are sampled from this family. We do not have access to 9, and instead get environment labels I, I2, ..., IN. The goal is to learn a latent space for the hyperparameters 6 We want to be smooth with respect to changes in dynamics from environment environment, which we can explicitly through the following objective:
[[)((1)) - (12))) = max SES
W/ (2(p5041155 ae, ((11),, p(s:41 on ((2))))
(2)
aEA
given environment labels I,, In and 2b - Rd, the encoder that maps from environment label, the set of positive integers, to 6.
3.2 LEARNING HIP-BMDPS
The premise of our work is that the HiP-BMDP formulation will improve sample efficiency and
generalization performance downstream tasks, We examine two settings, multi-task reinforcement
We overload notation here since the true state space is latent. We again overload notation here to refer the learned hyperparameters the ones are latent.
4

---

state distributions as Published as a conference paper at ICLR 2021
learning (MTRL) and meta-reiinfrrcemet learning (meta-RL). In both settings, we have access N training environments and a held-out set of evaluation environments, both drawn from a defined family. In the MTRL setting, 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 collect small number of samples from the evaluation environments learn each hidden parameter 0, and the zero-shot regime, where we average 0 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 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 0 for each new environment).
Learning HiP-BMDP approximation of a family of MDPs requires the following components: an encoder that maps observations from state space to learned, latent representation, é X - 2, ii) an environment encoder that maps an environment identifier to a hidden parameter 0, iii) a universal dynamics model I conditioned on task parameter 0. 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 the behavior policy. We train probabilistic universal dynamics model T output the desired G for which the 2- Wasserstein distance has a closed form:
W2(N(m1. Σ1), (m,,, Σz))² = m m2ll2 + +|2|/2 1//
where || |||= 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:
C(6, 7) = MSE
=(s()) ((2)))
learning ervor
+ MSE T(s (61)), oft))) + MSE af2, s
Model learning error
where red indicates gradients are stopped. Transitions {s] 41 s Li} and {sf² s I2] from two different environments (I1 & I2) sampled randomly from replay buffer. In practice, we scale the θ learning error, our task bisimulation metric loss, using scalar value denoted as og-
3.3 THEORETICAL ANALYSIS
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 <,, éT, eo-bisimulation abstraction,
R(11, (((2,,))R(22.a))
CR
sup A.A. 21 aa (a ((x)
[T((),))) (
DT
ET
sup
A.A. x 32 ( )
6 = 8||.
4f denotes the lifted version of T, where we take the next-step transition distribution from observa- tion space X' and lift it to latent space S. We can think of ER, ET as describing new MDP which close but not necessarily the same, if CR, cr > 0 - to the original Block MDP. These two error terms can be computed empirically over all training environments and are therefore not task-specific.
*This is not restrictive assumption to make, any distribution can be mapped to a Gaussian with an encoder of sufficient capacity.
5

---

Published as a conference paper at ICLR 2021
cø, 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-BMMP setting with a family of tasks.
Value Bounds. We first evaluate how the error in 0 prediction and the learned bisimulation repre- sentation affect the optimal LMa, of the learned MDP, by first bounding its distance from the optimal Q' of the true MDP for single-task.
Theorem (Q error), Given an MDP Mg built on (ER, ET, Eq) apprrxximaeeee bisimulation abstrac- tion of an instance of a RMP-BMDPP2 Ma. we denote the evaluation of the optimal Q function of Ma on M [Q'Aal Ma The value difference with respect the optimal Q'M upper bounded by
Rmax
1le'o - -Mollloollo sen + 7(r + 60)
2(1 71)
As in the HiP-MDP setting, we can measure the transferability of a specific policy π learned one task to another, taking into account enror from the learned representation.
Proof in Appendix C
Theorem 3 (Transfer bound). Given two MDPs Me, and Ma,, we can bound the difference in Q" between the two MDPs for a given policy π learned under an ER. cp, co, approximate abstraction of Mo, and applied to
Rmaxx
||Q'a, - 1Q2n IMo lloo SER + Eq, + +110, - 0,]11) 2(1 7)
This result clearly follows directly from Theorem 2 gives a bound on how far from optimal that policy is when applied to task j. Intuitively, the more similar in behavior tasks and j are, denoted by |10 0ylla, the better æ performs task j.
Given policy learned for task i, Theorem 3
Finite Sample Analysis. In MDPs (or families of MDPs) with large state spaces, it can be unreal- istic to assume that all states are visited at least once, 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) minggé(S)).e DDraa This sample complexity bound relies on a Hoeffding-style inequality, and therefore requires that the samples D be independent, which usually not the case when trajectories are sampled.
Theorem 4 (Sample Complexity). For any & which defines an (ER, eT, cg) approximate bisimulation abstraction on HiP-BMDP family Me, we define the empirical measurement of QiMg over D be
Drg.. Then, with probability 2 1- -,
10. 1QigglMo Ma snn Contae r(er co) Riaax
2(1 77) (1 >)- (m) log 21(((*)) 5
Pemax
(4)
This performance bound applies to all tasks in the family and has two terms that are affected by using state abstraction: the number of samples (D)) and the size of the state space (o(X)) We know that (x()) < |X| as behaviorally equivalent states are grouped together under bisimulation, and
ns (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.
4 EXPERIMENTS & RESULTS
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-reinforreement learning settings.

---

Published as a conference paper at ICLR 2021
We consider two setups for evaluation: i) 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 AAppeddix 1) and sample videos of policies at https://sites: http://g....e...olelll... wtt.ccmm/iiiiiee hwtm:/iiew/hbii wwweew/hhppbmaee bmdp
Environments. We create family of MDPs using the exist- ing environment-task pairs from DMC and change one environ- ment parameter to sample different MDPs, We denote this pa- rameter as the perturbation-parannn
We consider the follow- ing HiP-BMDPs: 1. Cartpole-Swi Sartppole-Swinuu--V.. the mass of the pole varies, 2. Cheetah-Run-Vo the size of the torso varies, 3. Malker-Run-VO: the friction coefficient between the ground and the walker's legs varies, 4. Malker-Run-- v1: the size of left-foot of the walker varies, and 5. Finger-Spinn Spren--Spi---O the size of the finger varies. We show an example of the different pixel observations for Cheetah- Ren---nnn- VO in Figure Additional environment details are in Appendix D
Figure 3: Variation in Cheetah-Run-V0 tasks.
We sample 8 MDPs from each MDP family by sampling different values for the perturbation- parameter. The MDPs are arranged in order of increasing values of the cetpurrbatio-ppaammett 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, F) are used for evaluating the model in the interpolation setup (i.e. the value of the perturbation -artrintt-attteeerre can be obtained by interpolation). MDPs (A, H are for evaluating the model in the extrapolation setup (i.e. the value of the crnrturrattio-ppaaameee can be obtained by extrapolation). We evaluate the learning agents by computing average reward (over 10 episodes) achieved by the policy after training for 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 multi-task setup where the agent is trained on four related, but different environments with pixel observations. We compare our method, HiP-BMDD with the following baselines and ablations: DeepMDP (Gelada et al. 2019) where we aggregate data across all training environments, it) JiP-BMDD-nubissim HiP-BMDP without the task bisimulation metric loss on task embeddings, iii) Distral, an ensemble of policies trained using the Distral algorithm (Th eet al. 2017) with SAC-AE (Yarats et al. 2019) the underlying policy, iv) PCGrad (Yu et al.) 2020). and v) GradNorm (Chen et al. 20188. For all models, the agent sequentially performs one update per environment. For fair comparison, we ensure that baselines have at least many parameters as HiP-BMDP. Distral has more parameters as it trains one policy per environment. Additional implementation details about baselines are in Apppendix D.I
In Figures 4 and TO (in Appendix), we observe that for all the models, performance deteriorates when evaluated on interpolaioonlexrraapltiion 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- the task-embeddings is not guaranteed to improve performance in all the environments (as can be seen the case of aaneetahh-u---- Distral, PCGrad, and GradNorm sometimes lag behind even the DeepMDP baseline, perhaps because they do not leverage a shared global dynamics model.
also uses task embeddings. Moreover, only incorporating
We also note that the multi-task learning baselines like
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 time- intensive to train pixel observations directly. Specifically, 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 qa (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 (bluc) converges faster to a
PEARL (Rakelly et al. 2019).
7

---

Published as a conference paper at ICLR 2021
ker-Bun--Vo
Walker-lun-V1
Finger-Soi--i
Ipleodi: Reward
6c0
Sphooon Rawward
4.u. 200-
430
200
60
0z
64
0.0
0.8
6.0
62
04
03
52
d4
Steps
(1006)
Stups
HiP-BMDP
HiP-BMDP-nobisim
DeepMDP
Distral
GradNorm
PCGrad
Figure 4: Multi- -ask Setting. Zero-shot generalization performance on the extrapolation tasks. We see that method, BMDP performs best against all basclines across all environments, Note that the environment steps (on the x-axis) denote the environment steps for each task. Since we are training over four environments, the actual number of steps is approx. 3.2 million.
threshold reward (green) than the baseline for Seartpole-Swingure and Wallker-Walk--i We provide additional results in Appendix E
Walker-Run-VI
Evaluating the Universal Transition Model. We investigate
Modal I Error (100 stop)
HP--MOP
how well the transition model performs in an unseen environment by only adapting the task parameter 6. 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
HP-3MDP-nolim
0 parameter, and the transition model is evaluated by unrolling the transition model for k-steps. We report the average, per- step model error in latent space, averaged over 10 environments. While we expect both the proposed setup and bascline setups to adapt the new environment, we expect the proposed setup adapt faster because of the exploitation of underlying structure. In Figure [6] indeed observe that the proposed HiP-BMDP model adapts much faster than the ablation JJP--MMDP-nobisim.
10
Number Updates ('10"3)
Figure 6: Average per-step model error (in latent space) after unrolling the transition model for 100 steps.
Relaxing the Block MDP Assumption. We incorporate sticky observations into the enviromment to determine how HiP-BMI -MP-BMDP behaves when the Block MDP assumption is relaxed. For some probability p (set to 0.1 in practice), the current observation 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.
RELATED WORK
Multi-task tast 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 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, 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 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 eet 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 constructed state-action abstraction that assumes
Caripols Swingup VO
Walker Walk V]
Cartool: Swhgup VI
Walker Walk VI
Filorde Rewwerd
250
25.3
3c0
Esioodi: Iteqeets
220
200
1so-
SVDP
I0o
MP BHEP
HP SND
30
Train Maps 30
Train Steps 10 (3)
Train steps 1 10"3)
Tran Steps
Figure 5: Fhw-shot generalization performance on the interpolation (2 left) and extrapolation (2 right) tasks.
Green line shows a threshold reward. 100 steps are used for adaptation to the evaluation environments.
8

---

Published as a conference paper at ICLR 2021
Finger-Spin-Vo
Finger-Spin-V0
800
Episodic Reward
600
Episodic Reward
600
p=0.01
400
p-0.02
400
p-0.05
200
p=0.1
p=0.2
0-5.5
0.0
0.2
0.6
0.8
0.0
0.1
Environment Steps (*10~6)
Environment Steps (*10^6)
HiP-BMDP
HiP-BMDP-nabiiim
DeepMDP
Distral
GradNorm
PCGrad
Figure Zero-shot generalization performance on the evaluation tasks in the MTRL setting with partial observability. HiP-BMDP (ours) consistently outperforms other baselines (left). We also show decreasing performance by HiP-BMDP as p increases (right). 10 seeds. 1 stderr shaded.
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 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 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 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. (2119) model and learn latent contexts upon which universal policy conditioned. However, no explicit assumption of universal structure is leveraged. Amit & Meir (2018); Yin et al.] (2020) give 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 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 means to gen- eralize across tasks with a common structure. Abbll et al.] (2018) learn transitive and PAC state abstractions for distribution over tasks, but they concentrate finite, tabular MDPs. One approach to form such abstractions is via bisimulation metrics (Givan et al. 2003) Ferns et al.) 2004) which formalize conerete way group behaviorally equivalent states. Prior work also leverages bisim- ulation for transfer (Castro & Precup 2010), but on the policy level. Our work instead focuses on learning latent state representation and established theoretical results for the MTRL setting. Recent work (Gelada et al., 2019) also learns latent dynamics model and demonstrates connections to bisimulation metrics, but does not address multi-task learning.
6 DISCUSSION
In this work, we advocate for new framework, HiP-BMDP, 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 6, which captures the behavioral similarity across tasks. We provide error and value bounds for the HiP-MDP (in appendix) and HiP- HiP-BMDP settings, showing improvements in sample complexity over prior work by producing bound that depends on the number of samples aggregate over tasks, rather than number of tasks seen training time. Our work relies on an assumption that we have access to 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 _ norms for measuring error and the value and transfer bounds. In future work will investigate tightening these bounds with I.P norms.

---

Published as a conference paper at ICLR 2021
REFERENCES
David Abel, Dilip Arumugam, Lucas Lehnert, and Michael Littman. State abstractions for lifelong reinforcement learning. In International Conference on Machine Learning, pp. 10-19, 2018.
Ron Amit and Ron Meir. Meta-learning by adjusting priors based on extended PAC-Bayes theory. In International Conference Machine Learning (ICMI., pp. 205-214, 2018.
Haitham Bou Ammar, Eric Eaton, Paul Ruvolo, and Matthew Taylor. Online multi-task learning for policy gradient methods. In International conference on machine learning, pp. 1206-12111 1214, 2014.
Richard Bellman. Dynamic Programming. Princeton University Press, Princeton, NJ, USA, edition, 1957.
Dimitri P. Bertsekas and John N. Tsitsiklis. Neuro-Dynamic Programming. Athena Scientific, 1st edition, 1996. ISBN 1886529108.
Emma Brunskill and Lihong Li. Sample complexity of multi-task reinforcement learning. Uncertainty in Artificial Intelligence Proceedings of the 29th Conference, UAI 2013, 09 2013.
Daniele Calandriello, Alessandro Lazaric, and Marcello Restelli. Sparse multi-task reinforcement learning. In Z. Ghahramani, M. Welling, C. Cortes, N. D. Lawrence, and K. Q. Weinberger (eds.), Advances in neural information processing systems 27, pp. 819-827. Curran Associates, Inc., 2014.
Pablo Samuel Castro and Doina Precup. Using bisimulation for policy transfer in mdps. In Twenty- Fourth AAAI Conference on Artificial Intelligence, 2010.
Zhao Chen, Vijay Badrinarayanan, Chen- Yu Lee, and Andrew Rabinovich. Gradnorm: Gradient normalization for adaptive loss balancing in deep multitask networks. In Jennifer G. Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning. ICML 2018, Stockholmasmässan, Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pp. 793-802, PMLR, 2018. URL httt:::/ppoooeeeiiiiiii ww.press//8) chen Ba htmm
Carlo D' Eramo, Davide Tateo, Andrea Bonarini, Marcello Restelli, and Jan Peters. Sharing knowledge in multi-task deep reinforcement learning. In International Conference on Learning Representa- tions, 2020.
Finale Doshi-Velez and George Konidaris. Hidden Parameter Markov Decision Processes: A Semi- parametric Regression Approach for Discovering Latent Task Parametrizations, 1308:3330 [cs], August 2013. arXiv: 1308.3513.
Simon S. Du, Akshay Krishnamurthy, Nan Jiang, Alekh Agarwal, Miroslav Dudík, and John Langford. Provably efficient RL with rich observations via latent state decoding. CoRR, abs/1901 100011,, 2019.
Norm Ferns, Prakash Panangaden, and Doina Precup. Metrics for finite markov decision processes. In Proceedings of the 20th Conference on Uncertainty in Artificial Intelligence, UAI '04, pp. 162-169, Arlington, Virginia, United States, 2004. AUAI Press. ISBN 0-9749039-0-6.
Norm Ferns, Prakash Panangaden, and Doina Precup. Bisimulation metrics for continuous markov decision processes. SIAM J. Comput., 4066):1662 December 2011. ISSN 0097-5397. doi: 10.1 10.1137/10080444X.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In Proceedings of the 34th International Conference on Machine Learning- Volume 70, pp. 1 1135. JMLR. org, 2017.
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Mare G Bellemare. Deep- mdp: Learning continuous latent space models for representation learning. arXiv preprint 12Xi::1906.02766. 2019.
Robert Givan, Thomas Dean, and Matthew Greig. Equivalence notions and model minimization in markov decision processes. Artificial Intelligence, 147(1-2):163-223 2003.
10

---

Published as a conference paper at ICLR 2021
Tuomas Haamoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. volume 80 of Proceedings of Machine Learning Research, pp. 1861-1870, Stockholmsmässan, Stockholm Sweden, 10-15 Jul
2018. PMLR. URI. htt.htttpp//pcceeeeeeeee
ngs ml wrress//80/haaanojalll... ht.m
Nan Jiang. Notes on State Abstractions. pp. 12, 2018. URL https::/nanjiant htt.a/ois.eddeeeeeeese//////////
Nan Jiang, Alex Kulesza, and Satinder Singh. Abstraction selection in model-based reinforcement learning. International Conference on Machine Learning, pp. 179-188, 2015.
Leslie Pack Kaelbling, Michael L. Littman, and Anthony Cassandra. Planning and acting in
partially observable stochastic domains. Artif. Intell., 101(1-2):99-33, May 1998. ISSN 0004- 3702.
Nicholas C. Landolfi, Garrett Thomas, and Tengyu Ma. model-based approach for sample-efficient multi-task reinforcement learning, 2019.
Lihong Li, Thomas J. Walsh, and Michael L. Littman. Towards a unified theory of state abstraction for mdps. In Proceedings of the Ninth International Symposium on Artificial Intelligence and Mathematics, pp. 531-539, 2006.
Yuping Luo, Huazhe Xu, Yuanzhi Li, Yuandong Tian, Trevor Darrell, and Tengyu Ma. Algorith- mic framework for model-based deep reinforcement learning with theoretical guarantees, In International Conference on Learning Representations, 2019.
Andreas Maurer, Massimiliano Pontil, and Bernardino Romera-Paredes. The benefit of multitask representation learning. J. Mach. Learn. Res., 17(1):283322884 January 2016. ISSN 1532-4435.
Rémi Munos. Error bounds for approximate value iteration. In Proceedings of the 20th National Conference on Artificial Intelligence Volume 2, AAAI'05, pp. 1006-1011. AAAI Press, 2005.
ISBN 157735236x.
Alfred Müller. Integral probability metrics and their generating classes of functions. Advances in Applied Probability, 29(2):429-44 443, 1997. doi: 10.2307/1428011.
Emilio Parisotto, Jimmy Ba, and Ruslan Salakhutdinov. Actor-mimic: Deep multitask and transfer
reinforcement learning. In ICLR, 2016
Christian F. Perez, Felipe Petroski Such, and Theofanis Karaletsos. Generalized Hidden Parameter MDPs Transferable Model- based RL. in Handful of Trials. arXiv:2002..3002 [cs, stat], February 2020. arXiv: 2002.03072.
Martin L Puterman. Markov decision processes: Discrete stochastic dynamic programming. Journal of the Operational Research Society, 1995.
Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient off-policy metta-reinforeement learning via probabilistic context variables. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.). Proceedings of the 36th International Conference on Machine Learning. volume 97, pp.
Long Beach, California, USA, 09-15 Jun 2019, PMLR.
Jonas Rothfuss, Dennis Lee, Ignasi Clavera, Tamim Asfour, and Pieter Abbeel. ProMP: Proximal meta-policy search. In International Conference on Learning Representations, 2019.
Yanchao Sun, Xiangyu Yin, and Furong Huang. Temple:
Learning template of transitions for sample
efficient multi-task rl, 2020.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, Timothy Lillicrap, and Martin Riedmiller. DeepMind control suite. Technical report, DeepMind, January 2018.
Yee Teh, Victor Bapst, Wojciech M. Czameckii John Quan, James Kirkpatrick, Raia Hadsell, Nicolas Heess, and Razvan Pascanu. Distral: Robust multitask reinforcement learning. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in neural information processing systems 30, pp. 1496-4500
Curran Associates, Inc., 2017.
11

---

Published as a conference paper at ICLR 2021
Andrea Tirinzoni, Riccardo Poiani, and Marcello Restelli. Sequential transfer in reinforcement learning with a generative model. In Proceedings of the 35th International Conference on Machine Learning, ICML 2020, Proceedings of Machine Learning Research. PMLR, 2020.
Denis Yarats, Amy Zhang, Ilya Kostrikov, Brandon Amos, Joelle Pineau, and Rob Fergus. Improving sample efficiency in model-free reinforcement learning from images. 2019.
Mingzhang Yin, George Tucker, Mingyuan Zhou, Sergey Levine, and Chelsea Finn. Meta-learning without memorization. In International Conference on Learning Representations, 2020.
Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn. Gradient surgery for multi-task learning, 2020.
Amy Zhang, Zachary C. Lipton, Luis Pineda, Kamyar Azizzadenesheli, Anima Anandkumar, Laurent Itti, Joelle Pineau, and Tommaso Furlanello. Learning causal state representations of partially observable environments. The Multi huuli--isciplinary Conference on Reinforcement Learning and Decision Making, 2019.
Amy Zhang, Clare Lyle, Shagun Sodhani, Angelos Filos, Marta Kwiatkowska, Joelle Pineau, Yarin Gal, and Doina Precup. Invariant causal prediction for block mdps. International Conference on Machine Learning (ICML), 2020a.
Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. arXiv preprint arXiv:2006. /0742, 2020b.
12

---

Published as a conference paper at ICLR 2021
A BISIMULATION BOUNDS
We first look at the Block MDP case only (Zhagg et al. 2020a), which can be thought of as the single-task setting in HiP-BMDP. We compute approximate error bounds in this setting by denoting & an (cn, (areerr)app >-approximate bisimulation abstraction, where
sup
[( a))2,,))
A ( (
sup A (m)))))())))))m)))))))))))) (((a))
TT denotes the lifted version of T, where we take the next-step transition distribution from observa- tion space X' and lift it to latent space S.
Theorem 5. Given an MDP M built on (ER,ET) -app-rx-ixi bisimulation abstraction of Block MDP M, denote the evaluation of the optimal Q function of M on M [Qy]M. The value
difference with respect to the optimal Q'M is upper bounded by
Rymax
|0M- - 11 SE + 7ET 2(1 77)
Proof. From Theorem 2 in Jiang (2018).
B THEORETICAL RESULTS FOR THE HIP-MDP SETTING
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 family of environments Me, we bound the difference in expected value between two sampled MDPs, Ma..-M E Me using ((0.,6,). Additionally, we make the assumption that we have behavior policy that is near both optimal policies zे. TØ.- We use KI. divergence to define this neighborhood for Tz.
KLL (T, = Es-p- [KL(T(-ls), T7. (21
(5)
We start with a bound for specific policy ग. One way to measure the difference between two tasks Ma., MM, 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 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 n, the difference in expected value between two MDPs drawn from the family of MDPs Ma,,Mo, MM bounded by
IV
Vē <
100 0,111.
(6)
Proof. We use telescoping sum to prove this bound, which is similar Luo et al. (2019). First, we let Z₁ denote the discounted sum of rewards if the first & steps are Ma,, and steps / are in Me,
Σ S R(ss an) 1=
Z : E
W20,a ( (s):a))
W2j (s),
By definition, we have Zoo Vi and Zo VoJ. Now, the value function difference can be written telescoping sum,
V
Σ
(7++1 2x).
(7)
k=0
13

---

Published as a conference paper at ICLR 2021
Each term can be simplified to
Zx+1 Zs = 7 k+t E.
Esy- 1 Toy (1ss,aa)) [V, (sk+1) - Voj (sk+1] sken Io, (llenaa))
as wxt To:
Plugging this back into Equation (7)
VØ V 0
7
E smp8 aw)((s)
Vo, Bo (1a.0) Vo, (5)) -uytTy (la,a) Vog
1
This expected value difference is bounded by the Wasserstein distance between To, Tog
? W W(To, To,)
IV
Voll
1-7
1 110: - 0z1l1 using Equation (1)
Another comparison to make how different the optimal policies in different tasks are with respect to the distance 100, -0yl.
Theorem 7. The difference in expected optimal value between two MDPs Mo. Mo, E MM is bounded by.
7 VV Vol S (1 (1-7)21166 - -0y112-
(8)
Proof.
[V) - Vo,(s)
max Q7, a) max Qū, as))
a
S max Q7, (s) - Qū, (s,aa)
We can bound the RHS with
sup 1Q6 a) - (s) S sup |ra, (8, -),, (s. a)) sup sup |Eg a,
To, (1l,,0) (5)) -Boy (11,,) Voj (6")
8,a
s,a
All MDPs MM have the same reward function, so the first term is O.
sup
@7 (s, a) - Qō, (s, a)l S y sup E EgeTo (1 Vo, (s') Ear (-la,)) VoJ (s")
sup
(1 (1 Vo'(')) - [s) + B,,'TTo (1),,)) Vo,(8') - [s)]
8,
s'nTo, (-1,,a)
suu
Ey (11.,0) (s'(s)) - [5)(1 + + sup E."TTo (11,,))
[Vo (8') - Vēj (s")]
8,
8,0
T' (-13,a)
Sysup
E. Ey (119,0) (((())) - (s']] + -- 0 0 7
7
8,a
S max
Vo (8) Vor (8) + '- 110 - 0z111
max
max Qa, (s,a) a) max Q2, (s, ) + 1- 1|: 0 ll2
?
Q7, (s, a) Qi, (s, + 1- 10: - 0
7
sup
saa
Solving for sups,a 10, (s,a) - Qū (s, a)):
sup
Q7, (s,a) a) - Qū (s, 0) - (1 (1-
7 11, 01111
8,0
Plugging this back in,
7
(T((s) VV (s)1 5
(1 7)2 10, -0y11.
14

---

Published as a conference paper at ICLR 2021
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.
B.2 EXPECTED ERROR BOUNDS
In MTRL, we are concerned with the performance over family of tasks. The empirical risk is typically defined as follows for I tasks (Maurer et al. 2016):
Earg (0)
Σ ((((F((((;;()))) Y)M
(9)
T
Consequently, we bound the expected loss over the family of environments & with respect to In particular, interested in the average approximation error and define it the absolute model eиror averaged across all environments:
8
Cavg (0)
Vo (8) - Vo,
(10)
E
Theorem 8. Given a family of environments Me, each parameterized with an underlying true hidden parameter 0,, .,.... Og, and let 0,, 02,.... 0g be their respective approximations such that the average approximation error across all environments bounded as follows:
67
Cang (0) S
(11)
7)2
where each environment's 's parameter 0, e-close to its approximation 0, i.e. d(O,, 0.) S e, where is the distance metric defined in Eq. 7
Proof. We here consider the approximation enror averaged across all environments as follows:
Eaug (0)
(x)
3
i=1
1
Eang (0)
max Q7 (8,a) max Qù, (8, a')l
3 i=1
(12)
max 10 (8,a) a) Qô, (s, a)
Let us consider environment 0, E Me for which we can bound the RHS with
sup (8, -a, 3.a
( (8, S sup Ira, (8, a) ro, (8, ())aalyysuu sup EN (-1s.)) Vo, (6')-EgrTTh, (Hl,a) Vo, (s")1 8.0
8.a
Considering the family of environments Me have the same reward function and is known, resulting in first term be O.
15

---

Published as a conference paper at ICLR 2021
sup 16. (8,a)
Qú, (8,
suup E B (1e,0) Vo((5') To (11..) (s")
7 sup s,a
Eg BeinTa (18,,a) (6)(()) [2) + ==,TTT, (13,a)) 5' s'~To, (11,0)
Vo (5)('))
V2 [s"]]
S sup
BunTT (1),a) (V((())) (oT + s To
[/((('))) [s)
5,8
sia
8 s''i' (1b,,))
S ssu E BernTy (11,,a) (6 - Vo ?
00 0.l
1-7
8,a
S y max
9. (s) Vi [s) + 1- 0 - 0,1
2
max 02 (s,a) max 2ú, (s, + --7 10: 0,1
7
max
a'
S y sup 0% (s. a) 26, (s, ?
10 -
1-
Solving for sups,a 14 (8, a) - Qš (8, a)l
sup 02 (s, a) Qa. (s. (s. 5 (1 (1- 10: 0.1
7
(13)
5,
Plugging Eq. 13 back in Eq. 12
8
Eang (0) <
10: 0:1
7)2
1=1
10-11 0i=1l + 10:2 01-21 + ... +.+ 10:-0 0z=el|
E(1 7)3
We now consider that the distance between the approximated 6 and the underlying hidden parameter E MM defined as in Eq. such that: d(O,, 0.) S ED
Plugging this back concludes the proof,
Eavg (0) 5 (1 7)
7/8
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 0 for each environment.
C ADDITIONAL RESULTS AND PROOFS FOR HIP-BMDP RESULTS
We first compute L norm bounds for Q error under approximate abstractions and transfer bounds.
Theorem (Q error), Given an MDP Má built on (cn, <r, (0) apprrxxmiee bisimulation abstrac- tion of an instance of a HiP- BMPD Mg. we denote the evaluation of the optimal Q function of M on M (QMs,MM. The value difference with respect the optimal Qu upper bounded by
Rona
l|euo 1Qi,, sen + 7 60)
2(1 7(1-))
Proof. In the HiP-BMDP setting, we have a global encoder & over all tasks, but the difference in transition distribution also includes 0. The reward functions are the same across tasks, there is
16

---

Published as a conference paper at ICLR 2021
no change to eg. However, now must incorporate difference in dynamics in ey. Assuming have two environments with hidden parameters 0,, 0g O we can compute CT 0:,03 across those two
environments by joining them into super-MDP.
TT
0.0,
sup
(TT,,,1,,)) TTo, (22, a))2
aEA. #1 22 EX (0(1)=
||0To, (21, 0) TTo, (z2, a)|(, + +To (22,
TTo, (12, (x).)))
sup aEA. #1 E (0(1))
sup
$To (II, &To, (12,
sup ac
(To, (22, a) TTo, (22, a))l
aEA4. 21,226X (((2))))))))))
67' + |00 03
This result intuitive 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) the HiP-BMDP setting by denoting approximation error of 0 as |10 - Ø|| < Eg.
Theorem 4 For any & which defines (en, er, ca) apprrxxmmaee bisimulation abstraction on HiP-BMDP family Me. we define the empirical measurement of Q'A, over to be Q Tap. Then,
with probability 2 1 5,
Rmax
Rmax
2) (14)
10 - -MinglMullo S en 7 t60)
+ 2(1 -7) (1 (1-7)2
log
7)2
2n²(D)
8
Proof.
l10', 1Q'ugl - $1Qin - 1QM 11 + 1100M,lm - -Qixg loo
Qiu. - -19 1l + 1192, Piuzlla
The first term is solved by Theorem 2, only need to solve the second term using McDiarmid's inequality and the knowledge that the value function of a bisimulation representation is - Lipschitz from (remreee I
First, write this difference to be deviation from expectation in order to apply the concentration inequality.
110²a, Diggll> = =10 - TBQTs, + +TQQi Tiquzllo s 110 - -T7Qin, 1100 + +1110mg - Vinglloo 100
1
< 1-7 -
-T0Q2,,1 100
Now we can apply McDiarmid's inequality.
Pp [ O'apl 2 < 2 exp
(- 2t Dostt)
R2 12 (11 -2)
Solve for the t that makes this inequality hold for all (0(z), a)) E X x A with a union bound over all |o(x)]|A) abstract states,
Raux t 1 - r V 2n)(D)
1 216(21)
log
Combine to get
Rmux
R
2))((')114
l1Quo - 1Qiggluolle S Sca + +c0)
+ (1 7)2 2n²(D)
log
2(1 7)
8
17

---

Published as a conference paper at ICLR 2021
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 vo to generated task-specific 0 parameters.
1: for each timestep 1..7' do
2: for each Ti do
3: aj (-(s)
4:
p' ('18,, at)
5:
D D (8); dz r(s at), ,
6:
UPDATECRITIC(D, (uses data only from ith task) 7: UPDATEACTOR(D. (uses data only from jth task)
8:
UPDATEUSIINGHI--MDDPP BMDPLOss(D,))
9: end for
10: end for
Algorithm 2 UpdlateWoodeUUiiggiii----M
Require: Batches of data for the different tasks (Ti)-=...T sampled from the Replay Buffer D,
learning rates and (2,, index of the current task i, Transition Model M, environment encoder
d).
1: for each batch of dataset t = L..T, t do
2: Compute L(vb, M) = (zb,,, M,i,t using Equation (3)
3:
d 4- 01 V Σ C
4: M + M 02 V 2,
5: end for
D ADDITIONAL IMPLEMENTATION DETAILS
In Figure X
show the variation in the left foot of the walker.
Figure 8: Variation in Walker (VI) across different tasks.
MTRL Algorithm The Multitask RL algorithm for the HiP-BMDP setting can be found in Al gorithm 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 algorithh for the HiP-MDP setting can found in Algorithm 3 We take the PEARL algorithm (Rakelly et al. 2019) and incorporate HiP-MDP objective (text shown in red color)
D. BASELINES
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 all the experiments. For GradNorm,
observe that the learned weights w, (for weighing tar-task loss) can become negative for some tasks (which means the model tries to unleam those tasks). In practice, we clamp the we values not become smaller than threshold.
18

---

Published as a conference paper at ICLR 2021
D.2 HYPER PARAMETERS
D.2.1 MTRL ALGORITIIM
All the hyper parameters (for MTRL algorithm) are listed in Table
Parameter name
Value
Actor learning rate
10
Actor update frequency
2
Actor log stddev bounds
[-10, 2]
Alpha learning rate
10-4
Batch size
128
Critic learning rate
10-3
Critic target update frequency
Critic Q-function soft-update rate TQ
0.01
Critic encoder soft-update rate Tene
0.05
Discount y
0.99
Decoder learning rate
10-3
Critic update frequency
Encoder feature dimension
100
Encoder learning rate
10
Number encoder layers
Number of filters in encoder Hidden Dimension
32
1024
Replay buffer capacity Optimizer
1000000
Adam
Temperature Adam's 's 31
0.5
Init temperature
0.1
Number of embeddings in vb Embedding dimension for z/
10
100
Learning rate for zb
10 10-3
ag (ic for (-) for Finger-Spin-Ve VO
0.1
Og (ic a for (/) for Cheetah-Run VO
0.1
Orp (ie for 1/) for Walker-Run-Vi
1.0
Otg (ie a for ()) for Walker-Run-Vi
0.01
Odissaal (ie for Distral) for Finger-Spin-VO
0.05
Odistaal (ie for Distral) for Cheetah-Run-Ve
0.01
Odissaal (ie for Distral) for Walker-Run VO
0.05
Odistall (ie for Distral) for Walker-Run,
0.01
Baistal (ie for Distral)
1.0
Table 1: A complete overview of used hyper parameters.
D.2.2 METARL ALGORITIM
For MetaRL, we use the same hyperparameters as used by PEARL Rakelly et al. (2019). We set
ag 0.01 for all environments, other than Walker-Stand environments where as 0.001.
E ADDITIONAL RESULTS
Along with the environments described in 4 we considered the following additional environments:
1. Wa Halker-Sttande -iker--ttaddoo: Walker-Stand task where the friction coefficient, between the ground and the walker's leg, varies across different environments.
2. Wa ker--elll k-v0: Mal ker-Wal k task where the friction coefficient, between the ground and the walker's 's leg, varies across different environments.
3. Walker-St. and-Vi1 Walker-Stand task where the size of left-foot of the walker varies across different environments.
19

---

Published as a conference paper at ICLR 2021
Algorithm 3 HiP-MDP training for the meta-RL setting.
Require: Batch of training tasks {Ti)i=....r from p(T), learning rates 01, 02, 0:3y 0g Initialize replay buffers B for each training task
2: while not done do
3: for each T. do
4: Initialize context C* (
5: for k = =, .,.., do
6:
Sample ~ (z))
7:
Gather data from No (als,, z) and add to B
8:
Update C {(s), aj. S/ rj)lj..... ~ B
9: end for
10: end for
11: for step in training steps do
12:
for each Ti do
13:
Sample context C Sa(Ba) and RL batch D' B
14:
Sample z ~ (z)))"
15:
gattor = Lactor (b',
Corite Coritie (bC
17:
LKL = 3DKI ([)((())))
18:
Sample a RL batch 6 from any other task j
19:
Compute Lisssim (), (,,5) using the equation end for
20:
21:
4 Φ 01 V Ex CCystii fKL + as Bisss)
0₂ + 0₂ a2 V Σ, Cactor
23:
0o 0Q a3 V E, Cerittic
24:
end for
end while
4. Walker-Walk-Vi: Walker-Walk task where the size of left-foot of the walker varies across different environments.
E.1 MULTI-TASK SETTING
In Figure 10 we observe that the HiP- -M-DP method consistently outperforms other baselines when evaluated on the interpolation environments (zero-shot transfer). As noted previously, effectiveness of our proposed model can not be attributed to task task-embeddings alone HiP-BMDP. model but does not include the task
nobisim model uses the same architecture as the bisimulation metric loss. We hypothesise that the Ensemble bascline behaves poorly because cannot leverage a shared global dynamics model.
Cheetah-Run-vn
Walker-Run-va
Walker-Run-VV
Finger-Spin-40
2000
Eppeste sumer
Sonn So sze- Se
Eppoodiic
480 200
22 04
d4
08
62
05
s's Enviroment Steps ("10"s)
62
54
Enviroment Steps (100b)
Envionmant Steps (1los)
Enveonment Stops ("iog)
HiP-BMDP
HHP-BMDD-nobisim
DeepMDP
Distral
GradNorm
PCGrad
Figure 9: Multi-Task Setting. Performance on the training tasks. that the environment steps (on the x-axis) denote the environment steps for each task. Since are training over four environments, the actual number of steps is approximately 3.2 million.
E.2 META-RL SETTING
We provide the Meta- RL results for the additional environments. Recall that extend the PEARL algorithm (Rakelly et al. 2019) by training the inference network (z(z)) with our additional HiP.
BMDP loss. The algorithm pseudocode can be found in Appendix D) In Figure 13] show the
20

---

Published as a conference paper at ICLR 2021
C Cheettah-Bun-Vil
Walcer- Run--V
Wilker-Run-Y
Inger-Spinvvl
Biizd Rovord
533
Eplcud Rovard
100
100
233
Enviromment Stops (10
Emvianment Steps (110n))
Enstronment Stops (i00g)
Envirahment Stees (1106))
32
HiP-BMDP
HIP-BMDP-nobisim
DeepMDP
Distral
GradNorm
PCGrad
Figure 10: Zero-shot generalization performance the interpolation tasks in the MTRL. setting. HiP-BMDP
(ours) consistently outperforms other baselines. 10 seeds, standard error shaded.
8 g
Moddl Erw (5 2000 5
MOdel Emrents sterl
Model EnT C stap)
10
Number of Updates (11
Number Uppittis (*
Number bet Updates (*10"3)
Number Updatees (*10
Figure 11: Average per-step model error (in latent space) after unrolling the transition model for steps.
results on the interpolation setup and in Figure 14 we show the results on the extrapolation setup. In some environments (eg Walker-Ward VI), the proposed approach (blue) converges faster threshold reward (green) than the baseline. In the other environments, the gains are quite small.
E.3 EVALUATING THE UNIVERSAL TRANSITION MODEL.
We investigate how well the transition model performs an unseen environment by only adapting the task parameter 8. We instantiate new MDP, sampled from the family of MDPs, and use behavior policy to collect transitions. These transitions are used to update only the 0 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 and 100 steps respectively. While 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 step and 100 step unrolls, the proposed HMP-BMDP model adapts much faster than the baseline HiP- BMDP-nobin (Figures [1 and 12)
21

---

Published as a conference paper at ICLR 2021
Veltore-uu-V1
Whike-Bun-20
Cheetah Cheetah-Rum-30 Rum-30
Modes Frre (100 orep
Model Enor (100 2000
Model Eree 0100 wla
20.
Model Enner (L000 stagt
1603
12
623
Hin eootteo
12.5
16
Nimber of Uplittes (110
Nu Rumber of Updatess 1110"31
N.mber Uptones -11033
kumber of Undates (110"3)
Figure 12: Average per-step model error (in latent space) after unrolling the transition model for 100 steps.
Walker
Walve
Run
Ep-ooit Rewtrr
Eprcott
Epccottc
SMOP
SMOF
#P BND
Tra Steps 1010 -3)
Train Steps ( -3)
Train Steps 10 *3)
Train Steps (* 10 "3)
Walcendiur--it
Waller Nalker/hke l-Vi
430
Walber andd-V
Eisoeele Rewwed
Editodlic Rewerd
3ND*
226 Troin Stops 1" 10 -31
280
Train Steps e 10 -3)
Tral Stops 10 "3)
Figure 13: Few-shot generalization performance the interpolation tasks Walker-Run- VO, Walker-Walk- VO Walker-Stand-Ve Cheetah-Run-Ve (top row), Walker-Run-Vi, Walker-Walk-Vi and Salker-Stand-VI (hottom row) respectively. note that for the Walker-Wall V1, the proposed approach (blue) converges faster to a threshold reward (green) than the baseline. In other environments, the gains are quite small.
Walkersun--v
400
Walcer-Sand-13
Chestat-Bun-a
Entodte Reward
302
130
100
Epcodii
4P
.400 ebsi
26 7a0
Train Stops 100 "3)
Train Stops 10 "3)
Train Stops 10
Train Train 10 "3)
Walcerdtune1
We WoleerWW loy!
Welbee-Sandaat
rpinodie Rewwnrs
Eattoel Roveers
430
337
200
150
RMAA
FMARI
PTARI
525
1050
1906
25.0
100
901
dic
165
466
061
853
1550
1209
S0eps
-3)
Train Slaps
-31
Staps
Figure 14: Few-shot generalization performance on the extrapolation tasks on Walker- Run--RRR--- Walker- Walk- VO, Walker-Stand-vo Cheetah-Run-V (top row), -un-er- Walker Walk-r- and Walker Stanker-Stand--ee V1 (bottom row) respectively. We note that for the Walker Walker-Walk-V. V1, the proposed approach (blue) converges faster a threshold reward (green) than the baseline. In other environments, gains are quite small.
22


Reviews:
[{"review_id": "k6JJrFGW0Gu", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\nThis paper combines representation learning via approximate bisimulation with the HiP-MDP framework for representing multiple tasks. The result is multi-task and meta-RL algorithms that operate from images and in the meta-RL case adapt to changes in the dynamics of the environment.\n\nPros:\nCombines HiP-MDP and block MDP formalisms\nEvaluates resulting method on a set of multi-task and meta-RL problems operating from image observations\nCons:\nNot clear (to me) if any insight is gained from the theoretical analysis (the derivation of value and sample complexity bounds for approximate bisimulation was performed in Gelada et al. 2019)\nEmpirical gains are modest and it’s not clear if they are due to the image representation learning component of the loss or other aspects of the method.\n\nDetailed Comments:\nI think this paper would be relevant to cite: Learning an Embedding Space for Transferable Robot Skills (Hausman et al. 2018).\n\nOne thing that stood out to me was the very small number of training tasks used - only 4! I wonder if using more training tasks improves the generalization of single-task methods like DeepMDP that might be able to simply generalize over the dynamics changes given more examples. How much overfitting to the training tasks do you observe for these baseline methods versus HiP-BMDP?\n\nIt seems like there might be a baseline missing - generic multi-task algorithm that conditions on the environment ID but does use the bisimulation loss to help process the image observations. From my understanding, all baselines have no image representation learning component except DeepMDP which is not a multi-task or meta-RL algorithm.  I am a bit concerned that the gains we are seeing here have more to do with image representation learning than the structural assumptions employed by the algorithm. \n\nI think that for the meta-RL setting, “environment ID” is replaced everywhere by experience seen so far in the task - is that correct? It would be good to make that clear, or if environment IDs are being used in meta-RL, how they are being used.\n\nI’m a bit confused why there’s a gap between your method in the multi-task experiments but not in the meta-RL ones. Do the meta-RL experiments use image observations? If so, how is the PEARL baseline adapted to handle images? If not, why not?\n\nWhy does HiP-BMDP handle the sticky observation setting better than the baselines? Isn’t it impossible for it to persist state information over time steps?\n\nRecommendation:\nBorderline\nI think the strong point of this paper would be the theoretical contribution of connecting bisimulation to HiP-MDPs. As I’m not very familiar with bisimulation theory, I cannot comment with confidence on the value of this contribution. It is not clear to me what new insights are gained from this extension.\nEmpirically, I am not convinced that the structural assumptions on the MDP used by the proposed algorithm yield performance improvement. For example, one assumption is that the reward function won’t change across tasks, but the method doesn’t really outperform a baseline that doesn’t make that assumption. I am concerned that the gains that are observed come from image representation learning via approximate bisimulation, and that this combined with any multi-task or meta-RL algorithm might achieve the same results.\n\n---**Update**---\nIncreased score 5 -> 7 thanks to clarifications from authors.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "RQmm50Xo4a5", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Hi,\n\nFirst I want to thank authors for putting this together, I enjoyed reading it. \n\n*Summary* Authors propose mixing Block MDP and Hidden Parameter MDP, and offer and algorithm (loss function) to learn this model that they claim to be useful for sample complexity and transferability among environments (that share state and action space). The proposed method showed promising results in experiments. \n\nIn general I enjoyed reading the paper, and I believe the idea is novel, a good step toward multi-task learning and well-written. My main concern is the validity of Block MDP assumption in real world. \n1. Authors have to be more upfront on why this assumption is needed, it's hard now to exactly find why do we need Block MDP assumption, what would break if we were to relax this?\n2. Authors have an experiment for this, but I would like to see how much p- probability of sticky observation would actually affect the performance, for example what would happen if we increase p, and what point the algorithm will break down? (Basically, I'm asking how sensitive the algorithm is to Block MDP assumption)\n\nIn general I enjoyed the paper, I'd like the theory part. \nIt seems like some notational clarity can help the paper, for example what is \\Phi in the definition of \\epsilon_T? Or what is Z in definition of \\phi ? \n", "extra_scores": {}, "extra_text": {}}, {"review_id": "Q5ZyXOYR507", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nThe authors then propose algorithms to learn parameters $\\theta$ of HiP-BMDP from observed trajectories of the system. The authors derive the error bounds over between the Q-value parametrized by the learned $\\theta$ and that of the actual parameter $\\theta^*$. The authors then apply these results to the transfer learning settings where one applies the learned parameter $\\theta$ to a different but somewhat similar environment parametrized by $\\theta_i$. The error bound over Q-values in such a transfer learning settings is also provided. These results, including Theorems 2 and 3, seem sensible, while I haven't checked the details of the proof. Finally, the authors validate the proposed method through extensive simulations.\n\nOverall, I believe this paper is well written and well organized. The assumptions of HiP-BMDP seem sensible and the derived error bounds look promising. These results are applicable in multi-taks learning settings. They explicate a set of parametric conditions under which one could accelerate the learning process of a future task from prior knowledge derived from previous tasks.", "extra_scores": {}, "extra_text": {}}, {"review_id": "NIQ82AGSfpI", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "### Summary\n\nThe authors combine hidden-parameter MDPs and state abstractions to model multi-task problems with dynamics parameterized by some latent variables and where the agent receives high-dimensional observations whose corresponding low-dimensional states, on which dynamics are defined, are unobserved (as in block MDPs). They provide both a theoretical analysis and an empirical evaluation of the proposed method, showing that it performs better than competitive baselines on complex continuous control domains.\n\n### Pros\n\n- The idea of combining HiP-MDPs with state abstractions seems interesting and relevant. One of the limitations of HiP-MDPs is indeed that they do not easily handle high-dimensional observations and the proposed approach overcomes this limitation. \n- The method seems also quite general, in the sense that it can be applied in multi-task/meta-RL settings and combined with existing algorithms.\n- Experiments are conducted on complex tasks and show convincing results.\n\n### Cons\n\n- I found the paper quite hard to read in its core parts (e.g., Sec. 3). This is in part due to a complex/confused notation, which overall made it hard for me to go through the theoretical part (and proofs). See detailed comments below.\n- The theoretical results focus on assessing the value-function errors for fixed abstractions/dynamics-structure rather than on the errors in learning the abstractions/structure themselves (and learning these components seems to be one of the primary concerns here).\n- The requirement that task IDs are known could be a potential limitation for applying this method\n\n### Detailed comments\n\n1. Are the environment labels/ids used only for communicating to the learner that the task changed or could they provide more information? For instance, if we face the same task twice in two non-consecutive episodes, are the corresponding labels equal?\n\n2. What makes these state abstractions \"robust\" (as in the title)? Is there any particular theoretical or empirical evidence that justifies this term?\n\n3. Existing analyses of multi-task settings, including the one of Brunskill and Li [1] and others [2,3,4], provide guarantees for task-structure (e.g., dynamics models) learned from data, while here the focus seems to be on guarantees for arbitrary abstractions/structures as a function of their errors, without considering how they are learned. This makes it difficult to interpret the results in the context of the proposed approach\n\n4. Regarding the statement: \"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\", the bounds derived in [4] depend only on a number of abstract \"transition templates\" that they define, while those derived in [3] depend on the minimum between the number of tasks and the number of states.\n\n5. To be clear, the notation \\| x - y \\|_1 is used for the Wasserstein distance and not l1-norm, right? In such a case, the sentence \"We omit d but use d(x, y) = \\|x − y\\|_1 in our setting\" is quite confusing.\n\n6. The definitions of the terms \\epsilon_R, \\epsilon_T, \\epsilon_\\theta could be moved to the theoretical analysis since they are introduced early and never used before Sec 3.3\n\n7. Sec. 3.2: the definition of \\phi maps S to Z. Isn't it X to Z?\n\n8. In Sec. 3.2, the 2-Wasserstein distance is used, while the 1-Wasserstein distance was introduced earlier. Which one is used?\n\n9. In Eq. 3, it was not clear to me why the \"model learning error\" term uses the MSE for two different environment labels (isn't one enough)?\n\n10. What is exactly the \"evaluation of the optimal Q function of \\bar{M} on M\" in Theorem 2? Do you mean that we take the optimal policy of \\bar{M} and test it on M?\n\n11. [if the above comment is correct] Theorem 2 was slightly weird to me: it takes the optimal policy of \\bar{M}, evaluates it on M_{\\theta_i}, and check how much the corresponding Q function differs from the optimal Q function of M_{\\theta_j}. If we wanted to figure out the transfer error of the optimal policy of \\bar{M} on M_{\\theta_j}, shouldn't we compare the optimal Q function of M_{\\theta_j} with the Q function of the policy on the same MDP (rather than on M_{\\theta_i})?\n\n12. How exactly is the \"empirical measurement of Q^*_{\\bar{M}} over D\" computed?\n\n13. In Th. 4, the last term bounds the concentration of expectations of Q functions (which are bounded by rmax/(1-\\gamma). Shouldn't there be a multiplicative rmax in front of the sqrt?\n\n14. Page 6: \"We compare the our method\" -> \"We compare our method\"\n\n15. In the multi-task experiments, it was not clear what method was used in combination with HiP-BMDPs to learn policies (though it is in appendix).\n\n16. How many adaptation steps for \"few-shot generalization\" were used in Figure 5?\n\n17. Below Figure 5: \"and use a behavior policy to collect transitions.\" What is the behavior policy exactly and how was it chosen?\n\nSome minor comments/questions:\n\n- Page 1: \"This additional structure gives us better sample efficiency, both theoretically and empirically.\" Better than what algorithm?\n- Page 2: \"don't\" -> \"do not\"\n- Page 2: \"which naturally leads to a gradient-based representation learning algorithm.\" Why is it natural to use a gradient-based algorithm in this setting?\n- Background: you could mention that the finite-MDP assumption is only for deriving the theoretical results, while the method can be applied in continuous domains (right?)\n- Page 2: \"a unobservable \"-> \"an unobservable\"\n- Page 3: \"where l2 distance corresponds to d\" Why? d isn't the l2 distance, right?\n- Sec. 3.1: \\phi in \\epsilon was never defined.\n\n### Overall comment\n\nOverall I believe that the paper is interesting and the results significant. However, at the present time, I have too many doubts to vote for acceptance. I will be happy to increase my score after the authors have clarified them.\n\n### Update\n\nI have increased my score to 7 after reading the authors' response and the updated paper.\n\n### References\n\n[1] Brunskill, Emma, and Lihong Li. \"Sample complexity of multi-task reinforcement learning.\" UAI (2013).\n[2] Liu, Yao, Zhaohan Guo, and Emma Brunskill. \"Pac continuous state online multitask reinforcement learning with identification.\" AAMAS (2016).\n[3] Tirinzoni, Andrea, Riccardo Poiani, and Marcello Restelli. \"Sequential transfer in reinforcement learning with a generative model.\" ICML (2020).\n[4] Sun, Yanchao, Xiangyu Yin, and Furong Huang. \"TempLe: Learning Template of Transitions for Sample Efficient Multi-task RL.\" arXiv (2020).", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper addresses the problem of learning and exploiting common (latent) task structure in multi-task reinforcement learning settings. The authors introduce a new formalism for capturing this type of structure and derive a gradient-based learning algorithm. They provide novel theoretical insights and strong empirical results.

Reviewers initially raised several concerns, regarding assumptions and especially accessibility of the paper (and in particular theoretical discussions). The majority of these concerns have been addressed in the detailed rebuttal. The resulting consensus is to accept the paper. Authors are encouraged to continue to improve accessibility of the paper for the camera ready submission.
Decision: Accept (Poster)