Variational Causal Dynamics: Discovering Modular World Models from Interventions
Authors: Anson Lei
Bernhard Schölkopf
Ingmar Posner
Venue: Submitted to ICLR 2023 (2023)

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.  

Under review as a conference paper at ICLR 2023
VARIATIONAL CAUSAL DYNAMICS: DISCOVERING MODULAR WORLD MODELS FROM INTERVENTIONS
Anonymous authors Paper under double-blind review
ABSTRACT
Latent world models allow agents 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 ex- ploits the invariance of causal mechanisms across environments to achieve fast and modular adaptation. By causally factorising transition model, VCD is able to identify reusable components across different environments. This is achieved by combining causal discovery and variational inference learn a latent repre- sentation and transition model jointly in an unsupervised manner. Specifically, we optimise the evidence lower bound jointly sition model structured as a causal graphical model. In evaluations on simulated environments with state and image observations, successfully identify causal variables, and to discover consistent causal structures across different environments. Moreover, given small number of observations a previously unseen, intervened environment, VCD is able to identify the sparse changes the dynamies and to adapt efficiently. In doing so, VCD significantly extends the capabilities of the current state-of-the-ar in latent world models while also comparing favourably in terms of prediction accuracy.
representation model and tran-
show that VCD able
INTRODUCTION
The ability to adapt flexibly and efficiently to novel environments one of the most distinctive and compelling fcatures 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 gen- eralisable, structural relationships within the environment (Behrens et al. 2018). It is conjectured that this latter aspect 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 and Schmidhuber] 2018) aim to capture the dynamics of an environment by distilling past experience into parametric predictive model. Advances in latent variable models have enabled the learning of world models in a compact latent space (Ha and Schmidhuber 2018 Watter et al.] 2015 Hafner let al. 2019b] Buussing et al., 2018 /Zhang et al.) 2019) from high-dimensional observations such as images. Whilst these models have enabled agents to act in complex environments via planning (e.g.
Hafner al. 2019b Sekar al. 2020) or learning parametric policies (c.g. Hafner al., 2019a Ha
and Schmidhuber, 2018), structurally adapting to changes in the environment remains a 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 tantalising prospect to discover such structure from observations.
Causality plays 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 causal generative process. This implies that many causal mechanisms remain invariant across domains (Schölkopf et al. 2012; Peters et al.) 2016 Zhang ett al., 2015). In this light, learning causal model of the environment enables agents reason about distribution shifts and to exploit the invariance of learnt causal mechanisms across different environments. Hence, we posit that world models with causal structure can facilitate modular

---

Under review as a conference paper at ICLR 2023
Training
Rollout
VCD
+
Loaruaie
Gragit
Figure I: Left: The general architecture of VCD. The dynamics model is trained using observation sequences from multiple environments by minimising the KI. divergence between the predicted state distribution and the encoded posterior state 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 treated as causal variable. Predictions are made using only the causal parents of each variable, according to a learnt causal graph.
transfer of knowledge. To date, however, methods for causal discovery (Spirtes et al. 2000 Pearl 2009 Peters et al., 2017; Brouillard et al., 2020: Ke et al. 2020) require access to abstract causal variables to learn causal models from data. These are not typically available in the context of world model learning, where we wish to operate directly on high-dimensional observations.
In order to benefit from the structure of causal models and the ability to represent high-dimensional observations, we propose Variational Causal Dynamics (VCD), which combines causal discovery with variational inference. Specifically, tion model using variational inference and sparsity regularisation from causal discovery. By jointly training representation and transition model, VCD learns a causally factorised world model that can modularly adapt to different environments. The key intuition behind our approach is that, since sparse causal structures can only be discovered on abstract causal variables, training the rep- resentation and the causal discovery module in an end-to-end manner acts as an inductive bias that encourages causally meaningful representations. By leveraging the learnt causal structure, VCD is able to identify the sparse mechanism changes in the environment and re-learn only the intervened mechanisms. This enables fast and modular adaptation to changes in dynamics.
train latent state-space model with structural transi-
2 RELATED WORK
Predictive models of the environment can be used to derive exploration- (Sekar et al. 2020) or reward-driven (Ha and Schmidhuber [018 Hafner et al./ [2019a]b] b) behaviours. In this paper, we fo- cus the learning of latent dynamics models. World models (Ha and Schmidhuber] 2018) train a representation encoder and a RNN-based transition model a two-stage process. Other ap- proaches (Hafner et al. 2019b: Zhang et al. [019] Watter et al. (2155) learn generative model by jointly training the representation and the transition via variational inference. PlaNet (Hafner et al. 2019b) parameterises the transition model with RNNs. E2C (Watter et al. 2015 Banijamali al. 2018) and SOLAR (Zhang et al.) 2019) use locally-lineat transition models, arguing that including constraints the dynamics model yields structured latent spaces that are suitable for control. Other approaches such as (Goyal et al. 2021b [BeckerEhhmck et al. (019) consider latent transition as discrete mechanisms. In similar vein, the use of latent prediction models have also been explored in the context of video prediction (Villegas et al.) 2022]. Our proposed approach shares the general principle that latent representations can be shaped by structured transition mechanisms (Ahuja et al.
2019 Denton and Fergus 2018 Assouel et al.)
2021). However, to the best of our knowledge, VCD is the first approach that implements causal transition model with high-dimensional inputs.
Causal discovery methods enable the learning of causal structure from data. Approaches can be cate- (e.g. (Spirtes et al. 2000) and score-based (e.g. (Hauser and Bühlmann) 2.12). The reader is referred to (Peters et al. 2017) for detailed review of causal discovery meth- ods. 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 to operate on lovel data, and causality (Schölkopf et al., 202I). Recent advances in this area include theoretical works exploring the conditions under which disentanglement of representation is possible (Yao et al. 2022
gorised as constraint-
2

---

Under review as a conference paper at ICLR 2023
Hyvarinen et al. 2019 Lachapelle et al. 2021 Lippe al. 2022a by Brehmer et al. 2022). Our
current work situates within this broader context of causal representation learning, and aims to iden- tify causal representations via the discovery of causal transition dynamics. (Lachapelle et al. 2021) discusses the theoretical identifiability of a latent state-space model similar to ours. Whilst we do not make identifiability claims, iet al., 2021). To the best of our knowledge, VCD is the first model that focuses on the adaptation capabilities of causal models and empirically shows their applicability to image observations.
work builds on the inductive biases developed in (Lachapelle
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 keypoints (Li al, 2020b) or object slots (Li et al. (2020a), and attribute changes in the dynamics to unobserved con- founders, 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 (Huang et al. 2022) and tem- porally intervened sequences (Lippe et al. 2022c). While the latter shares the general approach of learning sparse graphs in the latent space, considers 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.
Another branch of causality-inspired work leverages the invariance of causal mechanisms by learn- ing invariant predictors across environments (Tian and Pearl [00T]] Schölkopf et al.] [012 Peters iet al., 2016 Zhang et al.) 20177; Arjovsky et al., 2019). This invariance has been studied in the context of state abstractions MDPs (Zhang et al. 2020), and imitation learning from different environments (Bica et al. [22T)) In contrast, our approach models the full generative process of the data across different environments rather than learning discriminative predictors.
BACKGROUND
This section describes the prerequisite definitions and formulations for our proposed method. We highlight the strengths and weaknesses of latent state-space models, causal models and causal dis- covery methods, and motivate our approach that builds upon these learn causal world models.
Latent state-space models In complex enviromment with high- sigh-dimensional observations, such as images, learning a 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 (Buesing et al., 2018). Given a dataset {(o-T, (.TT) > with observations of and actions a' timestep f, a generative model of the observations can be defined using latent states 20: as
/ po (2))a/22))00())) (2=|22=111_1__ add"")dd22
pOo a (T)
where pe (o and po (2*112 aa" are the observation model and the transition model respec- tively. For simplicity, we do not learn the policy term p(a" ( (z*) and omit it throughout this paper as it constant with respect to the parameter O. The variational evidence lower bound can be written as
ELBO
Σ (=)0)))) (ogg (at (Ko
a 1]]}
(2)
where (2 learnable approximate posterior. See Appendix A for the derivation. RSSM (Hffne et al. 2019b) parameterise the transition model as a feed-forward recurrent neural network.
Despite their success, these models cannot reason about changes in the environment. Specifically, they are unable to structurally utilise prior knowledge 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 in modular and efficient manner.
Causal graphical models A causal graphical model (CGM) (Peters et al. 2017) is defined as set of random variables (X1,-., Xa), their joint distribution Px. and directed acyclic graph
3

---

Under review as a conference paper at ICLR 2023
(DAG), 9 (X, E), where each edge (i,j) E implies that X, is direct cause of X,. The joint distribution admits a causal factorisation such that
d
p(1,,-.. Fa)
II p PA,), 1=0
(3)
where PA, is the set of parents the variable X; in the graph. Each of the conditional distributions can be considered as an independent causal mechanism.
In contrast standard graphical models, CGMs support the notion of interventions, i.e., local changes in the causal distribution. Formally, an intervention on the variable X; is modelled as replacing the conditional distribution (PA.) while leaving the other mechanisms unchanged.
Given the set of intervention targets / C X, the interventional distribution can be written as
p' Ta) II p(z; (PA) II p'(x;(PA)) eg'
sE/
where (-I--) is the new conditional distribution corresponding to the intervention. The SMS hypoth- esis (Schölkopf et al. 202I) posits that naturally occurring distribution shifts tend to correspond to sparse changes in a causal model when factorized as 3 i.e., changes of few mechanisms only. Causal mechanisms thus tend to be invariant across environments (Schōlkopf et al.) [012]] Peters al.] 2016; Zhang et al. 2017). In this light. we argue that a causal world model can structurally inductive prior. In order to learn causal model the context of world models, we draw inspiration from recent advances in causal discovery which aim to learn causal structures from data.
leverage the invariance within distribution shifts
Differentiable causal discovery We focus on methods that formulate causal discovery as a con- tinuous optimisation problem (Brouillard et al. 22200) Ke et al.) [020] [Beniio et al. (019) as these can be naturally incorporated into the variational inference framework. Since the causal variables are learnt in our model, the causal discovery module is required to learn causal graphs from unknown intervention targets. In this work, we follow the formulation in Differentiable Causal Discovery with Interventional data (Brouillard et al. 2020) (DCDI), which optimises continuously parameterised probabilistic belief over graph structures and intervention targets. See Appendix for further detail.
In the context of learning world models from high-dimensional observations, the drawback of this approach, and of causal discovery methods in general, is 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, we present a 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-observation sequences from undisturbed environment, (o%) / ), and K intervened en- vironments, (k) ((k) ) Throughout this paper, we assume that 1) changes between the environments are due to 'soft' interventions where the structure of the causal graphs remains the same but individual parameters of the mechanisms change, 2) there are no instantancous causal de- pendencies in each timester!] and 3) the observation function does not change across environments.
The general approach of VCD follows the latent state-space model framework (Fig 2a where latent representation of the observations is jointly learnt with transition model by maximis- ing the ELBO. In contrast to existing approaches that parameterise the transition probability p(='|=--1,a'-) a1-1) as feedforward neural network (Fig 2b). VCD learns a causal transition model by utilising inductive biases from causal discovery. Importantly, our approach is grounded in the hypothesis that causal structures can only be discovered on semantically meaningful causal repre- sentations of the system. We therefore argue that training a representation and transition model jointly to optimise causal discovery objective can lead to latent representation that affords causal transition models and is therefore semantically meaningful Taking the view that changes in the dy-
This assumption can readily relaxed as DCDI can handle instantaneous edges.
Whilst we do not make theoretical identifiability claims for our model, the identifiability gaurantees of similar state-space model discussed in (Lachapelle al. 202T)
4

