Hierarchical Prototypes for  Unsupervised Dynamics Generalization in Model-Based Reinforcement Learning
Authors: Jiaxian Guo
Mingming Gong
Yali Du
Zhen Wang
Dacheng Tao
Venue: Submitted to ICLR 2023 (2023)
Keywords: Unsupervised Dynamics Generalization
Model-Based Reinforcement Learning

Abstract: By incorporating the environment-specific factor into the dynamics prediction, model-based reinforcement learning (MBRL) is able to generalise to environments with diverse dynamics.In the majority of real-world scenarios, the environment-specific factor is not observable, so existing methods attempt to estimate it from historical transition segments. Nevertheless,earlier research was unable to identify distinct clusters for environment-specific factors learned from different environments, resulting in poor performance.
To address this issue,
We introduce a set of environmental prototypes to represent the environmental-specified representation for each environment. By encouraging learned environment-specific factors to resemble their assigned environmental prototypes more closely, the discrimination between factors estimated from distinct environments will be enhanced. To learn such prototypes, we first construct prototypes for each sampled trajectory and then hierarchically combine trajectory prototypes with similar semantics into one environmental prototype. Experiments demonstrate that environment-specific factors estimated by our method have superior clustering performance and can consistently improve MBRL's generalisation performance in six environments consistently.

HIERARCHICAL PROTOTYPES FOR UNSUPERVISED DY- NAMICS GENERALIZATION IN MODEL-BASED REIN-
FORCEMENT LEARNING
Anonymous authors Paper under double-blind review
ABSTRACT
Generalization remains a central challenge in model-based reinforcement learning. Recent works attempt to model the environment--pecife factor and incorporate it as part the dynamic prediction to enable generalization different contexts. By estimating environment-specifie factors from historical transitions, earlier research was unable to clearly distinguish devironment-speciine factors from different envi- ronments, resulting in poor performance. To address this issue, introduce set representation factors
of environment prototypes to represent the dnvironmental-specifn for each environment. By encouraging learned environment-specci semble their assigned environmental prototypes more closely, the discrimination of factors between different environments will be enhanced. To learn such prototypes in the unsupervised manner, we propose a hierarchical prototypical method which first builds trajectory embeddings according the trajectory label information, and then hierarchically constructs environmental prototypes from trajectory prototypes sharing similar semantics. Experiments demonstrate that environment-speciffc factors estimated by our method have superior clustering performance and can improve MBRL's generalisation performance six environments consistently.
INTRODUCTION
Reinforcement learning (RL) has achieved great success in solving sequential decision-making prob-
lems, e.g., board games (Silver al. 2016 2017 Schrittwieser et al.) 2020), computer games (Mnih) et al.
2013 Silver et al.] 2018 Vinyals et al. 2019), and robotics (Levine & Abbeel 2014 Bousmalis ct al.
2018 but still suffers from the low sample efficiency problem, making challenging to solve
real-world problems, especially for those with limited or expensive data (Gottesman et al. 2018
Lu et al. 2018 2020 Kiran et al. 2020 In contrast, model-based reinforcement learning (MBRL)
(Janner et al.)
2019 Kaiser et al." 2019 Schrittwieser et al. 2020; Zhang et al. 2019 van Hasselt
et al. 2219;] Hafner et al. 2001bb a Lenz et 2015) has recently received wider attention, because it
explicitly builds predictive model and can generate samples for learning RL policy to alleviate the
sample inefficiency problem.
As sample-efficient alternative, the model-based RL method derives policy from the learned environmental dynamics prediction model. Therefore, the dynamics model's prediction accuracy is highly correlated with policy quality (Janner al. 2019). However, it has been evidenced that the learned dynamics prediction model is not robust to the change of environmental dynamics (Lee et al. 2020 Seo et al., 2020 Guo et al. 202T). and thus the agent model-based RL algorithms has a poor generalization ability on the environments with different dynamics. Such vulnerability to the change environmental dynamics makes model-based RI. methods unreliable in real- aal-world applications where the factors that affect dynamics are partially observed. For example, the friction cocfficient of the ground is usually difficult to measure, while the changes
can largely affect the dynamics when controlling robot walking the grounds, leading to the performance degradation of agent trained by model-based RL methods (Yang et al. 2209
Gu et al.) 2017 Vagabandi et al. 2018b).
Recent Studies (Seo et al.) 2020 Nagabandi et al. 2018 Sa Lee et al. 2020 Guo et al., 2021) have demonstrated that incorporating environmental factor Z into dynamics prediction facilitates the generalisation of model-based RI. methods to unseen environments. However, environmental factors

---

