Year: 2021

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 ẑi, ẑ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 Ẑ from the same environment. As a 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 Ẑ can significantly reduce dynamics prediction errors and improve the performance of model-based RL methods on zero-shot new environments with unseen dynamics.

heading: 1 INTRODUCTION
text: 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 et al., 2018; Vinyals et al., 2019), and robotics (Levine & Abbeel, 2014; Bousmalis et al., 2018). However, solving real-world problems with RL is still a 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 al., 2019; Schrittwieser et al., 2020; Zhang et al., 2019b; van Hasselt et al., 2019; Hafner et al., 2019b;a; Lenz et al., 2015), which explicitly builds a predictive model to generate samples for learning RL policy, has been widely applied to a variety of limited data sequential 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 a significant performance decline of MBRL methods (Lee et al., 2020; Nagabandi et al., 2018a; Seo 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 et al., 2018a; Yang et al., 2020; Rakelly et al., 2019; Gu et al., 2017; Bousmalis et al., 2018; Raileanu et al., 2020; 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 S and action space A but the transition dynamics between states p(st+1|st, at) varies across different environments. Following previous works (Petersen et al., 2018; Gottesman et al., 2018), we focus on the unsupervised dynamics generalization setting, i.e. the id or label information of dynamics function in training MDPs is not available. This setting appears in a 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|st, at) varies across patients. However, it is difficult to label which patients share similar dynamics.
To build a generalized dynamics prediction function that can generalize to different transition dynamics, the shift of the transition dynamics can be modelled as the change of unobserved factors across different environments, i.e. there are hidden environment-specified factors Z ∈ Z which 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). It is natural to assume that Z is the same in a single environment but varies across different environments. As such, these unobserved environment-specified factors do not violate the nature of MDP in a single environment, but their changes can affect the dynamics functions across environments. Therefore, the dynamics function f : S × A → S can naturally be augmented by incorporating Z to be f : S ×A×Z → S (Rakelly et al., 2019; Zhou et al., 2019; Lee et al., 2020; Seo et al., 2020). Learning the augmented dynamics function is difficult because the environment-specified factor Z is unobservable. Previous methods (Rakelly et 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 1a) . However, in the unsupervised dynamics generalization setting, the extracted information from historical transition segments inevitably contains redundant information unrelated to the dynamics in transition segments. The redundant information would cause the surrogate for Z to lose a crucial property that characterizes Z: Z should be similar in the same environment and dissimilar in different environments. As shown in Figure 1b, the learned environment-specified information of a state-of-theart method CaDM (Lee et al., 2020) does not form clear clusters for different environments. Because the learned Ẑ fails to represent environmental information, the learned dynamics prediction 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 f̂ can learn modals of transition dynamics in different environments, we need to cluster Z according to their belonging environment.
In this paper, we provide an explicit and interpretable description to learn Z as a vector Ẑ (i.e. the estimation of Z) from the history transition segments. To cluster Ẑ from the same environment, we introduce a relational head module as a learnable function to enforce the similarity between Ẑs learned from the same environments. However, because environment label is not available, we can only cluster the Ẑs from the same trajectory, so we then propose an interventional prediction module to identify the probability of a pair of ẑi, ẑj belonging to the same environment through estimating Ẑ’s direct causal effect on next states prediction by do-calculus (Pearl, 2000). Because Z from the same environment surely have the same causal effect, we can directly maximize the similarity of Ẑs with the similar causal effect using the relational head, and thus can cluster Ẑ 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 on a range of tasks in OpenAI gym (Brockman et al., 2016) and Mujoco (Todorov et al., 2012), and empirically show that Ẑ estimated by our method enjoy less redundant information than baselines, and 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 on a new environment, and even achieve comparable results with the method directly cluster Ẑ with the true environment label.

