A Relational Intervention Approach for Unsupervised Dynamics Generalization in Model-Based Reinforcement Learning
Authors: Jiaxian Guo
Mingming Gong
Dacheng Tao
Venue: ICLR 2022 Poster (2022)
Keywords: Model-Based Reinforcement Learning
Unsupervised Dynamics Generalization

Abstract: The generalization of model-based reinforcement learning (MBRL) methods to environments with unseen transition dynamics is an important yet challenging problem.
Existing methods try to extract environment-specified information $Z$ from past transition segments to make the dynamics prediction model generalizable to different dynamics. However, because environments are not labelled, the extracted information inevitably contains redundant information unrelated to the dynamics in transition segments and thus fails to maintain a crucial property of $Z$: $Z$ should be similar in the same environment and dissimilar in different ones. As a result, the learned dynamics prediction function will deviate from the true one, which undermines the generalization ability. To tackle this problem, we introduce an interventional prediction module to estimate the probability of two estimated $\hat{z}_i, \hat{z}_j$ belonging to the same environment.
Furthermore, by utilizing the $Z$'s invariance within a single environment, a relational head is proposed to enforce the similarity between $\hat{{Z}}$ from the same environment. As a result, the redundant information will be reduced in $\hat{Z}$. We empirically show that $\hat{{Z}}$ estimated by our method enjoy less redundant information than previous methods, and such $\hat{{Z}}$  can significantly reduce dynamics prediction errors and improve the performance of model-based RL methods on zero-shot new environments with unseen dynamics. The codes of this method are available at \url{https://github.com/CR-Gjx/RIA}.

Published as a conference paper at ICLR 2022
A RELATIONAL INTERVENTION APPROACH FOR UN- SUPERVISED DYNAMICS GENERALIZATION IN MODEL- BASED REINFORCEMENT LEARNING
Jiaxian Guot Mingming Gongt Dacheng Taotš
The University of Sydneył The University of Melbournet JD Explore A jquo59346uni. 40uni. sydney. edu. au mingming. gong@unime edu. au da cheng. taolgmai com
ABSTRACT
The generalization of model-based reinforcement learning (MBRL) methods to environments with unseen transition dynamics is an important yet challenging problem. Existing methods try to extract environment--pecified information Z from past transition segments make the dynamics prediction model generalizable different dynamics. However, because environments are not labelled, the extracted information inevitably contains redundant information unrelated the dynamics in transition segments and thus fails to maintain crucial property of Z: Z should be similar in the same environment and dissimilar in different ones. As result, the learned dynamics prediction function will deviate from the true one, which undermines the generalization ability. To tackle this problem, we introduce an interventional prediction module estimate the probability of two estimated 2, 2 belonging to the same environment. Furthermore, by utilizing the Z's s invariance within single environment, relational head is proposed to enforce the similarity between 2 from the same environment. As result, the redundant information will be reduced in Ž. We empirically show that ž estimated by our method enjoy less redundant information than previous methods, and such 2 can significantly reduce dynamics prediction errors and improve the performance of model-basee RI. methods on zero-shot new environments with unseen dynamics. The codes of this method are available at httttt:///git thtb. com CR- G
INTRODUCTION
Reinforcement learning (RL) has shown great success in solving sequential decision-making problems, such as board games (Silver et al. 2016 2017; Schrittwieser et al., 2020), computer games (e.g. Atari, StarCraft II) (Mnih et al. 2013 Silver al. [0188] Vinyals et al. 2019). and robotics (Levine) & Abbeel 2014: Bousmalis et al.] 2018). 88. However, solving real-world problems with RL is still challenging problem because the sample efficiency of RL is low while the data in many applications
is limited or expensive to obtain (Gottesman et al. 2018 Lu et al. 2018 2020:
Kiran et al. 2020).
Therefore, model-based reinforcement learning (MBRL) (Janner et al. 2019 Kaiser et 2019
Schrittwieser et al. 2020 Zhang et al.] 2019b; van Hasselt et al. 2019 Hafner el al. 2019b Lenz
al. 2015), which explicitly builds predictive model to generate samples for learning RL policy,
has been widely applied to variety of limited data sequential decision decision-making problems.
However, the performance of MBRL methods highly relies on the prediction accuracy of the learned environmental model (Janner et al. 2019). Therefore, a slight change of environmental dynamics may cause significant performance decline of MBRL. methods (Lee et al., 2020 Nagabandi et al.. 2018a; (Se et al.) 2020). The vulnerability of MBRL to the change of environmental dynamics makes them unreliable in real world applications. Taking the robotic control as an example (Nagabandi
2020 Rakelly et al.] 2019 Gu et al. 2017 Bousmalis et al.) 2018 Raileanu jt al. 2020 Y Yang et al.
et al. 2018a| Yang et al.
2019), dynamics change caused by parts damages could easily lead to the failure of MBRL algorithms. This problem is called the dynamics generalization problem in MBRL, where the training environments and test environments share the same state and action space AA but the transition dynamics between states (88++1 18t, (z) varies across different environments, Following

---

Published as a conference paper at ICLR 2022
previous works (Petersen et al.) 2018 Gottesman et al. 2018), we focus on the unsupervised dynamics generalization setting, L.e. the id or label information of dynamics function in training MDPs is not available. This setting appears in wide range of applications where the information of dynamics function is difficult to obtain. For example, in healthcare, patients may respond differently to the same treatment, i.e., p(st+1(88 0z) varies across patients. However, it is difficult to label which patients share similar dynamics.
CaDM
TMGL
Ownorn Broddr
Transition segments,
Trawiio Spcand
(a)
(d)
Figure I: (a) The illustration of why historical states and actions are encoded in environment-speciifed factor Z, (b)(c)d) The PCA visualizations of estimated context (environmental-sper -specffi) vectors in Pendulum task, where the dots with different colors denote the context vector (after PCA) estimated from different environments. More visualization results are given at Appendix A.13
To build a generalized dynamics prediction function that can generalize to different transition dynamics, the shift of the transition dynamies can be modelled the change of unobserved factors across different environments, i.e. there are hidden environment-specci can affect the environmental dynamics. This is analogous to the human intuition to understand the change of dynamics, e.g. patients may show different responses to identical treatments because the differences of their gene sequences can affect how well they absorb drugs (Wilke et al. 2007). is natural to assume that Z is the same in single environment but varies across different environments. As such, these unobserved wwvironment-speciffe environment, but their changes can affect the dynamics functions across environments. Therefore, the dynamics function f S x A - S can naturally be augmented by incorporating Z to be 2020)
factors Z E 2 which
factors do not violate the nature of MDP in a single
f: S A 2 - S (Rakelly al. 2019 Zhou et al. 2019 Lee al. 2020 Seo et al.
Learning the augmented dynamics function is difficult because the environment-specifee factor Z
unobservable. Previous methods (Rakelly al. 2019 Zhou et al., 2019 Lee et al. 2020) try
to extract information from historical transition segments and use it as a surrogate for Z (Figure
la However, in the unsupervised dynamics generalization setting, the extracted information from historical transition segments inevitably contains redundant information unrelated to the dynamics. The redundant information would cause the surrogate for Z to lose crucial property that characterizes Z: Z should be similar in the same environment and dissimilar in different environments. As shown in Figure Lb the mnvironment-speciffe. information 2 learned by CaDM (Lee et al. 2020) does not form clear clusters for different environments. Because the learned 2 fails to represent environmental information, the learned dynamics function will deviate from the true one, which undermines the generalization ability. To alleviate this problem, TMCL (Seo et al. 2020) directly clusters the environments by introducing multiple prediction heads, i.e. multiple prediction functions. However, TMCL needs to choose the proper prediction head for each new environment, making it hard to be deployed into the scenario with consistently changing environments, e.g. robots walk in the terrain which is constantly changing. To avoid adaptation at the deployment time, we thus need to learn a single generalized prediction function f. To ensure that ff can learn modals of transition dynamics in different environments, we need to cluster Z according to their belonging environments.
In this paper, we provide an explicit and interpretable description to learn Z as a vector 2 (i.e. the estimation of Z) from the history transition segments. To cluster 2 from the same environment, we introduce a relational head module as a learnable function to enforee the similarity between žs learned from the same environments. However, because environment label is not available, we can only cluster the Żs from the same trajectory, so we then propose an interventional prediction module to identify the probability of pair of E,, ε belonging the same environment through estimating Ets direct causal effect on next states prediction by do-calculus (Pearl, 2000). Because Zs from the same environment surely have the same causal effect, we can directly maximize the similarity of 2s with the similar causal effect using the relational head, and thus can cluster 2 according to the estimated environmental similarity and alleviate the redundant information that varies in an environment, e.g. historical states and actions. In the experiments, we evaluate our method
2

---

Published as a conference paper at ICLR 2022
range of tasks in OpenAl gym (Brockman et al. 2016) and Mujoco Todorov et al. 2012). and empirically show that 2 estimated by our method enjoy less redundant information than baselines. The experimental results show that our method significantly reduces the model prediction errors and outperforms the state-of-art model-based RL methods without any adaptation step new environment, and even achieve comparable results with the method directly cluster 2 with the true environment label.
2
RELATED WORK
Dynamics Generalization in MBRL Several meta-lcarming--ased MBRL methods are proposed Nagabandi et al. (201 8a b): Sæmundsson al. (0088: Huang et al. (202T) to adapt the MBRL into environments with unseen dynamics by updating model parameters via a small number of gradient updates Finn et al. (2017) or hidden representations of recurrent model [Doshi Velez & Konidaris (2016), and then Wang & van Hoof (2021) proposes a graph-structured model to improve dynamics forecasting. Ball et al. (2021) focuses on the offline setting and proposes an augmented model method to achieve zero-shot generalization. Lee et al. (2020); (Seo al. (2020) try to learn generalized dynamics model by incorporating context information or clustering dynamics implicitly using multi- choice learning, aiming adapt any dynamics without training. However, how explicitly learn the meaningful dynamics change information remains a big challenge.
Relational Learning Reasoning relations between different entities is an important way to build knowledge of this world for intelligent agents Kemp & Tenenbaum (2008). In the past decades, relational paradigm have been applied
application, e.g., reinforcement learning Zambaldi et al. (2018). arrestion-asswer Santoro et al. al.] (2017): Raposo et al. (2017). graph neural network Battaglia et al. (2018), sequential streams Santoro et al.] (2018). few-shot networkBBattaglia
wide range of deep learning-basee
learning Sung et al. (2018), object detection Hu et al. (2018) and self-supervised learning Patacchiola Storkey (2020). Different from previous methods that perform binary relational reasoning on enti- ties, method can also perform multiplies relations between entities through the learned similarity of entities, and thus can learn more compact and meaningful entity representation.
Causality in Reinforcement Learning Many works focus the intersection area of reinforce- ment learning and causal inference. For example, some works aims to alleviate the causal confusion problem in the imitation learning (de Haan al., 2019: Zhang et al.) [220cc] Kumor al. 2021).
batch learning (Bannon et al. 2020), and partial observability settings (Forney et al. [017
Kallus
& Zhou, 2018: Zhang al.] 2019a; Lu al. 2018) in the online environment (Lyle et 2022
Zhang & Bareinboim] 2018). (Wang et al.) (020a) also try to apply causal inference in the offline
setting, where the observational data always confounded. Lee & Bareinboim (2018) ; Bareiinboi
et al. 2015): Lattimore et al. () Mozifian et al. (2020); Volodin et al. (2020) also explore how to design an optimal intervention policy in bandits or RL settings. addition, (Zhang et al.) 2020a b)
improve the generalization ability of state abstraction. Different from these methods, we focus on the setting of unsupervised dynamics generalization, and measure the direct causal effect (Pearl 2013) between 2 and the next state to estimate the probability of them belonging to the same environment.
3
METHODS
In the section, we first introduce the formulation of the dynamic generalization problem. Then we present the details on how relational intervention approach learns the environmenntspeecffedd factors.
3.1 PROBLEM SETUP
The standard reinforcement learning task can be formalized as a Markov decision process (MDP) M = (S, r. P. 7: Po) over discrete time (Puterman) 2014, SSutton & Barto 2018), where S. A. 7 (0. Po are state space, action space, the reward discount factor, and the initial state distribution, respectively. The reward function r: S x A + R specifies the reward at each timestep t given and a,, and transition dynamics p( sr, a,) gives the next state distribution conditioned on the current state 8t and action a, The goal of RL is to learn a policy (-1s)) mapping from state S over the action distribution to maximize the cumulative expected return over timesteps E,,es.arcAl E 2²t r(s,, a;)]. In model- based RL, a model f is used to approximate the transition dynamics p. and then can provide training data to train policy ær or predict the future sequences for
3

---

Published as a conference paper at ICLR 2022
planning, Benefiting from data provided by learned dynamics model f. model-based RL has higher data efficiency and better planing ability compared with model-free RL.
Here consider the unsupervised dynamics generalization problem in model-based RL, where we are given K training MDPs {M j and test MDPs {M;"); that have the same state and action space but disjoint dynamics functions, and randomly sample several MDPs from training MDPs in each training iteration. We assume that all these MDPs have a finite number of dynamics functions, meaning that the MDPs can be categorized into finite number of environments and the MDPs in each environment share the same dynamics function but the environment id of MDPs is unavailable the training process. In the context of model-baset RL, how learn generalized dynamics model f is the key challenge to solve unsupervised dynamics generalization problem.
3.2 OVERVIEW
Positive: trajectory
Appregation X
Relational Head
12-
Factor Vector
Negative: From different trajectories
relation Loss
8y K
prediciim Loss
Prediction Head
Concetenation
Figure 2: An overview of our Relational Intervention approach, where Relational Encoder, Prediction Head and Relational Head are three learnable functions, and the eircles denote states (Ground-Trutie are with red boundary, and estimated states are with black boundary), and the rectangles denote the estimated vectors. Specifically, prediction Loss enables the estimated muvironmental-specfffe factor can help the Prediction head to predict the next states, and the relation Loss aims to enforce the similarity between factors estimated from the same trajectory or environments.
As analyzed in Section D we can incorporate the dnvironment-specifae factors Z E Z into the dynamics prediction process to generalize the dynamic functions on different environments, i.e. extending the dynamics function from f S x A + S to f S A x Z + S. Because Z is the same within an environment, we expect estimated 2s from the same environment are similar while those from different environments are dissimilar. Therefore, I models the commonalities of the transition dynamics in different environments and 7 models the differences. In the supervised dynamics generalization setting, where the environment id is given, one can easily learn Z by using metric losses, e.g., CPC (Oord et al. 2018) and relation loss (Patacchiola & Storkey 2020) to enforce that the estimated Zs are similar in the same environment and dissimilar in different environments. However, since the environment label unavailable the unsupervised setting, we have to simultaneously learn Z and discover the cluster structures. To this end, we propose an intervention module to measure the similarities between each pair of & and ž as the probability of them belonging to the same environment. Furthermore, aggregate Żs with high probability using relation loss. By simultancously updating the dynamics prediction and the relation loss, can cluster 2s from the same environment, and learn augmented dynamics prediction model f. Next, we will give details about our relational intervention approach.
then introduce relational head to
3.3 RELATIONAL CONTEXT ENCODER
To learn the pewironment-speciffe factor Z of each environment, we firstly introduce relational encoder 9 parameterized by é. Similar to previous methods (Nagabandi af. 20188| Rakelly et al.. 2019 Zhou et al.) 2019 Lee et al.) 2020, Seo et al.) 2020), we use the past transition segments
T: kit-1
(8t- k, 04 k) (8t-1, Ge- 1)) as the input to estimate its corresponding k
-- g(T
4

---

Published as a conference paper at ICLR 2022
After obtaining environment-specifee 2 timestep incorporate it into the dynamics prediction model f to improve its generalization ability different dynamies by optimizing the objective function following (Lee et al. 2020 Seo et al. 2020 Janner et al. 2019
log [(1441 9 (1
where k is the length of transition segments, t is the current timestep and N is the sample size. In practice, we sub-sample a mini-batch of data from the whole dataset to estimate [1 and use stochastic gradient descent update the model parameters.
However, as analyzed in Section 3.2 the vanilla prediction error (1) is not sufficient to capture devironment-speciffe Z of each environment, and even introduce redundant information into it. In order to eliminate the redundant information within transition segments and preserve the trajectory invariant information, function h to pull factors Z from the same trajectory together and push away those from different trajectories. Concretely, the estimated s2 -t- in mini-batch will be firstly aggregated as pairs, e.g. concatenate two factors as 1 2.), and the pairs having two factors from the same trajectory are seen as positives, and vice versa. Then the relational head h parameterized by &o takes a pair of aggregated factors as input to quantify the similarity of given two factors and returns similarity score y, To increase the similarity score & of positive pairs and decrease those negatives, we minimize the following objective:
introduce relational head (Patacchiola & Storkey 2020) learnable
Crelation P.
ΣΣ N(N 1=
y udd-log log h(E"t, +):: (1) ) (1 (1-A(12), (11l1,,,,)))) *1]+,)]] (2)
where nd 1 stands for positive pairs, and ndو 0 stands for negatives. Because the positive pairs have two factors belonging to the same trajectory, optimizing (2) can increase the similarity of žs estimated from the same trajectory, and push away those factors estimated from different trajectories in their semantical space. Therefore, by optimizing ( the information that is invariant within a trajectory will be encoded into 2 and the redundant information transition segments will be reduced. 22) can also be interpreted from the perspective of mutual information, if we regard the 2s from the same trajectory the positive pairs, optimizing 2 be seen maximizing the mutual information between 2s from the same trajectory (Please refer to Tssai et al. 2020) and Appendix A.5 and thus preserve the invariant information with the same trajectory. However, estimating trajectory invariant information is insufficient because the estimated 2s in the same environment will also be pushed away, which may undermine the cluster compactness for estimated Żs.
23
2k
S
S
S
S
Se+1
Sc-1
S
SELt
A--
A
A
A
Ae-
A
Randomize
Randomize
(a)
(b)
Figure 3: (a) The illustration of causal graph, and the red line denotes the direct causal effect from Z
8.41- (b) The illustration estimating the controlled causal effect.
3.4 INTERVENTIONAL PREDICTION
Because the environment id of trajectory unknown, we cannot directly optimize relational loss (2 to cluster 2 within an environment. We propose an interventional prediction method find the trajectories belonging to the same environment. Here we formalize the dynamics prediction model using graphical causal model, and the causal graph is illustrated as Figure (a), where the next state St+1 is caused by the current state S, action A, and 2, and the dynamics prediction model represents the causal mechanism between them. Because Z from the same environment should have
5

---

Published as a conference paper at ICLR 2022
the same influence on the states, and thus they should have the same causal effect on the next state St+1 if given S2 and A, under the causal framework. As such, we can find estimated 2 belonging to the same environment by measuring the similarity of their causal effect on S,41. As Figure (a) shows, there are multiple paths from Z to S1+1 and we roughly categorize them into two categories. The first category is the direct path between Z and S1+1 ( (shown as read in Figure 3). The second category contains all the indirect paths where Z influences S1+1 via previous states and actions. However, because the mediator in other paths e.g. S,, A,, may amplify or reduce the causal effect of Z, we only consider the direct path from Z to the next state(denote by the red line at Figure (a)), which means that we need to block all paths with meditors from 2 to S1+1. By means of do-calculus (Pearl 2000), we can estimate the direct causal effect of changing Z = 3 to Z 2* on S through calculating the controlled direct effect (CDE) (Pearl 2013) by intervening mediators and 2:
CDF (a = [T[[S 4 (S((S 81, A = ac), do Z =²]] - E[S
8t, A a.), do 7 = ==*]]
(3)
BB = St, Ar = ae, Z = 80 E| St st, A: - a, Z = 1*1,
(4)
where do is the do-calculus (Pearl 2000). There is no arrow entering 2, so the do operator on 2 can be removed. Also, since there is no confounder between the mediators (St.AA) and S1+1.
can remove the do operator of them as well, and the equation become as ( Because the direct causal effects may differ for different values of St and A,, we should sample St and A, independently of Z, Le. sampling St and A₂ (Pearl et al. (016) uniformly to get the average controlled direct causal effect from 2 to S1+1- However, if we use the uniformly generated S; and A,, the sampled distribution may differ from the training distribution, resulting in inaccurate the next state prediction. As such, we directly sample St and A, from the observational data. For the convenience of optimization, we only use a mini-batch of S. and A pairs (si, at), and concatenate them with á and g to calculate the average controlled direct effect under f:
N EICDESS (s),a)). N
1
ACDE S*
where N is the batch size, j and i are the id of 2 estimated from two transition segments. Specifically, because the factors 2 estimated from the same trajectory should be the same, and thus we minimize their controlled direct E as Ediat between them in the optimization process. Now we can use the calculated ACD Egs, g as the semantic distance di-A between estimated 2t and 21, and thus we can aggregate factors 2 estimated from similar trajectories by the proposed relational head h. As such, we apply transformation to convert distance metric di-k to a similarity metric w E (0, 1]. which is uddk = exp( (rpp(=""), ), where 33 is a factor controlling the sensitivity of the distance metric. Specifically, because the scale and size of state varies in different tasks, e.g. dims in Pendulum but 20 dims in Half-Cheetah, the optimal 33 may vary in different task. As such, we apply the normalization in the distance metric d, i.e., normalize d with batch variance, to convert it as a relative distance within single task, thus making the optimal stable in different tasks. Then we can directly aggregate similar trajectories by extending the loss function (2) as follows:
N
N
1
Ci ci-relation 2.
l + (1- y. w -log h([2', (12,,
N(N
j=1
+ +11 y )و (1 wid) log (1 h(l2t, (118t.21:4)) 21:41))]
(6)
where the first term indicates that the factors from the different trajectories can be aggregated with the similarity weight w and 1 those from the same trajectory, and the second term means that factors from different trajectories should be pushed with each other with weight w. Similar to the analysis in section 3.3, optimizing the loss function (6) can increase the similarity between Ž with weight w, and w. Because 2s estimated from the same environment have
push away from them with the weight similar effects, these factors will be assigned with high similarities (estimated by the intervention operation of our paper). By simultaneously updating the prediction loss ( and intervention relation loss 6 estimated Żs within the same environment will be aggregated, and the learned dynamics function f can learn the modals of transition dynamics according to the 2 in different clusters. The training procedure of our approach can refer Algorithm process in Appendix A.2

---

Published as a conference paper at ICLR 2022
4 EXPERIMENTS
In this section, we conduct experiment to evaluate the performance of our approach by answering the following questions:
Can our approach reduce the dynamics prediction emros in model-based RL? (Section 4.2.
Can our approach promote the performance of model-based RL on environments with unseen dynamics? (Section 4.2.2)
Can our approach learn the semantic meaningful dynamics change? (Figuer 1 and Ap- pendixA.13)
Is the similarity of w measured by the intervention module reasonable? (Appendix A.7
Can solely relational learning improve the performance model-based RL? (Section
4.1 ENVIROMENTAL SETUP
Implementation Details Our approach includes three learnable functions, including relational encoder, relational head and prediction head. All three functions are constructed with MLP and optimized by Adam (Kingma & Ba 2014) with the learning rate le-3. During the training procedure, the trajectory segments are randomly sampled from the same trajectory to break the temporal correla-
tions of the training data, which was also adopted by (Seo et al. 2020 Wang al., 2020b 2019). Specifically, the length of the transition segments, i.e., k. is 10. All implementation details can be found in Appendix A.I
Datasets Following the previous methods (Lee et al. 2020 Seo et al. 2020), we perform ex-
periments on classic control task (Pendulum) from OpenAl gym (Brockman al. 2016) and simulated robotics control tasks (HalfCheetah, Ceiphee-HalCCheetah, Ant, Hopper, Slim-Humanowe from Mujoco physics engine (Todorov et al. 2012)).
Dynamics Settings To change the dynamics of each environment, we follow previous methods (Zhou et al. 2019 Packer et al., 2019) Lee et al. [020 Seo et al.) (2020) to change the environmental parameters (e.g. length and mass of Pendulum) and predefine them in the training and test environ- mental parameters lists. At the training time, we randomly sample the parameters from the training parameter list train relational context encoder and dynamies prediction model. Then test model on the environments with unseen dynamics sampled from the test parameter list. Specifically, the predefined parameters in the test parameter list are outside the training range. The predefined training and test parameter lists for each task are the same with (Lee et al. 2020). and all details are given Appendix [.I
Planning Following (Lee et al. 2020 Seo et al. 2020), we use the model predictive model (MPC) (Maciejowskij 2002) to select actions based on learned dynamics prediction model, and assume that the reward functions of environments are known. Also, the cross-entropy method (CEM) (De Boer cct al. 2005) is used to optimize action sequences for finding the best performing action sequences. Baselines We compare our approach with following state-of-the-art model-based RL methods on dynamics generalization. Also, to show the performance gap between our method and supervised dynamics generalization, perform the method using true environment label to cluster Z.
Probabilistic ensemble dynamics model (PETS) (Kurutach et al. 2018): PETS employs an probabilistic dynamics models capture the uncertainty in modeling and planning.
Meta learning based model-based RL (ReBAL and GrBAL) (Nagabandi et al., 2018b a):
These methods train a dynamics model by optimizing a meta-objective (Finn et al. 2017). and update the model parameters by updating hidden with a recurrent model or by updating gradient updates at the test time.
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-wise Multiple Choice Learning (TMCL) (Seo et al. 2020): This method is the state-of-the-ar model-based RL method on dynamics generalization, which introduces the multi-choice learning to cluster environments. TMCL needs the adaptation in the test procedure, while our method does not, so we also report the performance of TMCL without adaptation in Figure 7 for the fair comparison.
7

---

Published as a conference paper at ICLR 2022
True Label: The method uses relational head to cluster Z with the true environment label (not the ground-trum truth of Z). All hyperparameters are same with our method for the fair comparison.
4.2 PERFORMANCE COMPARISONS WITH BASELINES
4.2.1 PREDICTION ERROR COMPARISONS
We first evaluate whether the dynamics model trained by our methods can predict next-states more accurately not. Figure 4 shows that the average dynamics prediction error of dynamics prediction models trained by three methods (CaDM (Lee et al. 2020). TMCL (Seo et al., (020) and ours). We can see that the dynamics model trained by our relational intervention method has superior prediction performance over other state-of-the -f-ther-tt- methods, achieving the lowest prediction errors on almost all six tasks. Specifically, the prediction errors of our model are lower than others by a large margin in Hopper and Pendulum, outperforming the state-of-the--e aht methods by approximately 10%.
Pendulum
Hopper
Halfchestah
Slim Humanoid
CaDM
CabM
TMCL
TMCE
TMCL
TMCL.
-M
Ours
-4
445
40
Trensips (T00))
Termsius (TP4))
Figure 4: The average prediction emors of dynamics models training environments during training process (over three times). Specifically, the x axis is the training timesteps and axis is the log value of average prediction prediction errors. More figures are given Appendix A.8
Table 1:
The average rewards of baseline model-based RL methods and ours on test environments with unseen dynamics. Here report the average rewards over three (ours ten). Specifically, the results of methods with * are from the paper (Lee et al. 2020).
PETS*
ReBAL
GrBAL
CaDM
TMCL
Ours
Ratio
Pendulum Ant 965,883.5
-1103
-943.6
-1137.9
-713.95:121.1
.691.2-1.3.4
-587.5116444
15.0%
63.0
44.7
1660=57.8
2994,9#243.88
3297.9:159.77
10.1%
Hopper
821.2
846.2
621
845.2120.41
999.35
1057.4137.2
5.8%
HalfCheetah
1720.9
52
-69.1
5876.61799.0
9039.6=1065
10859.2:146..1
20.1%
C_ C.HalfCheetah
1572
868.7
2814
3656.44856.22
3998.8++866.2
4819,3+-409.3
20.5%
Slim Humanoid
784.5
97.25
-480.7
859.1+24.01
2098.7+109.1
2432.61465.1
15.9%
4.2.2 PERFORMANCE COMPARISONS
Then we evaluate the generalization of model-based RL agents trained by our methods and baselines on test environments with unseen dynamics. Following the setting of (Seo et al. 2020), we perform experiments three runs (ours with 10 runs to reduce random errors), and give the mean of rewards at
We T better than vanilla PETS (Kurutach et al.) 2018), while methods (Lee et al. 2020: Seo et al. 20200 that aim to learn generalized dynamics prediction model superior others significantly. Among which our approach achieves the highest rewards on all six tasks among all methods. Figure 5 shows the mean and standard deviation average rewards during the training procedure, indicating that the performance of our methods is better than the other two methods consistently at the training time, which is sufficient to show the superiority of our method over other methods. fair comparison between TMCL (no adaptation) and our method can be found at Appendix [A.6] In addition, we observe that our method achieves comparable results with the method directly cluster 2 using the truth environment label, which indicates that our intervention module actually can assign high similarities into Zs estimated from the same environment in an unsupervised manner. We also observe the same results in the similarity visualization in the Appendix A.7] where we find that Żs from the same environment are assigned significant higher similarities than those pairs from different environments.
that the meta-learning based methods (Nagabandi et al. 2018b al do not perform
8

---

Published as a conference paper at ICLR 2022
Pendulum
Happer
Halfcheelah
Slim Humanoid
2000
CaDM
10:00
Oun
E00
True Label
1900
CaDM
600
CaDM
CaDM
TMCL
TMCL
Our
Ours
Clurs
Tror Label
Tiue
Figure 5: The average rewards of trained model- -ased RL agents unseen test environments. The results show the mean and standard deviation of returns averaged over three runs. The fair comparison between TMCL (no adaptation) and our method can be found in Appendix A.6
4.3 ABLATION STUDY
In this section, we evaluate the effect of the proposed relation head and intervention prediction on the generalization improvement, respectively. Because the intervention prediction is based on the relational head, we compare the performance of our approach with and without the intervention. As Figure 6a and [bb show, after incorporating the relational head and intervention prediction, the performance of model-based agents and the generalization of the dynamics prediction model are both improved. However, although the model without the intervention module has lower prediction errors in the Pendulum task, also has lower rewards than the whole model. One possible reason is that the Pendulum is simple for the dynamics prediction model learn, and thus the dynamics prediction model with the vanilla relational head is a little over- wer--fitting on the training environments (Please refer prediction errors on test environments are given in Appendix A.9.. limiting the performance improvement. This phenomenon confirms the importance of our intervention prediction on reducing the sereectorr-sppcciie redundant information.
Hallcheelah
Pendulum
Halfcheelas
Pendulum
1200
CaDM
CaDM
Chara Intervation)
Intervaation)
Cura
Dun
Dan (ko istervvtion)
CalMM
Our (un Int.er cation)
Curr
Theeno
2906
Yrustha 2904
(a)
(b)
Figure 6: (a) The average rewards of trained model-based RL agents unseen environments. The results show the mean and standard deviation of returns averaged over three runs. (b) The average prediction errors over the training procedure. Prediction errors on test environments given in Appendix A.9
5 CONCLUSION
In this paper, we propose a relational intervention approach to learn a generalized dynamics prediction model for dynamies generalization in model-based reinforcement learning. Our approach models the dynamics change as the variation of environment-specifee factor 2 and explicitly estimates Z from past transition segments. Because environment label is not available, it is challenging to extract Z from transition segments without introducing additional redundant information. We propose an intervention module to identify the probability of two estimated factors belonging the same environment, and a relational head cluster those estimated Żs are from the same environments with high probability, thus reducing the redundant information unrelated to the environment. By incorporating the estimated Ž into the dynamics prediction process, the dynamics prediction model has stronger generalization ability against the change of dynamics. The experiments demonstrate that our approach significantly reduce the dynamics prediction error and improve the performance of model-based agents on new environments with unseen dynamics.

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGE
Mr Jiaxian Guo is supported in part by Australian Research Council Projects FL-110110111 and LE- 200100049, Dr Mingming Gong supported by Australian Research Council Project DE210101624.
7 ETHICS STATEMENT
Our paper provides a method to generalize the agent trained by the model-based reinforcement learning into new environments with unseen transition dynamics, which may significantly improve the robustness of trained agents in the complex real-world environment, thus broadening the application scenarios of robots trained by reinforcement learning. Although far away apply such an algorithm to real-world applications, still need prevent the algorithm from being applied in military areas.
8 REPRODUCIBILITY STATEMENT
We have run our experiments over three runs (ours with 10 runs to reduce random errors) to reduce random errors, and public hyperparameters used our experiments. The codes of this method are available at ptttp::/gitthe com/CE
REFERENCES
Philip J Ball, Cong Lu, Jack Parker-Holder, and Stephen Roberts. Augmented world models facilitate zero-shot dynamics generalization from single offline environment. arXiv preprint arXiv:2104.05632, 2021.
James Bannon, Brad Windsor, Wenbo Song, and Tao Li. Causality and batch reinforcement learning: Complementary approaches to planning in unknown domains. arXiv preprint arXiv:2006.02579, 2020.
Elias Bareinboim, Andrew Forney, and Judea Pearl. Bandits with unobserved confounders: A causal approach. Advances in Neural Information Processing Systems, 28: 1342-11335 2015.
Peter W Battaglia, Jessica B Hamrick, Vietor Bapst, Alvaro Sanchez-Gonzaler Vinicius Zambaldi. Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, et al. Relational inductive biases, deep learning, and graph networks. arXiv preprint arXiv: 1806.0/261, 2018.
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 on 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 arXiv: 01540, 2016.
Pieter-Tjerk De Boer, Dirk P Kroese, Shie Mannor, and Reuven Rubinstein. tutorial on the cross-entropy method. Annals of operations research, 134(1):19-67, 2005.
Pim de Haan, Dinesh Jayaraman, and Sergey Levine. Causal confusion in imitation learning. arXiv preprint arXiv:1905. 11979, 2019.
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.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-leariig for fast adaptation of deep networks. In International Conference Machine Learning, pp. 11126-188 135. PMLR, 2017.
10

---

Published as a conference paper at ICLR 2022
Andrew Fomey, Judea Pearl, and Elias Bareinboim. Counterfactual data-fusion for online rein- forcement learners. In International Conference on Machine Learning, pp. 1156-1164, PMLR, 2017.
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: 12298, 2018.
Shixiang Gu, Ethan Holly, Timothy Lillicrap. and Sergey Levine. Deep reinforcement learning for robotic manipulation with asynchronous aff-policy updates. In 2017 IEEE international conference robotics and automation (ICRA), pp.
IEEE, 2017.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint /9Xiv:/9/2.01/0 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-2563 PMLR, 2019b.
Han Hu, Jiayuan Gu, Zheng Zhang. Jifeng Dai, and Yichen Wei. Relation networks for object detection. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 3588-3597, 2018.
Biwei Huang, Fan Feng, Chaochao Lu, Sara Magliacane, and Kun Zhang. Adarl: What, where, and how to adapt in transfer reinforcement learning. arXiv preprint arXiv:2107.02229, 2021.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When trust your model: Model-based policy optimization. arXiv preprint arXiv: /906: 08253, 2019.
Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, al. Model-based reinforcement learning for atari. arXiv preprint arXiv:/903..0 aee::/903000744 2019.
Nathan Kallus and Angela Zhou. Confounding-robust policy improvement. arXiv preprint 18X5:0880508595 2018
Charles Kemp and Joshua B Tenenbaumm The discovery of structural form. Proceedings of the National Academy of Sciences, 105(31):1088--066)) 2008.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv: 1412.6980 2014.
B Ravi Kiran, Ibrahim Sobh, Victor Talpaert, Patrick Mannion, Ahmad A Al Sallab, Senthil Yogamani, and Patrick Pérez. Deep reinforcement learning for autonomous driving: A survey. arXiv preprint arXiv:2002. 00444, 2020.
Daniel Kumor, Junzhe Zhang, and Elias Bareinboim. Sequential causal imitation learning with unobserved confounders. 2021.
Thanard Kurutach, Ignasi Clavera, Yan Duan, Aviv Tamar, and Pieter Abbeel. Model-ensemble trust-region policy optimization. arXiv preprint arXiv: 1802. /0592, 2018.
P. Langley. Crafting papers machine learning. In Pat Langley (ed.), Proceedings of the 17th International Conference on Machine Learning (ICML 2000), pp. 1207-1216, Stanford, CA, 2000. Morgan Kaufmann.
Finnian Lattimore, Tor Lattimore, and Mark D Reid. Causal bandits: Learning good interventions via causal inference. arXiv preprint 1rX6::1606.03203 2016.
Kimin Lee, Younggyo Sco, Seunghyun Lee, Honglak Lec, and Jinwoo Shin. Context-aware dynamics model for generalization in model-based reinforcement learning. In International Conference Machine Learning, pp. 5757-5766. PMLR, 2020.
Sanghack Lee and Elias Bareinboim. Structural causal bandits: where to intervene? Advances in Neural Information Processing Systems 31, 31, 2018.

---

Published as a conference paper at ICLR 2022
Ian Lenz, Ross A Knepper, and Ashutosh Saxena. Deepmpe: Learning deep latent features for model predictive control. In 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.
Chaochao Lu, Bernhard Schölkopf, and José Miguel Hernández-Lohato. Deconfounding reinforce- ment learning in observational settings. arXiv preprint arXiv: /8/2. 2018.
Chaochao Lu, Biwei Huang, Ke Wang. José Miguel Hernández-Lobaro, Kun Zhang, and Bernhard Schölkopf. Sample-efficient reinforcement learning via counterfactual-basel data augmentation. arXiv preprint arXiv:2012.09092 2020.
Clare Lyle, Amy Zhang, Minqi Jiang, Joelle Pineau, and Yarin Gal. Resolving causal confusion in reinforcement learning via robust exploration. In Self-Supervision for Reinforcement Learning Workshop-ICLA 2021, 2021.
Jan Marian Maciejowski. Predictive control: with constraints. Pearson education, 2002.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv:1312.5602, 2013.
Melissa Mozifian, Amy Zhang, Joelle Pineau, and David Meger. Intervention design for effective sim2real transfer. arXiv preprint S.Kiv:20/2.00055: 2020.
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-reinforcemmnt learning. arXiv preprint 1803: 1347, 2018a.
Anusha Nagabandi, Chelsea Finn, and Sergey Levine. Deep online learning via meta-learning: Continual adaptation for model- rl. arXiv preprint arXiv:1812.0767) 2018b.
Aaron den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv: 1807.03748 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 arXiv:2006. 05849, 2020.
J Pearl. Causality: Models, reasoning, and inference 47cambridge university presscambridge, united kingdom. pearl, j. 2000. Causality: models, reasoning, and inference, 47, 2000.
Judea Pearl. Direct and indirect effects. arXiv preprint arXiv:1301.2300, 2013.
Judea Pearl, Madelyn Glymour, and Nicholas Jewell. Causal inference in statistics: A primer. John Wiley & Sons, 2016.
Brenden K Petersen, Jiachen Yang, Will S Grathwohl, Chase Cockrell, Claudio Santiago, Gary An, and Daniel M Faissol. Precision medicine as a control problem: Using simulation and deep reinforcement learning to discover adaptive, personalized multi-cytokine therapy for sepsis. arXiv preprint arXiv: 10440, 2018
Martin L. Puterman. Markov decision processes: discrete stochastic dynamic programming. John
Wiley & Sons, 2014.
Roberta Raileanu, Max Goldstein, Arthur Szlam, and Rob Fergus. Fast adaptation to environ- ments via policy-dynamics value functions. In International Conference on Machine Learning, pp. 7930-7931. PMLR, 2020.
Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient off off-policy menffreeinorreement learning via probabilistic context variables. In International conference on machine learning, pp. 5331-5340. PMLR, 2019.
12

---

Published as a conference paper at ICLR 2022
David Raposo, Adam Santoro, David Barrett, Razvan Pascanu, Timothy Lillicrap, and Peter Battaglia. Discovering objects and their relations from entangled scene representations. arXiv preprint arXiv:1702.05506, 2017.
Steindór Sæemundsoon, Katja Hofmann, and Mare Peter Deisenroth. Meta reinforcement learning with latent variable gaussian processes. arXiv preprint arXiv:
2018.
Adam Santoro, David Raposo, David GT Barrett, Mateusz Malinowski, Razvan Pascanu, Peter Battaglia, and Timothy Lillicrap. A simple neural network module for relational reasoning. arXiv preprint arXiv: / 01000.0112,,,.. 2017.
Adam Santoro, Ryan Faulkner, David Raposo, Jack Rae, Mike Chrzanowski, Theophane Weber, Daan Wierstra, Oriol Vinyals, Razvan Pascanu, and Timothy Lillicrap. Relational recurrent neural networks. arXiv preprint 188.v:1806.01822. 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 609, 2020.
Younggyo Seo, Kimin Lee, Ignasi Clavera, Thanard Kurutach, Jinwoo Shin, and Pieter Abbeel. Trajectory- wise multiple choice learning for dynamics generalization in reinforcement learning. arXiv preprini arXiv:2010.1383 2020.
David Silver, Aja Huang, Chris Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panncershelvam, Marc Lanctot, et al. Mastering the game go with deep neural networks and tree search. nature, 529475877:4844488, 2016.
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):354-33) 2017.
David Silver, Thomas Hubert, Julian Schrittiessrr, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Mare Lanctot, 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.
Flood Sung, Yongxin Yang, Li Zhang, Tao Xiang, Philip HS Torr, and Timothy M Hospedales. Learning to compare: Relation network for few-shot learning. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1199-1208, 2018.
Richard Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
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-Philippe Morency, and Ruslan Salakhutdi- nov. Neural methods for point-wise dependency estimation. arXiv preprint arXiv:2006.0555: 2020.
Hado P Hasselt, Matteo Hessel, and John Aslanides. When parametric models in reinforcement learning? Advances in Neural Information Processing Systems, 32: 32:14322-14333, 2019.
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-554 2019.
Sergei Volodin, Nevan Wichers, and Jeremy Nixon. Resolving spurious correlations in causal models of environments via interventions. arXiv preprint arXiv:2002.05217, 2020.
Lingxiao Wang, Zhuoran Yang, and Zhaoran Wang. Provably efficient causal reinforcement learning with confounded observational data. arXiv preprint arXiv:2006. /23/000001331111 2020a.
13

---

Published as a conference paper at ICLR 2022
Qi Wang and Herke van Hoof. Model-based meta reinforcement learning using graph structured surrogate models. arXiv preprint arXiv:2102.06291, 2021.
Zhen Wang, Rui Zhang, Jianzhong Qi, and Bo Yuan. Dbsvec: Density-based clustering using support vector expansion. In 2019 IEEE 35th International Conference on Data Engineering (ICDE), pp. 280-291. IEEE, 2019.
Zhen Wang, Liu Liu, and Dacheng Tao. Deep streaming label learning, In International Conference on Machine Learning, pp. 9963-99722 PMLR, 2020b.
Russell A Wilke, Debbie W Lin, Dan M Roden, Paul В Watkins, David Flockhart, Issam Zinch, Kathleen M Giacomini, and Ronald M Krauss. Identifying genetic risk factors for serious adverse drug reactions: current progress and challenges. Nature reviews Drug discovery, 6(11 ((11):9049116, 2007.
Jiachen Yang, Brenden Petersen, Hongyuan Zha, and Daniel Faissol. Single episode policy transfer in reinforcement learning. arXiv preprint arXiv:1910007119, 2019.
Yuxiang Yang, Ken Caluwaerts, Atil Iscen, Tingnan Zhang, Jie Tan, and Vikas Sindhwani. Data efficient reinforcement learning for legged robots. In Conference on Robot Learning, pp. 1-10. PMLR, 2020.
Vinicius Zambaldi, David Raposo, Adam Santoro, Victor Bapst, Yujia Li, Igor Babuschkin, Karl Tuyls, David Reichert, Timothy Lillicrap, Edward Lockhart, et al. Deep reinforcement learning with relational inductive biases. In International Conference on Learning Representations, 2018.
Amy Zhang, Zachary Lipton, Luis Pineda, Kamyar Azizzadenesheli, Anima Anandkumar, Laurent Itti, Joelle Pineau, and Tommaso Furlanello. Learning causal state representations of partially observable environments. arXiv preprint arXiv: 1906. 10437, 2019a.
Amy Zhang, Clare Lyle, Shagun Sodhani, Angelos Filos, Marta Kwiatkowska, Joelle Pineau, Yarin Gal, and Doina Precup. Invariant causal prediction for block mdps. In International Conference on Machine Learning. pp. 1 11214-111 11224. PMLR, 2020a.
Amy Zhang, Shagun Sodhani, Khimya Khetarpal, and Joelle Pineau. Learning robust state abstrac- tions for hidden-- block mdps. In International Conference on Learning Representations, 2020b.
Junzhe Zhang and Elias Bareinboim. Fairness in decision-making- the causal explanation formula. In Thirty-Second AAAI Conference on Artificial Intelligence, 2018.
Junzhe Zhang, Daniel Kumor, and Elias Bareinboim. Causal imitation learning with unobserved confounders. Advances in neural information processing systems, 33, 2020c.
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. 7444-7453. PMLR, 2019b.
Wenxuan Zhou, Lerrel Pinto, and Abhinav Gupta. Environment probing interaction policies. arXiv preprint arXiv:1907. 001119977777771 17400 2019.
14

---

Published as a conference paper at ICLR 2022
A APPENDIX
We promise that we will public all codes after the acceptance of this paper.
A.1 ENVIRONMENTAL. SETTINGS
We follow the environmental settings of Lee et al. (2000) dynamics generalization. The details of settings are given as follows:
Pendulum We modify the mass and the length / of Pendulum to change its dynamics.
Half-Cheetah We modify the mass of regid link m and the damping of joint d of Half- Cheetah agent to change its dynamics.
Crppled_Cheetann We cripple the of leg of Half-Cheetab agent to change its dynamics.
Ant We modify the mass of ant's leg to change its dynamics. Specifically, modify two legs by multiplying its original mass with m, and others two with
Slim Humanoid We modify the mass 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.
The training and test modified parameter list can be found at the Table 2
Table 2: The environmental settings in our paper.
Training Parameter List
Test Parameter List
Episode Length
m E (0.75.008..885,0990000.33 1.1.05.1111115:121.2215
m
(0.2,0.4,0.5,007
1.3.1,5
Pendulum
200
1 (0.75508.0088,0090000055 I.1.0...L1111551221..25...
Le (0.2.0.4.0.5.007,
1.3.1.5.166.188)
m (0.2.0.3,0..40055,
E (0.755.088.10001155..2533
1.5.1.6.1.7.188
1000
d E (0.75,0.85, 1.0011111155
d (0.2.0.3,0.4.0055
1.5.1.6.1.7.188))
C Cheetah
e (0.1,2.3)
C E [4.5]
1000
m E 11.20,0.5..033,0.35,,00
Ant
m (0.85,0.90.099511000)
1000
0.45,0.50.0.55,0.66)
E (0.40.0.50.0.60.007700
m E {0. (0.8000.90110000115511..3
1,0,,.,0,,..7,,,,
Slim Humanoid
1000
d [0.80.0.900100011111...25
d E (0.40.0.50.0.60000.77
1.50...60.1.70,..88)
Hopper
m (0.5, 0.75, 1.0, 1.25, 1.5)
m (0.25, 0.375, 1.75, 2.0)
500
A.2 ALGORITIM
The training procedure is give at Algorithm
A.3 TRAINING DETAILS
Similar to the Lee al. (2020), train our model-based RL agents and relational 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, B is 6e- 1. All module are learned by Adam optimizer with 0.001 learning rate.
15

---

Published as a conference paper at ICLR 2022
Algorithm 1 The training algorithm process of our relational intervention approach
Initialize parameters of relational 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 - = -( x)
Collect (St, Qz. S141, t(- from M with dynamics prediction model 0
Update B <- B U (8t, dt, Sc-1: Tt, -
end for
for Each Dynamics Training Iteration do
P Update 6,0 and
for k = 1 to K do Sample data P.b.P L.b.K
f1- . T..M and j.b.P K.b.K
-it-t-1 T.M with batch size B,from B B..B..P
-(T and Get the estimation of the environment-specifice factor kt-
st..B..P s(T..B.P -(((-kt-i;d)
2t-ka-1
Estimate the probability w of ,.. tt-kit-1 and B.....P -- belonging to the same environment.
Compute the total loss
Ctot= - 0.0 (T:M 144 -
B.B.. si...,P +
L.B.K ci-relation (3) + 0.0 (TeM -
Update 0 0 4 Vo.0 ftot
end for
end for
end for
A.4 NETWORK DETAILS
Similar to the Lee et al. 2020), the relational encoder is constructed by a simple hidden-layer MLP, and the output dim of eavironmentallspeeiff vector 10. The relational head is modelled single FC layer. The dynamics prediction model is a 4 hidden-layer FC with 200 units.
A.5 CONNECTION BETWEEN RELATION LOss AND MUTUAL INFORMATION
Given pair of data (z, y) E X V, we donote the joint distribution of and Y are Pxy, and their marginal distributions are Px and Py, respectively. By definition, the mutual information between X and Y is:
T(X;Y) Epxy [log( p(z, v)
p(z)p(y)
To estimate mutual information between X and Y. (Tsai et al. 2020) proposes a probabilistic classifier method. Concretely, we can use Bernoulli random variable C to classify one given data pair (I, y) from the joint distribution Pxy (C = 1) or from the product of marginal distribution P(Y) (C 0). Therefore, the mutual information I(X:Y) between X and Y be rewrite as:
p(z,z) v) p(z)p(y)
I(X;Y) Epsy [log(
p(z,yl -1)1-11)) =1)
Exxy [log(
p(x, y 0)
PC = P)P(C = 11x,
EExx
(8)
p(C 1)P(C 0|2,
P(C=0 can be measured by classifier hi(C(z;, y), and it can be learned by our relation loss with relational head h:
Obviously, p(C=0)
P(C= can be approximated by the sample size, i.e. "PxPy, while P(C=1la.))
"Pxy
crelation = - C log h[I vl: () + (1 - C) log (1 - h( vi: v))+
p.
where C' = 1 if the given pair (x, y) is from the joint distribution Pxy, and C' = 0 if the given pair (z,, y) is from the product of the marginal distributions Px Because p(C=0) tend be constant, p(C=1)
16

---

Published as a conference paper at ICLR 2022
optimizing relation loss is actually estimating the mutual information I(X;Y) Y) between X and Y. As such, if we regard the pairs (2) from the same trajectory/environment as positive pairs, and others negative pairs. optimizing 2 actually maximizing the mutual information between (2) from the same trajectory/envioomment and thus preserve the trajectory/enviromment invariant information. If the readers are interested in the concrete bound about this method to estimate mutual information, please refer to (Tsai et al. 2020).
A.6 FAIR COMPARTSON WITH TMCLL
Because TMCL needs an adaptation process when deploying it into the real world while our method does not. For the fair comparison and show the significance of our method over TMCL, we test the performance of TMCL with no adaptation, and show the results below:
Pendulum
Hopper
Halfcheelah
Slim Hummnoid
1200
Cublee
20.0
2010
1000
6008 Cant
Calo1
Mar)
-000
Trentigss (70/4)
Ttrestion
Figure 7: The estimated similarities between 2s (learned by the model without Intervention module) from different environments and anchors (mass=1) on different tasks, where the box represent the range of 50% samples, the line in the middle of a box denotes the average similarity and the top/bottom lines denote the max/min similarities.
We can see that the average returns of TMCL without adaptation are significantly lower than ours, especially for the classic control task Halfcheetah. The experimental comparison with TMCL without adaptation is direct evidence to support our claim: environmeent-separate generalization of dynamics functions, and our method can significantly outperform baselines in zero-shot unseen test environments with different dynamics. Specifically, the performance of TMCL with no adaptation is still superior to the CaDM, this is because TMCL uses the invariance of Z within a trajectory (Z should predict other states within a trajectory in TMCL), which is similar to our paper with no intervention module.
7s is important for the
A.7 SIMILARITIES VISUALIAZTION
To evaluate the correctness of the estimated similarity of our intervention module, we use a ga estimated from the environment where the mass is 1 as the anchor, and randomly sample 200 2 estimated from different environments (including mass = 1). Then we calculate the similarity between anchor 2t and 21, and visualize the similarities according to their environments. As Figure 8 shows, ss belonging to the same environment with the anchor á² have significant higher similarities than those belonging other environments, and even higher than 0.8 in some tasks (all are higher than 0.6), which shows that our intervention module can successfully identify whether two és from the same environment or not.
To study the role of the intervention module, we also visualize the similarity of 2¹ learned by the model without the intervention module, and the results are given as Figure 9 Figure 9 shows that many contexts from different environments still have high similarities. This indicates that the existing relational learning cannot separate environment-specifeed factors Zs. By contrast, after incorporating the intervention module, the contexts from different environments have significantly smaller similarities than those from the same environments. The comparison between Figures & and g directly shows that our intervention module is valuable to predict whether two contexts are from the same environment or not.
A.8 PREDICTION ERRORS ON TRAING ENVIRONMENTS
The prediction emrrr of each method training environment are given Figure 10
17

---

Published as a conference paper at ICLR 2022
Slim Humanoid
Pendulum
HalfCheetah
Hopper
03
00
Figure 8: The estimated similarities between 2s (learned by the model with Intervention module) from different environments and anchors (mass=1) different tasks, where the box represent the range of 50% samples, the line in the middle of box denotes the average similarity and the top/bottom lines denote the max/min similarities.
Slim Humanoid
Pendulum
HalfCheelah
Hopper
1.00
1.52
04
0.00
maaid 75HMAAOOSS mamrrt Staaprt.3
Figure 9: The estimated similarities between Żs (learndd by the model without Intervention module) from different environments and anchors (mass=1) different tasks, where the box represent the range of 50% samples, the line in the middle a box denotes the average similarity and the top/bottom lines denote the max/min similarities.
Hopper
Ant
Cripple_ Halkcheelhh
-25
TMCL
TMCL
-30
TMCL
Prodicion Eroo
-2.5
+2.0
-4.0
-3.
Procuction
-45
-40
-50
-65
-95
-60
15
10
15
20
13
15
20
Timestops ('104)
Timestops (1104)
Timestops (110%)
Perdulum
Humsndid
-25
CaDM
CaDM
CaDM
-0.0
-3.5
Ours
Ours
Ours
Predition Errr
PreditionE Emar
-22
15.0
-44
-66
-26
-28
-30
Timeslops 10144)
Timeslops (1104)
Timestops (1104)
Figure 10: The average prediction errors of dynamics models on training environments during training process (over three times). Specifically, the axis is the training timesteps and axis is the log value of average prediction prediction errors. More figures are given at Appendix A.8
A.9 PREDICTION ERRORS ON TEST ENVIRONMENTS
The prediction emrors of each method test environments given at Table Specifically, test each test environment 10 times, and plot the average prediction error to reduce random errors (Figure
11
18

---

Published as a conference paper at ICLR 2022
Pendulum
Hopper
CaDM
CaDM
3
TMCL (No Adapt)
TMCL (No Adapt)
2
Our (No Intervention)
Our (No Intervention)
Prediction Erro
Prediction Error
-1
-2
10
15
20
Timesteps (*1044)
Timesteps (M004)
Figure 11:
The average prediction errors dynamics models on environments during training process (over three times). Specifically, the x axis is the training timesteps and y axis is the average log value of prediction prediction errors.
Table 3: The prediction errors of methods on test environments
CaDM (Lee et al. 2020)
TMCL (Seo et al. 2020)
Ours
Hopper Ant
0.0551
0.0316 0.0138
0.0271 = 0.0011
0.3850 + 0.0256
0.1560 + 0.0106
0.1381 + 0.0047
Calffcheettah
0.0815 + 0.0029
0.0751 +0.0123
0.0525 + 0.0061
HalfCheetah
0.6151 = 0.0251
1.0136 - 0.6241
0.4513 10.2147
Pendulum
0.0160 #0.0036
0.01301 0.0835
0.0030 + 0.0012
Slim Humanoid
0.8842 + 0.2388
0.3243 + 0.0027
0.3032 + 0.0046
A. PREDICTION ERRORS ON SPECIFIED ENVIRONMENT
The prediction errors of each method on specified environment are given at Table and
Table 4: The prediction errors of methods on specified environment of Hopper Task.
CaDM (Lee et al.
2020
TMCL (Seo et al. 2020
Ours
0.25
0.0443 + 0.0049
0
0.0120 = 0.0025
1.75
0.0459 0.0006
0.0131 + 0.0138
0.0132 0.0013
Table 5: The prediction errors of methods specified environment of Ant Task.
CaDM (Lee 2020)
TMCL (Seo al. 2020)
Ours
0.30
0.0928 + 0.0019
0.0910 1 0.0200
0.0669 + 0.0040
0.50
0.1013 0.0057
0.0887 + 0.0212
0.0671 0.0034
Table 6: The prediction errors of methods specified environment of Slim Humanoid Task.
CaDM (Lee et al.
2020
TMCL (Seo et al. 2020)
Ours
0.50
0.1614 0.0165
0.1860 + 0.0040
0.1282 + 0.0295
0.70
0.1512 0.0152
0.1550 + 0.0186
0.1236
0.0162
1.50
0.1601 + 0.0202
0.1873 + 0.0087
0.1444 = 0.0233
1.70
0.1439 0.02029
0.1688 = 0.01032
0.1217
0.0206
19

---

Published as a conference paper at ICLR 2022
A.11
THE AVERAGE RETURNS ON TEST ENVIRONMENTS DURING TRAINING PROCESS
The average returns on test environments during training process are given at Figure 12
Happer
Art
Croppic Haltcheetah
CaDM
CaDM TMCL
CaDM
TMCL
5000
TMCL
3000
2500
Rewards
Rewards
Rowards
2000
3000
800
1800
500
13
22
15
20
10
20
Timestops (*184)
Timestops (*1044)
Timestoos (1104)
Hofchootth
Panduluu
Humanaid
CaDM TMCL
CaDM
-000
TMCL
Ours
2000
Ours
Rowards
8000
Flewards
-500
Rewards
1200
0000
-1200
500
-1200
CaDM
2020
TMCL
1400
-1000
10
10
20
Timestops (*10%)
Timestage (11044)
T Timeslege (1104)
Figure 12: The average rewards of trained model-based RL agents on unseen environments. The results show the mean and standard deviation of returns averaged over three runs.
A.12 QUANTITATIVE CLUSTERING PERFORMANCE COMPARISON
To quantitatively evaluate the 2s' clustering performance,
K-means algorithm to predict each Z's environment id, and compare them with the true environment id. The details are provided in demo of K-means and evaluation metrics The results given below. Specifically, TMCL has lower clustering performances than CaDM, but TMCL still has higher returns on test environments than CaDM. This because TMCL clusters environments via multiplying dynamics functions rather than separating Zs.
Table 7: Quantitatively clustering evaluation results of 2 on Pendulum.
homo
compl
V-meas
ARI
AMI
CaDM TMCL
1
0.655
0.627
0.516
0.599
0
0.298
0.217
0.088
0.165
Ours (no Intervention)
0
0.768
0.762
0.760
0.653
Ours
1
0.932
0.932
0.937
0.931
Table 8: Quantitatively clustering evaluation results of Ž on Halfcheetah.
homo
compl
V-meas
ARI
AMI
CaDM
0.262
0.260
0.203
0.257
TMCL
0
0.239
0.165
0.051
0.126
Ours (no Intervention)
0.368
0.362
0.265
0.353
Ours
0.416
0.411
0.312
0.405
According to the quantitative clustering performance measures, can that the clustering performance of our method is superior to baselines by a large margin, and the results are consistent with the performance the test environments.
20

---

Published as a conference paper at ICLR 2022
Table 9: Quantitative clustering evaluation results of 2 on Slim Humanoid.
homo
compl
V-meas
ARI
AMI
CaDM
0.046
0.045
0.027
0.042
TMCL
0
0.002
0.002
0.000
0.000
Ours
0
0.055
0.052
0.037
0.058
Table 10: Quantitative clustering evaluation results of 2 on Cripple Halfcheetah.
homo
compl
V-meas
ARI
AMI
CaDM
0.733
0.716
0.686
0.701
TMCL
0
0,253
0,000
0.000
0.000
Ours
1
0.853
0.851
0.860
0.849
Table 11: Quantitative clustering evaluation results of 2 on Hopper.
homo
compl
V-meas
AMI ARI
CaDM
0
0.019
0.018
0.010
0.015
TMCL
0
0.023
0.008
0.000
0.003
Ours
0
0.130
0.108
0.049
0.089
A.13 VISUALIZATION
A.14
T-SNE VISUALIZATION
CaDM
TMCL
maxe-iB-
-20
mess:
maxentt!
Ours (No Intervention)
Ours
maass:00
mese:1p.j g]
massero mg
maserja.11
maxeeja 15:5
mess-111
#1a0s-1 16)
mess-11-19 (515
massett
mese-11-3)
-10
-20
-30
40
20
40
Figure 13: The T-SNE visualization of estimated context (environmental-specfffee vectors in the Pendulum task, where mass = 0.5 and mass =1.3 are from test environments.
A.15 PCA VISUALIZATION
21

---

Published as a conference paper at ICLR 2022
CaDM
TMCL
masem0.
masert
80
nasssng
nasset1 18
40
rasee11 (15)
mess-(o.a
mess-1o.5
-40
mass-1oas
massrt
-00
mass-11:15
-40
mass-11.3]
Ours (No Intervention)
Ours
manam]1.3]
-5
-15
-10
Figure 14: The T-SNE visualization of estimated context (environmental-specfitt) vectors in the Halfcheetah task, where mass 0.5 and mass =1.5 are from test environments.
CaDM
TMCL
mase-..5 0.5
maxe
mase
maxem|B.75
rasse'
masen[!]
musse' 1X
maser[1 15
muser
maxer[1.3]
0.5
-1.5
Ours NO Intervention
Ours
0.5
02
0.1
00
.01
Figure 15: The PCA of estimated context (environmentaalsseecffe)) vectors in Pendulum task, where
mass 0.5 and mass = 1.3 are from test environments.
22

---

Published as a conference paper at ICLR 2022
CaDM
TMCL
0.75
muss-18
muss-
0.30
0.5
0.2
-0.75 -0.00
40.29
0.02
0.25
1.00
Ours NO Intervention
Ours
maxempo 5]
masamjo.5)
mass:10 7g
mese-16.81
masen10 88
mase-jpas
maseet))
masue
mass-11-15
(1099-(1.)
mess-11.3)
-100
-50
50
100
150
Figure 16: The PCA of estimated context (envirommental-specinn vectors in HalfCheetah task,
where mass = 0.5 and mass =1.5 are from test environments.
Ours
CaDM
TMCL
1.5
max 5
mard277 7g
1.0
mare 25
6.5
C44
6/9
04
C3
22
30
02
C.4
055
0.3
3.50
3.25
0.30
025
0.50
375
1.0
1.25
1e12
Figure 17: The PCA of estimated context (environmeentalsppecfff)) vectors in the Hopper task.
CaDM
TMCL
Ours
11)
01891
-35
101899
-12
00877
-10
-1.
-65
6.0
Figure 18: The PCA of estimated context (environmental-spenn vectors in the Crip- ple Halfcheetah task.
23

---

Published as a conference paper at ICLR 2022
CaDM
TMCL
Ours
-1
-2
-04
-02
06
Q4
30
Figure 19: The PCA of estimated context (environmennta-specfffe vectors in the Slim Humanoid task.
24


Reviews:
[{"review_id": "f_DGQFWfKf", "reviewer": "Reviewer_q6e8", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "Strengths\n- novel idea for using CDE for clustering trajectories\n- I like the separation of the two objectives (1) and (6). (1) enforces that the z's are useful for predicting dynamics. (6) enforces that the z's within the same trajectory _and_ the same environment are clustered together. The authors show that in practice the z's do end up being predictive of the environment id.\n- Figure 1 shows that the learned context vectors are clearly distinguishable\n\nWeaknesses\n- the results in Figures 4 and 5 show that TMCL very similarly to the proposed method, which seems to slightly refute the authors claim that \"However, because environments are not labelled, the extracted information inevitably contains redundant information unrelated to the dynamics in transition segments and thus fails to maintain a crucial property of Z: Z should be similar in the same environment and dissimilar in different ones. As a result, the learned dynamics prediction function will deviate from the true one, which undermines the generalization ability.\" Essentially, why is it necessary to \"explicitly learn the meaningful dynamics change information\" instead of doing it implicitly? The empirical results do not seem to suggest that being explicit is necessarily much more helpful for generalization and performance.\n- I was interested to see that in Figure 6, not having the intervention does not seem to hurt performance or generalization that much, and in some cases even improves it. Would the authors be able to conduct an analysis similar to Figure 7 for the case where the interventional prediction is not used and test whether the contexts can still be predictive of the environment? ", "summary_of_the_paper": "This paper considers the unsupervised dynamics generalization problem. In this problem there are a set of train MDPs and a set of test MDPs, all with the same state and action spaces, but with different dynamics functions. The authors build an approaches that use past transition segments to estimate an environment-specific vector $Z$, which is then intended to act as a contextual input to the dynamics function. The problem the authors identify is that a naive application of inferring $Z$ from past transition segments undermines the ability for Z to be similar in the same environment and dissimilar in different environments. The authors claim that generalization ability is hurt by the inability to preserve this property. The authors propose a loss (equation 6) that pull $Z$'s from the same trajectory together and push away $Z$'s from different trajectories together. However, to avoid not pushing away different trajectories from the same environment, the authors propose to estimate the controlled direct effect (CDE) of the $Z$ on the next state $S_{t+1}$: if $z^i$ and $z^j$ from trajectories $i$ and $j$ have a similar CDE, then the authors declare $i$ and $j$ to be likely from the same environment. This likelihood is used as a weighting coefficient on the loss: the authors use this weighting coefficient as a \"soft\" indicator variable to distinguish examples from the same environment (positive examples) or examples from different environment (negative examples). Figure 1 shows that a PCA visualization of the learned context vectors using the proposed method separate the context vectors into distinct clusters. Empirical evaluation on standard benchmarks shows that the proposed method does slightly better than the next best baseline, TMCL from (https://arxiv.org/abs/2010.13303).", "summary_of_the_review": "It's not clear how much the interventional prediction aspect of this work actually helps (see Figure 6), even though this seems to be the main proposed contribution. It is also not clear what advantage that the proposed method has over TMCL, other than the ability to qualitatively visualize distinct cluster's in the z's. But if this does not hurt performance or generalization (Fig 4 and 5), what would be the reason that the interventional prediction is crucial?\n\n== After Rebuttal ==\nThe quantitative results seem to still indicate a marginal improvement over the baselines, but perhaps what is really required is to identify environments that the community cares about that would more drastically distinguish the proposed method from those that do not use the proposed intervention prediction method. However, given the current experiments, the authors have adequately addressed my concerns and therefore I have raised my score to a 6.", "strength_and_weaknesses": "", "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": {}}, {"review_id": "e858gdRzcN", "reviewer": "Reviewer_db3k", "summary": "", "questions": "", "limitations": "", "rating": 8, "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": "#### **Strengths**\n- The paper tackles a very relevant and important problem in the RL community, where generalization to changing dynamics is one of the key aspects to enable RL algorithms for real-world applications.\n- To the best of my knowledge, the novelty of the proposed method is significant. While other works have used similar ideas to characterize the environment dynamics using a vector $Z$, the paper presents an interesting idea to reduce the redundant information in $Z$, or in other words, learn a better representation of $Z$. The interventional prediction, while being simple, is elegant and works well in practice. I believe this interventional approach would open some new directions to improve current unsupervised representation learning methods, which use instance-based discrimination objectives (similar to (2) in the paper).\n- The improvements of the proposed method are significant compared to the existing baselines. Visualizations such as Figures 1,7,10-13 empirical show that the proposed method actually learns semantical information of $Z$.\n- The paper is well written and organized.\n\n#### **Questions**\n- What happens to the representation of $Z$ when we omit the interventional module? I would love to see a visualization similar to Figure 1. My wild guess is $Z$ will be spread out?", "summary_of_the_paper": "This paper proposes a new method to tackle the problem of generalization to unseen dynamics for model-based RL. Previous methods learn to predict a vector $Z$ that characterizes a particular environment dynamics from past transitions. However, as the environment id or label is not available, this vector inevitably contains redundant information, which might hurt the generalization of the model. The paper therefore proposes an interventional approach to estimate the probability that two vectors $\\hat{z}_i$ and $\\hat{z}_j$ belong to the same environment, and then uses a relational head to force similarity between them. As a result, the redundancy in $Z$ is reduced which leads to improved generalization.", "summary_of_the_review": "Overall, I like the idea of the paper and believe it has significant contributions to the RL community. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "A4DfVTRnqn", "reviewer": "Reviewer_3e12", "summary": "", "questions": "", "limitations": "", "rating": 8, "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": "# Strengths\n- The problem statement identified by the paper is clear and compelling, which motivates the paper well. \n- The method is interesting and novel, and seems well-designed to tackle the problem being addressed\n- The empirical results and visualisations both support the claims that this method is performing better than previous SOTA methods due to the better clustering of contexts.\n\n# Weaknesses\n- The notation in the last paragraph of page 2 is confusing - what's the difference between Z and \\hat Z?\n- In general, I found the explanation of the method difficult to parse at first - it took me several passes to put all the pieces together, understand the justification for each piece and how it fulfilled it's role. I think that's partly due to slightly confusing notation (see above point) which isn't well defined, and also due to the complexity of the method. In some sense this can't be avoided, but it's still worth mentioning here.\n- The biggest weakness of the paper (although I don't think it prevents acceptance at all) is that the method is a fairly incremental improvement to previous methods, and hence isn't highly novel. There's not much that can be done for this, unless additional benchmarks were created and assessed on (which is clearly a lot of work).\n\n# Questions\n- In the PCA visualisations, how are the different dots of the same colour generated? Is from different subsequences of transitions from a single episode, or multiple episodes, or both? It would be interesting to see how the context prediction evolves over an episode, and whether it converges over time as more information is gathered within the episode (as we might expect it to).\n- The relational encoder take in a fixed number of past transitions, is that correct? How does this work for time-steps near the beginning of the episode, where not enough transitions are available? Have you investigated using a recurrent encoder which processes all the information in the episode?\n- What's L_dist in the loss function in Algorithm 1?\n\n# Suggestions for Improvement\n- AugWM missing from related work: https://arxiv.org/abs/2104.05632, and seems relevant\n- It would be interesting to visualise clustering of masses which are in testing distribution as well, to see how well the relational embedding generalises.\n- The justification that their method improves through better clustering is only supported qualitatively through visualisation. It would be great to see a quantitative measure of how good the clustering is in the latent space for different methods, and then seeing if this correlates with evaluation performance or prediction error across the range of methods and environments (or even across hyperparameter options for each method).\n- It would be good to see a plot of average prediction error on test environment during training, to see how that changes as the methods train.\n- You provide some intuition as to why the ablation with no intervention has lower prediction error but lower performance in 4.3, but you could test this explicitly by plotting the test prediction error and seeing whether it's higher for the no-intervention ablation\n- PCA is quite a limited visualisation technique for clustering - have you considered other techniques such as t-SNE or UMAP?", "summary_of_the_paper": "This work tackles the problem of unsupervised dynamics generalisation in model-based reinforcement learning, improving approach introduced by additional of inferring a latent context which conditions the dynamics model and captures the variation in dynamics between environments. They introduce a set of auxiliary losses based on relational intervention and causal reasoning to encourage the inferred context to be the same in trajectories from the same environment (even when the environment identification is unknown) through inferring which environments are the same and weighting the context-similarity loss by this environment similarity. The method shows improved prediction error and test reward on a range of continuous control tasks compared to SOTA baselines from MBRL, Meta-RL and dynamics generalisation MBRL specifically. Visualisations qualitatively show that their method clusters transition segments from the same trajectory together better than previous methods, which could explain their improved performance.", "summary_of_the_review": "I think the paper is well-motivated, presents and interesting and novel method and demonstrates that method's improved performance on the problem being targeted, and hence deserves acceptance. I think the visualisations and ablations are useful qualitative evidence for the benefit of the method, but additional visualisations, or more quantitative investigation of the clustering hypothesis would be beneficial. Also, the explanation of the method is sometimes difficult to understand. The method is only an incremental improvement on previous work, which is necessary to perform but limits it's impact and novelty.\n\nEDIT: After seeing the revisions and additional experiments, I have raised my confidence from 3 to 4.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "3HoQTrhl7sn", "reviewer": "Reviewer_RE3u", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "Novelty\n\nTo the best of the reviewer’s knowledge, the paper is fairly novel. The authors use CaDM and TMCL, other model-based approaches that aim at generalization, as baselines. Figures 10-13 suggest that TMCL fails to learn separable encodings Z, while Figure 9 suggests that CaDM achieves lower returns. The proposed approach combines the best of both worlds achieving comparable returns and learning Zs that are clustered for the same domain.\n\nSignificance\n\nThe reviewer is unsure about the effect size of the proposed method. First of all, the experiments were repeated using 3 random seeds only. Combined with the overlapping confidence intervals of returns on a fair amount of test environments (e.g. Fig 5), the significance of the results is unclear. Increasing the number of runs would improve the credibility of the conclusions. For a discussion about the importance of rigorous experimentation in RL, see [1].\n\nOn a positive side, the paper follows the setup from previous papers and performs an ablation study of the proposed modification. The reviewer also appreciates qualitative findings about separability of learned Z in Figures 1 and 7 and quantitative evaluations of model prediction errors in training environments in Figures 4 and 8.\n\nDetailed comments and questions\n\n- Figure 1 demonstrates the linear projection of context vectors Z obtained by different methods and serves (at least partially) as a motivation for the paper. Would the conclusion that Zs for TMCL are not separable hold for a non-linear method for visualization (say, t-SNE)?\n- In Table 3, C_HalfCheetah row, CaDM has a lower error but the number for the proposed method is bold.\n- Conditioning $p(s’|s,a)$ on $Z_{t-k:t-1}$ makes the dynamics non-markovian. Would it be possible to use an RNN-like baseline in this setting?\n- Could authors present the plots for prediction errors on *testing* environments?\n- What would be a relevant model-free baseline in this setting?\n\n[1] Henderson, Peter, Riashat Islam, Philip Bachman, Joelle Pineau, Doina Precup, and David Meger. \"Deep reinforcement learning that matters.\" In Proceedings of the AAAI conference on artificial intelligence, vol. 32, no. 1. 2018.\n\n-------------\n\nPOST-REBUTTAL UPDATE:\n\nI update the score to 6 and confidence to 4", "summary_of_the_paper": "The paper studies multi-task generalization in model-based reinforcement learning (MBRL). One of the standard ways to approach the problem is given by inferring a latent variable Z encoding each task and then conditioning the dynamics model on it. The authors propose to encode segments of state-action trajectories into Z vectors and maximize similarity between Zs from the same trajectory. Using the mechanism, the paper then studies the ability of the agents to generalize to unseen variations of the training environments. For example, if the agent was trained on MuJoCo HalfCheetah with body mass 1, it is asked to generalize to body masses 0.5 and 1.5. The paper compares the method against the baselines across 3 axes: in terms of dynamics prediction error, in terms of returns on testing environments, and in terms of separability of the inferred Z for different environments.", "summary_of_the_review": "The reviewer leans towards recommending the submission for rejection. The ultimate objective of an RL system is maximizing the returns and it is unclear from the empirical evaluation to which extent the proposed method improves the performance of the agent on unseen tasks. However, the reviewer has limited familiarity with the related work on causality and it is possible that they missed something. Addressing the outlined concerns might improve the overall score.\n", "strength_and_weaknesses": "", "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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: Description of paper content:

The authors propose a dynamics model that can generalize to novel environments. The train and test MDPs have the same state and action spaces but different dynamics. Environment specific inference is achieved by estimating latent vectors Z that describe the non-stationary or variable part of the dynamics. These Z-s are inferred from trajectory segments in unlabeled environments. The Z-s are learned contrastively: Z-s from the same trajectory are pulled together, and Z-s from separate trajectories are pushed apart. However, to mitigate the error of distancing Z-s from different trajectories but the same environment, Z-s on trajectories with similar transitions are also pushed together using a soft clustering penalty. These losses are justified based on ideas from Pearl’s causal inference.

Summary of paper discussion:

The reviewers concluded that the contributions are conceptually interesting and “somewhat” novel. The reviewers felt that the empirical performance gains of the method over baselines were demonstrated but not extremely impressive.
Decision: Accept (Poster)