unobservable in the majority of applications; for instance, the friction coefficient is not available for robots. Therefore, estimating semantical meaningful Z for each environments is the first step for generalization of model-based RL. However, it is not easy implement, because the environment is hard to label. For example, it is impractical to measure the friction coefficient of every road.
Without the label information of environments, Zs estimated from previous methods (Seo et al. [2020)
Nagabandi et al., 20118a Lee et al. 20200 Guo et al. (022)) cannot form clear clusters for different environments Figure 3 Shows. These entangled Zs cannot represent the distinct environmental
specific information, and thus may deviate the learned dynamics prediction function from the true
one, resulting in the poor generalization ability.
In this paper, we propose a hierarchical prototypical method (HPM) with the objective of learning an invironment-specifie representation with distinct clusters. By representing environment-speciin information semantically meaningfully, HPM learns more generalizable dynamics prediction function. To achieve this, our method propose to construct a set of environmental prototypes to capture environment--peccfie information for each environment. By enforcing the estimated 2 to be more similar to its respective environmental prototypes and dissimilar to other prototypes, the estimated 2s can form compact clusters for the purpose of learning a generalizable dynamics prediction function. Because environmental labels are not available, we cannot construct environmental prototypes directly. To address this issue, we begin by developing easily-learned trajectory prototypes based on the trajectory label. Then, environmental prototypes can be created by merging trajectory prototypes with similar semanties, as suggested by the natural hierarchical relationship between trajectory and environment.
With the built hierarchical prototypical structure, we further propose prototypical relational loss to learn Z from past transitions. Specifically, we not only aggregate the 2s with similar causal effects by optimizing the relational loss (Guo et al. 2021) but also aggregate 2 with its corresponding trajectory and environmental prototypes via the relational loss. In addition, to alleviate the over- wver-penalization of semantically similar prototypes, we propose to penalize prototypes adaptively with the intervention similarity. In the experiments, we evaluate
method on range of tasks in OpenAI gym (Brockman iet al. 2016) and Mujoco (Todoroy et al. 2012). The experimental results show that our method can form more clear and tighter clusters for Zs, and such 2s can improve the generalization ability of model-based RL methods and achieve state-of-art performance in new environments with different dynamics without any adaptation step.
2 RELATED WORK
Model-based reinforcement learning With the learned dynamics prediction model, Model-based Reinforcement Learning (MBRL) takes advantage of high data efficiency. The learned prediction
model can generate samples for training policy (Du & Narasimhan 2019) Whitney et al., 2019)
or planning ahead in the inference (Atkeson & Santamaria 1997; Lenz al. 2015]] TTssa al.) 2012). Therefore, the performance of MBRL highly relies on the prediction accuracy of the dynamics predictive model. To improve the predictive model's l's accuracy of MBRL, several methods were proposed, such as ensemble methods (Chua et al.) 2018), latent dynamics model (Hafner et al. 2019b at Schrittwieser et al.) 2020), and bidirectional prediction (Lai al. 2020). However, current predictive methods are still hard to generalize well on unseen dynamics, which hinders the application of MBRL methods in the
problems.
Dynamics generalization in model-based reinforcement learning known dynamics, meta-learning methods (Nagabandi et al.) 2018a bj Sæmundsson al.) 2018 attempted to update model parameters by updating a small number of gradient updates (Finn et al. 2017) or hidden representations of a recurrent model (Doshi- Velez & Konidaris (016). Then, using multi-choice learning, (Lee et al. 2020 Seoo et al. 2020) attempted to learn generalised dynam- ics model by incorporating with the goal of adapting any dynamics without training. Through relational learning and causal effect estimation, RIA (Guo et al. (021) aims to explicitly learn meaningful environmental-spesn -pecificc information. However, the dynamics change learned by RIA still suffer from high variance issue.
To adapt the MBRL to un-
information or clustering dynamics implicitly,
Prototypical methods By learning an encoder to embed data in a low-dimensional representa- tion space, prototypical methods gain a set of prototypical embeddings, which are referred to as
2

---

Context Encoder
Estimation of
Enwrionments
Historical Segments
Latent Space of
Trajectory Prototype
Environment Prototype
Prototypical Relational Learning
Figure I:
An overview of our Hierarchical Prototypical Method, where the context encoder estimates the sivironmental-specificc factor 2 and environments includes four ants with different destroyed leg with red color. Items extracted from different environments are different colors. We construct prototypes for each trajectory and environment, and denote them as diamond and star, respectively. Each estimated & are optimized with its corresponding trajectory and environment prototype using our prototypical relational learning as dotted line shows.
prototypes (Asano al. 2020 Caron et al. 2020b) that form the basis of this representation space. Prototypical methods aim to derive compact data representations gathering around corresponding prototypes [Li et al. 2021 Oord al. 2018 Wang et al., 2021). which captures some basic semantic structures. Therefore, prototypical methods have been applied into many areas, e.g. self-supervised
learning [i et 2020; Caron al. 2020a), few-shot learning (nelll al., 2017 Bateni et al., 2020
Simon et al. (020), domain adaptation (Tanwisuth et al. 2021) and continue learning (De Lange & Tuytelaars] 2021 Yu al. 2020). In the RI. area, (Yarats al. (021) ties representation learning with exploration through prototypical representations for mage--base RI., while our method focuses on the unsupervised dynamics generalization problem in model-based RL, aiming to learn seman- tical meaningful dynamics change using prototypical method. Specifically, our method propose hierarchical method to construct environmental prototypes from trajectory prototypes.
3 METHODS
this section, we first introduce the formulation of the unsupervised dynamic generalization problem in model-based reinforcement learning. Then present the details of how our hierarchical prototype method learns the environment-specifcc factors.
3.1 PROBLEM SETUP
We formulate the standard reinforcement learning as a markov decision process (MDP) M (S,A, r, ff Pu) over discrete time (Puterman) 2014 Sutton & Barto 2018). where S, A. (0. and Po are state space, action space, the reward discount factor, and the initial state distribution, respectively. Dynamics function ʃ S A + S gives the next state 81+1 conditioned state 8; and action a,, and reward function r : A t given 8; and ag. The goal of RL is learn policy T (--ss) mapping from state E S over the action distribution to maximize the cumulative expected return E.,,ES,. 7f r(st, at)] over timesteps. In model-based RI., we aim to learn prediction model fto approximate the dynamics function f. and then f can generate training data to train policy predict the future sequences for planning. With the data provided by learned dynamics model f. model-basee RL has higher data efficiency and better planing ability compared with model-free RI.
the current specifies the reward each timestep
In this paper, consider the unsupervised dynamics generalization problem in model-based RL. Different from the standard reinforcement learning, there exists an unobserved variable Z that can affect the dynamics prediction function f in the dynamics generalization problem. The goal of dynamics gemeralization is to derive a generalizable policy from given K training MDPs (M J and expect the policy can generalize well on L test MDPs [M*eeL J==0- Without losing generality, we assume all MDPs share the same state and action space but preserve different factor Z.
3

---

In the context of model-based reinforcement learning, need to learn the dynamics function before learning policy. In order to generalize the dynamic functions on different environment, we need to incorporate the unobserved variable Z into dynamics prediction process, i.e., extending the dynamics function from f S x A + S to S x A x 2 + S. Since Z is not available, we should estimate it from past transition segments 7t---t1 { k; at=k); -.. (st-1, ((-1)) (Seo et al. 2020
Lee
it al. 2020 Guo et al., (021).
Next, will present how our hierarchical prototypes method estimates Z, and enable it learn the dynamics function f that can generalize to environments with unseen dynamics. In Section 3.2 we present how our method hierarchically constructs prototypes as a representative embedding to represent enviromennaal-peciin we update prototypes dynamically and how to estimate environmental-speciff factors Z from past transition segments using prototypes. Once Z are estimated, we describe how they enable dynamics function f to generalize well environments with different dynamics.
information for each environment. In Section 3.3 we describe how
3.2 HIERARCHICAL ENVIRONMENT PROTOTYPES CONSTRUCTION
The objective of our method is to construct a set of prototypes to represent the environmental-speciic information for each environment, and guide the context encoder to estimate environmental-specffic variable Z from historical transition segments. In each training iteration, trajectory from a subset of MDPs in the training MDPs. Because labels of MDPs are not available, cannot estimate environmental prototypes directly. Furtunately, we still have the trajectory label information, and thus we can construct the prototypes for each sampled trajectory first. Specifically,
randomly sample
we denote the prototype for j-th trajectory as Cira. Because different trajectories may be sampled from single environment, the trajectory prototypes from the same environment should share similar semantics for dynamics prediction. Therefore, we can construct environmental prototypes hicrarchi- cally from trajectory prototypes sharing similar semantics. In this way, environmental prototypes and trajectory prototypes form a natural hierarchical structure, and environmental prototypes can be constructed utillisig trajectory label information even if no environmental label is available.
If we denote the wira as the semantical similarity between the trajectory prototypes clya and arras we can construct a trajectory similarity matrix w Figure 2 (b) shows, where each row of such as w represents the similarity between ctr, and all other trajectory prototypes. Because it is unknown how many environments are in the sampled trajectories, we directly construct environmental prototypes cene for each trajectory prototype cira. Specifically, each environmental prototype cene is the mean of its corresponding trajectory prototype cira and ctraa 's top k similar trajectory prototypes.
Σ Cra
KE(T)}
where T. denotes the index set of the top-K K similar trajectory prototypes with cira. In this way, we can obtain the i-th environmental prototypes, but before that, we need to calculate the semantic similarity matrix w.
(o).a))
Fina
(ol.og]
Cero (s). of
f.
(ol.nn)
2111
1/15
1/13
1018
:
1P23
1222
1028
(ef, ot)
2
<...
fo
10b1
1ik
1000
d
....
(oL mb)
(s) Colculating Similartty between Prototypes via Cousal Direct Effect
(b) Marging Top. k Simiaar Trajectory Prototypes into Enironment Prototypes
Figure 2: (a) The illustration of causal direct effect estimation between two trajectory prototypes, where we calculate the mean difference over a batch of predicted next states as the similarity w of the two prototypes. (b) The illustration constructed the similarity matrix between trajectory prototypes. We use the mean of top_k (denoted by the red color) similar trajectory prototypes as the corresponding environmental prototypes.
Normally, we can directly use the euclidean distance to discriminate the similarity between different trajectory prototypes. However, this ignores the semantic effect of trajectory prototypes on dynamics
4

---

prediction. If two trajectories prototypes are from single environment, their trajectory prototypes should share the same semantics, i.e., and their effects on the dynamics function should be the same. Therefore, similarity estimation. However, it is challenging to estimate the effects of the trajectory prototype on the dynamics function because Z is not the only factor that can influence the dynamies function. To remove the effects of other factors, e.g. states and actions, on the dynamics function, our method draws inspiration from the recently proposed RIA method (Guo al., 202T) calculate the direct causal effects (CDE) of trajectory prototypes. By controlling all factors that have effects on the dynamics function over a mini-batch, we can solely estimate average CDE between different
consider take account the semantic effect the dynamics prediction into
trajectory prototypes as their semantic difference d. Concretely, we compute d between two trajectory prototypes using a mini-batch of S, and A, pairs (st, aj) as Figure 2 (a) shows:
N d, Σ
CDE di (st,at)).
where N is the batch size, i and are the id of trajectory prototypes. Please refer to Appendix A.6 for the details of CDE. With semantic difference d, we can convert it as the semantic difference w via = exp( = ) where 3 is factor that controls the sensitivity of w. With the calculated similarity w, can construct environmental prototypes via Eq.
Next, we will describe how to update the built trajectory and environmental prototypes to ensure that hierarchical prototypes representative for each trajectory and environment, and how they help learn the context encoder.
3.3 PROTOTYPICAL RELATIONAL LEARNING
As Figure shows, we introduce a context encoder g parameterized by á to estimate environmental- specific factor 2t from the past transition segments Tt-k:t-1 { ky at dt-k),..,, (5t-1, 4.-1)} following previous methods:
= =
In order to learn the context encoder and encourage the estimated environmental-specifee factor #f to be semantically meaningful, we optimize via the proposed prototypical relational loss to form a clear cluster for Zs from the same environments. Concretely, & Storkey (020) as a learnable function h to derive the environmental-specifee estimation 2 closely surrounded its associated cluster prototypes. To achieve this, we concatenate the 2t and its assigned prototypes, e-g., ctr, as the positive pair, and the concatenation of other prototypes are negative pairs. Then we use the relational head h parameterized by o quantify the similarity score of ğ. To increase the similarity score ý of positive pairs and decrease those of negatives, we can regard it as simple binary classification problem to distinguish positive and negative pairs. This can be regarded as maximizing the mutual information between Zs and its corresponding prototypes (Please refer to (Tsai et al. 20200] Guo et al. 202T) and Appendix A.3). However, it neglects the semantic comelation among different prototypes, and so it may excessively penalize some semantically relevant prototypes. To alleviate such over- pecr-peaalizatin, we propose to penalize prototypes adaptively with the intervention similarity (Guo et al. 2021) through the following objective:
introduce relational head (Patacchiola
N
N
C p--p-a prreationn
I + (1 (1- w w log c'='),']::)))
-NNN - 1)
i=1 j=1
+(1 - 1.) (1 - (1-w^d)
log (1 - M(=*t,e):))))) cl:4))))
where w ranges from 0 to 1, and we use it as the similarity between different prototypes. In addition, the first term of Eq. (3) clusters 2 with prototypes et with the similarity weight wh.., and the second term push them away with weight 1 - -d To maintain the hierarchical prototypes structure (Li et al. 2020 Guo et al. 2022), we simultaneously update the context encoder by optimizing the objective Eq. between with trajectory and environmental prototypes. Specifically, the calculation of
similarity Wene between environmental prototypes and 2 is same with Section 3.2 as In addition,
also optimize the relation loss among different Zs following (Guo et al., 2021 Li et al. 2020) because Z itself can be regarded as an instance prototype, and thus can retain the property of local smoothness and help bootstrap clustering.
5