heading: 2 RELATED WORK
text: Dynamics Generalization in MBRL Several meta-learning-based MBRL methods are proposed Nagabandi et al. (2018a;b); Sæmundsson et al. (2018) 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 updating hidden representations of a recurrent model Doshi-Velez & Konidaris (2016), and then Wang & van Hoof (2021) proposes the graph structured model to model the dynamics accurately to improve dynamics forecasting. Ball et al. (2021) focus on the offline setting and proposes an augmented model method to achieve zero-shot generalization. Lee et al. (2020); Seo et al. (2020) try to learn a generalized dynamics model by incorporating context information or clustering dynamics implicitly using multi-choice learning, aiming to adapt any dynamics without training. However, how to explicitly learn the meaningful dynamics change information remains a big challenge. Relational Learning Reasoning relations between different entities is a way to build knowledge of this world for intelligent agents Kemp & Tenenbaum (2008). In the past decades, relational paradigm have been applied to a wide range of deep learning-based application, e.g., reinforcement learning Zambaldi et al. (2018), question answer Santoro et al. (2017); Raposo et al. (2017), graph neural network Battaglia et al. (2018), sequential streams Santoro et al. (2018), few-shot 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 entities, our 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 on the intersection area of reinforcement learning and causal inference. For example, some works aims to alleviate the causal confusion problem in the imitation learning (de Haan et al., 2019; Zhang et al., 2020c; Kumor et al., 2021), batch learning (Bannon et al., 2020), and partial observability settings (Forney et al., 2017; Kallus & Zhou, 2018; Zhang et al., 2019a; Lu et al., 2018) in the online environment (Lyle et al., 2021; Zhang & Bareinboim, 2018), (Wang et al., 2020) also try to apply causal inference in the offline setting, where the observational data is always confounded. Lee & Bareinboim (2018); Bareinboim et al. (2015); Lattimore et al. (2016); Mozifian et al. (2020); Volodin et al. (2020) also explore how to design an optimal intervention policy in bandits or RL settings. In 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 Ẑ and the next state to estimate the probability of them belonging to the same environment.

heading: 3 METHODS
text: 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 environment-specified factors.

heading: 3.1 PROBLEM SETUP
text: The standard reinforcement learning task can be formalized as a Markov decision process (MDP) M = (S,A, r, p, γ, ρ0) over discrete time (Puterman, 2014; Sutton & Barto, 2018), where S, A, γ ∈ (0, 1], ρ0 are state space, action space, the reward discount factor, and the initial state distribution, respectively. The reward function r : S ×A → R specifies the reward at each timestep t given st and at, and transition dynamics p(st+1|st, at) gives the next state distribution conditioned on the current state st and action at. The goal of RL is to learn a policy π(·|s) mapping from state s ∈ S over the action distribution to maximize the cumulative expected return over timesteps Est∈S,at∈A[ ∑∞ t=0 γ
t r(st, at)]. In model-based RL, a model f is used to approximate the transition dynamics p, and then f can provide training data to train policy π or predict the future sequences for 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 we consider the unsupervised dynamics generalization problem in model-based RL, where we are given K training MDPs {Mtri }Ki=1 and L test MDPs {Mtej }Lj=1 that have the same state and action space but disjoint dynamics functions, and we 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 a finite number of environments and the MDPs in each environment share the same dynamics function but the environment id of MDPs is unavailable in the training process. In the context of model-based RL, how to learn a generalized dynamics model f is the key challenge to solve unsupervised dynamics generalization problem.

heading: 3.2 OVERVIEW
text: As analyzed in Section 1, we can incorporate the environment-specified factors Z ∈ Z into the dynamics prediction process to generalize the dynamic functions on different environments, i.e. extending the dynamics function from f : S × A → S to f : S × A × Z → S. Because Z is the same within an environment, we expect estimated Ẑs from the same environment are similar while those from different environments are dissimilar. Therefore, f models the commonalities of the transition dynamics in different environments and Z 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 Ẑs are similar in the same environment and dissimilar in different environments. However, since the environment label is unavailable in 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 of pair of ẑi and ẑj as the probability of them belonging to the same environment. Furthermore, we then introduce a relational head to aggregate Ẑs with high probability using relation loss. By simultaneously updating the dynamics prediction and the relation loss, we can cluster Ẑs from the same environment, and learn an augmented dynamics prediction model f̂ . Next, we will give details about our relational intervention approach.

heading: 3.3 RELATIONAL CONTEXT ENCODER
text: To learn the environment-specified factor Z of each environment, we firstly introduce a relational encoder g parameterized by φ. Similar to previous methods (Nagabandi et al., 2018a; Rakelly et al., 2019; Zhou et al., 2019; Lee et al., 2020; Seo et al., 2020), we use the past transition segments τt−k:t−1 = {(st−k, at−k), ..., (st−1, at−1)} as the input of g to estimate its corresponding ẑt−k:t−1:
ẑt−k:t−1 = g(τ i t−k:t−1;φ).
After obtaining environment-specified ẑt−k:t−1 at timestep t, we incorporate it into the dynamics prediction model f̂ to improve its generalization ability on different dynamics by optimizing the objective function following (Lee et al., 2020; Seo et al., 2020; Janner et al., 2019):
Lpredθ,φ = − 1
N N∑ i=1 log f̂(sit+1|sit, ait, g(τ it−k:t−1;φ); θ), (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 to update the model parameters.
However, as analyzed in Section 3.2, the vanilla prediction error (1) is not sufficient to capture environment-specified 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, we introduce a relational head (Patacchiola & Storkey, 2020) as a learnable function h to pull factors Ẑ from the same trajectory together and push away those from different trajectories. Concretely, the estimated ẑit−k:t−1 in a mini-batch will be firstly aggregated as pairs, e.g. concatenate two factors as [ẑi, ẑj ], and the pairs having two factors from the same trajectory are seen as positives, and vice versa. Then the relational head h parameterized by ϕ takes a pair of aggregated factors as input to quantify the similarity of given two factors and returns a similarity score ŷ. To increase the similarity score ŷ of positive pairs and decrease those negatives, we minimize the following objective:
Lrelationϕ,φ = − 1
N(N − 1) N∑ i=1 N∑ j=1 [ yi,j ·log h([ẑi, ẑj ];ϕ)+(1−yi,j)·log (1−h([ẑi, ẑj ];ϕ)) ] , (2)
where yi,j = 1 stands for positive pairs, and yi,j = 0 stands for negatives. Because the positive pairs have two factors belonging to the same trajectory, optimizing (2) can increase the similarity of Ẑs estimated from the same trajectory, and push away those factors estimated from different trajectories in their semantical space. Therefore, by optimizing (2), the information that is invariant within a trajectory will be encoded into Ẑ and the redundant information in transition segments will be reduced. (2) can also be interpreted from the perspective of mutual information, if we regard the Ẑs from the same trajectory as the positive pairs, optimizing (2) can be seen as maximizing the mutual information between Ẑs from the same trajectory (Please refer to (Tsai 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 Ẑs in the same environment will also be pushed away, which may undermine the cluster compactness for estimated Ẑs.

heading: 3.4 INTERVENTIONAL PREDICTION
text: Because the environment id of a trajectory is unknown, we cannot directly optimize relational loss (2) to cluster Ẑ within an environment. We propose an interventional prediction method to find the trajectories belonging to the same environment. Here we formalize the dynamics prediction model using a graphical causal model, and the causal graph is illustrated as Figure 3 (a), where the next state St+1 is caused by the current state St, action At and Ẑ, and the dynamics prediction model f represents the causal mechanism between them. Because Z from the same environment should have
the same influence on the states, and thus they should have the same causal effect on the next state St+1 if given St and At under the causal framework. As such, we can find estimated Ẑ belonging to the same environment by measuring the similarity of their causal effect on St+1. As Figure 3 (a) shows, there are multiple paths from Z to St+1 and we roughly categorize them into two categories. The first category is the direct path between Z and St+1 ( shown as read in Figure 3). The second category contains all the indirect paths where Z influences St+1 via previous states and actions. However, because the mediator in other paths e.g. St, At, 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 3 (a)), which means that we need to block all paths with meditors from Ẑ to St+1. By means of do-calculus (Pearl, 2000), we can estimate the direct causal effect of changing Z = ẑj to Z = ẑk on St+1 through calculating the controlled direct effect (CDE) (Pearl, 2013) by intervening mediators and Ẑ :
CDEẑj ,ẑk (st, at) =E[St+1|do(St = st, At = at), do(Z = ẑ j)]− E[St+1|do(St = st, At = at), do(Z = ẑk)]
(3)
=E[St+1|St = st, At = at, Z = ẑj ]− E[St+1|St = st, At = at, Z = ẑk], (4)
where do is the do-calculus (Pearl, 2000). There is no arrow entering Ẑ, so the do operator on Ẑ can be removed. Also, since there is no confounder between the mediators (St, At) and St+1, so we can remove the do operator of them as well, and the equation become as (4). Because the direct causal effects may differ for different values of St and At, we should sample St and At independently of Z, i.e. sampling St and At (Pearl et al., 2016) uniformly to get the average controlled direct causal effect from Ẑ to St+1. However, if we use the uniformly generated St and At, the sampled distribution may differ from the training distribution, resulting in inaccurate the next state prediction. As such, we directly sample St and At from the observational data. For the convenience of optimization, we only use a mini-batch of St and At pairs (sit, a i t), and concatenate them with ẑ
j and ẑk to calculate the average controlled direct effect under f̂ :
ACDEẑj ,ẑk = 1
N N∑ i=1 |CDEẑj ,ẑk(sit, ait)|, (5)
whereN is the batch size, j and k are the id of Ẑ estimated from two transition segments. Specifically, because the factors ẑ estimated from the same trajectory should be the same, and thus we minimize their controlled direct effect 5 as Ldist between them in the optimization process. Now we can use the calculated ACDEẑj ,ẑk as the semantic distance dj,k between estimated ẑi and ẑj , and thus we can aggregate factors Ẑ estimated from similar trajectories by the proposed relational head h. As such, we apply a transformation to convert distance metric dj,k to a similarity metric w ∈ (0, 1], which is wj,k = exp(−d j,k
β ), where β is a factor controlling the sensitivity of the distance metric. Specifically, because the scale and size of state varies in different tasks, e.g. 3 dims in Pendulum but 20 dims in Half-Cheetah, the optimal β may vary in different task. As such, we apply the L2 normalization in the distance metric d to convert it as a relative distance within a 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:
Li−relationϕ,φ = − 1
N(N − 1) N∑ i=1 N∑ j=1 [ [yi,j + (1− yi,j) · wi,j ] · log h([ẑi, ẑj ];ϕ)
+ (1− yi,j) · (1− wi,j) · log (1− h([ẑi, ẑj ];ϕ)) ] , (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 1− w. Similar to the analysis in section 3.3, optimizing the loss function (6) can increase the similarity between Ẑ with weight w, and push away from them with the weight 1−w. Because Ẑs estimated from the same environment have similar effects, these factors will be assigned with high similarities (estimated by the intervention operation of our paper). By simultaneously updating the prediction loss (1) and intervention relation loss 6, estimated Ẑs within the same environment will be aggregated, and the learned dynamics function f̂ can learn the modals of transition dynamics according to the Ẑ in different clusters. The training procedure of our approach can refer to Algorithm process in Appendix A.2.

heading: 4 EXPERIMENTS
text: 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 errors in model-based RL? (Section 4.2.1) • 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 AppendixA.13)
• Is the similarity of w measured by the intervention module reasonable? (Appendix A.7) • Can solely relational learning improve the performance of model-based RL? (Section 4.3)

heading: 4.1 ENVIROMENTAL SETUP
text: 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 1e-3 learning rate. During the training procedure, the trajectory segments are randomly sampled from the same trajectory to break the temporal correlations of the training data, which was also adopted by (Seo et al., 2020). Specifically, the length of the transition segments is 10, and the hyper-parameters are the same for all experiments, and details can be found in Appendix A.1. Datasets Following the previous methods (Lee et al., 2020; Seo et al., 2020), we perform experiments on classic control algorithm (Pendulum) from OpenAI gym (Brockman et al., 2016) and simulated robotics control tasks (HalfCheetah, Cripple-HalfCheetah, Ant, Hopper, Slim-Humanoid) 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., 2020; Seo et al., 2020) to change the environmental parameters (e.g. length and mass of Pendulum) and predefine them in the training and test environmental parameters lists. At the training time, we randomly sample the parameters from the training parameter list to train our relational context encoder and dynamics prediction model. Then we test our 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 details about training and test parameter lists for each task are the same with (Lee et al., 2020), and all details are given in Appendix A.1. Planning Following (Lee et al., 2020; Seo et al., 2020), we use the model predictive model (MPC) (Maciejowski, 2002) to select actions based on learned dynamics prediction model, and assume that the reward function of environments is known. Also, the cross-entropy method (CEM) (De Boer et 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, we 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 to 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 a 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-art model-based RL method on dynamics generalization, which introduces the multi-choice learning to cluster environments. For the fair comparison, we also report the performance of TMCL without adaptation in Figure 7.
• True Label: The method uses our relational head to cluster Ẑ with the true environment label (not the ground-truth of Z),. All hyperparameters are same with our method for the fair comparison.

heading: 4.2 PERFORMANCE COMPARISONS WITH BASELINES

heading: 4.2.1 PREDICTION ERROR COMPARISONS
text: We first evaluate whether the dynamics model trained by our methods can predict next-states more accurately or 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., 2020) and ours). We can see that the dynamics model trained by our relational intervention method has superior prediction performance over other state-of-the-art 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-art methods by approximately 10%.