---

Under review as a conference paper at ICLR 2023
Tramition Midel
gisl
2++
K
K
(a) Latent state space model
(b) Typical transition model
(c) VCD transition model
Figure 2: (a) The structure latent state space models in one timestep. (b) The structure of latent transition models where the probability distribution is modelled fully connected neural network. (c) The causal transition model in VCD. Each individual variable has a separate conditional distri- bution p:- Each mechanism is conditioned on 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 the black lines, which are shared across environments.
namics can be attributed to sparse causal interventions, we posit that a causally factorised transition model facilitates modular adaptation new environments by explicitly leveraging the invariance of causal mechanisms. We propose algorithm for such 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 CGM (Eq. 4p. In the following paragraphs, motivate and describe the three main architectural features VCD that enable causal discovery: independent mechanisms, sparse causal dependencies and sparse interventions. Access the causal graph G 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 4.3
Independent mechanisms In contrast to a fully-connected transition model, the transition proba- bility is factorised into individual conditional distribution as
p(2'121-1 ( II ad'-))
(5)
where is the dimension of the latent space, to be set as a hyperparameter. z; denotes the ith dimension of the latent state, and each conditional distribution p; is one-dimensional normal dis- tribution with mean and variance given by separate neural networks. This separation of parameters motivated by the Independent Causal Mechanism principle, which states that the causal generative process of system is composed of modules that do not inform each other (Schölkopf et al.) 2021) This explicit modularity of the model structure enables the notion interventions, where individual conditional distributions are locally changed without affecting the other mechanisms.
Sparse causal dependencies Following the structure of CGM, we 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 MG is 1 if and only if is causal parent of z/. This is consistent with the intuition that, in physical sys-
[z'-1, :---1aa-11, a tems, 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
p1 aa = II p o 0 aat-1)),
(6)
5

---