---

In order to improve its generalization ability different dynamics, incorporate the estimated environment-specifee &, into the dynamies prediction model and optimize the objective function following (Lee et al. 2020 Seo et al.) 2020 Janner et al. 2019:
log f( 16 g
(4)
where & is the length of transition segments, is the current timestep, and N is the sample size. In addition, we also enable the built prototypes to optimize Eq. 4 to ensure that the learned prototypes semantically meaningful. Overall, our method simultaneously optimize the prediction loss Eq. 4
and prototypical relational loss Eq. (3) with prototypes in different levels to learn context encoder and semantic meaningful prototypes, which encourage the estimated environmentallspeccin 2 can form clear clusters, and thus can learn a generalizable prediction function f.
3.4 DIFFERENCE TO RIA
Our paper refers the idea of RIA (Guo al. 2021) to estimate semantic similarities between different prototypes. However, our method differs RIA from three aspects: 1) RIA estimates the semantic similarities between different instance estimation while our method estimates the semantic similar- ities between different prototypes. Considering the number of prototypes procedure are faster and more stable than RIA. 2) Our method fully takes the advantage the hierarchy between trajectory and environments, and construct environmental prototype based on trajectory label information while RIA ignores it. Thus our method can achieve better performance than RIA. 3) RIA only pulls 2 and other estimations with similar semanties, but our prototypical relational learning further pulls the ? and its corresponding trajectory Ctra and environmental prototypes Ceno.
limited, the training
4 EXPERIMENT
In this section, we perform experiments to evaluate the effectiveness of our approach by answering the following questions: 1) Can our method encourages the learned to form a clear cluster? (Section 4.2 2) Can the learned 2 with the clear cluster reduce the dynamics prediction emros in model-based RL? (Supplementary Material A.2); 3) Can the learned 2 with the clear cluster promote the performance model- based RL in environments with unseen dynamics? (Section 4.3): 4) Is our method sensitive to hyperparameters? (Section
4.1 ENVIRONMENTAL SETUP
Implementation details Our method includes three learnable functions and set of learnable trajectory prototypes. The learnable functions are context encoder, relational head and prediction head, and they all are constructed with MLP and optimized by Adam (Kingma & Baj (014)) with 1e-3 learning rate. During the training procedure, the trajectory segments are randomly sampled from the same trajectory break the temporal correlations of the training data, which also adopted by (Seo et al. 2020 Guo et al. 2021). Specifically, we combine A: = 3 similar trajectory embedding into environmental embedding, and the length of the transition segments is 10, and the hyper-paramett are the same for all experiments, and details can be found in supplementary material A.
Tasks Following the previous methods (Lee et al.) 2020 Seo et al., 2020). perform experiments on the classic control algorithm (Pendulum) from OpenAl gym (Brockman et al.) (2016) and simulated robotics control tasks (HalfCheetah, Swimmer, Ant. Hopper, Slim-Humanoid from Mujoco physical
engine (Todoroy et al. 2012).
Dynamics settings
construct different dynamics of environments, we change the environmental parameters (e.g., length and mass of Pendulum) and predefine them in the training and test envi- ronmental parameters lists following previous methods (Zhou et al.] 22019] Packere Packer et al. 2019 Lec et al. 2020 Seo et al. 2020 Guo et al. 2)21). Specifically, for the convthe training environmental parameters lists for all tasks are {0.75. 0.8, 0.85, 0.90, 0.95, 1, 1.05. 1.1, 1.15, 1.2, 1.25). and test en- vironmental parameters lists are {0.2, 0.4, 0.5, 0.7, 1.3, 1.5, 1.6, 1.8). We
that the parameters
in test list are out of range of the parameters in the training set. At the training time, we randomly sample the parameters from the training parameter list to train context encoder and dynamics

---

TMCL
CaDM
RIA
Ours
TMCL
CaDM
RIA
Ours
Figure 3: The PCA visualization of environmental-specn factors estimated by TMCL (Seo et al. 2020), CaDM (Lee al. 2020). RIA (Guo et 2021 and ours on the Half-cheetah (upper part) and Pendulum (lower part) task.
prediction model. Then we test our model on the environments with unseen dynamics sampled from the test parameter list. All details are given in supplementary material|A.I
Planning Following (Lee et al. 2020; [Soo et al., 2020), we use the model predictive model (MPC) (Maciejowski] 2002) to select actions based on learned dynamics prediction model, and assume that reward functions are known. In addition, we use the cross-entropy method (CEM) (De Boer et al.) 2005) to find the best action sequences.
Baselines In this paper, we compare our approach with the following state-of-the-art model-based RL methods on dynamics generalization:
Context-aware dynamics model (CaDM) (Lee et al. 2020): This method design several auxiliary loss, including backward and future states prediction to learn the context from transition segments.
Trajectory-winn Multiple Choice Learning (TMCL) (Seo et al. 2020): TMCL introduces multi-choice learning to adapt to different environments. For a fair comparison, we use the adaptation version of this method.
Relation Intervention Approach (RIA) (Guo et al., 2021): This method proposes to relational intervention loss to cluster Zs from the same environments.
It has been clearly evidenced that Probabilistic ensemble dynamics model (PETS) Kurutach et al. 2018) and Meta learning based model-based RL methods, hidden- hidden-parameter parameter model GrBAL (Nagabandi et al. 2018b a), perform worse than CaDM (Lee et al. 2020). TTCL (Seo et al.) 2020) and RIA (Guo et al. 2021 so we do not consider them as baselines in our paper.
Recurrent model ReBAL and
4.2 CLUSTER VISUALIZATION AND ANALYSIS
Table 1: The quantitative evaluation results of estimated environmental-specifee factors.
ARI
AMI
V-means
TMCL.
CaDM
RIA
Ours
TMCL
CaDM
RIA
Ours
TMCL.
CaDM
RIA
Ours
HalfCheetah
0.006
0.128
0,212
0,570
0.058
0.175
0.333
0.681
0.06
0.176
0.314
0.680 0.975
Pendulum
0.060
0.471
0.754
0.971
0.054
0.529
0.838
0.967
0.051
0.531
724
We perform PCA visualization of estimated 2s from baselines and our method as Figure 3 to evaluate the cluster performance of estimated Zs. We can see that our method can achieve better cluster performance qualitatively. Specifically, most Żs estimated by RIA (Guo et al. 2021) have good cluster performance in general, but the outliers decrease the cluster performance. By contrast,
7

---

that there are fewer outhiers in method than them in RIA because the built prototypes and the proposed prototypical relational loss can enforce constraints into estimated Žs. More qualitatively cluster comparisons can be found in Supplementary Material A.8
We also quantitatively evaluate the cluster performance 2s estimated by baselines and our method. Here we firstly perform k-means (MacQucen et al. 1967) on the estimated Žs, and then use the ground-truth environmental label calculate the cluster performance. Here we use the popular mutual information-base metric AMI (Vinh et al. 2010), random-index metric ARI(Hubert & Arabie 1985) and V-means (Rosenberg & Hirschberg 2007) as the evaluation metrics. The results are shown in Table 3 we can see that 7s estimated by our method achieves the highest cluster performance. More quantitatively cluster comparisons can be found in Supplementary Material A.8
4.3
PERFORMANCE COMPARISONS
Ant
Swimmer
3500
CaDM
CaDM
CaDM TMUL
3000
TMCL
1000
RIA
RIA
RIA
2000
Curs
900
Rewards
Rewards
Rewards
700-
15
400-
0-
Timesteps ('1044) Hallcheelah
('10%4)
Timestopa (M104) Pendukum
Humianod
CaDM
2500
CaDM
10000
8000
Ours
RIA
1900
-000
RIA
Ours
1000
-1000
C4DM
TMCL
2000
-000
RIA
-1200
-1000
Ous
10
Timestops (*1014)
Timestaps (*1044)
Timastops (*1 (104))
Figure 4: The average returns of model-based RL agents unseen test environments. The results show the mean and standard deviation of returns averaged over five runs. Specifically, we use the no adaptation version of TMCL for a fair comparison. The performance comparisons dynamics prediction errors are given at Appendix A.4 A.4
Then, we evaluate the generalization of model-based RL agents trained by our methods and baselines test environments with unseen dynamics. Following the setting of (Seo et al. (2020), perform experiments across five runs, and show the test returns on the test environments in Figure 6 Note that the results are slightly different from the results in RIA and TMCL paper since we change the parameter lists that change the environmental dynamics. Specifically, we change the parameter lists of all environments the
for the convenience of performing environments.
As Figure 6 shows, we can see that our method can achieve significantly better performance than baselines Ant, Halfcheetah, and Pendulum. Specifically, we can see that our method outperforms the second- -bst method RIA by 20% in Ant and Halfcheetah environments, which indicates that the changing parameter can largely change their dynamics. In addition, we can see that our method achieves only slightly better performance than baselines in Hopper, Swimmer, and Slim iamn_uumaddddd problems. For Hopper and Slim Humanoid environment, we observe that both RIA and our method can achieve comparable results in all test environments, which indicates that the change of dynamics for Hopper is easy to model and solve. For the Swimmer environment, we observe that TMCL (Seo) al. 2020) sometimes may have significant performance decline at the final training iteration. This may be because that TMCL may fail learn the modalities dynamics function in the no adaptation version. Also, method still achieves better performance than RIA the Swimmer task.
8

---

4.4 ABLATION STUDY
Pendulum
Halfcheeteh
Pendulum
4400
-003
-800
Rewards
Rewerds
k 3
43%
1193
1200
Causal Effict Puelidean
Causal Elffet
-140
10
10.0
10.0
(*104)
(11044)
Figure 5: Left Image: The sensitivity analysis about how many trajectory prototypes should be com- bined into environmental prototypes. Right Two Images: The similarity metries used combining trajectory prototypes into environmental prototypes.
In this section, first perform sensitive analysis of how many trajectory prototypes should be combined into environmental prototypes. The experiments are conducted at the Pendulum task, and the results shown as the left image of Figure 5 we can that matter what & is, our method consistently outperforms the baseline CaDM (Lee et al. (020), which indicates that our method is robust the selection of value. Specifically, k means that there no hierarchical prototypes because one trajectory prototype can decide one environmental prototype, and thus environmental prototypes are the same as trajectory prototypes. We can see that all experimental results with k > 1 are better than the experimental result with k: 1, which shows the effectiveness of our proposed hierarchical prototypes method and the necessity of the built environmental prototypes. The results of l: = 1 achieve the best performance on the Pendulum task, we use it as the default parameter in all experiments.
We also perform ablation study about the similarity metric used to calculate the similarity among trajectory prototypes. For most cluster methods, e.g. k-means (MacQueen et al.) 1967), they usually calculate the similarity among entities using the Euclidean distance, while our method uses the direct causal effect as the similarity metric. To evaluate the effectiveness of the similarity metrics based on direct causal effect (Pearl 2013), we perform experiments the Halfcheetah and Pendulum tasks, and we can see that using the causal effect to calculate the similarities among trajectory prototypes can achieve better performance than using Euclidean distance on both tasks.
LIMITATION
Our paper only considers the unsupervised dynamics generalization in model-hased reinforcement learning, but model-free RL also suffers from this problem, and we will apply our method to model- free RL in future work. In addition, there are many other generalization problems in reinforcement learning area, e.g. observation generalization (Wang et 2020; Kirk et al., [202]] Ghosh al. 2021) and action generalization (Jain et al. 2020). and it would be interesting to extend our method into other generalization settings and train generalizable agents.
CONCLUSION
In this paper, we focus on the unsupervised dynamics generalization problem in model- reinforcement learning, and propose a hierarchical prototypical method to construct environmental prototypes in an unsupervised manner. With the learned environmental prototypes, we further propose a prototypical relational loss to learn a context encoder to estimate auvironmental-spccife factors from past transition segments, which enables the dynamics prediction function in model-based reinforcement learning to generalize well on environments with unseen dynamics. The experiments demonstrate that our method can form clearer and tighter clusters for 2s from the same environment and improve the performance of model-based agents in new environments with unseen dynamics.

---

7 REPRODUCIBILITY STATEMENT
We acknowledge the importance of reproducibility for research work and try whatever we can to ensure the reproducibility of our work. As for the implementation of our method, details such as hyperaarameters are provided in Section 4.1 and Appendix A.1 We will publicly release all codes after the acceptance of this paper.
REFERENCES
Yuki M. Asano, Christian Rupprecht, and Andrea Vedaldi. Self-labelling via simultaneous clustering and representation learning. In International Conference on Learning Representations (ICIR), 2020.
Christopher G Atkeson and Juan Carlos Santamaria. A comparison of direct and model- reinforcement learning. In Proceedings of international conference on robotics and automation, volume 4, pp. 3557-3564. IEEE, 1997.
Peyman Bateni, Raghav Goyal, Vaden Masrani, Frank Wood, and Leonid Sigal. Improved few-shot visual classification. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 14493-14502, 2020.
Konstantinos Bousmalis, Alex Irpan, Paul Wohlhart, Yunfei Bai, Matthew Kelcey, Mrinal Kalakrish- nan, Laura Downs, Julian Ibarz, Peter Pastor, Kurt Konolige, et al. Using simulation and domain adaptation to improve efficiency of deep robotic grasping. In 2018 IEEE international conference robotics and automation (ICRA), pp. 4243 4250. IEEE, 2018.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. Openai gym. arXiv preprint wxXiv:6606.01500. 2016.
Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. Unsupervised learning of visual features by contrasting cluster assignments. Advances in Neural Information Processing Systems, 33:9912 9924, 2020a.
Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. Unsupervised learning of visual features by contrasting cluster assignments. 2020b.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learn- ing in a handful of trials using probabilistic dynamics models. arXiv preprint 1805 2018.
Pieter- Tjerk De Boer, Dirk P Kroese, Shie Mannor, and Reuven Y Rubinstein. A tutorial on the cross-entropy method. Annals of operations research, 134(1):19667 2005.
Matthias De Lange and Tinne Tuytelaars. Continual prototype evolution: Learning online from non- stationary data streams. Proceedings of the IEÉE/CVF International Conference on Computer Vision, pp. 8250 8259, 2021.
Finale Doshi Velez and George Konidaris. Hidden parameter markov decision processes: A semi- parametric regression approach for discovering latent task parametrizations. IJCAI: proceedings of the conference, volume 2016, pp. 1432. NIH Public Access, 2016.
Yilun Du and Karthic Narasimhan. Task-agnostic dynamics priors for deep reinforcement learning. In International Conference on Machine Learning, pp. 1696-1705, PMLR, 2019.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-leamning for fast adaptation of deep networks. International Conference on Machine Learning, pp. 1135. PMLR, 2017.
Dibya Ghosh, Jad Rahme, Aviral Kumar, Amy Zhang, Ryan P Adams, and Sergey Levine. Why generalization in rl is difficult: Epistemic pomdps and implicit partial observability. Advances in Neural Information Processing Systems, 34, 2021.
10

---

Omer Gottesman, Fredrik Johansson, Joshua Meier, Jack Dent, Donghun Lee, Srivatsan Srinivasan, Linying Zhang, Yi Ding, David Wihl, Xuefeng Peng, et al. Evaluating reinforcement learning algorithms in observational health settings. arXiv preprint arXiv:1805. 12298, 2018.
Shixiang Gu, Ethan Holly, Timothy Lillicrap, and Sergey Levine. Deep reinforcement learning for robotic manipulation with asynchronous off-policy updates. In 2017 IEEE international conference on robotics and automation (ICRA), pp. 3389-3396 IEEE, 2017.
Jiaxian Guo, Mingming Gong, and Dacheng Tao. A relational intervention approach for unsupervised dynamics generalization in model-based reinforcement learning. International Conference on Learning Representations, 2021.
Yuanfan Guo, Minghao Xu, Jiawen Li, Bingbing Ni, Xuanyu Zhu, Zhenbang Sun, and Yi Xu. Hesc: Hierarchical contrastive selective coding, arXiv preprint arXiv:2002.00555, 2022.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv:1912.01603 2019a.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International Conference on Machine Learning, pp. 2555 2565. PMLR, 2019b.
Lawrence Hubert and Phipps Arabie. Comparing partitions. Journal of classification 2(1):193-218. 1985.
Ayush Jain, Andrew Szot, and Joseph Lim. Generalization new actions in reinforcement learning. arXiv preprint arXiv:2011.01928, 2020.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When trust your model: Model-based policy optimization. arXiv preprint arXiv:I906 1906.08253, 2019.
Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, et al. Model-based reinforcement learning for atari. arXiv preprint arXiv: (:03::/903.00777 2019.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint 4rXhv:1412.6980, 2014.
Ravi Kiran, Ibrahim Sobh, Victor Talpaert, Patrick Mannion, Ahmad Al Sallab, Senthil Yogamani, and Patrick Pérez. Deep reinforcement learning for autonomous driving: A survey. arXiv preprint arXiv:2002.0044: 2020.
Robert Kirk, Amy Zhang, Edward Grefenstette, and Tim Rocktäschel. A survey of generalisation in deep reinforcement learning. arXiv preprint arXiv:2111.0979, 2021.
Thanard Kurutach, Ignasi Clavera, Yan Duan, Aviv Tamar, and Pieter Abbeel. Model-ensemble trust-region policy optimization. arXiv preprint arXiv: 10592, 2018.
Hang Lai, Jian Shen, Weinan Zhang, and Yong Yu. Bidirectional model- based policy optimization. In International Conference on Machine Learning, pp. 5618-5627. PMLR, 2020.
Kimin Lee, Younggyo Seo, Seunghyun Lee, Honglak Lee, and Jinwoo Shin. Context-aware dynamics model for generalization in model-based reinforcement learning. In International Conference Machine Learning, pp. 5757-5766. PMLR, 2020.
Ian Lenz, Ross A Knepper, and Ashutosh Saxena. Deepmpe: Learning deep latent features for model predictive control. Robotics: Science and Systems. Rome, Italy, 2015.
Sergey Levine and Pieter Abbeel. Learning neural network policies with guided policy search under unknown dynamics. In NIPS, volume 27, pp. 1071-1079. Citeseer, 2014.
Junnan Li, Pan Zhou, Caiming Xiong, and Steven CH Hoi. Prototypical contrastive learning of unsupervised representations. arXiv preprint arXiv:2005.044: 049666 2020.

---

Yunfan Li, Peng Hu, Zitao Liu, Dezhong Peng, Joey Tianyi Zhou, and Xi Peng. Contrastive clustering. In 2021 AAAI Conference on Artificial Intelligence (AAAI), 2021.
Chaochao Lu, Bernhard Schölkopf, and José Miguel Hemmndezz-Lbbao. Deconfounding reinforce- ment learning observational settings, arXiv preprini arXiv: 1812. 10576, 2018.
Chaochao Lu, Biwei Huang, Ke Wang, José Miguel Hemánndez-Lobato, Kun Zhang, and Bernhard Schölkopf. Sample-efficiee reinforcement learning via counterfactual-basel data augmentation. arXiv preprint arXiv:2012.09092 2020.
Jan Marian Maciejowski. Predictive control: with constraints. Pearson education, 2002.
James MacQueen et al. Some methods for classification and analysis of multivariate observations. In Proceedings the fifth Berkeley symposium mathematical statistics and probability, volume 1, pp. 281-297. Oakland, CA, USA, 1967.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, loannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv:1312.5602, 2013.
Anusha Nagabandi, Ignasi Clavera, Simin Liu, Ronald S Fearing, Pieter Abbeel, Sergey Levine, and Chelsea Finn. Learning to adapt in dynamic, real- world environments through meta-reinforcementt learning. arXiv preprint arXiv: 1347, 2018a.
Anusha Nagabandi, Chelsea Finn, and Sergey Levine. Deep online learning via meta-learning: Continual adaptation for model-based arXiv preprint 181.v:111..0767/ 2018b.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv: 037443378, 2018
Charles Packer, Katelyn Gao, Jernej Kos, Philipp Krähenbühl, Vladlen Koltun, and Dawn Song. Assessing generalization in deep reinforcement learning, 2019.
Massimiliano Patacchiola and Amos Storkey. Self-supervised relational reasoning for representation learning. arXiv preprint Ttl::2006.058443 2020.
J Pearl. Causality: Models, reasoning, and inference 47cambridge university presscambridge, united kingdom. pearl, 2000. Causality: models, reasoning, and inference, 47, 2000.
Judea Pearl. Direct and indirect effects. arXiv preprint 1:Xivv1301.2333 2013.
Martin L Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 2014.
Andrew Rosenberg and Julia Hirschberg. V-measure: A conditional entropy-based external cluster evaluation measure. In Proceedings of the 2007 joint conference empirical methods in natural language processing and computational natural languuge learning DEMNLP-CONL.N pp. 410-420, 2007.
Steindór Sæmuundsoon, Katja Hofmann, and Mare Peter Deisenroth. Meta reinforcement learning with latent variable gaussian processes. arXiv preprint arXiv: 1803.07551, 2018.
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature. 588(7839):604-600 2020.
Younggyo Seo, Kimin Lee, Ignasi Clavera, Thanard Kurutach, Jinwoo Shin, and Pieter Abbeel. Trajectory-wise multiple choice learning for dynamies generalization in reinforcement learning. arXiv preprint arXiv:2010.13300 2020.
David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, loannis Antonoglou, Veda Panneershelvam, Mare Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):484 489, 2016.
12

---

David Silver, Julian Schrittwieser, Karen Simonyan, Ioannis Antonoglou. Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, et al. Mastering the game of go without human knowledge. nature, 550(7676)3354 2017.
David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Mare Lanctor, Laurent Sifre, Dharshan Kumaran, Thore Graepel, et al. A general reinforcement learning algorithm that masters chess, shogi, and go through self-play. Science, 362(6419): 1140-1144, 2018.
Christian Simon, Piotr Koniusz, Richard Nock, and Mehrtash Harandi. Adaptive subspaces for few-shot learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 1136-4145, 2020.
Jake Snell, Kevin Swersky, and Richard Zemel. Prototypical networks for few-shot learning, Advances in neural information processing systems, 30, 2017.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Korawat Tanwisuth, Xinjie Fan, Huangjie Zheng, Shujian Zhang, Hao Zhang, Bo Chen, and Mingyuan Zhou. A prototype-oriented framework for unsupervised domain adaptation. Advances in Neural Information Processing Systems, 34, 2021.
Yuval Tassa, Tom Erez, and Emanuel Todorov. Synthesis and stabilization of complex behaviors through online trajectory optimization. 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pp. 4906-4913, IEEE, 2012.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pp. 5026-5033, IEEE, 2012.
Yao-Hung Hubert Tsai, Han Zhao, Makoto Yamada, Louis-Philipee Morency, and Ruslan Salakhutdi- nov. Neural methods for point-wise dependency estimation. arXiv preprint arXiv:2006.05553 2020.
Hado P van Hasselt, Matteo Hessel, and John Aslanides. When to use parametric models in reinforcement learning? Advances Neural Information Processing Systems, 32:14322--14333 2019.
Nguyen Xuan Vinh, Julien Epps, and James Bailey. Information theoretic measures for clusterings Variants, properties, normalization and correction for chance. The Journal of Machine Learning Research, 11:2837-2854, 2010.
comparison:
Oriol Vinyals, Igor Babuschkin, Wojciech M Czarnecki, Michaël Mathieu, Andrew Dudzik, Junyoung Chung, David H Choi, Richard Powell, Timo Ewalds, Petko Georgiev, al. Grandmaster level in starcraft ii using multi-agent reinforcement learning. Nature, 575(7782)3350-354 2019.
Kaixin Wang, Bingyi Kang, Jie Shao, and Jiashi Feng. Improving generalization in reinforcement learning with mixture regularization. Advances Neural Information Processing Systems, 33: 7968-7978, 2020.
Xudong Wang, Ziwei Liu, and Stella X Yu. Unsupervised feature learning by cross-level instance- group discrimination. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 12586-12595, 2021.
William Whitney, Rajat Agarwal, Kyunghyun Cho, and Abhinav Gupta. Dynamics-awart embeddings. arXiv preprint arXiv: 09357,, 2019.
Jiachen Yang, Brenden Petersen, Hongyuan Zha, and Daniel Faissol. Single episode policy transfer in reinforcement learning. arXiv preprint arXiv:1910.07719, 2019.
Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Reinforcement learning with prototypical representations. In Proceedings of the 38th International Conference on Machine Learning, volume 139, pp. 11920-11931 2021.
13

---

Lu Yu, Bartlomiej Twardowski, Xialei Liu, Luis Herranz, Kai Wang, Yongmei Cheng, Shangling Jui, and Joost van de Weijer. Semantic drift compensation for class incremental learning. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 6982-6991, 2020.
Marvin Zhang, Sharad Vikram, Laura Smith, Pieter Abbeel, Matthew Johnson, and Sergey Levine. Solar: Deep structured representations for model-based reinforcement learning. In International Conference on Machine Learning, pp. 9444-7453. PMLR, 2019.
Wenxuan Zhou, Lerrel Pinto, and Abhinav Gupta. Environment probing interaction policies, arXiv preprint arXiv: 2019.
A APPENDIX
We promise that we will public all codes after the acceptance of this paper and we public all training details Appendix A.1 and A.3
A. ENVIRONMENTAL SETTINGS
We follow the environmental settings of Lee et al. (2020); Guo et al. (202T) and give the details of settings as follows:
Pendulum We modify the mass and the length / of Pendulum to change its dynamics,
Half-Cheetah We modify the mass of rigid link m and the damping of joint of Half- Cheetah agent to change its dynamics.
Swimmer We modify the mass of rigid link m and the damping of joint Swimmer agent to change its dynamics.
Ant We modify the mass of ant's leg m change its dynamics. Specifically, modify two legs by multiplying its original mass with m, and others two with
Slim Humanoid We modify the mass of rigid link m and the dampling of joint d of the Slim Humanoid agent to change its dynamics.
Hopper We modify the mass of m of the Hopper agent to change its dynamics.
Specifically, all training and test parameter lists are set as {0.75, 0.8, 0.85, 0.90, 0.95, 1, 1.05, 1.1, 1.15, 1.2, 1.25) and {0.2, 0.4, 0.5, 0.7, 1.3, 1.5, 1.6, 188,, respectively.
A.2
ALGORITHM
The training procedure is give at Algorithm] T
A.3 TRAINING DETAILS
Similar to the Lee et al. 2020); Guo al. (2021), train our model-based RL agents and context encoder for 20 epochs, and we collect 10 trajectories by a MPC controller with 30 horizon from environments each epoch. In addition, the cross entropy method (CEM) with 200 candidate actions is chosen as the planing method. Specifically, the batch size for each experiment is 128, /3 is 6e-1. All module are learned by Adam optimizer with 0.001 learning rate.
A.4 PREDICTION ERROR
A.5 NETWORK DETAILS
Similar to the Lee al. (2020), the context encoder is constructed by simple 3 hidden-layer MLP, and the output dim of envirommental-specine vector is 10. The relational head is modelled as single FC layer. The dynamics prediction model is 4 hidden-layer FC with 200 units.
14

---

Algorithm The training algorithm process of our relational intervention approach
Initialize parameters of context encoder é, dynamics prediction model 0 and relational head Initialize dataset B for Each Iteration do
sample environments M from training environments {M, Ji=D
Collecting Data
for I 1 to TaskHorizon do
Get the estimation of the environment-specified factor -t- = -( +)
Collect (St, Qz. S141, T( from M with dynamics prediction model 0
Update B +- B U (8t, dr, Sc-1: Tt, ((-k--)))
Initialize trajectory prototype Ctra for each sampled trajectory
end for
for Each Dynamics Training Iteration do
Update 0.0 and y
for k = to K do
J.b.P Sample data ....P f=t-t-1- Cira and T---t-1 Crra with batch size B.from B s.B.P R.B.P Get the estimation of the environment-specifee factor ft-kt-1 -( and
B....P
-[-ka-1 = g( g(P.B.P....
Estimate the similarity between Cira and C
Construct between Cý and C
Combing top k similar Ciyo into environmental prototypes C
Chrrd B.B stB..P'
(TEM ( 2.0
ci-relation gi.
(1-
C²) with prototypes in different
levels.
Update 0, Φ, + VoNo crat
end for
end for end for
Hopper
0.5
-2.5
CaDM
CaDM
TMCL
TMCL
TMCL
-3.0
RIA
RIA
OURS
Predicioon Enor
OURS
Prodimion Eroo
OURS
-1.0
Prodiction
+4.0
-1.5-
-3.0
-30
10
15
20
10
16
20
Timesteps ('10M)
Timedleps ('104)
Timestaps (1104)
Half Cheelah
Pendulum
Slim Humenoid
-1.5
CIDM TMCL
C2DM
CLDM
-20
TMCL
TMCL
Prediction Ereer
-2.5
OURS
Prediction Erao
OURS
OURS
PredictionEn
-3.5-
4.0
.50
-55-
-25
10
10
15
20
20
Timestips ("1044)
Timestops (0104)
Timestops (11044
Figure 6: The average prediction errors of mrodel-based RL agents on training environments. The results show the mean and standard deviation of prediction errors averaged over five runs. Specifically, we use the no adaptation version TMCL for a fair comparison.
15

---

Table The average prediction errors of model-based RL agents on test environments. The results show the mean and standard deviation of prediction errors averaged over five runs. Specifically, we use the no adaptation version of TMCL for fair comparison.
TMCL
CaDM
RIA
Ours
HalfCheetah Pendulum
0.025=0.011
0.027=:0.033
0.021#0.031
0.013=0.022
0.013+00084
0.016++0.004
0.003+00001
0.002 0.001
Ant
0.069++00013
0.054+00.026
0.051+00024
1.037=0.012
Hopper
0.061+0.013
0.069 +0.028
0.032++00233
0.026+ 0.011
Slim Humanoid
0.145 0.101
0.105:t0.012
0.100
Swimmer 0.078
0.067:00.085
0.037:200055
0.023
A.6 DIRECT CAUSAL EFFECT BETWEEN TRAJECTORY PROTOTYPES
Concretely, the direct causal effect difference between two trajectory prototypes e/ra and ctra can be calculated through the controllable causal effect (Pearl 2013) given as following:
CDE
(st, az) -
= st, Ae at) do(Z = cira]]
(5)
E[ |do(S: = St, A: = az), do(2 2 c(ra)l
(6)
=B|S = St, A, = 0z, d.all - E|S St 81, A: = ac, Z
ct ctral, (7)
where do the do-calculus (Pear 2000). Because we control all variables that can influence on
the S141. and there is no other confounder between the mediators (S,, A²) and S(+1 except Z (Guo
jet af., (022), we can remove all do operators in Eq. (6). Therefore, the intervention distribution of controlling 7, and (Sr, Ar) Eq. (6) is equal to the conditional distribution Eq. 7 In addition, the
should sample St and A, independently Z to calculate the average controlled direct effect.
direct causal effects between dir and cxya may differ for different values of S. and A,,
Concretely, directly a mini-batch of S₁ and As pairs (s), a;) to calculate the average controlled direct effect of them as Figure (a) shows:
Wik
N D
CDE Genn (i)) at)
(8)
where N is the batch size, j and i are the id of trajectory prototypes.
A.7 CONNECTION BETWEEN RELATION Loss AND MUTUAL INFORMATION
We denote the environmentaalspeeiffc factor as Z and its prototypes as C. By definition, the mutual information between Z and C should be:
-Epc) llog( c(z,c) p(=)p(c)
where Pzc is he joint distribution of Z and C, and Pz and Pc are their marginal distributions. To estimate mutual information between Z and C, we can use the probabilistic classifier method proposed by Tsai al. (020)). Concretely, we can one given data pair (z, c) from the joint distribution Pzc (Y = 1) or from the product of marginal distribution P(C) (Y 0). Therefore, the mutual information I [ZZ C*) between Z and C can be rewrite as:
Bernoulli random variable Y classify
I(Z:C) Ep [log( c(z,c)
p(2)p(c)
D/2 (12,cY=11,, =1)_1
Eze [log(
p(z,c) cYY =0)
P(Y O)P(Y = 1|2,
EEze Dlog
(10)
p(Y 1)P(Y = 0
Obviously, p(Y=0) IT can be approximated by the sample size, i.e. "PZPc while P(Y=Ils.c) npze
can be
P[Y=0
measured by a classifier (( (2, c) with the below our relational loss:
crelation - Y log A([z, c): s) + (1 Y) log (1 1((2,,);)))
(11)
16

---

where Y == 1 if the given pair (2, c) is from the joint distribution Pxy, and Y 0 if the given pair (z, c) is from the product of the marginal distributions Pz Pc. Because p(Y=0) tend be a constant, p(Y=1)
optimizing our relation loss is actually estimating the mutual information [(2; C*) between Z and C. Therefore, optimizing Eq. 11 actually maximizing the mutual information between (2) and its corrsponding prototype which represents the semantics of trajectorey or environment. If the readers are interested in the concrete bound about this method to estimate mutual information, please refer to
Tsai et al. 20200; Guo al. (021).
17

---

A.8
VISUALIZATION AND ANALYSIS
Table 3: The quantitative evaluation results of estimated environmennal-specifee factors.
ARI
AMI
V-means
TMCL
CaDM
Ours
TMCL
CaDM
RIA
Ours
TMCL.
CaDM
RIA
Ours
HaltChcetah
0.006
0.128
0.212
0.570
0.058
0.175
0.333
0.681
0.06
0.176
0.314
0.680
Pendulum
0.060
0.471
0.754
0.971
0.054
0.529
0.838
0.967
0.051
0.531
0.724
0.975
Slim_ Humanoid
0.001
0.139
0.212
0.472
0.004
0.121
0.245
0.613
0.004
0.139
0.213
0.612
Swimmer
0.052
0.583
0.586
0.615
0.052
0,582
0.597
0.638
0.012
0.528
0.595
0.637
18

---

TMCL
CaDM
RIA
Ours
Figure 7: The PCA visualization of anvironmental-specinn factors estimated by TMCL Seo et al. (202)) CaDM Lee et al. (2020), RIA Guo et al. (2021) and ours on the Pendulum task.
TMCL
CaDM
RIA
Ours
Figure 8: The PCA visualization of environmental-specifn factors estimated by TMCL Sco et al. (2020), CaDM Lee et al. (2020), RIA Guo et al. (2021) and ours the Halfcheetah task.
TMCL
CaDM
RIA
Ours
Figure 9: The PCA visualization of environmental-specif factors estimated by TMCL Seo et al. (2020). CaDM Lee et al. (2020), RIA Guo et al. (2021) and ours on the slim humanoid task.
TMCL
CaDM
RIA
Ours
Figure 10: The PCA visualization of spviron factors estimated by TMCL Seo et al. (002)), CaDM Lee et al. (2020), RIA Guo et al. (202T) and ours the swimmer task.
19


Reviews:
[{"review_id": "nBTQeDTkE_x", "reviewer": "Reviewer_wQ75", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "The paper is designed to address the issue of model-based reinforcement learning generalizing to new environments with slightly different transition dynamics. Specifically, the case where unobserved environment factors like the friction coefficient of the floor can differ between the training environments and the test environments. The proposed approach estimates latent environmental factors by clustering together trajectories from different environments. The results show that the learned clusters achieve good separation and more closely match the ground truth, and that the technique matches or exceeds the performance of closely related baselines in OpenAI Gym and Mujoco environments. ", "summary_of_the_review": "The clarity issues with the paper make it hard to evaluate the central technical contribution that makes this work different from the closest related prior work (Guo et al.). While the results appear promising, the paper is not yet at the level of an ICLR conference paper. ", "strength_and_weaknesses": "The biggest strength of the paper is that it achieves good results in both clustering and model-based RL performance, as evidenced by Table 1 and Figure 4. To the extent that these results are reproducible, they are of interest to the community. \n\nThe biggest issue with the paper is clarity. Highly critical points needed to understand the contributions of the paper and evaluate it are not made clear. For example, Section 3.2 makes conflicting statements about what kind of trajectory and environment labels are available. It first states that no environment labels are available, \"so we cannot estimate environmental prototypes directly. Fortunately, we still have the trajectory label information\". What is a trajectory label? What information does it provide? Does it indicate in which environment the trajectory was obtained? Wouldn't this be like an environment label? The section then goes on to show that each environment prototype is constructed as the \"mean of its corresponding trajectory prototype\". How do we know which trajectory prototypes correspond to which environments without environment label information? \n\nThe weakness with clarity also impacts the ability to evaluate the novelty of the work. This work seems to be largely based on prior work by Guo et al., including the main losses used (Eq 2 and 3). Therefore, it is not clear what this work contributes beyond the prior work by Guo et al. The authors include a section (3.4) on the \"Difference to RIA\" (Guo et al.), but this section seems to depend heavily on the point that the proposed approach uses trajectory label information, when (as stated above) it is not clear what a trajectory label *is*. Further, Section 3.4 states that unlike Guo et al., the proposed method estimates the similarity between different prototypes (which I take to mean different trajectories). But, according to the implementation details (Section 4.1), in Guo et al. \"the trajectory segments are randomly sampled from the same trajectory\", which seems to imply that Guo et al. applies the proposed losses to trajectories as well. Therefore the difference from Guo et al. is very unclear to me. For the authors to improve this point, it would be good to state how their application of the losses proposed by Guo et al. (Eq 2 & 3) differs for their specific use case. This could be done within sections 3.2 and 3.3, rather than having a high-level post-hoc explanation in section 3.4. Evidently there is some difference from RIA, since the results indicate their method can outperform RIA, but how this difference is achieved is not clear. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "deBHpdIfDXn", "reviewer": "Reviewer_M73e", "summary": "", "questions": "", "limitations": "", "rating": 3, "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": "The paper discusses the problem of learning from multiple MDPs, each corresponding to a different value for hidden factors such as physical properties (friction, gravity etc.) The authors propose estimating an embedding based on a similarity matrix of trajectories, and then use a contrastive loss to learn a $z$ that matches with the estimated embeddings. The learnt $z$ is fed back into a dynamics model and used in a MBRL algorithm. ", "summary_of_the_review": "The paper suffers from readability issues, quite critically in parts where the main contribution of the method is introduced. I found myself going back and forth on multiple occasions to try and understand the actual algorithm. From what I could gather, the method is essentially learning an encoding based on how similar trajectories coming from the different MDPs are. I do not see why the problem setup includes using model-based RL as a part of it. That should be part of the solution approach, and not the problem. Similarly, the experimental comparisons are only shown for model-based methods, which does not make sense given the problem setup should have nothing to do with model-based RL in particular. I think the paper needs a heavy revision in terms of the method description, even before going to the experimental results and hence I am advocating for a rejection. ", "strength_and_weaknesses": "\nStrengths\n\n- Important and Relevant problem\n- Promising general direction for the solution approach\n\nWeakness\n\n- Unclear writing\n- Bad paper structure\n- Lack of novelty in the proposed approach", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "OkpTcros43I", "reviewer": "Reviewer_JcMG", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "This work studies the problem of generalization in (model-based) RL by learning dynamics models that disentangle generalizable factors from environment-specific factors. Specifically, this paper proposes a hierarchical prototypical method (HPM) with the objective of learning to cluster different environments with similar environment-specific factors, thereby facilitating generalization or fast adaptation in new environments by associating it with a learned cluster.", "summary_of_the_review": "While the proposed algorithmic approach is interesting, the problem setting and assumptions (no simulator; only trajectories from diverse environments) are not well motivated. Important baselines are also missing.", "strength_and_weaknesses": "*Strengths*\n- The proposed clustering method is interesting and intuitive.\n- Learning prototypes of environments as a concept is again interesting, intuitive, and to my knowledge novel.\n- Proposed method shows stronger results than the baselines considered.\n\n*Weaknesses*\n- The authors do not provide a clear motivation for the problem setting. In practice, how is it possible to generate trajectories from a large collection of environments with different factors? What are some practical applications that have these characteristics? \n- In my mind, the most plausible way to generate trajectories from diverse environments is through simulation. In this case, there are several sim2real papers that also attempt to learn from simulation and generalize to real world. Such lines of work have not been discussed in related work.\n- In Fig 4, what would be the asymptotic performance of an oracle RL agent trained on the unseen test environment. Without this information, it is unclear how good is the performance of proposed method in absolute terms (i.e. relative to an oracle).\n- I feel several crucial baselines are missing. Examples are \n  - (a) Methods that perform some form of domain randomization, like EPOpt[1], that show that a single policy can be successful despite variations in the environment on several of the tasks considered in this paper.\n  - (b) A simple recurrent policy [2] which has been recently shown to be competitive with methods that explicitly perform adaptation.\n\n\nReferences\n- [1] Rajeswaran et al. 2016. EPOpt: Learning Robust Neural Network Policies Using Model Ensembles. \n- [2] Ni et al. 2021. Recurrent Model-Free RL Can Be a Strong Baseline for Many POMDPs.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject