Variational Causal Dynamics: Discovering Modular World Models from Interventions
Authors: Anson Lei
Bernhard Schölkopf
Ingmar Posner
Venue: Accepted by TMLR ()

Abstract: Latent world models allow agents to reason about complex environments with high-dimensional observations. However, adapting to new environments and effectively leveraging previous knowledge remain significant challenges. We present Variational Causal Dynamics (VCD), a structured world model that exploits the invariance of causal mechanisms across environments to achieve fast and modular adaptation. By causally factorising a transition model, VCD is able to identify reusable components across different environments. This is achieved by combining causal discovery and variational inference to learn a latent representation and transition model jointly in an unsupervised manner. Specifically, we optimise the evidence lower bound jointly over a representation model and a transition model structured as a causal graphical model. In evaluations on simulated environments with state and image observations, we show that VCD is able to successfully identify causal variables, and to discover consistent causal structures across different environments. Moreover, given a small number of observations in a previously unseen, intervened environment, VCD is able to identify the sparse changes in the dynamics and to adapt efficiently. In doing so, VCD significantly extends the capabilities of the current state-of-the-art in latent world models while also comparing favourably in terms of prediction accuracy.

Published in Transactions on Machine Learning Research (11/2023)
Variational Causal Dynamics: Discovering Modular World Models from Interventions
Anson Lei Applied AT Lab University of Ouford
onson@robots.oz..co
Bernhard Schölkopf
be@tuuemmp mpg..d
MPI for Intelligntt Systems Tübingen
Ingmar Posner
ingmar@robots.or.ccm
Applied Al Lab University of Oxford
Reviewed on OpenReview: /ttps: opw.reveew.eeeee..eeeee net/ forum? V9t0KYYNK1
Abstract
Latent world models allow agents to reason about complex environments with high- dimensional observations. However, adapting to new environments and effectively leveraging previous knowledge remain significant challenges. We present variational causal dynamics (VCD), a structured world model that exploits the invariance of causal mechanisms across environments achieve fast and modular adaptation. By causally factorising a transition model, VCD is able to identify reusable components across different environments. This is achieved by combining causal discovery and variational inference to learn latent represen- tation and transition model jointly in an unsupervised manner. Specifically, we optimise the evidence lower bound jointly over a representation model and a transition model struc- tured as a causal graphical model. In evaluations on simulated environments with state and image observations, we show that VCD able to successfully identify causal variables, and to discover consistent causal structures across different environments. Moreover, given small number of observations in previously unseen, intervened environment, VCD is able to identify the sparse changes in the dynamics and to adapt efficiently. In doing so, VCD significantly extends the capabilities of the current statc-of-the-art in latent world models while also comparing favourably in terms of prediction accuracy.
1 Introduction
The ability to adapt flexibly and efficiently to novel enviromments is one of the most distinctive and compelling features of the human mind. It has been suggested that humans do so by learning internal models which not only contain abstract representations of the world, but also encode generalisable, structural relationships within the environment (Behrens al. 2018). This latter aspect, is conjectured, is what allows humans to adapt efficiently and selectively. Recent efforts have been made to mimic this kind of representation in machine learning. World models (e.g. Ha & Schmidhuber 2018) aim to capture the dynamics of an environment by distilling past experience into a paramctric predictive model. Advances in latent variable models have enabled the learning of world models in a compact latent space (Ha & Schmidhuber 2018 Watter et al., 2015 Hafner et al. 2018 Buesing et al.) 2018 Zhang et al, 2019) from high-dimennsinnl obscrvations such as images. Whilst these models have enabled agents act in complex environments via planning (c.g. HHffer al. 2018 Sekar et al. 2020)) or learning parametric policies (e.g.
Hafner al. 2019 Ha & Schmidhuber) 2018), structurally adapting to changes in the environment, remains

---

Published in Transactions on Machine Learning Research (11/2023)
Training
Rollout
VCD
A
VCB
KDD
VCD
VcD
ACD
Leamable
2443
Causel
Figure 1: Left: The general architecture of VCD. The dynamics model is trained using observation sequences from multiple environments by minimising the KL divergence between the predicted state distribution and the encoded posterior stale distribution. Rollouts in the latent space can be performed by recursively applying the learnt transition model. Right: The structure of the causal transition model. Each dimension of the latent space is treated as a causal variable. Predictions are made using only the causal parents of each variable, according to a learnt causal graph.
significant challenge. The consequence of this limitation is particularly pronounced when deploying learning agents to environments, where distribution shifts occur. As such, we argue that it is beneficial to build structural world models that afford modular and efficient adaptation, and that causal modeling offers a tantalising prospect to discover such structure from observations.
Causality plays a central role in understanding distribution changes, which can be modelled as causal interventions (Schölkopf et al. 2021). The Sparse Mechanism Shift hypothesis (Schölkopf et al. 2021: Bengio et al. 2019) (SMS) states that naturally occurring shifts in the data distribution can be attributed to sparse and local changes in the generative process. This implies that many causal mechanisms remain invariant across domains (Schölkopf et al., 2012 Peters et al., 2016 Zhang et al., 2015). In this light, learning a causal model of the environment enables agents to reason about distribution shifts and to exploit the invariance of learni mechanisms across different enviromments. Hence, we posit that world models with a causal structure can facilitate modular transfer of knowledge. To date, however, methods for causal discovery (Spirtes et al.] 22000 Pcarl 2000; Peters et al., 2017: Brouillard et al.] 2020 Ke et al.] (020) require access to abstract causal variables to learn cansal models from data. These are not typically available in the context of world model learning, where wish to operate directly high-dimensional observations.
In order to benefit from the structure of causal models and the ability to represent high-dimensional obser- vations, we propose Variational Causal Dynamics (VCD), which combines causal discovery with variational inference. Specifically, we train a latent state-space model with a structural transition model using varia- tional inference and sparsity regularisation from causal discovery. By jointly learning a latent representation of the observations and a transition model, VCD learns a causally factorised world model that can modularly adapt to different environments D The key intuition behind our approach is that, since sparse structures can only be discovered on semantically mcaningful variables, training the representation and the causal discovery module in an end-to-end manner acts as an inductive bias that encourages well disentangeed representations. By leveraging the learnt causal structure, VCD is able to identify the sparse mechanism changes in the envi- ronment and re-learn only the intervened mechanisms. This enables fast and modular adaptation to changes in dynamics. The code for the experiments are available at https: //githu/ cow com/applied- ai ai-lab/VCn
2 Related Work
Predictive models of the environment can be used to derive exploration- (Sekar et al. 2020) reward-driven
(Ha & Schnidhuber 2018: Hafner et al. of latent dynamics models. World models (Ha & Schmidhuber 2018 train a representation encoder and a
2019 2018) behaviours. In this paper, we focus on the learning
RNN-based transition model in a two-stage process. Other approaches [Hffner et al. 2018 Zhang et al.
The assumption under which the model is valid discussed section
2

---

Published in Transactions on Machine Learning Research (11/2023)
Watter et al. 2015| learn generative model by jointly training the representation and the transition via variational in Inference PlaNet [Hafner et al. 2018 parameterises the transition model with RNNs. F2C (Watter et al.) 2015 Banijamali et al.) 20088) and SOLAR (Zhang et al.) 2019) use locally-linear transition models, arguing that including constraints in the dynamics model yields structured latent spaces that are suitable for control problems. In a similar vein, the use of latent prediction models have also been explored in the context of video prediction (Villegas et al., 2019 Denton & Fergus 2018: Assouel et al. 2022). Our proposed approach shares the general principle that latent representations can be shaped by structured transition mechanisms (Ahuja et al.) 2 However, to the best of our knowledge, VCD is the first approach that implements causal transition model given high-dimensional inputs in this context.
2019
Causal discovery methods enable the learning of causal structure from data. Approaches can be cat- egorised as constraint-based (e.g. (Spirtes et al., 2000)) and score-based (e.g. (Hauser & Bühlmann 2012) The reader referred to (Peters et al. 2017) for detailed review of causal discovery methods. Motivated by the fact that these methods require access to abstract causal variables, recent efforts have been made to reconcile machine learning, which has the ability operate on low-level data, and causality (Schölkopf et al., 2021). Recent advances in this area include theoretical works independent component analysis exploring the conditions under which disentanglement is possible (Yao et al. 2022 Hyvarinen et al. 2019 Lachapelle ett al. 2021 Our current work situates within this broader context of causal representation learning, and aims to identify causally meaningful representations via the discovery of causal transition dynamics. [Lachapelle et al.] 2021 Lachapelle & Lacoste-Julien, 2022) discusses the theoretical identifiability of a latent state-space model similar to ours. Whilst we do not make identifiability claims, our work builds on the inductive biases developed in these works. Building on this theoretical foundation, VCD focuses on the adaptation capabilities of causal models and empirically shows their applicability to image observations.
There has been significant interest in leveraging causal reasoning in world model learning (Li et al. 2020a b). These approaches typically rely on predefined representations such as keypoints (Li et al. (0200)) or object slots (Li el. al., 2020a), and attribule changes the dynamics to unobserved confounders, which can be estimated from observations. Our approach differs in that we couple the learning of the representation and the model, which enables dynamics-aware discovery of causal representations. Similar approaches of jointly learning representation and model structure have also been explored in different contexts such as reinforcement learning. AdaRT, (Huang et al. 2022) considers few-shot adaptation under sparse mechanism shift (SMS) by attributing changes in the environments to low-dimensional environment-specifee parameters. In contrast, by explicitly drawing from causal discovery, VCD models the dynamics of the scene as independent causal mechanisms. This, in turn, allows our approach to directly identify localised mochanism shifts between environments in an interpretable manner. Another branch of work considers temporally intervened sequences (Lippe et al. 2022bb 2023 Brehmer et al., 2022). While these share our general approach of learning sparse graphs in the latent space, they consider interventions direct changes to the causal variables in each time step, whereas VCD formulates interventions as shifts in the dynamics across environments, which enables the capability of domain adaptation. In the model-free RL literature, Raileanu et al. (2020); Sang et al.] (2022) investigate a similar setling where agents are trained in diverse environments and adapt to new environments with limited data.
Another branch of causalityinnpireed work leverages the invariance of causal mechanisms by learning invariant predictors across environments (Tian & Pearl 2001 Schölkopf al., 2012 Peters et al.) 2016 Zhang et al.) 2017 Arjovsky et al. 2019).. This invariance has been studied in the context of state abstractions in MDPs (Zhang et al. (020), and invariant policies can be learnt via imitation learning from different environments [Bica et al. 2021 Tn contrast, our approach models the full generative process of the data across different environments rather than learning discriminative predictors.
3 Background
This section describes the preroquisite definitions and formulations for our proposed method. We highlight the strengths and weaknesses of latent state-space models, causal models and causal discovery methods, and motivate our approach that builds upon these to learn causal world models.
3

---

Published in Transactions on Machine Learning Research (11/2023)
Latent state-space models In complex environment with high-dimensional observations, such im- ages, learning compact latent state space that captures the dynamics of the environment has been shown to be more computationally efficient than learning predictions directly in the observation space (Hafner et al, 2018 Buesing et al. 2018). Given a dataset of sequences {(o a. Q:7)) )) with obscrvations of and actions a* at discrete timesteps I, a generative model of the observations can be defined using latent states 20:7 as
T / II (2)(()))))((((11>>)))))))))))))1 (2(())))a)))p))) (p(2")2)))))2)))1)) a' )d22
p(ouT, OTT) =
(1)
where 22( and po(="12*-1, a are the observation model and the transition model respectively. For simplicity, we do not learn the policy term p(a"|z²) and omit it throughout this paper as it is constant with respect to the parameter 0. The variational evidence lower bound can be written as
T
EILBO(O, 6) =
Σ 2
E (o (o - [xt [KL[ga(> (oP) 10
(2)
where (=") is learnable approxinate posterior of the observations See Appendix A for the derivation.
RSSM (Hafner et al. 2018) employs a flexible transition model parameterised as a fully connected recurrent neural network, where the transition probability is split into a stochastic part and a deterministic recurrent part,
~' po(="1n²)
ht ='= fo 2 4
where (()) is instantiated as GRU Cho al. 2014 and ht is the associated hidden state. Intuitively, this provides path through which information be passed on over multiple timesteps.
Despite their success, these models cannot reason about changes in the environment. Specifically, they are unable to structurally utilise prior knowledge learnt from different environments under distribution shift. To this end, we argue that imposing a causal structure on the transition model equips the learning agent with the ability to adapt to changes a modular and efficient manner.
Causal graphical models A causal graphical model (CGM) (Peters et al. 2017) defined a set of random variables (X.,....aa"" their joint distribution Px, and directed acyelic graph (DAG), (X, E), where each edge (i.j) E E implies that X, is direct cause of Xj. The joint distribution admits causal factorisation such that
p xa)
II (P.)), ==
where PA is the set of parents to the variable X; in the graph. Each of the conditional distributions can be considored as an independent causal mechanism.
In contrast to standard graphical models, CGMs support the notion of interventions, i.e., local changes the causal distribution. Formally, an intervention on the variable X; is modelled as replacing the conditional distribution p(z, while leaving the other mechanisms unchanged. Given the set of intervention targets I C X, the interventionss distribution can be written as
p' Ia) = P(.))A)) II p' (a,)A)) EET
(5)
egT
where p'(-l-) is the new conditional distribution corresponding to the intervention. The SMS hypothesis (Schölkopf et al. 2021) posits that naturally occurring distribution shifts tend to correspond to sparse changes
"In the current work, focus environments without rewards. However, the proposed method can be readily extended to include reward prediction.
The approximate postcrior, which only depends on the current observation, simplification of the full posterior which is conditioned on all observations. This consistent with prior the experiment. However, the algorithm presented in this work can be readily adapted use the full posterior instead.
and is found to sufficient for the dynamical systems
4

---

Published in Transactions on Machine Learning Research (11/2023)
Transition Model
P1
#
Inea
4
P2
#/
#
K
(a) Latent state space model
(b) Typical transition model
(c) VCD transition model
Figure 2: (a) The structure of latent state space models in one timestep. Solid lines denote the generative process of the environment, dashed lines denote approximate posteriors, shaded nodes denote observed variables, and the bi-directional arrow denotes the KL divergence in the ELBO term. In this work, we propose novel architecture for the transition model. (b) The structure of latent transition models where the probability distribution is modelled as fully connected neural network. (c) The causal transition model in VCD. Each individual variable has a separate conditional distribution pi. Each mechanism is conditioned on a subset of the previous state and action. The blue lines highlight the intervened mechanism which is specific to the environment k, as opposed to the mechanisms corresponding to the black lines, which are shared across environments.
in causal model when factorized as 4 i.e. changes of few mechanisms only. Causal mechanisms thus
tend to be invariant across environments (Schölkopf ct al.) 2012 Peters et al. 2016 Zhang et al. 2017 In this light, we argue that causal world model can structurally leverage the invariance within distribution shifts as an inductive prior. In order to learn causal model in the context of world models, we draw inspiration from recent advances in causal discovery which aim to learn causal struetures from data.
Differentiable causal discovery We focus on mothods that formulate causal discovery as a continuous optimisation problem [Brouillard et al., 2020 Ke et al. 2020 Bengio et al. 2019 Lippe et al.) 2022a) as these can be naturally incorporated into the variational inference framework. Furthermore, since the causal variables are learnt in our model, the causal discovery module is required learn causal graphs from unknown intervention targets. In this work, follow the formulation in Differentiable Causal Discovery with Interventional data (Brouillard et al. 2020) (DCDI), which optimises continuously param- eterised probabilistic belief over graph structures and intervention targets. See Appendix B for further detail.
In the context learning world models from digh-dimensional observations, the drawback this approach, and of causal discovery methods in general, that it requires access to semantically meaningful causal variables, much like classical AI required symbols in terms of which algorithms could be formulated. In the next section, present method to perform causal discovery and learn causally meaningful representations jointly.
4 Variational Causal Dynamics
Similar to causal discovery with interventional data, variational causal dynamics (VCD) learns from action- 0:T
obsorvation sequences from an undisturbed environment, (o0.) 2 ) and K intervened environments,
{(o97] ( ) Throughout this paper, we assume that 1) changes between the environments are due 'soft' interventions where the structure of the causal graphs remains the same but individual parameters of the mechanisms change, 2) there are no instantaneous causal dependencies in each timestop, and 3) the observation function does not change across environments.
5

---

Published in Transactions on Machine Learning Research (11/2023)
The general approach of VCD follows the latent state-space model framework (Fig 2a where a latent repre- sentation of the observations is jointly learnt with transition model by maximising the ELBO. In contrast to existing approaches that parameterise the transition probability ((z"|24-1,a'-)) as a feedforward neural network (Fig 2b). VCD learns a structured transition model by utilising inductive biases from causal discov- ery. Importantly, our approach is grounded in the hypothesis that causal structures can only be discovered on meaningfully disentangled representations of the system. We therefore argue that training a represcntation and a transition model jointly optimise a causal discovery objective can lead to latent representation that affords sparse transition models and is therefore semantically meaningful. We discuss this approach in relation to existing theoretical works on identifiability in Section 4.5 Taking the view that changes in the dynamics can be attributed to sparse causal interventions, we posit that a causally factorised transition model facilitates modular adaptation to new environments by explicitly leveraging the invariance of causal mechanisms. We propose an algorithm for adaptation in Section 4.4
4.1 Causal transition model
In order to perform causal discovery on the transition dynamics, the transition model in VCD is designed to mimic the structure of a CGM (Eq. 5 In the following paragraphs, we motivate and describe the three main architectural features of VCD that enable causal discovery: independent mechanisms, sporse causal dependencies and sparse interventions. Access to the causal graph 9 and the intervention targets for each environment I& is assumed throughout this subsection. We describe the method of learning these jointly with the model in Section
Independent mechanisms In contrast to fully-connected transition model, the transition probability is factorised into individual conditional distribution as
"(2912"-1''"" aa-1) II [=platla-1 aa-
(6)
where d is the dimension of the latent space, to be set as a hyperparameter. 2 denotes the ith dimension of the latent state, and each conditional distribution p; is a one-dimensional normal distribution with mean and variance given by separate neural networks. This factorisation and separation of parameters motivated by the Independent Causal Mechanism principle, which states that the causal generative process of a system's 's variables is composed of autonomous modules that do not inform or influence each other (Schölkopf et al., 2021). This explicit modularity of the model structure enables the notion of interventions, where individual conditional distributions are locally changed without affecting the other mechanisms.
Sparse causal dependencies Following the structure of a CGM, condition each variable only on its causal parents according to the learnable causal graph G. rather than the full state, Given a graph G, we define the binary adjacency mask M where the entry M is 1 if and only if (3---_a²=1) is a causal parent of z).. This is consistent with the intuition that, in physical systems, states interact with each other in a sparse manner (Goyal et al. 2021a), and actions tend to have a direct effect on only a subset of the states. Under this parameterisation, the causal transition probability can be written as
d
p(2'12'11, al = II pr(c o aa')),
where @ denotes the element-wise product, [+:]] denotes the concatenation of vectors, and MS is the binary mask that selects only the causal parents of 2; under the graph g.
Sparse interventions Following the SMS hypothesis (Schölkopf et al. 2021 distributions across the K intervened environments are due to sparse interventions in the ground truth causal gencrative process. In order to incorporate sparse interventions in VCD, given the set of learnt intervention targets I& for each environment, we define the binary intervention mask RD where the entry R is 1 if and only if the variable 2, is in the set of intervention targets in environment k. For each variable 2,, R2 acts
we assume that changes in
6

---

Published in Transactions on Machine Learning Research (11/2023)
switch between reusing shared observational model and an environment-spoeffic interventional model. The full interventional causal model of the transition probability in the environment. & can be written as
= II a
(4 @ [2- 1)- P: (0) p= 1
(1) (4 @ [2- 1 t p
(8)
where p0)) the observational mechanism that shared across all environments and p p(k) is the intervened distribution specific to environment k. Intuitively, this model ensures that all environments reuse the same conditional distributions pob unless it is deemed that particular mechanism has been intervened on (i.e. R == 1). This modular adaptation bet environments facilitates structural transfer of knowledge be- tween different environments. Note that the structure is that of a causal graphical model. The overall architecture of the VCD transition model is summarised Fig. 2c
4.2 Recurrent module
Similar to RSSM (Hafner et al. 2018), we augment the model with a deterministic recurrent path to enable long-term predictions. To ensure that each conditional distribution only has access the causal parents, in the same way that each conditional distribution modelled by a separate network, each conditional distribution keeps separate recurrent unit and a corresponding hidden activation:
- p()(=)))
= f(hhh-,, Mº o ,a
where is a recurrent module specific to the variable 24, instantiated a GRU CCoo al. 2014).
4.3
Training
The task of model training is determine the model parameters 0, the approximate posterior parameters well the causal graph 9 and the intervention targets I. These be jointly trained in way similar to DCDI. We parameterise the belief over the causal adjacency matrix M as a random binary matrix. Each entry My follows a Bernoulli distribution with success probability aaa,), where Oij is a scalar parameter and a-) is the sigmoid function. Similarly, a random binary matrix R is parameterised using the scalar variable Bes for each entry. Unlike DCDI, due to the existence of latent variables, we cannot directly maximise the data likelihood. Instead, we maximise the expected ELBO across all environments over causal graphs and intervention masks,
6(0,0,033)
Σ
Eg, ELBO 0:T O.T 0,,,9) - A,
(10)
kE...I......K
where
0:T D.T ELBO
Σ E Bost |=)()
-E (=t- lot-1) KL[ (2 (0 (2 a
(11)
(8) (x'12"-1, ,,a-1) is further factorised as in Equation The gradients through the outer expectation and the expectation term in ELBO are estimated using the straight-through Gumbel-max trick (Jang et al. 2017) and the reparameterisation trick (Kingma & Welling 2014) respectively. For further implementation details, derivation of the lower bound, and model architectures, see Appendices A and
Po
4.4 Adaptation
Due to the modular nature of the transition model, VCD can naturally adapt to new, unseen environments by jointly inferring the intervention targets and the new model parameters for the intervened mechanisms. Specifically, the transition model in new environment can be written
new (2 a a11 = II p (0) (a o 12--1 a - p1< -
o [2t-1 a/)
(12)

---

Published in Transactions on Machine Learning Research (11/2023)
where R' is the intervention mask for the new environment and pl is the environment-speecfic distribution. The parameters of (0) are carried over from the training environments, meaning that the model reuses trained mechanisms unless the variable is intervened on in the new environment. Under the SMS hypothesis, only a small subset of the mechanisms need to be adapted in a new environment. As such, by leveraging past experience, VCD can adapt more quickly to environment change.
One way to implement modular transfer is to train VCD on trajectories in the new environment while reusing the trained parameters for the causal graph and the mechanisms in the undisturbed environ- ment. This can be done by identifying the set of intervention targets the now environment and learning the model parameters for the corresponding conditional distributions 0'. These can be jointly trained in an analogous way Eq. I0
Ladapt (0' 0'' = Eg,r ELBO
GOT Ree neww 0', o, -
(13)
where 3' are the Bernoulli parameters for the distribution over intervention targets. Under this framework, VCD has the capability to identify the sparse mechanism changes in new environment and learn the transition mechanisms only for the intervened variables. In Section 5 we show that VCD can identify sparse interventions from a small amount of data and is able to refit. the world model with improved data efficiency.
4.5 Identifiability of Causal Graph and Causal Representations
VCD's ability to adapt in a modular fashion hinges on the discovery of the correct, or at least approximately correct, causal structure as well as on causally meaningful representation space. The conditions under be learnt from data are studied under the term identifiabilily. While make
which such structures identifiability of causal graph and representation in the paper, approach is formal claims about
grounded in recent theoretical results which makes use of the sparsity of causal relationships (Brouilardd
et al. 2020 Lachapelle el. al. 2021). In this subsection, we foreground the relationship between these results and our approach.
[Brouillard al. 2020) shows how the correct casual graph can be identified in a setting similar ours. Given observed causal variables X in & different environments with intervention targets I' and ground-truth casual graph G*, the following score function is proposed,
K
I) := sup Σ
Ex (
- Az|ZI,
(14)
where (()) is the likelihood of the data under the model environment k parameterised by o, and other definitions are the same as eq 10 The main theoretical results in (Brouillard al.) 2020) shows that with small enough Àg and Az, maximising this score function will result in the correct intervention targets and causal graph that is I*-markov equinalent to the ground-truth graph. In other words, with suitable choice of hyperparameters, maximising the score function will lead to the correct causal graph and intervention largets. The main idea of this theorem is that if the model learns the underlying distribution perfectly, any learnt graph that is not equivalent to the ground-truth truth graph is necessarily less sparse and hence have a lower score. VCD uses the same score function as the training objective with the log likelihood replacod with the evidence log lower bound since we do not have access the causal variables.
To this end, (Lachapelle et al. 2021 considers the identifiability of causal variables from high di- mensional observations, i.c. whether an encoder can learn latent representations that correspond to causally meaningful variables. Given a dataset of observattion-cctoon sequences, (o', a T generated by the transition model:
o' = f(='),
I | (a))
(15)
p(2/|=<1, ( = he he(2))expp
( < G 2 a 1/ ,,a)))))
(16)
8

---

Published in Transactions on Machine Learning Research (11/2023)
where f is invertable function that maps the latent causal variables to high dimensional observations. The theory holds conditional distributions within the exponential family, where T is the sufficient statistios, X is a transition function that outputs the natural parameters of the distribution, h and vf are the base measure and the normalising constant respoctively. G and G are the binary masks corresponding to the causal graph of the observed system, analogous to the learnt causal masks for VCD. In this framework, the =(f
learnable parameters are 0 == (f, X, G). The main theorem states that, if the distribution induced by the learnt model 0 matches the ground-truth model 8, then under certain technical conditions concerning the structure of the problem, ||G||o < ||G||o implies that 0 and 0 are equivalent up to permutations. In other words, if the learnt model perfectly captures the data distribution with causal graph that is at least as sparse as the ground-truth causal graph, then the learnt casual variables must be completely disentangled. The intuition behind this theorem can be demonstrated with a simple example. Consider system with transition model that only depends on one causal variable. With an entangled representation, i.c. information about the causal variable is spread across several latent dimensions, the model must take into account all of these latent dimensions in order to make predictions. As such, the model will require multi- ple edges the learnt causal graph. In contrast, only one edge needed if the representation is disentangled.
In sum, these results show that sparsity plays a central role in the identification of both the causal graph and causal representation. With this in mind, the architecture of VCD is designed to afford a sparse transition model that mimics the structure of CGM. Whilst, in practice, the condition of perfectly matched distributions might not hold, this general principle of sparsity regularisation forms the basis of our approach. In section 5 we empirically show that VCD is able to learn a disentangled latent representation and identify most of the causal relationships.
5 Experiments
In this section, we demonstrate that VCD is able to learn from multiple environments with different dynamics by learning a causal world model that explicitly captures the changes between the envi- ronments as interventions. We evaluate, qualitatively and quanti- tatively, the learnt representations and causal structures and show that VCD is able to capture the dynamics of the scene using sparse dependencies between states, as well as identify sparse changes be- tween the environments. Moreover, we illustrate that, in an unseen environment, VCD can leverage past experience to perform mod- ular adaptation, resulting in significantly improved data efficiency over the baselines.
2
4
Dataset We evaluate VCD on a simulated dataset a 2-D multi- body system which contains four particles that affect each other via a spring or an electrostatic-lik force. The configuration of the environment is shown in 3ig The action 2² is an external force that applics to particle 4. The environment is designed such that there is an unambiguous ground-truth causal graph between the causal variables and well-defined changes in the dynamics. We consider changes such as strengthening or removing one of the springs, increasing or decreasing the mass of a particle, or constraining the position of a particle along the æ or the y axis. Each of these changes can be considered as an intervention on one or more causal variables. See Appendix C for further details of the environment and the full list of interventions. We evaluate VCD in two experiments: Mixed-state, where the observation is given by applying affine transformation to the ground-truth positions of the particles, and Image, where the observation is given by rendering the environment as a RGB image.
Figure 3: The configuration of the multi-body environment. The red and the blue lines represent an attractive and repulsive force respectively.
Basclines We compare the performance of VCD against RSSM Hafner et al. 2018), a statc-of-tho-art latent world model that served as inspiration for VCD. As RSSM does not support learning from multiple environments, we consider two adaptations of RSSM with different levels of knowledge transfer wetween
9

---

Published in Transactions on Machine Learning Research (11/2023)
Mixed-state
Image
RSSM
MultiRSMM
VCD
State Enor
Pixel Error
50
100
125
150
175
20
40
60
timestep
timestep
Figure The rollout error measured on validation trajectories. The models receive observations on each timestep for the first half of cach trajectory (shadod), after which the model is rolled out based on latent predictions. The reported error in the mixed-state environment (left) and the image environment (right) are squared error in the ground-truth state space and squared error in pixel values respectively. VCD outperforms the baselines in both modes of observation. Evaluations on the image experiments using non-pixel-basee errors are available in Appendix
Table 1: The modified MCC scores for RSSM and VCD. Higher score mcans better discntanglement.
Expriment
RSSM
VCD
Mixed-state
0.728
0.975
Image
0.715
0.908
environments: (1) RSSM, where one transition model is trained over all environments, i.., maximum parameter sharing across environments; and (2) MultiRSSM, where individual transition models are trained on each environment, with shared encoders and decoders. This corresponds to the case where no knowledge about dynamics is transferred, i.c., each model is local expert. Compared to those two extremes of knowledge sharing, VCD captures environment-specde
behaviours whilst reusing invariant mechanisms.
Prediction performance VCD and the baselines are trained on dataset composed of 2000 trajectories from each of the undisturbed environment and five intervened environments. The models evaluated trajectories from a validation set drawn from the training environments. Fig. 4 shows the rollout error for each of the models with state and image observations. Unlike VCD and MultiRSSM, RSSM can only capture the average dynamics the environments due to parameter sharing. As such, it is not able to capture environment-speeific behaviours. This is reflected the prodiction accuracy of the model. In the mixed-state experiment, VCD is able to identify changes between the envirouments and performs as well as MultiRSSM in terms of prediction error. similar trend is shown in the image space, where VCD outperforms the baselines. See Appendix D for image rollout examples that demonstrate the difference in behaviour between RSSM and VCD.
Causal discovery and Disentanglement One of the key hypotheses of this work is that jointly learning representation and transition model using causal discovery leads to causally meaningful representations. Here we examine the quality of the learnt latent space and the causal structure of the transition model. Fig. 5 shows the image reconstructions of points drawn from straight line along dimension 6 of the latent space. Since both models are initialise with the same encoder, this provides a qualitative intuition as to how the causal discovery inductive bias shapes the latent space. Whilst this dimension of the latent space in RSSM is able to encode only the blue particle, VCD able learn an anis-aligned coordinate (æ coordinate) of the blue particle with respect the environment. Note that the motion of bouncing off the boundaries is only separable in the x, y frame, implying that the dynamies in axis-aligned coordinates is sparser. In table 1 we
10

---

Published in Transactions on Machine Learning Research (11/2023)
RSSM
VCD
along dimension 6
Figure 5: Reconstructed images from points sampled along dimension 6 of the latent space. Compared to the latent space learnt in RSSM, VCD is ablc to learn an axis-aligned representation where the blue particle moves horizontally. See Appendix D for more examples.
Sparsity of learnt causal graphs and intervention largets. We compare the learnt graph with the ground truth causal graph by mapping each ground-trute state dimension to a latent dimension by picking the dimension with the highest correlation score.
Table 2:
Observation
# of edges
Correct Edges
Missed Edges
False Positives
Mixod state
causal graph intervention targets
42/288
8
18/80
10
5
Image
causal graph
73/288
22
13
intervention targets
31/80
17
provide the mean absolute correlation coefficient scord" for the representations for both RSSM and VCD, which measures the quality of representation disentanglement by comparing the learnt latent space with the ground-truth states. We see that both cases, training with the sparsity regularisation in VCD achieves a higher MCC score, which confirms that VCD learns & more disentangled representation. More qualitative exploration of the learnt latent space is provided in Appendix [ By considering the correlation between cach dimension of the learnt representation and the ground-truth states, it is possible to compare the learnt causal graph with the ground-truth (label 2 As desired, the learnt causal graphs in both the mixed-state and the image environment are found to be sparse and encode most of the causal edges the graph. Taken together with the prediction accuracy of the model, this means that VCD is able retain causal edges that are important for predicting the dynamics of the environment while pruning out spurious edges. A more in-depth analysis of the learnt causal graph is provided in Appendix [) showing that the edges in the learnt causal graph do correspond to meaningful causal relationships.
Adaptation We provide empirical evidence that VCD can adapt a new environment with less data compared to RSSM and MultiRSSM by reusing learnt mechanisms in modular fashion. We collect datasets of different. sizes in a previously unseen intervened environment. where particle RSSM and MultiRSSM adapt to the new environment by optimizing the ELBO (Eq, 2 with the difference that MultiRSSM instantiates a new transition model randomly and RSSM initialises the transition model using pre-trained parameters. Following Eq. 13 VCD performs adaptation by jointly estimating the inter- vention targets and the parameters of the intervened mechanisms. The models are evaluated on a validation set of trajectories in the new environment. Fig. shows the rollout error of the models trained on datasets
constrained horizontally.
*Due the fact that the ground-truth state dimension (8) less than the latent state dimension (16), modify the MCC score by caleulating the of the top information.
of the 16, effectively ignoring the latent variables that do capture the
11

---

Published in Transactions on Machine Learning Research (11/2023)
Mixed-States
VCD
15 trajectory
26
R75
59
25 100 ds 24 al trajectories
24 i5 us 2o trajectories
as 1o 1is 64 trajectories
15 as 15 de 1h
256 traje trajectories
10
Images
302
100
200
209
24
280
160
100
178
38
160
A 160 trajectories
20
10 trajectories
trajectories
640 trajectories
2560 trajectories
Figure 6: Rollout errors in the mixed-state and image experiments where the models are trained on datasets of varying sizes. The models receive observations for the first half of the trajectory (shaded), and perform latent space prediction for the rest. VCD significantly outperforms the other models with little data. The bar below each plot shows the learnt intervention targets in the new environment. VCD is able to reuse most of the previously learnt mechanisms, as indicated by the sparsity in the learnt intervention targets. In the mixed-stale, verify that the learnt intervention target (dimension 8) corresponds to the ground-truth target (y1), see Appendix D Evaluations using non non-pixel-basee errors available in Appendix F
of varying size. Across all models, performance improves the dataset grows. However, in contrast to RSSM and MulGiRSMM, which overfits the dataset when the number of trajectories is small, VCD is able to predict significantly more accurately. This is because VCD estimates the intervention targets and reuses trained modules that remain invariant. Remarkably, in the mixed-state experiment, VCD able to converge to a single latent dimension that encodes the y position of particle with just one training trajectory. Sec A Appendix for more experiments with different interventions.
6 Conclusion and discussion
In this paper, we propose VCD, predictive world model with a causal structure that is able to consume high-dimensional observations. This is achieved by jointly training a representation and a causally structured transition model using modified causal discovery objective. In doing so, VCD is able identify causally meaningful representations of the observations and discover sparse relationships in the dynamics of the system. By leveraging the invariance of causal mechanisms, VCD is able to adapt to new environments efficiently by identifying relevant mechanism changes and updating in modular way, resulting in significantly improved data efficiency.
Limitations and future directions Our experiments highlight VCD's ability to discover causal variables and mechanisms from image observations. Future work should investigate the efficacy of the proposed method in more challenging environments such as [Baradel et al. 2019 Ahmed et al. 2021). One possible drawback of the proposed framework is that. the causal graph is assumed to be static across all timesteps and all environments. In more complex environments, could be more difficult for the model to discover temporally sparse interactions such as collisions, and adapt to scenes with different numbers of objects.
such, exciting directions for future research also include exploring the synergy between causal world models and object- Kügelgen et al. 2020). temporally-local
centric generative models (Wu et al. 2021 Engeleke al. 2021
12

---

Published in Transactions on Machine Learning Research (11/2023)
causal inftuence detection (Sitzeer et al. 2021 Pitis et al. Kipf et al. 2019
2020) and relational reasoning Goyal et al. 2021b
Broader Impact Statement
While the present work significantly advances the current stato-of-the-art in world-modelling and causal representation learning, we expect its immediate impact outside of the machine learning community to be low as current methods can not yet deal effectively with real-world scenarios.
Acknowledgments
This research was supported by an EPSRC Programme Grant (EP (EP//V00748/1), The authors would like to acknowledge the use of the University of Oxford Advanced Research Computing (ARC) facility in carrying out this work. http://dxxx http://x..oi.org//// (0..522/zzoodoo 22558.. The authors thank Ceri Ngai, Jack Collins, Oiwi Parker Jones, Frederik Nolte and Jun Yamada for useful comments.
References
Ossama Ahmed, Frederik Träuble, Anirudh Goyal, Alexander Neitz, Manuel Wüthrich, Yoshua Bengio, Bernhard Schölkopf, and Stefan Bauer. Causalworld: A robotic manipulation benchmark for causal structure and transfer learning. In International Conference on Learning Representations, 2021.
K. Ahuja, J. Hartford, and Y. Bengio. Propertics from mechanisms: An equivariance perspective on identi- fiable representation learning. arXiv preprinl arXim:2110.1572 15796, 2021.
M. Arjovsky, L. Bottou, I. Gulrajani, and D. Lopez-Paz. Invariant risk minimization. ar preprinl arXiv arXiv:1907.08800 2019.
Rim Assouel, Lluis Castrejon, Aaron Courville, Nicolas Ballas, and Yoshua Bengio. Vim: Variational in- dependent modules for video prediction. In Conference on Causal Learning and Reasoning, pp. 70 89, PMLR, 2022.
E. Banijamali, R. Shu, H. Bui, A. Ghodsi, et al. Robust locally-linear controllable embedding. In Interna- tional Conference on Artificial Intelligence and Statistics, pp. 1751-1759. PMLR, 2018
Fabien Baradel, Natalia Neverova, Julien Mille, Greg Mori, and Christian Wolf. Cophy: Counterfactual learning of physical dynamics. arXiv preprint or Xiu: 100..999000000 12000, 2019.
T. Behrens, T. Muller, J. Whittington, Mark, A. Baram, K. Stachenfeld, and Z. Kurth-Nelson. What
is cognitive map? organizing knowledge for flcxible behavior. Newron, 100(2):490- 509, 2018. ISSN 0896-6273. doi: https://doi.org//0001////// http:///....erroo...111.00.00000 20.8.0 URL http:::/www.acieeceeiretttt com/ science/artien /ttaa////....ccee//////999622233
Y. Bengio, T. Deleu, N. Rahaman, R. Ke, S. Lachapelle, O. Bilaniuk, A. Goyal, and C. Pal. A meta-transfer objective for learning to disentangle causal mechanisms, 2019. URL 1rxiv:
I. Bica, D. Jarrett, and M. der Schaar. Invariant causal imitation learning for generalizable policies. Aduances in Newral Information Processing Systems, 34, 2021.
Johann Brehmer, Pim De Haan, Phillip Lippe, and Taco Cohen. Weakly supervisod causal representation learning. In ICLR2022 Workshop on the Elements of Reasoning: Objects, Structure and Causality, 2022. URL htttps://openrrviiiii wwwpaarrviiieeeeeffrroeeeee www.://oorummmddrrroUUU.........
P. Brouillard, S. Lachapelle, A. Lacoste, S. Lacoste-Julien, and A. Drouin. Differentiable causal discovery from interventional data. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (ods.), Advances in Neural Information Processing Systems, volume 33, pp. 21865- 21877. Curran Associates, Inc., 2020. URL http:::///rooeedgg.norrrrrpp neurips. httprrpps.c/pppp///////L./. f8b7 F5577aaaaatt00000022200000222222
13

---

Published in Transactions on Machine Learning Research (11/2023)
L. Buesing, T. Weber, S. Racanière, S. M. Ali Eslami, D. Rezende, D. Reichert, F. Viola, F. Besse, K. Gregor, D. Hassabis, and D. Wierstra. Learning and querying fast generative models for reinforcement learning. ArXiu, 0bs/1802.030006 2018.
K. Cho, B. van Merrienboer, C. Gulcehre, F. Bougares, Н. Schwenk, and Y. Bengio. Learning phrase representations using rmn encoder-docoder for statistical machine translation. In Conference on Empirical Methods in Natural Language Processing (EMNLP 2014), 2014.
Emily Denton and Rob Fergus. Stochastic video generation with a learned prior. In International conference on machine learning, pp. 1174-1183. PMTR, 2018.
M. Engelcke, O. Parker Jones, and I. Posner. Genesis-v2: Inferring unordered object representa- tions without iterative refinement. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and
J. Wortman Vaughan (eds.), Advances in Newral Information Processing Systems, volume 34, pp. 8085-8094. 8094. Curran Associates, Inc., 2021. URL httpps:/pproceedings www.eeeeerrrirccceeeeeeee w/tp:/ps.c//ppppe///////iii pdf
A. Goyal, A. Lamb, J. Hoffmann, S. Sodhani, S. Levine, Y. Bengio, and B. Schölkopf. Recurrent independent mechanisms, In 9th International Conference on Learning Representations (ICLR), May 2021a. URL https://openreviers net/ htt.eeenet///ddpommmmmmmmmml.l. Fi==ml.cmm1UUxx-
Anirudh Goyal, Aniket Didolkar, Nan Rosemary Ke, Charles Blundell, Philippe Beaudoin, Nicolas Heess, Michael C Mozer, and Yoshua Bengio. Neural production systems. Advances in Neural Information Processing Systems, 0125673-256877 2021b.
D. Ha and Schmidhuber. Recurrent world models facilitate policy evolution. In Adoances Neu- ral Information Processing Systems 31, pp. 2451 2463. Curran Associates, Inc., 2018 URL https:
1/ //papers... ceapers.iis..o /worldmodels..iiieee 10
https:
Hafner, T. Lillicrap, J. Ba, and M. Norouzi. Dream to control: arXiv preprint Xrxi 2106 2019.
Learning behaviors by latent imagination.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. arXiv preprint arXit:1811.04300 2018.
A. Hauser and Bühlmann. Characterization and greedy learning of interventional Markov equivalence classes of directed acyelic graphs. Journal of Machine Learning Research, 13:2409 -4644, 2012.
Biwei Huang, Fan Feng, Chaochao Lu, Sara Magliacane, and Kun Zhang, AdaRL: What, where, and how to adapt in transfer reinforcement learning. In International Conference on Learning Representations, 2022. URL huttps://openreviiii hetffforrmiiii---tllblpyy
Aapo Hyvarinen, Hiroaki Sasaki, and Richard Turner. Nonlinear ica using auxiliary variables and generalized contrastive learning. In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 859 868. PMLR, 2019.
E. Jang, S. Gu, and B. Poole. Categorical reparameterizzttin with gumbel-softmax. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings, 2017. URL
R. Ke, O. Bilaniuk, A. Goyal, S. Bauer, H. Larochelle, B. Schölkopf, M. C. Mozer, C. Pal, and Y. Bengio. Learning neural cansal models from unknown interventions, 2020. arXivv1910.01075x2.
D. Kingma and M. Welling. Auto-cncoding variational bayes In 2nd International Conference on Learning Representations, ICLR, 2014. URL http:://arxiv http://arriv.orr/aab/11222 6114
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arxiv preprint
14 2014.
14

---

Published in Transactions on Machine Learning Research (11/2023)
Thomas Kipf, Elise van der Pol, and Max Welling. Contrastive learning of structured world models. arXiv preprint 1911. 12247, 2019.
S. Lachapelle, P. López, Y. Sharma, K. Everett, R. Le Priol, A. Lacoste, and S. Lacoste-Julien. Dis- entanglement via mechanism sparsity regularization: A new principle for nonlinear ica, 2021. URL
10098
Sebastien Tachapelle and Simon Lacoste-Julien. Partial disentanglement via mechanism sparsity. In UAT 2022 Workshop on Causal Representation Learning, 2022. URL http://p.///eeeerreiiee////////s 53fqtUWF
Minne Li, Mengyue Yang, Furui Liu, Xu Chen, Zhitang Chen, and Jun Wang. Causal world models by unsupervised deconfounding of physical dynamics. arXiv preprint arXiv:2012. 2.1220,,.,,,28,, 2020a.
Yunzhu Li, Antonio Torralba, Anima Anandkumar, Dieter Fox, and Animesh Garg. Causal discovery in physical systems from videos. Advances in Neural Informalion Processing Systems, 33:9180 00009992222111 2020b.
Phillip Lippe, Taco Cohen, and Eistratios Gavves. Efficient neural causal discovery without acyclicity constraints. In International Conference on Learning Representations, 2022a. URL https://openreviies
Phillip Lippe, Sara Magliacane, Sindy Löwe, Yuki M Asano, Taco Cohen, and Stratis Gavves. Citris: Causal identifiability from temporal intervened sequences. In International Conference on Machine Learning, pp. 13557 13603. PMLR, 2022b.
Phillip Lippe, Sara Magliacane, Sindy Löwe, Yuki M Asano, Taco Cohen, and Efstratios Gavves. Causal representation learning for instantancous and temporal effects in interactive systems. In The Eleventh International Conference on Learning Representations, 2023. URL http::///pppnrrreeeeeeefft/ttrne itZ6ggvMnzS
J. Pearl. Causality: Models, Reasoning, and Inference. Cambridge University Press, New York, NY, 2nd edition, 2009.
J. Peters, P. Bühlmann, and N. Meinshausen. Causal inference using invariant prediction: identification and confidence intervals. Journal of the Royal Statistical Society, Series B: Statistical Methodology (with discussion), 78(5):947-100 1012, 2016.
J. Peters, D. Janzing, and B. Schölkopf. Elements of Causal Inference Foundations and Learning Algorithmss. MIT Press, Cambridge, MA, USA, 2017.
Silviu Pitis, Elliot Creager, and Animesh Garg. Counterfactual data augmentation using locally factored dynamics. Advances in Neural Information Processing Systems, 33:3976 3990, 2020.
Roberta Raileanu, Max Goldstein, Arthur Szlam, and Facebook Rob Fergus. Fast adaptation to new envi- ronments via policy-dynamia value functions, In Proceedings of Machine Learning and Systems 2020, pp. 7078-7089. 2020.
Tong Sang, Hongyao Tang, Yi Ma, Jianye Hao, Yan Zheng, Zhaopeng Meng, Boyan Li, and Zhen Wang. Pandr: Fast adaptation to new environments from offiine experiences via decoupling policy and environ- ment representations, 2022.
B. Schölkopf, D. Janzing, J. Peters, E. Sgouritsa, K. Zhang, and J. Mooij. On causal and anticausal learning, J. Langford and J. Pincau (eds.), Proceedings of the 29th International Conference on Machine Learning (ICMT), pp. 1255-1262, New York, NY, USA, 2012. Omnipress. URL httt:://icmm.cc/2011222 htttm/m/.///1122r.r.////
6df
B. Schölkopf, F. Locatello, S. Bauer, N. R. Ke, N. Kalchbrenner, A. Goyal, and Y. Bengio. Toward causal representation learning. Proceedings of the JEEE Advances in Machine Learning and Deep Neural Net- works, (09(5):612-634 2021. doi: (0.1109/JPROC.2000 300133854. URL httpps:/iieexpploree wew...reeeeee..r.... haamp//staampjppppaannbeeereee
15

---

Published in Transactions on Machine Learning Research (11/2023)
Maximilian Seitzer, Bernhard Schölkopf, and Georg Martius. Causal inftuence detection for improving efficiency in reinforcement learning. Advances in Newral Information. Processing Systems, 34:22905-222 229188 2021.
R. Sekar, O. Rybkin, K. Daniilidis, P. Abbeel, D. Hafner, and D. Pathak. Planning to explore via self- supervised world models. In ICML, 2020.
P. Spirtes, C. Glymour, and R. Scheines. Causation, Prediction, and Search. MIT Press, Cambridge, MA, 2nd edition, 2000.
J. Tian and J. Pearl. Causal discovery from changes. In Proceedings of the 17th Annual Conference on Uncertainly in Arlificial Intelligence (UAI), pp. 512 522, 2001.
Ruben Villegas, Arkanath Pathak, Harini Kannan, Dumitru Erhan, Quoc V Le, and Honglak Lec. High fidelity video prediction with large stochastic recurrent neural networks. Advances in Neural Information Processing Systems, 32, 2019.
J. von Kügelgen, 1. Ustyuzhaninov, P. Gehler, M. Bethge, and Schölkopf. Towards causal generative scene models via competition of experts. arXiv preprint arXiv:2004. 12906, 2020.
M. Watter, J. Springenberg, J. Boedecker, and M. Riedmiller. Embed to control: A locally lin- ear latent dynamics model for control from raw images. In C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Garnett (eds.), Advances in Neural Information Processing Systems, vol- ume 28. Curran Associates, Inc., 2015. URL https://proceediing wttpeeeednnnnnnnnussss.om httprrssss.cpppppr//01////////// pdf
alafc58c6 1lar68866a440000339992200001-6-3
Wu, O. Parker Jones, M. Engelcke, and Posner. Apex: Unsupervised, object-centric scene segmentation and tracking for robot manipulation. In IEEE/RS: RRSJ International Conference on Intelligent Robots and Systems (IROS), 2021.
Weiran Yao, Yuewen Sun, Alex Ho, Changyin Sun, and Kun Zhang. Learning temporally causal latent processes from general temporal data. In International Conference on Learning Representations, 2022.
URL https: net worffroummmdiinii -dDDLN/--DDL
A. Zhang, C. Lyle, S. Sodhani, A. Filos, M. Kwiatkowska, J. Pineau, Y. Gal, and D. Precup. Invariant causal prediction for block mdps. In International Conference on Machine Learning, pp. 11214-11224. PMLR, 2020.
K. Zhang, M. Gong, and B. Schölkopf. Multi-source domain adaptation: A causal view. In Proceedings of the 29th AAAI Conference on Artificial Intelligence, pp. 3150-30 3157, 2015.
K. Zhang, B. Huang, J. Zhang, C. Glymour, and B. Schölkopf. Causal discovery from nonstation- ary hheehheeroggnouus data: Skeleton estimation and orientation determination. In Proceedings of the Twenty- Sixth International Joint Conference Artificial Intelligence (IJCAI 2017), pp. 1347 1353, 2017.
M. Zhang, S. Vikram, L. Smith, Abbocl, M. Johnson, and S. Levine. Solar: Deep structured representations for model-based reinforcement learning. In International Conference on Machine Learning, pp. 7444 7453. PMLR, 2019.
16

---

Published in Transactions on Machine Learning Research (11/2023)
A
Derivations
Using the approximate posterior q(z" Jo TT = II. a(>")o"), the variational log lower bound for the latent state-space model (Eq. 2) can be derived from importance weighting and Jenson's inequality:
T II
log p(o 0:T (.TT) = log
plo (paaa))"))((""") a' 1 Idz
(17)
t=
log /
II p(o't )l(a')aa))))))))) p(
(=21)0') (1=1)0)dd""
(18)
T log ([ p( ,at-1)
B ((2))00)
(19)
(22100')
Σ Bote [logr(o'|e" logp(a'|=')]
+ 1o(==-(/=--) (1a()100 [logp(="|>''),, (
(20)
[ott)( [logp(o'(2") (2*) + logp(a'|<"")
(B(=t-11o--1) [KL
(21)
Since the policy p(a"l=") is constant with respect to the model parameters, we omit this term and write the ELBO objective as
T
(L,
Σ Baols (Log(p(("))=)) 1=
- Beo(("-1 ( [""_ []]]
(22)
where 0 is the model parameter and o is the parameter for the approximate posterior. and ø are omitted henceforth to simplify notation. In VCD, the KL divergence term can be further decomposed by exploiting the structure of the transition model (Eq. S and the assumption that variables within each timestep are independent:
KLL <<=*)o"))lp (2912t (
(23)
f log
(00 (2)0)))))) ((21)0')
(1'""o')de"
(24)
/ log
Pi (2 @ [2t-1 a² -1--1)11 1 Ri (1) (+)
1
(>/ MS @ [x--11 (2'-_1,a'-11 a
((23)o')d2)
(25)
(<2510t)
Σ (1 Rig / log
Pi
(0) (= 2 [3--1 111
((25)o))d2
d l0))
Pi (2 o 1 a1-11)
+ R1: log
((=5)o))dt)
al
(26)
a(2/10')
(1
JLL JKL
((atlo')llz? M2 @ (a--)aa-111)
R KIL
( (K) (2f o 1 a -
(27)
The KL terms can be computed analytically since the conditional distributions in the last expression are univariate Gaussian distributions. In training time, the gradients through the expectation terms in the ELBO is estimated by drawing a sample from the posterior distribution using the reparameterisation trick Kingma & Welling 2014).
17

---

Published in Transactions on Machine Learning Research (11/2023)
B Implementation Detail
B.1 DCDI and Graph Learning
This section covers the formulation of DCDI (Brouillard et al. 2020 and the graph learning method. These are subsequently used in the learning of VCD.
Given samples from an observed data distribution Px and K intervened distributions p(k) DCDI optimises a probabilistic belief over causal graphs G and intervention targets I. Specifically, these are encoded as random binary matrices, Mg and Rz, where My = implies that the edge (f.j) is the causal graph, and R = 1 implies that the variable is in the intervention targets in environment k. Each entry in M follows an independent Bernoulli distribution, parameterised by matrix a where P(M 1j = 1) = a(ss)). R is similarly parameterised by B. Under this parameterisation, causal discovery can be formulated maximising the expected data log likelihood with sparsity regularisation,
L(0, a. Ea,
Ac|s1 Σ loglps (8) (xf,..., - -
(28)
where p(() is the data likelihood under causal graph G and intervention targets I in the kth environment, as factorised in (5. The conditional distributions are parameterised as feedforward neural networks with parameter 8; λG are hyperparameters to control sparsity. In the original DCDI framework, this is also subject to an acyclicity constraint. However, this is not neccessary in the context of our work as we assume there are no instantaneous causal effects (i.e., within a timestep).
The training objective for VCD can be viewed as modified version of the DCDI objective, where the likelihood term is replaced with the ELBO (Eq. 22
K T Σ
LCD (0,0, = Ea,
(LBO(0,0:9,0) - Aolgi - X [
(29)
t=
Note that the expected number of edges in & and I given and ,3 is simply the sum of the probability of each entry being one. Therefore, the training objective can be computed as:
LCD (0.00, = En,
ΣΣ k=
ELBO((,6:G,I) A Eo(ay)) - -xr (BKs). oB(as)
(30)
The gradients through the outer expectation can be estimated using the Gumbel-Softmax trick (Jang et al. 2017). To implement this, the ELBO term is evaluated with a sample of the causal graph using the following expression for each entry,
M (((a(a Lúj > 0.5) + o(aiy + Lij) - stop gradient(oraa + Liz)),
(31)
where II(-) is the indicator function, Lij a sample from the logistic distribution, and stop gradient is a function that does not change the value of the argument but sets the gradient zero. Samples for the intervention targets similarly acquired. Note that the sample used throughout cach trajectory, i.c. the same sample graph and intervention largets are used for all of T timesteps.
B.2 Model Architecture
Mixed-state In the mixed-state experiment, all conditional distributions (including encoders, decoders and transition models) are parameterised by foodforward MLPs with two hidden layers of hidden units each. The recurrent modules are implemented as GRUs (Cho al. 2014) with 64 hidden units. Distributions in the latent space 16-dimensional diagonal Gaussian distributions with predicted mean and log variance.
18

---

Published in Transactions on Machine Learning Research (11/2023)
Table 3: List of interventions
ID
Intervention
Intervention targets
1
Remove spring between 1 and 2 Remove spring between 2 and 3
31 M1, 22, M/2
2
22, V2, 23: M8
3
Increase mass 1
11 U1, V4
4
Increase mass 2
5
Increase mass 3
03, M3, 244,
Decrease mass 1
$1, M1. 24. V/1
Decrcaso mass 2
22. M2
8
Decrease mass 3
33. M3. 24: 1/2
9
Tncrease spring constant between and 2 31, 22, 1/2
10
Increase spring constant between 2 and 3 Constrain movement of to vertical only
22, 1/2 23, M/
11
21
12
Constrain movement. to horizontal only
V/
13
Constrain movement of to vertical only
J22
14
Constrain movement to horizontal only
1/2
15
Constrain movement of to vertical only
Z3
16
Constrain movement. of to horizontal only
y3
17
Constrain movement of 4 to vertical only
T4
18
Constrain movement of 4 to horizontal only
y/
Image In the image experiment, the encoders and decoders are parameterised as convolutional and deconvolutional networks from (Ha & Schmidhuber 2018)). In the RSSM models, the transition models are parameterised feedforward MIPs with two hidden layers of 300 hidden units. The recurrent module is a GRU with 300 hidden units. In VCD, to compensate for the fact each dimension in the latent space has separate model, the number of hidden units in the GRU and MLP reduced to 32 to avoid wver-paraameterisation. We found that initialising the encoders and decoders by pretraining them as variational autoencoder helped with training stability for both RSSM and VCD.
both experiments, the training objective maximised using the ADAM optimiser (Kingma & Ba 2014) with learning rate 10-3 3 for mixed-state, and 10-1 for images. In both environments, we clip the log variance to 3, with a batch size of two trajectories from each of six environments with T' 50. In VCD, the hyperparameters Àg, dr are both set to 0.01. All models trained on single Nvidia Tesla V100 GPU.
C Experiment Detail
Interventions The ground truth states of the multi-body dynamics environment the æ and y coordinates of each particle. The full list of possible interventions the environment provided in Table 3 Note that the forces between particle 1, 3 and 4 are proportional to their masses. Hence intervening the mass of particle and 3 also affect the dynamics of particle 4. In the experiments, all models are trained in the undisturbed environment and intervened environments 1, 5, 11, 14, 17.
In the mixed-state experiment, the observation function is a mixing matrix where each entry drawn from unit Gaussian distribution. the image experiment, the observation given by rendering the system 128 128 image. In both experiments, the models are trained on training of 2000 trajectories from each of the six environments and evaluated on validation set of 400 unseen trajectories.
19

---

Published in Transactions on Machine Learning Research (11/2023)
Ground Truth
VCD
RSSM
time
Figure 7: Sample rollouts (without observations) from an intervened environment. VCD successfully captures the constraint on the yellow particle (vertical movement only) over a long time horizon. The frames are sampled 5 timesteps apart from each other.
ground truth state
ground state
3-
10
14
12
RSSM latent dimension
VCD latent dimension
Figure 8: The average magnitude of the derivative of the ground-truth states x; w.r.t. the learnt latent states Zj. VCD is able learn lalent representation where each dimension captures information about state only, ie. only one state variable changes when perturb the latent representation along one dimension. On the other hand, the RSSM latent space encodes state information in more entangled manner, i.c. multiple ground-truti states are affected when the latent representation is perturbed in one dimension.
D Qualitative Exploration for Learnt Causal Structure
D.1 Rollout prediction
Section 5 shows the rollout accuracy of the baselines and VCD. Here we demonstrate qualitatively that VCD is able to capture auvironment-specif behaviours that RSSM cannot learn due to maximal parameter sharing. Figure shoos sample rollouts the image space from RSSM and VCD intervened environment 11, i.c. the yellow particle is constrained horizontally and only moves vertically. The yellow particle in the VCD rollout stays along vertical line whereas RSSM fails to capture this environment-specffcc constraint.
D.2 Representation quality
Mixed-state Since the observation mixing function is linear and invertable in the mixod-state environment, directly access the level of disentanglement of the latent space with respect to the ground-truth state
20

---

Published in Transactions on Machine Learning Research (11/2023)
ground truth state
ground-rruth state
Image RSSM atent state
Image
Figure 9: The R² a linear regression between the ground-truth states and the learnt latent variables. This shows a similar pattern as the mixed-state experiments where the VCD representation general only capture information about one ground-truth state.
Mixed-state
Image
3-
causal parents
5
causal parents
8-
a
11-
12
12
16
17
1011 14
10
14
causal variables
causal variables
Figure 10: The learnt causal graphs for the mixed-state and image experiments, obtained by binarising the learnt edge probabilities such that blue square at (i,j) implies a(aij) > 0.5, Lee. is causal parent of
of the environment, i.e. the and y coordinate of the particles. Fig. S shows the average magnitude of the entries of the Jacobian matrix between the ground-truth state and the learnt latent state. Each entry measures the changes in each ground-truth state variable when the latent representation is perturbed along each dimension. By using the discovery of causal transition models an inductive bias, VCD is able to learn disentangled representation where each ground-truth state is captured by only one latent dimension. In contrast, RSSM learns a representation that is not sparse.
Image The quality of the learnt representation discussed Section 5 Here we show reconstruction samples along all dimensions of the latent space in the RSSM and VCD representation. Note that since both encoders initialised from the same pretrained VAE, the difference in the latent space arise because of the causal discovery objective in VCD. Fig. 13 shows the changes in the reconstruction when the RSSM reprosentation is perturbed in cach dimension of the latent space. Fig. 12 shows the same for VCD. As discussed in the main lext, VCD learns axis-aligned representation that aflords a sparse causal graph. For example, dimension 1 and 3 captures the y and x coordinates of the green particle respectively. contrast, RSSM learns a representation that not axis-aligned. Fig. plots the R² linear regression scores between the learnt latent variables and the ground-truth variables, showing that the VCD representation in general disentangle individual states, as indicated by the more salient squares.
21

---

Published in Transactions on Machine Learning Research (11/2023)
Mixed-state
Image
Environment
environment
intervention targets
intervention targets
Figure 11: The learnt intervention targets for the mixed-state and image experiments, obtained by binarising the learnt probabilities such that a blue square at (k: implies ((3r)) > 0.5, ie. is an intervention target in environment k.
D.3 Learnt causal graphs and intervention targets
In this subsection, we explore the quality of the learnt causal graph and intervention targets. Fig. 10 shows the learnt causal graphs for the mixed-state and image experiments. Fig. [1 shows the learnt intervention targets for each environment. The sparsity of the learnt graph and targets is summarised in Table VCD has identified 42 and 73 causal edges in the mixed-state and image experiments respectively, out of 288 possible edges. Viewed in conjunction with the prediction performance results, this shows that VCD is able to learn world model that is sparsely connected and affords modular parameter sharing wilhout compromising on prediction accuracy.
Mixed-state In the mixed-state experiments, each ground-truth state can be mapped to a latent. dimension using the average magnitude Jacobian matrix (Fig, We use this mapping to compare the learnt causal graph with the ground truth causal structure of the environment and report the number of correctly identified edges. Table summarises the quality of the learnt graph. VCD is able identify majority of the correct causal dependencies and all of the intervention targets. Upon further inspection of the learnt causal graph, find that all of the seven missed edges correspond to the 1/ terms that scale the electrostatic- forces. We hypothesise that the model cannot capture these dependencies as they are not as significant as the other forces.
Image In the image space, a similar comparison is made using the R² scores. Each ground-truth state is mapped to latent dimension by choosing the dimension with the highest R² score. This leaves unused latent states, which are ignored in the analysis provided in Table We also qualitatively explore the learnt causal relationships. Focusing dimension for example, which encodes the coordinate of the green particle, the ground-truth causal parents of this variable is the x coordinates of the yellow and the blue particles, In the learnt causal graph, on column three, the learnt causal parents are dimensions 3, 6, 8, 14 and 15. The visualisation of the latent space (Fig. 12] suggests that dimensions and 8 captures the & coor- dinates of the yellow and blue particles respectively, meaning that VCD has learnt the correct causal parents.
A similar analysis can be carried out on the intervention targets. Focusing on intervened environ- ment 0, for example, the learnt intervention targets are dimension 1, 2, 3, 7, 8, 11, 13, 15. The ground-truth intervention ID is 1, iee the intervention targets 2, coordinates of the yellow and green particles (see. table 3) By inspecting the VCD latent space (Fig. [2), dimensions and 3 encodes the position of the green particle and dimensions 7, 8 and 11 encodes the position of the yellow particle. This shows that VCD is able to identify the changes in the environments.
In summary, while VCD identifies some false positive edges, it is able capture the causal parents and the intervention targets in each environment.
22

---

Published in Transactions on Machine Learning Research (11/2023)
dm
dm1
am7
4m
din
dm
SimG
dn
dms
4m 10
4m 11
am01 12
di
4m 24
dim 25
WCB
Figure 12: Visualisation of the learnt latent space in VCD.
23

---

Published in Transactions on Machine Learning Research (11/2023)
4mo
dm
4m
dm
dm
em.
4n
di
dh
dm
dm 10
dn.33
dm 12
4mi1
4lm 15
Figure 13: Visualisation of the learnt latent space RSSM.
24

---

Published in Transactions on Machine Learning Research (11/2023)
Mixed-Stan States
RZGM NUNRSSN VCD
i do 65 IE
LE 1M
A trajectories
DI
trajectory
1
Images
58
13c
10 trajectories
40 trajectories
100 trajectories
640 trajectories
2560 trajectories
Figure 14: Adaptation results for intervention 4.
Mixed States
VCD trajectory
trajectories
16 trajectories
64 trajectories
256 trajectories
r'
10
Images
02
40
40
30
12
19 trajectories
&
trajectories
trajectories
trajectories
trajectories
T
Figure 15: Adaptation results for intervention 9.
25

---

Published in Transactions on Machine Learning Research (11/2023)
Mixed- States
RSSH NUERSSN VCD
1E DF trajectory
10 4 trajectories
21
16 trajectories
64 trajectories
256 trajectories
Images
250
120
100
180
180
10 trajectories
40 trajectories
160 trajectories
640 trajectories
2560 trajectories
Figure 16: Adaptation results for intervention 13.
Hit @ 5 scores
100
RSSM
Multi-RSSM
VCD
@ 5 (%)
10
20
30
40
50
rollout timesteps
Figure 17: The H@5 score (higher is better) over rollout timesteps, i.e. unshaded regions in the main plots. This shows consistent result the main text, where VCD outperforms the baselines.
E Further Experiments
This section provides extra adaptation experiment results similar Fig. 0 the main text, where the
models adapt to intervention number 12 (see table 3). We present experiment results for adaptation to different types of interventions (intervention numbers 9, 13). Fig 14
15
and 16 show the adaptation plots.
The results exhibit consistent pattern where VCD significantly outperforms the baselines in the low data regime by identifying sparse mechanism changes.
F Non-pixel-based error
In the image experiments, while pixel error is indicative of model performance in the short term, it provides limited clarity in long term predictions as two non-overlapping balls leads the same pixel error regardless of their distance. In this section, we present an alternative latent space distance- based evaluation metric, Hit at (H@5). This metric is defined by the proportion of episodes where the predieted latent representation
26

---

Published in Transactions on Machine Learning Research (11/2023)
Mulld RSSM
VCD
22
H85 (%6
41
20-
21-
25
10 trajectory
40 trajectories
160 trajectories
640 trajectories
2560 trajectories
Figure 18: The adaptation plots with the H scores.
Learnt causal graph with no interventions
Rollout Error with 40 adaptation episodes
800
RSSM
700
No Intervention
VCD
600
500
L2 error
400
10
300
200
14
100
16
0
20
40
60
80
100
time
10
12
14
(a) Rollout curves compared to the RSSM baseline and VCD with full interventional training.
(b) Learnt causal graph.
Figure 19: Training VCD without interventions results a model that shows some causal structure. Com-
pared to the VCD with full interventional training, the adaptation performance of the model is slightly degraded, but still significantly better than the RSSM baseline.
is in the top-5 nearest neighbour set of the encoded ground-truth image. A H@5 score of 1 means all predicted latent states lie within close distance to the ground-trth representation. Fig 17 and I8 shows the prediction errors and the adaptation plots respectively. These are consistent with the results shown in the main text, where VCD outperforms the baselines in prediction accuracy as well as adaptation speed.
G Importance of Interventional Training
One of our hypothesis is that explicitly training the model to exhibit minimal changed across multiple interventional environments aids causal discovery and hence improves adaptation efficiency, Fig. [9 presents results where VCD is trained without interventions. can be seen that the model is still able to learn causal structures, albeit. less sparse compared to the full model. The adaptation performance is slightly worse than the full model but still much better than the baseline. This corroborates our claim that interventional training improve structure learning and adaptation efficiency.
27


Reviews:
[{"review_id": "vfSTu7hinf", "reviewer": "Reviewer_Kqkr", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "* Strengths:\n1. The paper was written very clearly and with helpful diagrams. Apart from minor notational suggestions I think I can follow the ideas pretty easily. \n2. The idea of combining a dynamical model, training on causally intervened data, and using the variational approach to serve to improve data efficiency seems nontrivial and interesting. Each of the components are well justified by either theory of hypothesis, which I think helps convincing the readers. \n3. All assumptiosn are clearly stated, including the crucial ones at the beginning of Section 4. \n\n* Weaknesses\nThe main weakensses of the paper are not technical but rather on writing, which can be easily addressed in a revision, I think. In general, the writing of the paper could be toned downed a little given the quite restrictive assumptions made in the begining of Section 4. The clarity of the paper naturally leads to a few questions out of curiosity, which hopefully could be answered by the authors. The experiments could be enriched by an ablation study I proposal below, but other than those I think the paper is of very high quality.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper combines techniques of causal discovery and variational inference to enhance data efficiency during learning / parameter estimation (adaptation in the paper), leveraging a dynamical model. The key contributions in the generative framework seem to be:\n1. It extends the work of (Peters et al., 2017) by using a dynamical model and intervening on random subsets of the latent variables, eqn (7)\n2. It implements the sparse changes in the intervened environments motivated by (Schölkopf et al., 2021), eqn (8) \nThere are also other design specifics of the dynamical model motivated by simplicity assumptions. \nThe authors then propose to optimise the ELBO of this model using a simple variational posterior. Crucially, the training data consists of intervened distributions that aid the model to learn modularised causal components. The validity of this approach is tested in simple datasets.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Critical:\n1. The variatioanl posterior in (2) and (11) depends only on the current observation, which can be insufficient for certain dynamical systesm, because it is a drastic simplification of the full posterior that depends on all observations. I understand that this came from a previous work, but some discussions on this should is necessary. \n2. Please put disclaimers in the introduction that various assumptions are made in a later section.\n3. I find the overloaded of terminologies very confusing: for example, there are in total 200 mentions of the word \"causal\", followed by a word in the set {discovey, dynamics,  transition model, graph, mechanism, variable, structure,...}, more than article “a\" (188 times) Do the word \"causal\" here mean the same thing? Please simplify the terminologies. These make reading very difficult, especially the introduction and the second paragraph of Section 4. \n\nMinor:\nLast suggestion is to run an ablation study to support the benefit of interventioanl training. For example, could train the models with fewer intervened environments, but keep the total datasize controlled.", "broader_impact_concerns": "Don't think this is mentioned. I suggest the author comment on possible broader impact.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "YSleayzf3k", "reviewer": "Reviewer_aecq", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n\n* The submission is clear and easy to follow.\n* The considered problem is important in practice.\n\n### Weaknesses\n\n* Although the motivation is clear, I'm not convinced that such objective can really learn the causal graph. Specifically, as the causal graph and the intervention target are learnt from the data, I don't believe that the learnt results are necessarily the targets we want.\n* The experimental results are not so satisfactory. Specifically, I don't see strong evidence that the model really learns the causal graph and intervention target as we want.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This submission proposed one new objective for training world models that can potentially enforce the causal discovery as well as the domain generalization. Experiments partially support the claim.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "No", "audience_explanation": "", "requested_changes": "Can the authors elaborate more on the identifiability issue as I mentioned in the weakness part? Specifically, can the authors make some discussions on when and how we can identify the causal graph and the intervention target, even in some simple scenarios? I hope to see more theoretical and experimental results on this part.", "broader_impact_concerns": "Not applicable at this time.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "oIjpAioQh2", "reviewer": "Reviewer_a2UN", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n\n1. The motivation of this work is clear, which is a problem worth studying under reinforcement learning.\n\n2. Various experiments as well as relevant analyses are performed in this work, which illustrates the efficacy of the method.\n\n3. This paper is well written.\n\nWeaknesses\n\n1. Some experimental details of this work are not clear enough.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes Variational Causal Dynamics(VCD), which learns a causal world model with a causal structure. The proposed VCD is able to identify causally meaningful representations of the observations and discover sparse relationships in the dynamics of the system. Thus, through prior knowledge obtained from different environments, VCD can quickly adapt to a new environment using less data. The experiments demonstrate that the proposed method is superior to the current SOTA.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "There are some questions about this paper:\n\n1. The problems and solutions proposed in this work are novel and interesting. Whether the author can supplement related toy experiments to illustrate the problem intuitively.\n\n2. In the introduction, the author mentions \"jointly training a representation and a transition model\". Here, what does representation specifically refer to? I recommend that the authors state it in the Introduction section.\n\n3. These works[1-2] are related to this work, and I recommend authors to cite them appropriately.\n\n[1] Sang, T.; Tang, H.; Ma, Y.; Hao, J.; Zheng, Y.; Meng, Z.; Li, B.; and Wang, Z. 2022. PAnDR: Fast Adaptation to New Environments from Offline Experiences via Decoupling Policy and Environment Representations. arXiv preprint arXiv:2204.02877.\n\n[2] Raileanu, R.; Goldstein, M.; Szlam, A.; and Fergus, R. 2020. Fast Adaptation to New Environments via Policy-Dynamics Value Functions. In ICML, volume 119, 7920–7931. PMLR.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes VCD, a model that aims to discover causal dynamics to enable fast adaptation to new environments via modular knowledge transfer. While previous approaches to causal discovery necessitate access to abstract causal variables, the proposed approach operates directly on high-dimensional observations. This is achieved via a training objective that combines causal discovery with variational inference.

Reviewers mentioned that the proposed method is clearly written (Reviewer Kqkr, Reviewer aecq, Reviewer a2UN), the motivation is clear (Reviewer a2UN) the idea is interesting (Reviewer Kqkr), the problem studied is important (Reviewer aecq). The authors have provided sufficient empirical evidence to support their claims (Reviewer a2UN). While Reviewer aecq raised concerns about unsatisfactory empirical evidence, the authors have since updated their paper during the rebuttal to address those concerns, and have clarified the claims made and relationship to theoretical identifiability results. Therefore, I think the paper provides convincing evidence for the claims made, and it is interesting and relevant for TMLR. 

As per comments made by reviewers, I encourage the authors when preparing the camera ready, to further emphasize (e.g. in the Abstract of Introduction) the fact that VCD may not always discover the exact causal graph and is not accompanied with theoretical gurarantees about identifiability; the claims about VCD successfully identifying causal variables are made via empirical investigations.
Decision: Accept (poster)