heading: 4.2.2 PERFORMANCE COMPARISONS
text: 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 Table 1. We can see that the meta-learning based methods (Nagabandi et al., 2018b;a) do not perform better than vanilla PETS (Kurutach et al., 2018), while methods (Lee et al., 2020; Seo et al., 2020) aiming to learn a generalized dynamics prediction model are superior to others significantly. Our approach achieves the highest rewards on all six tasks among all methods. Figure 5 shows that the mean and standard deviation of 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. A fair comparison between TMCL (no adaptation) and our method can be found at Appendix A.6. Also, we observe that our method achieve comparable results with the method directly cluster Ẑ using the truth environment label, which indicates that our intervention module actually can assign high similarities into Ẑs estimated from the same environment in an unsupervised manner. The claim is also evaluated by the similarity visualization in the Appendix A.7, where we find that Ẑs from the same environment will be assigned significant higher similarities than those pairs from different environments.

heading: 4.3 ABLATION STUDY
text: 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 6b 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, it also has lower rewards than the whole model. One possible reason is that the Pendulum is simple for the dynamics prediction model to learn, and thus the dynamics prediction model with the vanilla relational head is a little over-fitting on the training environments (Please refer to 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 trajectory-specified redundant information.

heading: 5 CONCLUSION
text: In this paper, we propose a relational intervention approach to learn a generalized dynamics prediction model for dynamics generalization in model-based reinforcement learning. Our approach models the dynamics change as the variation of environment-specified factor Z and explicitly estimates Z from past transition segments. Because environment label is not available, it is challenging to extractZ from transition segments without introducing additional redundant information. We propose an intervention module to identify the probability of two estimated factors belonging to the same environment, and a relational head to cluster those estimated Ẑs are from the same environments with high probability, thus reducing the redundant information unrelated to the environment. By incorporating the estimated Ẑ into the dynamics prediction process, the dynamics prediction model has a stronger generalization ability against the change of dynamics. The experiments demonstrate that our approach can significantly reduce the dynamics prediction error and improve the performance of model-based agents on new environments with unseen dynamics.

heading: 6 ETHICS STATEMENT
text: 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 it is far away to apply such an algorithm to real-world applications, we still need to guard against the possible harm caused by its application in military robotics.

heading: 7 REPRODUCIBILITY STATEMENT
text: We have run our experiments over three runs (ours with 10 runs to reduce random errors) to reduce random errors, and public hyperparameters used in our experiments. We promise that we will public codes after the acceptance of this paper.

heading: A APPENDIX
text: We promise that we will public all codes after the acceptance of this paper.

heading: A.1 ENVIRONMENTAL SETTINGS
text: We follow the environmental settings of Lee et al. (2020) in dynamics generalization. The details of settings are given as follows:
• Pendulum We modify the mass m and the length l of Pendulum to change its dynamics.
• Half-Cheetah We modify the mass of regid link m and the damping of joint d of HalfCheetah agent to change its dynamics.
• Crppled Cheetah We cripple the id of leg c of Half-Cheetah agent to change its dynamics.
• Ant We modify the mass of ant’s leg m to change its dynamics. Specifically, we modify two legs by multiplying its original mass with m, and others two with 1m .
• Slim Humanoid We modify the mass of rigid link m and the dampling of joint d of the Slim Humanoid agent to change its dynamics.
• Hopper We modify the mass of m of the Hopper agent to change its dynamics.
The training and test modified parameter list can be found at the Table 2.

heading: A.2 ALGORITHM
text: The training procedure is give at Algorithm 1.

heading: A.3 TRAINING DETAILS
text: Similar to the Lee et al. (2020), we 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 at 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 64, β is 5e-3. All module are learned by a Adam optimizer with 0.001 learning rate.
Algorithm 1 The training algorithm process of our relational intervention approach Initialize parameters of relational encoder φ, dynamics prediction model θ and relational head ϕ Initialize dataset B ← ∅ for Each Iteration do
sample environmentsMi from training environments {Mtri }Ki=0 . Collecting Data for T = 1 to TaskHorizon do
Get the estimation of the environment-specified factor ẑit−k:t−1 = g(τ i t−k:t−1;φ) Collect (st, at, st+1, rt, τ it−k:t−1) fromMi with dynamics prediction model θ Update B ← B ∪ (st, at, st+1, rt, τ it−k:t−1)
end for for Each Dynamics Training Iteration do . Update φ,θ and ϕ
for k = 1 to K do Sample data τ i,b,Pt−k:t−1 , τ i,b,K t:M and τ j,b,P t−k:t−1 , τ j,b,K t:M with batch size B,from B
Get the estimation of the environment-specified factor ẑi,B,,Pt−k:t−1 = g(τ i,B,P t−k:t−1;φ) and
ẑij,B,,Pt−k:t−1 = g(τ j,B,P t−k:t−1;φ)
Estimate the probability w of ẑi,B,,Pt−k:t−1 and ẑ j,B,,P t−k:t−1 belonging to the same environment. Compute the total loss Ltot=Lpredφ,θ (τ i,B,,K t:M , ẑ i,B,,P t−k:t−1) + L i−relation φ,ϕ (ẑ i,B,,P t−k:t−1) + Ldistφ,θ (τ i,B,K t:M , ẑ i,B,P t−k:t−1) Update θ , φ , ϕ← ∇θ,φϕ 1BL tot
end for end for
end for

heading: A.4 NETWORK DETAILS
text: Similar to the Lee et al. (2020), the relational encoder is constructed by a simple 3 hidden-layer MLP, and the output dim of environmental-specific vector ẑ is 10. The relational head is modelled as a single layer MLP with 10 units (same with the dim of ẑ) and an output layer. The dynamics prediction model is a 4 hidden-layer MLP with 200 units.

heading: A.5 CONNECTION BETWEEN RELATION LOSS AND MUTUAL INFORMATION
text: Given a pair of data (x, y) ∈ X × Y , we donote the joint distribution of X and Y are PXY , and their marginal distributions are PX and PY , respectively. By definition, the mutual information between X and Y is:
I(X;Y ) = EPXY [log( p(x, y)
p(x)p(y) )] (7)
To estimate mutual information between X and Y , (Tsai et al., 2020) proposes a probabilistic classifier method. Concretely, we can use a Bernoulli random variable C to classify one given data pair (x, y) from the joint distribution PXY (C = 1) or from the product of marginal distribution P (X)P (Y ) (C = 0) . Therefore, the mutual information I(X;Y ) between X and Y can be rewrite as:
I(X;Y ) = EPXY [log( p(x, y)
p(x)p(y) )]
= EPXY [log( p(x, y|C = 1) p(x, y|C = 0) )]
= EPXY [log( p(C = 0)P (C = 1|x, y) p(C = 1)P (C = 0|x, y) )] (8)
Obviously, p(C=0)p(C=1) can be approximated by the sample size, i.e. nPXPY nPXY , while P (C=1|x,y)P (C=0|x,y) can be measured by a classifier h(C|x, y), and it can be learned by our relation loss with relational head h:
Lrelationϕ,φ = − [ C · log h([x, y];ϕ) + (1− C) · log (1− h([x, y];ϕ)) ] , (9)
where C = 1 if the given pair (x, y) is from the joint distribution PXY , and C = 0 if the given pair (x, y) is from the product of the marginal distributions PXPY . Because p(C=0) p(C=1) tend to be a constant, optimizing our relation loss is actually estimating the mutual information I(X;Y ) between X and Y . As such, if we regard the pairs of (ẑ) from the same trajectory/environment as positive pairs, and others are negative pairs, optimizing 2 is actually maxmizing the mutual information between between (ẑ) from the same trajectory/environment, and thus preserve the trajectory/environment invariant information. If the readers are interested at the concrete bound about this method to estimate mutual information, please refer to (Tsai et al., 2020).

heading: A.6 FAIR COMPARISON WITH TMCL
text: Because TMCL need 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:
We can see that the average returns of TMCL without adaptation are significantly lower than ours, especially for classic control task Halfcheetah, which is a direct evidence to support our claim: environment-separated Zs is important for the 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 use the invariance of Z within an trajectory (Z should predict other states within a trajectory in TMCL), which is similar to our paper with no intervention module.

heading: A.7 SIMILARITIES VISUALIAZTION
text: To evaluate the correctness of the estimated similarity of our intervention module, we use a ẑi estimated from the environment where the mass is 1 as the anchor, and randomly sample 200 ẑj estimated from different environments (including mass = 1). Then we calculate the similarity between anchor ẑi and ẑj , and visualize the similarities according to their environments. As Figure 8 shows, ẑjs belonging to the same environment with the anchor ẑi have significant higher similarities than those belonging to 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 ẑs from the same environment or not.
To study the role of intervention module, we also visualize the similarity of ẑi learned by the model without 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-specified 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 8 and 9 directly shows that our intervention module is valuable to predict whether two contexts are from the same environment or not.

heading: A.8 PREDICTION ERRORS ON TRAING ENVIRONMENTS
text: The prediction errors of each method on training environment are given at Figure 10.

heading: A.9 PREDICTION ERRORS ON TEST ENVIRONMENTS
text: The prediction errors of each method on test environments are given at Table 3. Specifically, we test each test environment 10 times, and plot the average prediction error to reduce random errors (Figure 11).

heading: A.10 PREDICTION ERRORS ON SPECIFIED ENVIRONMENT
text: The prediction errors of each method on specified environment are given at Table 4, 5 and 6.

heading: A.11 THE AVERAGE RETURNS ON TEST ENVIRONMENTS DURING TRAINING PROCESS
text: The average returns on test environments during training process are given at Figure 12.

heading: A.12 QUANTITATIVE CLUSTERING PERFORMANCE COMPARISON
text: To quantitatively evaluate the Ẑs’ clustering performance, we use 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 are given at below. Specifically, TMCL has lower clustering performances than CaDM, but TMCL still has higher returns on test environments than CaDM. This is because TMCL clusters environments via multiplying dynamics functions rather than separating Zs.
According to the quantitative clustering performance measures, we can see that the clustering performance of our method is superior to baselines by a large margin, and the results are consistent with the performance on the test environments.
A.13 VISUALIZATION

heading: A.14 T-SNE VISUALIZATION

heading: A.15 PCA VISUALIZATION

Review summaries:
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).

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.

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.

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.