Under review as a conference paper at ICLR 2023
where @ denotes the element- product, [+ denotes the concatenation of vectors, and M, 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 we assume that changes in distributions across the K intervened environments are due to sparse interventions in the ground truth causal generative process. In order to incorporate sparse interventions in VCD, given the set of learnt intervention targets Za for each environment, we define the binary intervention mask R where the entry R is if and only if the variable 2; is in the set of intervention targets in environment k. For each variable 2;, R acts switch between reusing shared observational model and an environment--pecific interventional model. The full interventional causal model of the transition probability in the environment & can be written as
D² p"(="12"-1 =
II p (=:] o (2""1,a''" a'' ( (2!1M, @
(7)
where p() is the observational mechanism that is shared across all environments and p()) is the intervened distribution specific to environment k. Intuitively, this model ensures that all environ- ments reuse the same conditional distributions yd unless it is deemed that particular mechanism has been intervened on. This modular adaptation between environments facilitates structural transfer of knowledge between different environments. Note that the structure that of causal graphical model. The overall architecture of the VCD transition model is summarised in Fig. 2c
4.2 RECURRENT MODULE
Similar to RSSM (Hafner et al. 2019b), we augment the model with a deterministic recurrent path to enable long-term predictions. To ensure that each conditional distribution only has access to the causal parents, in the same way that each conditional distribution is modelled by a separate network, each conditional distribution keeps separate recurrent unit and a corresponding hidden activation:
2 (z)
h f((ht .M. o I=!"",a, ");
where is recurrent module specific the variable 2,, instantiated GRU (Cho et al. 2014
4.3 TRAINING
The task of model training is to determine the model parameters 0, the approximate posterior pa- rameters ø. as well as the causal graph and the intervention targets I. These can be jointly trained in a way similar to DCDI. We parameterise the belief over the causal adjacency matrix 1/9 as random binary matrix. Each entry M follows Bernoulli distribution with success probability (aiy). where oi is a scalar parameter and a()) is the sigmoid function. Similarly, a random binary matrix RD is parameterised using the scalar variable 3ri 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,
L(0, a)
Σ kE(O,I.....KK]
Eg.z ELBOO( 0:7 ((k) ;0, I. - Ac|gi - [Z),
(9)
( in which the transition model, Po p0 (29|=--1, al _'-1), ), is further factorised as in Equation (7). The gradients through the outer expec- tation and the expectation term in ELBO are estimated using reparameterisation tricks (Kingma and Welling. 2014 Dang et al., 2017). For further implementation details, derivation of the lower bound,
where the ELBO term is given by the expression in Equation
and model architectures, see Appendices A and
4.4 ADAPTATION
Due to the modular nature of the transition model, VCD can naturally adapt to new, unseen environ- ments 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 as
p new (2 1-1 = II P: (0) (a/) @ 1 a 11)2 1112-1111pp 26 pl o
(10)

---

Under review as a conference paper at ICLR 2023
where R' is the intervention mask for the new environment and p. is the environment- specific distri- bution. One way to implement modular transfer is to train VCD on trajectories in the new environ- ment while fixing the trained parameters for the causal graph and the mechanisms in the undisturbed environment. The new parameters can he jointly trained in an analogous way to Eq. 2 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. Following the SMS hypoth- esis, only small subset of the mechanisms need to be adapted in new environment. Hence, by leveraging past experience, VCD can adapt quickly to environment change. In Section [5] we show that VCD can identify sparse interventions from small amount of data and is able refit the world model with less data than baseline approaches.
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 quantitatively, 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 between the environments. More- over, we illustrate that, in an unseen environment, VCD can leverage past experience to perform modular adaptation, resulting in significantly improved data efficiency over the baselines.
Dataset We evaluate VCD on simulated dataset of a 2-D multi-body system which contains four particles that affect each other via spring or an electrostatic- like force. 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 removing one of the springs, increasing or decreasing the mass of a particle, or constraining the position of a particle along the x or the y axis. Each of these changes can be considered as intervention causal variables. We evaluate VCD in two experiments: Mixed-state, where the observation is given by applying an 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.
more
Baselines We compare the performance of VCD against RSSM (Hafner et al. 2019b), a state-of- the-art latent world model that served inspiration for VCD. As RSSM does support learning from multiple environments, we consider two adaptations of RSSM with different levels of knowl- edge transfer between environments: (1) RSSM, where one transition model is trained over all environments, i.e., maximum parameter sharing across environments; and (2) MultiRSSM, where individual transition models are trained on each environment. This corresponds to the case where no knowledge about dynamics is transferred, i.e., each model is local expert. We hypothesise that, compared to these two extremes of knowledge sharing, VCD is able to capture environment-specific behaviours whilst reusing invariant mechanisms via modular transfer.
Prediction performance VCD and the baselines trained on dataset composed of 2000 tra- jectories from each of the undisturbed environment and five intervened environments. The models evaluated trajectories from a validation 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 of the environments due to parameter sharing. As such, it is not able to capture environment-specife behaviours. This is reflected in the prediction accuracy of the model. In the mixed-state experiment, VCD is able to identify changes between the environments and performs as well as MultiRSSM in terms of prediction error. A sim- ilar 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.
drawn from the training environments. Fig.
We note that in the image space, MultiRSSM does not perform well compared to RSSM and VCD. We hypothesise that this is due to the fact that the environment speciifi transition models are only trained on data from one environment each, compared to RSSM which has access to data from all environments. VCD does suffer from this due modular parameter sharing.
7

---

Under review as a conference paper at ICLR 2023
Mixad-state
Image
23.0
830
MUJIRSSM
730
VCD
500
13.0
030
330
5.-
3.0
130
175
80
timestep
timestop
Figure 3: The rollout error measured on validation trajectories. The models receive observations each timestep for the first half of each trajectory (shaded), 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) 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- con-pixel-bar based errors are available in Appendix
RSSM
VCD
along dimension
Figure 4: Reconstructed images from points sampled along dimension 6 of the latent space. Com- pared the latent space learnt RSSM, VCD able to learn an axis-aligned representation where the blue particle moves horizontally. See Appendix D for more examples.
Causal discovery One the key hypotheses this work that jointly learning a representation and transition model using causal discovery leads causally meaningful representations. Here we examine the quality of the learnt latent space and the causal structure of the transition model. Fig. 4 shows the image reconstructions of points drawn from straight line along dimension 6 of the latent space. Since both models are initialised with the same encoder, this provides qualitative intuition 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 is able to learn an axis- aligned coordinate (xr coordinate) of the blue particle with respect to the environment. Note that the motion of bouncing off the boundaries is only separable in the y frame, implying that the dynamics in axis-aligned coordinates is sparser. We present further analysis of the learnt causal representation in Appendix [ As desired, the learnt causal graphs in both the mixed-state and the image environment are found to be sparse.
Adaptation We provide empirical evidence that VCD can adapt to a new environment with less data compared to RSSM and MultiRSSM by reusing learnt mechanisms in modular fashion. We collect datasets of different sizes previously unseen intervened environment where particle is constrained horizontally. RSSM and MultiRSSM adapt to the new environment by optimizing the ELBO (Eq. 2) with the difference that MultRRSM instantiates a new transition model randomly and RSSM initialises the transition model using pre- pre-trained parameters. VCD performs adaptation by jointly estimating the intervention targets and the parameters the intervened mechanisms. The models are evaluated validation set of trajectories in the environment. Fig. shows the
8

---

Under review as a conference paper at ICLR 2023
Mixeed-States
14
Lo
30
traectorie
Images
00
122
10
10 octooiss
40 trajectonies
100 trajectories
640 trajectories
2560 trajecto fee
Figure 5: 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 reuse most of the previously learnt mechanisms, as indicated by the sparsity in the learnt intervention targets. In the mixed-state, we verify that the learnt intervention target (dimension 8) corresponds to the ground-truth target (11). see Appendix D Evaluations the image experiments using non-pixel-base based errors are available in Appendix
rollout error of the models trained on datasets of varying size. Across all models, performance improves as the dataset grows. However, contrast to RSSM and MultiRSSM, 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 is able to converge to a single latent dimension that encodes the y position of particle with just one training trajectory. See Appendix D for more experiments with different interventions.
CONCLUSION AND DISCUSSION
In this paper, we propose VCD, predictive world model with causal structure that is able to consume high-dimensional observations. This is achieved by jointly training a representation and causally structured transition model using modified causal discovery objective. In doing so, VCD is able to 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 a 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.) 2.21.. One possible drawback of the proposed framework that the causal graph is assumed to be static across all timesteps and all environments. In more complex environments, it could be more difficult for the model to discover temporally sparse interactions such as collisions, and adapt to scenes with different numbers of objects. As such, exciting directions for future research also in- clude exploring the synergy between causal world models and object-centric generative models (Wul et al.) 2021 Engelcke et al. 2021 von Kügelgen et al.) 2020). temporally-local causal influence detection (Seitzer et al. 2021 Pitis et al. (020) and relational reasoning (Goyal et al. (0216) Kipf
al. 2019).

---

Under review as a conference paper at ICLR 2023
Reproducibility Statement In order to ensure reproducibility of the results, derivation of the vari- ational lower bound, details of the model architectures used in the experiments, implementation de- tails for the proposed algorithms, and details of the experiment setup is included in App. A B and C The code for the experiments as well as the data generation process will be made avai available for the reviewing process and intend to release the code for the camera-ready version of the paper.
While the present work significantly advances the current state-of-the-art in world-modelling and causal representation learning, we expect its immediate impact outside of the machine learning community to be low current methods can not yet deal effectively with real- world scenarios.
Ethics Statement
REFERENCES
Ossama Ahmed, Frederik Träuble, Anirudh Goyal, Alexander Neitz, Manuel Wüthrich, Yoshua Bengio, Bern- hard 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. Properties from mechanisms: An equivariance perspective on identifiable representation lcarning. arXiv preprint
2021.
M. Arjovsky, L. Bottou, I. Guirajani, and D. Lopez-Paz. Invariant risk minimization.
arXiv preprint
arXiv:/90.00289) 2019.
Rim Assouel, Lluis Castrejon, Aaron Courville, Nicolas Ballas, and Yoshua Bengio. Vim: Variational indepen- dent modules for video prediction. In Conference on Causal Learning and Reasoning, pages 70-89 89. PMLR, 2022.
Banijamali, R. Shu, Bui, Ghodsi, Robust locally -inear linear controllable embedding, In International Conference on Artificial Intelligence and Statistics, pages 1751-1759. PMLR, 2018.
Fabien Baradel, Natalia Neverova, Julien Mille, Greg Mori, and Christian Wolf. Cophy: Counterfactual learn- ing of physical dynamics. arXiv preprint 19Xiv: 12000, 2019.
Philip Becker-Ehmck, Jan Peters, and Patrick Van Der Smagt. Switching linear dynamics for variational Bayes filtering. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference Machine Learning, volume 97 of Proceedings of Machine Learning Re- search, pages 553 562. PMLR, 09 15 Jun 2019. URL https:/pprocccol.ggg mmlr www..ppress/v11111 becker- www.eeehks.... 9a.html
T. Behrens, T. Muller, J. Whittington, S. Mark, A. Baram, K. Stachenfeld, and 7. Kurth-Nelson. What is cognitive map? organizing knowledge for flexible behavior. Newron, 100(2):490-500 2018. ISSN 0896- 6273. doi: Ceterce//stttccee////////666//13
www.::00000...erooo0000011111110 2018. 10.000. URL httppp::/wwww...eeeeeeeeeeeeeeee science//rticle//
Y. Bengio, T. Deleu, N. Rahaman, R. Ke, S. Lachapelle, O. Bilaniuk, A. Goyal, and C. Pal. A mela-transfer objective for learning to disentangle causal mechanisms, 2019. URL a 1901. 10912
I. Bica, D. Jarrett, and M. der Schaar. Invariant causal imitation learning for generalizable policies. Ad- vances Neural Information Processing Systems, 34, 2021.
Johann Brehmer, Pim De Haan, Phillip Lippe, and Taco Cohen. Weakly supervised causal representation learning. In ICI,22022 Workshop the Elements of Reasoning: Objects, Structure and Causality, 2022. wRR/httpss//openreee
forum? d-UUXUURUUGSS
Brouillard, S. Lachapelle, A. Lacoste, S. Lacoste-Julien, and A. Drouin. Differentiable causal discov- from interventional data. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 21865-21877, 21877. Cur- Associates, Inc., 2020. URL htttps:// lhttps://proceedingg.. www.iiigg..uurss....mm https/aas....cc///////////...... [8b72a3a0d339d9562bb22
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. ArXiv, abs/1 802 03006, 2018.
10

---

Under review as a conference paper at ICLR 2023
K. Cho. B. Merrienboer, C. Gulcehre, F. Bougares, H. Schwenk, and Y. Bengio. Learning phrase represen- tations using rin encoder-deeoder 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, pages
PMLR, 2018.
M. Engeleke. O. Parker Jones, and I. Posner. Genesis-v2: Inferring unordered object representations with- out iterative refinement. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances Neural Information Processing Systems, volume 34, pages Curran Associates, Inc., 2021. URL
43cc51
A. Goyal, A. Lamb. J. Hoffimann. S. Sodhani, S. Levine, Y. Bengio, and B. Schölkopf. Recurrent independent mechanisms. In 9th International Conference Learning Representations (ICI.), May 2021a. URI.
https://cpenrevview /et/tett//tttte wwt.//dfff// dwd.mmmmmmmll... EUXY-
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 Pro- cessing Systems, 34:25673 256677, 2021b.
D. Ha and J. Schmidhuber. Recurrent world models facilitate policy evolution. In Advances in Neural Informa- fion Processing Systems 3/, pages 2451-2463. Curran Associates, Inc., 2018. UtL/httsss:/ppaerrs
woeentwwoollcce dordd- faci -oce https://wcrl http:://wwdddloddelss gttlls....tttu thub.io
nips. cips.cc/ppaper/5511 www5511--eecueeee
D. Hafner, T. Lillicrap. J. Ba, and M. Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint 191iv:/112.0100000 2019a.
D. Hafner, T. Lillicrap. T. Fischer, T. Villegas. D. Ha, H. Lee, and J. Davidson. Learning latent dynam- ics for planning from pixels. In Proceedings of the 36th International Conference Machine Learn- ing, volume 97 of Proceedings of Machine Learning Research. PMLR, 09-15 Jun 2019b. URL https: //proceedings mlr. l ptw.s/resssss... httpes/////////errr/ee 19.. html
A. Hauser and Bühlmann. Characterization and greedy learning of interventional Markov equivalence classes of directed acyclic graphs. Journal of Machine Learning Research, 13:2409 13:2409-2466, 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 Learning Representations, 2022. URL http:httt///ppereeetiit
net/ forum? d155m15ppv....
Aapo Hyvarinen, Hiroaki Sasaki, and Richard Tumnrr Nonlinear ica using auxiliary variables and generalized contrastive learning. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 859 868, PMLR, 2019.
E. Jang, S. Gu. and B. Poole. Categorical reparameterization 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, Bilaniuk, A. Goyal, Bauer, Larochelle, Schölkopf, M. C. Mozer, C. Pal, and Y. Bengio. Learning neural causal models from unknown interventions, 2020. 110.0191...007552.
D. Kingma and M. Welling. Auto-encoding variational bayes. In 2nd International Conference on Learning Representations, ICLR, 2014. httth://p//axxt wwwpaaawww.rrrraaaaaaaaaaa/ 6114
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization.
arXiv preprint
arXiv: 14XX::4112.6980 2014.
Thomas Kipf, Elise van der Pol, and Max Welling. Contrastive learning of structured world models. arXiv preprint 191 1. 122247, 2019.
Lachapelle, López, Y. Sharma, K. Everctt, R. Le Priol, A. Lacoste, and Lacoste-Julien. Disentanglement via mechanism sparsity regularization: A new principle for nonlinear ica, 2021. /tt::s/s:///axx
org/ /a////b//at / /2070 ..2...055 10098
Minne Li, Mengyue Yang, Furui Liu, Xu Chen, Zhitang Chen, and Jun Wang. Causal world models by unsu- pervised deconfounding of physical dynamics. arXiv preprint arXiv:20/2. 14228, 2020a.
11

---

Under review as a conference paper at ICLR 2023
Yunzhu Li, Antonio Torralba, Anima Anandkumar, Dieter Fox, and Animesh Garg. Causal discovery in physi- cal systems from vidcos. Advances Neural Information Processing Systems, 33:9180-911 -19099999,2--001 2020b.
Phillip Lippe. Sara Magliacane, Sindy Löwe, Yuki Asano. Taco Cohen, and Efstratios Gavves. iCITRIS:
Causal representation learning for instantaneous temporal effects. In UAI 2022 Workshop on Causal Repre- sentation Learning, 2022a. httpstts:: http::::://epeneveee.... net oo um? www.oooooommmmiiiiieeae 1d=xeDKT2s2222
Phillip Lippe, Sara Magliacane, Sindy Löwe.
Yuki M Asano, Taco Cohen, and Efstratios Gavves. Intervention design for causal representation learning In UA/ 2022 Workshop on Causal Representation Learning, 2022b. /
net/ for um wwwummmdmii- /h42hh
Phillip Lippe, Sara Magliacane, Sindy Löwe, Yuki Asano, Taco Cohen, and Stratis Gavves. Citris: Causal identifiability from temporal intervened sequences. In International Conference on Machine Learning, pages 13557-13603 PMLR, 2022c.
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 Statistical Methodology (with discus- sion). 7855)994-11122, 2016.
J. Peters. D. Janzing, and B. Schölkopf. Elements of Causal Inference Foundations and Learning Algorithms. MIT Press, Cambridge, MA, USA, 2017.
Silviu Pitis, Elliot Creager, and Animesh Garg. Counterfactual data augmentation using locally factored dy- namics. Advances Neural Information Processing Systems, 33:3976-3990, 2020.
B. Schölkopf, D. Janzing, J. Peters, E. Sgouritsa, K. Zhang, and Mooij. On causal and anticausal Icarning. In Langford and J. Pincau, editors, Proceedings of the 29th International Conference on Machine Learning (ICML), pages 1255-1262, New York, NY, USA, 2012. Omnipress. URL http::/iieee //cml. www.m...cc.....mmmmmmmm..mm...me 2C12/
papers/ / 625.pdd.
B. Schölkopf, F. Locatello, S. Bauer, N. R. Ke, N. Kalchhrenner, A. Goyal, and Y. Bengio. Toward causal representation learning. Proceedings of the IEEE Advances in Machine Learning and Deep Neural Net- works, 109(5)6612 634, 2021. doi:
20W./PPRRC..221335858883 leee. st amp/ st jap? jswparnunber- 9363924
URL /ttpp::// Settps://ieeexplee /eeeeeeeeeepeeeeeeeee
Maximilian Seitzer, Bernhard Schölkopf, and Georg Martius. Causal influence detection for improving effi-
ciency in reinforcement learning. Advances in Neural Information Processing Systems, 34:22905-22911 2021.
R. Sekar, O. Rybkin, K. Daniilidis, P. Abbeel, D. Hafner, and D. Pathak. Planning explore via self-supervised world models. In ICMI., 2020.
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 Uncer- tainty in Artificial Intelligence (UAI), pages 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 Ktigelgen, I. Ustyuzhaninov, P. Gehler, M. Bethge, and B. 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- latent dynamics model for control from raw images. In C. Cortes, N. Lawrenee, D. Lec, M. Sugiyama, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 28. www.seeerriiccoee ptt../pp.cccpppar///////111..
Curran Associates, Inc., 2015. URL htLps: http:httpes//pppceeddleeeeeeeeee
Y. Wu, O. Parker Jones, M. Engeleke, and 1. Posner. Apex: Unsupervised, object-centric scene segmenta- tion and tracking for robot manipulation. In JEFF/RSI International Conference on Intelligent Robots and Systems (IROS), 2021.
12

---

Under review as a conference paper at ICLR 2023
Weiran Yao, Yuewen Sun, Alex Ho, Changyin Sun, and Kun Zhang. Learning temporally causal latent pro- cesses from general temporal data. In International Conference Learning Representations, 2022. URL https://openrevi.cs net/ Lor.mnnddR
Zhang, Lyle, Sodhani, A. Filos, M. Kwiatkowska, Pineau, Gal, and Precup. Invariant causal prediction for block mdps. International Conference on Machine Learning, pages 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, pages 3150- -150,,117,0000 2015.
K. Zhang, B. Huang. J. Zhang. C. Glymour, and B. Schölkopt. Causal discovery from nonstation- ary/heterogencous data: Skeleton estimation and orientation determination. In Proceedings of the Twenty- Sixth International Joint Conference on Artificial Intelligence (IJCAI 2017), pages 1347-1353, 2017.
M. Zhang. S. Vikram, L. Smith. P. Abbeel, M. Johnson, and S. Levine. Solar: Deep structured representations for model-based reinforcement learning. In International Conference on Machine Learning, pages 7444- 7453. PMLR, 2019.
13

---

Under review as a conference paper at ICLR 2023
A DERIVATIONS
Using the approximate posterior a(2"-T)e 1o (aaTTo00 latent state-space model (Bq. 2 can be derived from importance weighting and Jensen's 's inequality:
( ((="'o") the variational log lower bound for the
log plo 0: aa log / II p(o' (aaa 1=')<=="|>>> a add
(11)
log /II I='ool=")aaaa)) p( aat1)
T
2(2%)0t g(2 dd))
(12)
Jogg [Iooo
> "paaall' p( al
2 [l
(13)
a(2'10t)
Botet (ee) (logp(o'l=") + log
+ Bogee Poot 1) (att [logp(> 12 122 ,a'=)) 1 1ogg(="10"))
(14)
Byler (o)) [logp(o'(>") logp(a'
- -Bg(rr-1jot-1) KL |a(="1o I[x((") a []]
(15)
Since the policy p(a'l=") is constant with respect to the model parameters, omit this term and write the ELBO objective
ELBO((O 6) = Σ B [Logepo)
( KLLL (LLigs((' (2 (z'oo') (2oo "aa"aaaa" [)], (16)
where is the model parameter henceforth to simplify notation. In VCD, the KL divergence term can be further decomposed by exploiting the structure of the transition model (Eq. 77. and the assumption that variables within each timestep are independent:
the parameter for approximate posterior. and Φ are omitted
TTL (g==100 o') ) (x'12" "laa=",aa"" b
(17)
p (k) (x'k"
log
a(= o(='\o")dd"*
(18)
((27(00)
(0) (2 @ 13"--1,,00 _a"",a"" aa-11"1 Rh (6)
(2f 1. @ a1
log
(19)
q(2:(0')
Σ
d
(1 REii / log
Pi (0) (2 o 13t aaaa
a(=))o')d2)
{(<((o')
(8) (2 o [3 a1
Ri: log
((=)10")dz)
(20)
{(<((o')
Σ ((1 RE.) (((2) (=
d
(=! @F 3='"-(aa'-11) 1
0
+ RTL (g(=)lo')1le (2 2 0|="_1,a" -a-")
(21)
The KL terms can be computed analytically since the conditional distributions in the last expression are uni- variate Gaussian distributions. In training time, the gradients through the expectation terms in the ELBO estimated by drawing a sample from posterior distribution using the reparameterisation trick (Kingma and Welling 2014).
B IMPLEMENTATION DETAIL.
B. DCDI AND GRAPH LEARNING
This section covers the formulation of DCDI (Brouillard et al. 2020 and the graph learning method. These suhsequently used in the leamning of VCD.
14

---

Under review as a conference paper at ICLR 2023
Given samples from an observed data distribution 2 and intervened distributions p DCDI op- timises a probabilisticc belief over causal graphs 9 and intervention targets I. Specifically, these are encoded as random binary matrices, Mg and Rz, where M, = implics that the edge (i,j) is in the causal graph, and R = implics that the variable 2; is in the intervention targets in environment k. Each entry in M follows independent Bernoulli distribution, parameterised by matrix where P(M = 1) ( Rig is similarly parameterised by Under this parameterisation, causal discovery be formulated maximising the expected data log likelihood with sparsity regularisation.
K Σ Kinnn
L(O, 3) Ea,8
Logbss (a" - - >r|Z]|
(22)
where p(() is the data likelihood under causal graph G and intervention targets I in the kth environment, factorised in eq.44. The conditional distributions are parameterised as feedforward neural networks with parameter 0; λG are hyperparameters to control sparsity. 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 timestep).
original DCDI framework, this is also suhject
The training objective for VCD can be viewed modificd version of the DCDI objective, where the likelihood term is replaced with the ELBO (Eq. T6].
K Σ>
L LVD (0,6,,, B.,
ELEO
-a|s >r|I
(23)
Note that the expected number of edges & and Z given a and |3 is simply the sum of the probability of cach entry being onc. Therefore, the training objective can be computed as:
K T
L VCD (0,0,0,8) =a.
BLBO(0,0;9,1)) do Ea(ow) -4 Σ (Bxs). LAEO
(24)
de.
The gradients through the outer expectation can be estimated using the Gumbel-Softmax trick (Jan et al. 2017). To implement this, the ELBO term is evaluated with sample of the causal graph using the following expression for each entry.
My = I(a(as L2) 0.5) + alays + Laj) - stop stop-gradient(((auu Lez));
(25)
where I(-) is the indicator function, Laj is sample from the logistic distribution, and stop-gre gradient is a func- tion that does not change the value of the argument but sets the gradient zero. Samples for the intervention targets are similarly acquired. Note that sample is used throughout each trajectory, i.e. same sample graph and intervention targets 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) parameterised by feedforward ML.ss with two hidden layers of 64 hidden units each. The recurrent modules implemented as GRUs (Cho al. (014) with 64 hidden units. Distributions in the latent space are 16-dimensional diagonal Gaussian distributions with predicted mean log variance.
Image In the image experiment, the encoders and decoders are parameterised as convolutional and deconvolutional networks from (Ha and Schmidhuber 2018). In the RSSM models, the transition models are parameterised as feedforward MLPs with two hidden layers of 300 hidden units. The recurrent module is GRU with 300 hidden units. In VCD, compensate for the fact each dimension the latent space has sepa- rate model, the number of hidden units in the GRU and MLP are reduced to 32 avoid over-parameteriisaion We found that initialising the encoders and decoders by pretraining them as a variational autocncoder helped with training stability for both RSSM and VCD.
both experiments, the training objective maximised using ADAM optimiser (Kingma and Ba 2014) with learning rate 10 for mixed-state, and 10 for images. In both environments, we clip the log variance to -3, with batch size of two trajectories from each of six environments with T = 50. In VCD, the hyperparameters Äg, Az are both set 0.01. All models trained on single Nvidia Tesla V100 GPU,
C EXPERIMENT DETAIL
15

---

Under review as a conference paper at ICLR 2023
Table 1: List of interventions
ID
Intervention
Intervention targets
Remove spring between and
21, y1, 422 1/2
Remove spring between 2 and
22, U/, 4:3,
Increase mass
11 U1 JT4 94
4 Increase mass
22, y/2
Increase mass
T3s M3. J.4s M/4
Decrease mass
01 U1, F. 7/4
Decrease mass
T2, 92
Decrease mass
33, U3, JF4, 1/
Increase spring constant between and 2 01, y1, 222,
10
Increase spring constant between 2 and 3 322 92 23 M3
11
Constrain movement of I to vertical only
12
Constrain movement of I to horizontal only
8/1
13
Constrain movement of 2 to vertical only
32
Constrain movement of 2 horizontal only
5/2
15
Constrain movement of 3 to vertical only
J73
Constrain movement of 3 to horizontal only
1/3
17
Constrain movement of 4 to vertical only
Constrain movement of 4 to horizontal only
y/
Table 2: The modified MCC scores for RSSM and VCD both experiments
Expriment
RSSM
VCD
Mixed-state
0.728
0.975
Image
0.715
0.908
Interventions The ground truth states multi-body dynamics environment is æ and y coordinates of each particle. The full list of possible interventions on the environment provided in Table Note that the forces between particle 3 and 4 are proportional to their masses. Hence intervening the mass of particle and also affect the dynamics of particle 4. In the experiments, all models trained in the undisturbed environment and intervened environments 1, 14, 17.
mixed-state experiment, observation function a mixing matrix where each entry drawn a unit Gaussian distribution. In the image experiment, the observation given by rendering the system to a 128 x 128 x 3 image. In both experiments, the models are trained on training set of 2000 trajectories from each of the six environments and evaluated validation of 400 unseen trajectories.
D QUALITATIVE EXPLORATION FOR LEARNT CAUSAL STRUCTURE
D.1 ROLLOUT PREDICTION
Section 5 shows the rollout accuracy of the baselines and VCD. Here demonstrate qualitatively that VCD ceonmmett--ppeiffe behaviours RSSM cannot learn due to maximal parameter sharing. Figure shows sample rollouts in the image space from RSSM and VCD in intervened environment 11, i.c. the ycllow particle is constrained horizontally and only moves vertically. The yellow particle in the VCD rollout stays along vertical linc whereas RSSM fails to capture this cuvironment-specite constraint.
to capture
D.2 REPRESENTATION QUALITY
In this subsection, we explore the quality of the learnt latent space. The key hypothesis of our work is that training the representation jointly with the transition model to maximise causal discovery objective serves as inductive bias that helps structure the latent space causally meaningful way. For hoth experiments,
16

---

Under review as a conference paper at ICLR 2023
Ground Truth
VCD
RSSM
Figure 6: Sample rollouts (without observations) from an intervened environment. VCD success- fully captures the constraint on the yellow partiele (vertical movement only) over long time hori- zon. The frames are sampled timesteps apart from each other.
ground truth state
ground truth state
[
RSSM latent dimension
VCD latent dimension
Figure 7: The average magnitude of the derivative of the ground-truth states z; w.r.t the learnt latent states 2j. VCD is able to learn latent representation where each dimension captures information about one state only, i.c. only one state variable changes when we perturb the latent representation along one dimension. On the other hand, the RSSM latent space encodes state information in more entangled manner, i.e. multiple ground-truth states are affected when the latent representation perturbed in one dimension.
we provide the mean absolute correlation coefficient score representations RSSM VCD We see that in both cases, training with the sparsity regularisation VCD achieves a higher MCC score, which confirms that VCD learns more disentangled representation.
Mixed-state Since observation mixing function linear and invertable the mixed-stan state environment, we can directly access the level of disentanglement of the latent space with respect to the ground truth state of the environment, i.c. the æ and coordinate of the particles. Fig. [hsows 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 each ground-truth state variable when the latent representation perturbed along each dimension. By using the discovery causal transition models as an inductive bias, VCD is able learn disentangled representation where each ground-truth state captured by only one latent dimension. In contrast, RSSM learns representation that is not sparse.
The quality the learnt representation is discussed Section Here we reconstruction sam- ples along all dimensions of the latent space the RSSM and VCD representation. Note that since both encoders are initialised from the same pretrained VAF, the difference in the latent space arise hecause of the causal discovery objective in VCD. Fig. 12 shows the changes in the reconstruction when the RSSM repre-
Image
'Due to the fact that the ground-truth truth state dimension (8) is less than the latent state dimension (16), we modify the MCC score by calculating the mean of the top 8 scores out of the 16, effectively ignoring the latent variables that do not capture the state information.
17

---

Under review as a conference paper at ICLR 2023
Doundoo
R2
image RSSM bntnt state
Figure 8: The R2 of linear regression between the ground-truth states and the learnt latent vari- ables. This shows a similar pattern as the mixed-state experiments where the VCD representation in general only capture information about one ground-truth state.
Mixed-state
Image
causal parents
causa parents
9
10
10
12
15
34
causal variables
causal variables
Figure 9: 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 o(ass) >> 0.5, i.e. is causal parent of j.
sentation is perturbed in each dimension of the latent space. Fig. shows the same for VCD. As discussed the main VCD learns axis-aligned representation that affords sparse causal graph. For example, dimension and 3 captures the y and coordinates the green particle respectively. contrast, RSSM learns a representation that is not axis-aligned. Fig. [ plots the H² linear regression scores between the learnt la- tent variables and the ground-trume variables. showing that the VCD representation can in general disentangle individual states, as indicated by the more salient squares.
D.3 LEARNT CAUSAL GRAPHS AND INTERVENTION TARGETS
In this subsection, explore the quality of the learnt causal graph and intervention targets. Fig. 9 shows the learnt causal graphs for the mixed-state and image experiments. Fig. shows the learnt intervention targets for each environment. The sparsity the learnt graph and targets summarised Table VCD identified 42
Mixed-state
Image
Envinnment
envilanmenn
14
15
intervention targets
intervention targets
Figure 10: The learnt intervention targets for the mixed-state and image experiments, obtained by binarising the learnt probabilities such that blue square at (k, i) implics a (Bur) > 0.5, i.c. i is an intervention target in environment h.
18

---

Under review as a conference paper at ICLR 2023
Table Sparsity of learnt causal graphs and intervention targets. We compare the learnt graph with the ground truth causal graph by mapping each latent dimension to a ground- state shown in Fig. 7 and Fig. Note that the number of correct edges and the false positives do not sum to the number of edges in the graph because the edges corresponding to 'dummy' variables are ignored.
Observation
# of edges
Correct Edges
Missed Edges
False Positives
Mixed state
causal graph
42/288
19
intervention targets
18/80
Image
causal graph
73/288
22
4
13
intervention targets
31/80
10
0
17
and 73 causal edges in the mixed-staaee 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 sparsely connected and affords modular parameter sharing without compromising on prediction accuracy.
Mixed-state In the mixed-state experiments, each ground-truth state can be mapped to latent dimension using the average magnitude Jacobian matrix (Fig. 7 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. quality
learnt graph. VCD able to identify a majority dependencies and all of the intervention targets. Upon further inspection of the learnt causal graph, we find that all of the seven missed edges correspond to the 1/|] || terms that scale the electrostatic -ikkettii--kkeeeeeee forces. We hypothesise that the model cannot capture these dependencies as they not significant the other forces.
Table summarises
correct causal
Image In the image space, a similar comparison is made using the scores. Each ground-truth state is mapped to a latent dimension by choosing the dimension with the highest R² score. This leaves unused latent states, which ignored in the analysis provided in Table 3 We also qualitatively explore the Ieant causal relationships. Focusing on dimension 3, for example, which encodes the x coordinate of the green particle, the ground-truth causal parents this variable is the æ coordinates 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. suggests that dimensions and 8 captures the coordinates of the yellow and blue particles respectively, meaning that VCD has learnt the correct causal parents.
similar analysis can be carried out on the intervention targets. Focusing on intervened environment 0, for example, the learnt intervention targets are dimension 1, 2, 3, 7, 8, 11, 13, 15. The ground-truth intervention ID is 1, i.c. the intervention targets are y coordinates of the yellow and green particles (see. table T By inspecting the VCD latent space (Fig. [T), dimensions I and encodes the position of the green partiele and dimensions 8
encodes the position 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.
E FURTHER EXPERIMENTS
This section provides extra adaptation experiment results similar to Fig. Bj in the main text, where the models adapt to intervention number 12 (sce table [T. We present experiment results for adaptation to different types of interventions (intervention numbers 4, 9, 13). Fig [ig[[3] 14]]14aad and show the adaptation plots. The results exhibit a consistent pattern where VCD significantly outperforms the baselines in the low data regime by identifying sparse mechanism changes.
F NON-PIXEL--BABI ERROR
In the image experiments, while pixel error indicative of model performance in the short term, it provides limited clarity long term predictions as two non-overlapping balls leads to same pixel error regardless of their distance. In this section, we present an alternative latent space distance- based evaluation metric, Hit at 5 5). This metric is defined by the proportion of episodes where the predicted latent representation is in the top-5 nearest neighbour set of the encoded ground-truth image. A H@5 score of means all predicted latent states lie within close distance the ground-truth representation. Fig T6 and T7 shows the prediction errors
19

---

Under review as a conference paper at ICLR 2023
am
4m
dm3
dm
dm
am
4m
am
dm 20
6m
dr 12
am 23
am 14
um
WED
Figure 11: Visualisation of thg learnt latent space in VCD.

---

Under review as a conference paper at ICLR 2023
din
dm
6m
0m
RESH
Figure 12: Visualisation of the learnt latent space RSSM.
21

---

Under review as a conference paper at ICLR 2023
Mixed-Stan States
VCD McRKSSH
Erajedary
Enjecter
Images
traectories
tojectares
Injecturee
Figure 13: Adaptation results for intervention 4.
Mixed States
Douectory
Egjedotis
Dajecories
trancDoRR
Images
sajecarris
Urajcdriis
100
Unjeedoris
2300 Dalctor
Figure 14: Adaptation results for intervention 9.
Mixed States
ROSM MUBIRSSV
di in
Images
petories
Injectates
Injecttries
trajetores
Figure 15: Adaptation results for intervention 13.
22

---

Under review as a conference paper at ICLR 2023
Hit
100
RS5M Multi-RSSM
80
@ 5 (%6
20
0
rollout timesteps
Figure 16: 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.
RS-M
123N VCD
Hg5 180
X 10 trajectory
Erajectories
160 trajectories
Trajectorios
2560 trajectories
Figure 17: The adaptation plots with the H@S scores.
and the adaptation plots respectively. These consistent with the results shown in the main text, where VCD outperforms the baselines in prediction accuracy as well as adaptation speed.
23


Reviews:
[{"review_id": "jIgEIUY0Z8", "reviewer": "Reviewer_qfCJ", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors propose a method for causal representation learning from interventional data and demonstrate how such a model can be adapted to new settings. They consider time-series data generated by a causal graphical model and a rendering function (e.g. to pixels); key assumptions include that there are no instantaneous causal effects and the causal structure is invariant over time steps. They then assume access to data from the observational and multiple interventional distributions. In this setting, a VAE with latent causal structure is trained. The authors argue that such a model should be able to adapt efficiently to data from a new environment provided the sparse mechanism shift hypothesis holds. This is then demonstrated empirically on toy data.", "summary_of_the_review": "This work has a lot of potential: learning latent causal structure from interventional data on the pixel level could be very useful. The VAE setup makes sense. It is commendable that the authors do not stop at the identification of causal structure, but show that learning it is actually useful for predicting the future of a roll-out and for adapting to new settings.\n\nUnfortunately, in its current form the paper does not demonstrate that the approach works. There should be either a theoretical discussion (like an identifiability result) or a convincing empirical demonstration that the model actually learns the correct causal structure. Right now, the paper has neither, and I can't recommend it for acceptance yet.\n\nI encourage the authors to expand either of these two angles (or, ideally, both). If they can add such a demonstration, this could become a great paper and an impactful contribution to the field.\n\n**Update after rebuttal**: Thanks to the authors for clear answers to my question and for addressing some of my criticisms in the upadted paper version. While the paper is improved, I still believe that it does not sufficiently demonstrate that the approach works. My score remains unchanged.", "strength_and_weaknesses": "**Strengths**:\n1. The authors tackle an interesting and largely unsolved problem. Both learning latent causal structure from a different representation and demonstrating the benefits of that for rapid adaptation may be very impactful.\n2. The model setup (a VAE with latent causal structure) makes sense.\n3. It is great that the authors consider a weak form of supervision: the model is trained only on interventional data without any additional labels. This seems much more easily attainable than labelled data [e.g. 2] or counterfactual data [e.g. 3]\n4. It is also great that the authors do not stop at the learning of causal structure, but demonstrate the benefits of that causal structure for more efficient finetuning on new settings.\n5. The paper is clearly written and a pleasure to read.\n6. The figure design is also good.\n6. The authors have provided clean and well documented code along with the paper.\n\n**Weaknesses**:\n1. The setup makes some strong model assumptions, in particular that the causal graph is invariant over time steps and that there are no instantaneous causal effects. While the authors point this out, it would be good to stress these assumptions in the introduction and conclusions a bit more and to discuss them: why are they needed? Are they realistic? How could we relax them?\n2. Relatedly, I have a hard time thinking of example systems that are compatible with all the assumptions here. It would be great if the authors could provide some motivating examples.\n3. Probably my biggest criticism of the paper is that it is not clear what exactly the method should be able to achieve – there is essentially no theoretical discussion. In particular, the authors do not make any claims about identifiability. Up to which equivalence class should the causal variables, the causal graph, and the rendering function be identified? What is the role of the assumptions on the data-generating process for identifiability? What properties do the rendering function and the distributions need to satisfy? Footnote 1 points the reader to Ref. [1], but if the identifiability result of that paper applies here, that should really be discussed more explicitly (which is also important to judge the novelty of this paper).\n4. The empirical evaluation also leaves something to be desired. It is okay that there is just a single toy dataset, but the performance in terms of causal representation learning should really be discussed quantitatively. There are many adequate metrics of disentanglement and graph similarity that could be applied here. The numbers that the authors do provide are not particularly impressive. In the appendix, the authors claim that it is \"not trivial to obtain a way to map the [ground-truth factors] to a latent dimension\". I would politely encourage the authors to have a look at how other representation learning papers deal with that – for instance, they could use a feature importance matrix like in Figure 7 to determine a map between ground-truth factors and learned latents and then compute graph edit distances.\n5. The authors cite many relevant works, but missed Refs. [3-5] given below. In particular, it seems that Ref. [5] considers a very similar problem. This is clearly concurrent, so don't take this as a criticism of novelty, but it would make the paper stronger if you could mention and contrast their and your findings.\n\n**References**:\n- [1] Lachapelle et al, \"Disentanglement via Mechanism Sparsity Regularization: A New Principle for Nonlinear ICA\", CLeaR 2022, 2107.10098\n- [2] Lippe et al, \"CITRIS: Causal Identifiability from Temporal Intervened Sequences\", ICML 2022, 2202.03169\n- [3] Brehmer et al, \"Weakly supervised causal representation learning\", 2203.16437\n- [4] Lippe et al, \"iCITRIS: Causal Representation Learning for Instantaneous Temporal Effects\", 2206.06169\n- [5] Ahuja et al, \"Interventional Causal Representation Learning\", 2209.11924", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "ggUKB3p0WPw", "reviewer": "Reviewer_xUbv", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "In this work a model for sequences with the following properties is proposed:\n- Modular description of the transition function\n- Handling of interventions, defined as the modification of a target of the transition function\n- Learning of the dependency structure\n- Use of latent variables\n  \nThe authors build on top of the formulation from \"Differentiable Causal Discovery with Interventional data\" (DCDI), which exhibits the above properties, except for the latent variables. The addition of the latent variables introduces intractability that is solved through the use of a variational inference.\n", "summary_of_the_review": "A step in a good direction, drawing most of the insights from existing work, limited empirical evaluation.", "strength_and_weaknesses": "Strengths:\n- The addressed problem is highly relevant. The above-mentioned list of properties seems necessary for artificial intelligence and few works check all these boxes. \n- The experiments show the advantages of the current proposal and even show the emergent property of axis alignment.\n\nWeaknesses:\n- I'd like to see more justification as to why this model should be called \"causal\", as opposed to simply \"predictive and modular, while encouraging sparsity\". While it seems reasonable that a predictive, modular, sparse model would capture causal dependencies, is there a formal reason why it couldn't capture consistent temporal correlations? What would happen if I were to provide as evidence a video in which entropy increases, but run backwards (so that entropy decreases)? What if I were to show any other video with consistent correlations but lack of causation?\n- The experiments are very limited (just one toy example), and the proposed model is not studied in detail:\n  - In particular, the dynamics seem pretty deterministic. How does it handle uncertain actions? Can it handle actions in which the target state is multimodal given a source state?\n  - How would this model handle situations in which the uncertainty gives rise to multimodality? For instance, imagine that there's a very small wall in the middle of a room and a ball can bounce off it only in a concrete trajectory, but won't if the trajectory deviates slightly. In this case, when predicting in open loop (observations are not fed back into the model), will it be able to handle the multimodality of the predictions properly?\n- Other\n  - Format violation: More than 9 pages\n  - I believe |G| is never defined", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "RJ4vw2QCcv", "reviewer": "Reviewer_Tfgu", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper describes variational causal dynamics (VCD), an approach that replaces a latent state-space model with a causal transition model that is learnt via a combination of differentiable causal discovery (Brouillard et al., 2020) and a set of independent deep networks that are defined the inferred causal graph. Additionally, they have a simple masking strategy to define the interventional distribution. The whole pipeline is trained via variational inference (optimizing the ELBO). They evaluate the approach of a simple physics-based task involving multiple coloured balls connected by spring forces.", "summary_of_the_review": "Overall, I liked the paper - it proposes a natural latent variable model that takes advantage of the recent advances in our understanding of causal representation learning, but I felt like it was let down by an experimental section that didn't do a good enough job of testing the limits of the approach. ", "strength_and_weaknesses": "Strengths:\n - A very clear and well written paper that combines most of the key ideas from the recent progress in causal representation learning and causal discovery into a latent state-space model. None of the individual components were particularly surprising, but the combined architecture and train procedure is appealing and very well motivated.\n\nWeakness:\n - The experiments were conducted on a very toy dataset. This style of data generating process is common in the recent papers in causal representation learning (although even there, richer datasets like 3DIdent, etc. are becoming more common), but in those papers, the experiments are primarily to ask whether identifiable models are learnable in practice (and to evaluate finite sample / optimization error / etc issues that are hidden in identifiability proofs). Because this paper isn't claiming any theoretical advances (and to be clear, I don't think it needs to!), it leans more heavily on experiments to show the value of these approaches. Given the really strong setup and motivation, I was hoping that the experiments would either provide strong evidence for the value of this approach, or alternatively show where there is need to further improvements. \n\n   The experiments do provide evidence that there is a benefit to taking advantage of causal structure, but it is not clear whether this generalizes beyond these environments that are specifically designed to show the difference between these approaches. And given that the image based experiments showed relatively modest differences when a large number of trajectories are used (Figure 3), I'm not sure that there will be any significant different in a more natural environment (e.g. an RL-style task). In the adaption experiments with low data regimes, the differences are far more dramatic. But that raises the question: is the advantage of imposing causal structure only about sample complexity, such that if we see enough complexity, the differences are washed out?\n\n - Causal discovery section: I would have liked to see a quantitative metric like the mean absolute correlation coefficient to measure disentanglement. In this setting you'll need to do a matching based on covariance between the true latents and the model's latents and drop the least correlated latents (to account for differences in size). Dropping the least correlated latents will bias the scores upwards, but they will still be comparable across models that use the same architecture. I would also include a linear disentanglement score ($r^2$ of a linear regression between the true latents and the predicted latents). \n\n - The paper draws on a number of areas to compose its architecture  (causal discovery / identifiable representation learning /  independent mechanisms / etc), so I think the experimental section would be significantly improved if it gave guidance on what components need to improve in order to see bigger differences with RSSM and to generalize to richer environments. I.e. were do we need breakthroughs to see the biggest gain in performance? Is it causal discovery that's letting the method down? Or the encoder? Or the parameterization of the independent mechanisms? Something else?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "Ce9IMuSSPI", "reviewer": "Reviewer_gksh", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors propose a method for latent dynamics modelling. In latent space, the latent variables are assumed to affect variables in the next time step via a sparse graph interaction. The model allows for soft interventions on the mechanisms. The model is trained via the ELBO and the graph structure, causal mechanisms and interventino targets are trained via the method from Brouillard et al. (2020). The method is evaluated on a toy image dataset of moving balls.", "summary_of_the_review": "The paper in limited in technical novelty and has neither theoretical guarantees, nor has been convincingly shown to work well in practice.", "strength_and_weaknesses": "Strength:\n- The paper is slightly more applicable than [1], as it doesn't require each environment to have labelled intervention targets.\n- The paper is clearly written.\n- The authors have provided code with their paper.\n\nWeakness:\n- The empirical evaluation is only done on a toy dataset of coloured circles. This is insufficient to assess the performance of the proposed model.\n- The paper has no identifiability guarantees.\n- The paper is very similar to [1], but this is not fairly reflected in the text.\n- In table 2, it appears like the graph and intervention targets aren't learned very accurately.\n\n[1] Lippe, Phillip, et al. “CITRIS: Causal Identifiability from Temporal Intervened Sequences.” 2022", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "-YHmA0blkY6", "reviewer": "Reviewer_LfHe", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper presents VCD, a world model that can modularly transfer knowledge between different interventional environments. It causally factories the transition model in the latent space and jointly learns an approximate posterior model, transition model, causal graph, and the intervention targets together with a derived ELBO loss. VCD learns both observational and interventional mechanisms which enable modular adaptation between different interventional environments, as it taking advantage of the invariances of causal mechanism.", "summary_of_the_review": "Given the above strengths and shortcomings of this paper, I recommend this paper be considered slightly below the acceptance threshold.", "strength_and_weaknesses": "Strengths:\n\t\n1. Modularity: When encountered with new and unseen environments, VCD can quickly adapt by training the VCD on trajectories in the new environment while fixing the trained parameters for the causal graph and the mechanisms in the undisturbed environment.\n\n2. Sample efficiency: During adaptation, VCD is more sample efficient as compared to the baselines RSSM and MultiRSSM.\n\n3. Extensive experiments are done in the multi-body environment.\n\t\nWeakness:\n\t\n1. Typos in equation 14: p(z^t|o^t)  should be q(z^t|o^t)\n\n2. The experiment is only tested on one simple environment.  Lack of complex environment to demonstrate its claimed effect on higher-dimensional raw image input.\n\n3. The causal graph is learned in the latent space and depends on the dimension d of the latent space. It would be interesting to provide comments on the choice of d, and the relationship between d and the actual number of edges between entities in the multi-body system. \n\n4. The multi-entity environment only has 4 entities. How will the model scale with a higher number of interacting objects in the task and a higher number intervene objects?\n\t\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject