A R ELATIONAL I NTERVENTION A PPROACH FOR U NSUPERVISED D YNAMICS G ENERALIZATION IN M ODELB ASED R EINFORCEMENT L EARNING
A BSTRACT
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-speciﬁed 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 ˆ
z i , ˆ 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 ˆ Z from the same environment.
As a result, the redundant information will be reduced in ˆ Z .
We empirically show that ˆ Z estimated by our method enjoy less redundant information than previous methods, and such ˆ Z can signiﬁcantly 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 https://github.com/CR-Gjx/RIA .
I NTRODUCTION
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 efﬁciency 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 signiﬁcant 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 ( s t +1 | s t , a t ) 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 difﬁcult to obtain.
For example, in healthcare, patients may respond differently to the same treatment, i.e. , p ( s t +1 | s t , a t ) varies across patients.
However, it is difﬁcult to label which patients share similar dynamics.
[Figure]
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-speciﬁed 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-speciﬁed 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 difﬁcult because the environment-speciﬁed 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.
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 environment-speciﬁed information ˆ Z learned by 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 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 environments.
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 s 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.
The experimental results show that our method signiﬁcantly 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.
R ELATED W ORK Dynamics Generalization in MBRL
Several meta-learning-based MBRL methods are proposed
Nagabandi et al.
(2018a;b); Sæmundsson et al.
; Huang et al.
(2021) 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 a 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 ofﬂine setting and proposes an augmented model method to achieve zero-shot generalization.
Lee et al.
Seo et al.
(2020) try to learn a generalized dynamics model by incorporating context information or clustering dynamics implicitly using multichoice 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 an important 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.
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., 2020a) also try to apply causal inference in the ofﬂine setting, where the observational data is always confounded.
Lee & Bareinboim (2018); Bareinboim et al.
(2015); Lattimore et al. ; Moziﬁan et al.
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.
M ETHODS
In the section, we ﬁrst introduce the formulation of the dynamic generalization problem.
Then we present the details on how relational intervention approach learns the environment-speciﬁed factors.
3. P ROBLEM S ETUP
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 speciﬁes the reward at each timestep t given s t and a t , and transition dynamics p
( s t +1 | s t , a t ) gives the next state distribution conditioned on the current state s t and action a t .
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 E s t ∈S ,a t ∈A [ (cid:80) ∞ t =
γ t r ( s t , a t )] .
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.
Beneﬁting from data provided by learned dynamics model f , model-based RL has higher data efﬁciency 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 {M tri } Ki =1 and L test MDPs {M tej } 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 ﬁnite number of dynamics functions, meaning that the MDPs can be categorized into a ﬁnite 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.
3. O VERVIEW
[Figure]
As analyzed in Section 1, we can incorporate the environment-speciﬁed 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 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.
3. R ELATIONAL C ONTEXT E NCODER
To learn the environment-speciﬁed factor Z of each environment, we ﬁrstly 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 = { ( s t − k , a t − k ) , ..., ( s t − 1 , a t − 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-speciﬁed ˆ 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):
[Equation]
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 sufﬁcient to capture environment-speciﬁed 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 ﬁrstly 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:
[Equation]
where y i,j = 1 stands for positive pairs, and y i,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 insufﬁcient because the estimated ˆ Z s in the same environment will also be pushed away, which may undermine the cluster compactness for estimated ˆ Z s.
[Equation]
[Table] [Table]
+
[Table] [Table]
+
[Table]
/sF/
[Table] [Table]
=</latexit
[Table]
Randomize <latexit sha1_base64="DGg75FA1fO/Y3mrIued6VccyOZc=">AAACdXicbVHLbtNAFJ2YVwmvli4RkoUDYhXsCESXBTYsi9S0FXEUXY
[Table]
+
[Table] [Table]
[Table] [Table]
+
[Table] [Table]
[Table]
+jc11JSuySVj7HcpQOgk2jHpy1JVqiU5PVBv6dxfN8TeKm7u/DY4m4+LDOP/2Ptv/vPmXLfaCvWJvWcE+sn32lR2wKeNMsJ/sF/s9uEheJqPkzWVqMth4dtk/kbz7C0XzwUc=</latexit
[Table]
Randomize <
[Table]
sha
[Table]
=
[Table] [Table]
=
[Table]
+v21HmpZlrqmjkb2ALn8aXdMs4jdTScqUZnTnnHt3HVFYKT3n+Z5DcuXvv/oOth8NHj588fba98/zIm9ZxnHIjjTupwKMUGqckSOKJdQiqknhcLb/0+vEPdF4YfUgri3MFp1o0ggNFavppEahbbGf5OF9HehsUG5CxTRwsdgbfy9rwVqEmLsH7WZFbmgdwJLjEbli2Hi3wJZziLEINCv08rLvt0teRqdPGuHg0pWv2uiOA8n6lqpipgM78Ta0n/6fNWmr25kFo2xJqflmoaWVKJu1HT2vhkJNcRQDcidhrys/AAae4oGGp8ZwbpUDXpdA1WoyXplBaZyj6yr6gBYlEeEVeZR52kUVnuzh+jc11JSuySVj7HcpQOgk2jHpy1JVqiU5PVBv6dxfN8TeKm7u/DY4m4+LDOP/2Ptv/vPmXLfaCvWJvWcE+sn32lR2wKeNMsJ/sF/s9uEheJqPkzWVqMth4dtk/kbz7C0XzwUc=</latexit
[Figure]
3. I NTERVENTIONAL P REDICTION
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 ﬁnd 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 S t +1 is caused by the current state S t , action A t and ˆ Z , and the dynamics prediction model f represents the causal mechanism between them.
Because Z from the same environment should have the same inﬂuence on the states, and thus they should have the same causal effect on the next state
S t +1 if given S t and A t under the causal framework.
As such, we can ﬁnd estimated ˆ Z belonging to the same environment by measuring the similarity of their causal effect on S t +1 .
As Figure 3 (a) shows, there are multiple paths from Z to S t
+
and we roughly categorize them into two categories.
The ﬁrst category is the direct path between Z and S t +1 ( shown as read in Figure 3).
The second category contains all the indirect paths where Z inﬂuences S t +1 via previous states and actions.
However, because the mediator in other paths e.g. S t , A t , 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 S t +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 S t +1 through calculating the controlled direct effect (CDE) (Pearl, 2013) by intervening mediators and ˆ Z :
[Equation]
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 ( S t , A t ) and S t +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 S t and A t , we should sample S t and A t independently of Z , i.e. sampling S t and A t (Pearl et al., 2016) uniformly to get the average controlled direct causal effect from ˆ Z to S t +1 .
However, if we use the uniformly generated S t and A t , the sampled distribution may differ from the training distribution, resulting in inaccurate the next state prediction.
As such, we directly sample S t and A t from the observational data.
For the convenience of optimization, we only use a mini-batch of S t and A t pairs ( s it , a it ) , and concatenate them with ˆ z j and ˆ z k to calculate the average controlled direct effect under ˆ f :
[Equation]
where N is the batch size, j and k are the id of ˆ Z estimated from two transition segments.
Speciﬁcally, because the factors ˆ z estimated from the same trajectory should be the same, and thus we minimize their controlled direct effect 5 as L dist between them in the optimization process.
Now we can use the calculated ACDE ˆ z j , ˆ z k as the semantic distance d j,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 d j,k to a similarity metric w ∈ (0 , 1] , which is w j,k = exp ( − d j,k β ) , where β is a factor controlling the sensitivity of the distance metric.
Speciﬁcally, 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 normalization in the distance metric d , i.e. , normalize d with batch variance, 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:
[Equation]
where the ﬁrst 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.
E XPERIMENTS
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)
4. E NVIROMENTAL S ETUP 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 1e-3.
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; Wang et al., 2020b; 2019).
Speciﬁcally, the length of the transition segments, i.e. , k , is 10.
All implementation 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 a classic control task (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 predeﬁne 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.
Speciﬁcally, the predeﬁned parameters in the test parameter list are outside the training range.
The predeﬁned 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 functions of environments are known.
Also, the cross-entropy method (CEM) (De Boer et al., 2005) is used to optimize action sequences for ﬁnding 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.
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.
• 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.
4. P ERFORMANCE C OMPARISONS WITH B ASELINES 4.2. P REDICTION E RROR C OMPARISONS
We ﬁrst 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.
Speciﬁcally, 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%
[Figure]
i c
[Figure]
[Table]
P ERFORMANCE C OMPARISONS
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) that aim to learn a generalized dynamics prediction model are superior to others signiﬁcantly.
Among which our approach achieves the highest rewards on all six tasks among all methods.
Figure 5 shows 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 sufﬁcient 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.
In addition, we observe that our method achieves 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.
We also observe the same results in the similarity visualization in the Appendix A.7, where we ﬁnd that ˆ Z s from the same environment are assigned signiﬁcant higher similarities than those pairs from different environments.
[Figure]
4. A BLATION S TUDY
In this section
[Figure]
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
[Figure]
after incorporating the relational head and intervention prediction
[Figure]
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
[Figure]
and thus the dynamics prediction model with the vanilla relational head is a little over-ﬁtting on the training environments (Please refer to prediction errors on test environments are given in Appendix A.9), limiting the performance improvement.
This phenomenon conﬁrms the importance of our intervention prediction on reducing the trajectory-speciﬁed redundant information.
[Figure]
C ONCLUSION
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-speciﬁed factor Z 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 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 signiﬁcantly reduce the dynamics prediction error and improve the performance of model-based agents on new environments with unseen dynamics.
A CKNOWLEDGE
Mr Jiaxian Guo is supported in part by Australian Research Council Projects FL-170100117 and LE200100049.
Dr Mingming Gong is supported by Australian Research Council Project DE210101624.
E THICS S TATEMENT
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 signiﬁcantly 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 prevent the algorithm from being applied in military areas.
R EPRODUCIBILITY S TATEMENT
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.
The codes of this method are available at https://github.com/CR-Gjx/RIA
A A PPENDIX
We promise that we will public all codes after the acceptance of this paper.
A. E NVIRONMENTAL S ETTINGS
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.
Speciﬁcally, we modify two legs by multiplying its original mass with m , and others two with 1 m .
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 modiﬁed parameter list can be found at the Table
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
A LGORITHM
The training procedure is give at Algorithm 1.
A. T RAINING D ETAILS
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.
Speciﬁcally, the batch size for each experiment is 128, β is 6e-1.
All module are learned by a Adam optimizer with 0.001 learning rate.
Algorithm
Initialize parameters of relational encoder φ , dynamics prediction model θ and relational head ϕ
Initialize dataset B ← ∅ for Each Iteration do sample environments M i from training environments {M tri }
Ki = (cid:46)
Collecting Data for T = 1 to TaskHorizon do
Get the estimation of the environment-speciﬁed factor ˆ z it − k : t − 1 = g
( τ it − k
t − 1 ; φ )
Collect ( s t , a t , s t +1 , r t , τ it − k : t − 1 ) from M i with dynamics prediction model θ
Update B ← B ∪ ( s t , a t , s t +1 , r t ,
τ it − k : t − 1 ) end for for Each Dynamics Training Iteration do (cid:46) Update φ , θ and ϕ for k = 1 to K do
Sample data τ i,b,Pt − k : t − 1 , τ i,b,Kt : M and τ j,b,
Pt − k : t − 1 , τ j,b,Kt : M with batch size B,from B
Get the estimation of the environment-speciﬁed factor ˆ z i,B,,Pt − k : t − 1 = g
( τ i,B,Pt − k : t − 1 ; φ ) and ˆ z ij,B,,Pt − k : t − 1 = g ( τ j,B,Pt − k : t − 1 ; φ )
Estimate the probability w of ˆ z i,B,,Pt − k : t − 1 and ˆ z j,B,,Pt − k : t − 1 belonging to the same environment.
Compute the total loss
L tot
[Equation]
predφ
[Equation] [Equation]
[Equation]
Update θ , φ , ϕ ← ∇ θ,φϕ 1 B L tot end for end for end for
A. N ETWORK D ETAILS
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-speciﬁc vector ˆ z is 10.
The relational head is modelled as a single FC layer.
The dynamics prediction model is a 4 hidden-layer FC with 200 units.
A. C ONNECTION BETWEEN R ELATION L OSS AND M UTUAL I NFORMATION
Given a pair of data ( x, y ) ∈
X × Y , we donote the joint distribution of X and Y are P XY , and their marginal distributions are P X and P Y , respectively.
By deﬁnition, the mutual information between
X and Y is:
[Equation]
To estimate mutual information between X and Y , (Tsai et al., 2020) proposes a probabilistic classiﬁer method.
Concretely, we can use a Bernoulli random variable C to classify one given data pair ( x, y ) from the joint distribution P XY ( 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:
[Equation]
[Equation]
Obviously, p ( C =0) p ( C =1) can be approximated by the sample size, i.e. n PXPY n PXY , while P ( C =1 | x,y ) P ( C =0 | x,y ) can be measured by a classiﬁer h ( C | x, y ) , and it can be learned by our relation loss with relational head h :
[Equation]
where
C = 1 if the given pair ( x, y ) is from the joint distribution P XY , and C
= 0 if the given pair ( x, y ) is from the product of the marginal distributions P X P Y .
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 maximizing the mutual information between (ˆ z ) from the same trajectory/environment, and thus preserve the trajectory/environment 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)
[Figure]
A. F AIR C OMPARISON WITH TMCL
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 signiﬁcance of our method over TMCL
[Figure]
we test the performance of TMCL with no adaptation
[Figure]
and show the results below:
[Figure]
We can see that the average returns of TMCL without adaptation are signiﬁcantly lower than ours, especially for the classic control task Halfcheetah.
The experimental comparison with TMCL without adaptation is direct evidence to support our claim: environment-separated Z s is important for the generalization of dynamics functions, and our method can signiﬁcantly outperform baselines in zero-shot unseen test environments with different dynamics.
Speciﬁcally, 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.
A. S IMILARITIES V ISUALIAZTION
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 j s belonging to the same environment with the anchor
ˆ z i have signiﬁcant 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 the intervention module, we also visualize the similarity of ˆ z i 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-speciﬁed factors Zs.
By contrast, after incorporating the intervention module, the contexts from different environments have signiﬁcantly 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.
A. P REDICTION E RRORS ON T RAING E NVIRONMENTS
The prediction errors of each method on training environment are given at Figure 10.
[Figure] [Figure]
[Figure]
A. P REDICTION E RRORS ON T EST E NVIRONMENTS
The prediction errors of each method on test environments are given at Table 3.
Speciﬁcally, we test each test environment 10 times, and plot the average prediction error to reduce random errors (Figure 11).
[Figure] [Figure]
Hopper
[Figure]
[Table]
A. P REDICTION E RRORS ON S PECIFIED E NVIRONMENT
The prediction errors of each method on speciﬁed environment are given at Table 4, 5 and 6.
[Figure]
[Table] [Table]
[Table]
A. T HE A VERAGE R ETURNS ON T EST E NVIRONMENTS DURING T RAINING P ROCESS
The average returns on test environments during training process are given at Figure 12.
[Figure]
A. Q UANTITATIVE C LUSTERING P ERFORMANCE C OMPARISON
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.
Speciﬁcally, TMCL has lower clustering performances than CaDM
[Figure]
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.
[Table] [Table]
According to the quantitative clustering performance measures
[Table]
we can see that the clustering performance of our method is superior to baselines by a large margin
[Table]
and the results are consistent with the performance on the test environments
[Table]
[Table] [Table]
[Table]
A. V ISUALIZATION A. T-SNE V ISUALIZATION
[Figure]
Ours (No Intervention)
[Figure]
A. PCA V ISUALIZATION
[Figure]
[Figure] [Figure]
Ours_NO_Intervention
[Figure] [Figure]
[Figure]
CaDM mass=[0.8]mass=[0.9]mass=[1.0] mass=[1.15]
TMCL mass=[0.8]mass=[0.9]mass=[1] mass=[1.15]
Ours mass=[0.8]mass=[0.9]mass=[1] mass=[1.15]