1975). An example image for each of the objects is shown in Figure 11.

####  C.1.2 Dataset Generation

The datasets are generated by starting at a random sample of all causal
factors. For each next time step, we generate a sample according to the
conditional distributions of each causal factor (see Section C.1.3 for details
on the distributions). Additionally, we sample intervention targets
$I^{t+1}_{1},...,I^{t+1}_{7}$ for all 7 causal factors. For the datasets in
Section 6, we sample the targets from $I^{t+1}_{i}\sim\text{Bernoulli}(0.1)$,
and show additionally intervention settings on this dataset in Section D.1.3.
For each causal factor for which the intervention target is one, we replace
its previously sampled value with a new value randomly sampled from a uniform
distribution. For angles and hues, the distribution is $U(0,2\pi)$, while for
the positions, we use $U(-2,2)$. For the object shape, we uniformly sample one
out of the seven shapes. After performing the interventions, the sampled
vector of causal factors is used to generate an image of a resolution of
$64\times 64$ using Blender. Note that for visualization purposes, the
depicted images in this section are shown in higher resolution ($256\times
256$). This makes it easier to recognize the object shapes and their
rotations. However, we use a resolution of $64\times 64$ in experiments to
keep the computational cost of the experiments in a reasonable range.

Repeating this generation procedure for several steps results in one long
sequence, which we use as a dataset. For the experiments on Temporal-
Causal3DIdent Teapot, we generate a sequence of $150$,$000$ images. For the
experiments on Temporal-Causal3DIdent 7-shapes, the chosen dataset size was
$250$,$000$ images. The large dataset sizes were chosen to prevent any
sampling bias and focus the experiments on general identifiability. We noticed
that smaller dataset sizes such as $50$,$000$ images still gave good scores on
the correlation metrics, but a decrease in the triplet evaluation was
noticeable for most causal factors, especially the position and rotation.

####  C.1.3 Temporal Causal Relations

pos_xpos_ypos_zrot_$\alpha$rot_$\beta$rot_shue_shue_bhue_oobj_spos_xpos_ypos_zrot_$\alpha$rot_$\beta$rot_shue_shue_bhue_oobj_s

Figure 13: Causal relations between all dimensions of the causal factors of the Temporal Causal3DIdent dataset. The causal graph of Figure 5 summarizes pos_x, pos_y, pos_z into pos_o and rot_$\alpha$, rot_$\beta$ into rot_o. See Section C.1.3 for details on the conditional distributions. |  |  |  |
---|---|---|---|---
(1) - |  (2) rot-$\alpha$ | (3) pos-x | (4) pos-y, hue-b |  (5) rot-$\alpha$
|  |  |  |
(6) pos-y | (7) - | (8) - | (9) - | (10) hue-b, obj-s
|  |  |  |
(11) rot-$\beta$ | (12) - |  (13) rot-$\beta$ |  (14) rot-$\beta$ | (15) -
|  |  |  |
(16) pos-y, rot-$\beta$ | (17) - | (18) - | (19) - | (20) rot-s
Figure 14: An example sequence with 20 frames in the Temporal Causal3DIdent
7-shapes dataset with higher resolution (from left to right, top to bottom).
The causal variables denoted below each image indicate the variables which
were intervened on at this time step. For instance, when transitioning from
the first to the second image, all variables were sampled according to their
temporal dependency except rot-$\alpha$.

Below, we define the transition functions used in the causal graph of Temporal
Causal3DIdent dataset, as shown in Figure 5 (see Figure 13 for the relations
on individual causal dimensions). The chosen function forms are inspired by
the ones defined by von Kügelgen et al. (2021) for the Causal3DIdent dataset.

For the position, rotation, and hue values, we sample new values over time
with the following functions:

$\displaystyle\text{pos\\_x}^{t+1}$ $\displaystyle=$ $\displaystyle
f\left(1.5\cdot\sin(\text{rot\\_}\beta^{t}),\text{pos\\_x}^{t},\epsilon^{t}_{x}\right)$
(32) $\displaystyle\text{pos\\_y}^{t+1}$ $\displaystyle=$ $\displaystyle
f\left(1.5\cdot\sin(\text{rot\\_}\alpha^{t}),\text{pos\\_y}^{t},\epsilon^{t}_{y}\right)$
(33) $\displaystyle\text{pos\\_z}^{t+1}$ $\displaystyle=$ $\displaystyle
f\left(1.5\cdot\cos(\text{rot\\_}\alpha^{t}),\text{pos\\_z}^{t},\epsilon^{t}_{z}\right)$
(34) $\displaystyle\text{rot\\_}\alpha^{t+1}$ $\displaystyle=$ $\displaystyle
f\left(\text{hue\\_b}^{t},\text{rot\\_}\alpha^{t},\epsilon^{t}_{\alpha}\right)$
(35) $\displaystyle\text{rot\\_}\beta^{t+1}$ $\displaystyle=$ $\displaystyle
f\left(\text{hue\\_o}^{t},\text{rot\\_}\beta^{t},\epsilon^{t}_{\beta}\right)$
(36) $\displaystyle\text{rot\\_s}^{t+1}$ $\displaystyle=$ $\displaystyle
f\left(\text{atan2}(\text{pos\\_x}^{t},\text{pos\\_y}^{t}),\text{rot\\_s}^{t},\epsilon^{t}_{rs}\right)$
(37) $\displaystyle\text{hue\\_s}^{t+1}$ $\displaystyle=$ $\displaystyle
f\left(2\pi-\text{hue\\_b}^{t},\text{hue\\_s}^{t},\epsilon^{t}_{hs}\right)$
(38) $\displaystyle\text{hue\\_b}^{t+1}$ $\displaystyle=$
$\displaystyle\text{hue\\_b}^{t}+\epsilon^{t}_{b}$ (39)

where $f(a,b,c)=\frac{a-b}{2}+c$, and all $\epsilon$-variables being
independent samples from a Gaussian distribution with standard deviation $0.1$
for positions, and $0.15$ for angles and hues. Intuitively, the function $f$
represents that we create a ’goal’ position for each variable based on its
parents, and move towards the goal by taking the average between goal and
current position, with additive noise. This gives us a simulation of a moving
system, which, however, also permits large changes without interventions.

The function of the object hue depends on the categorical object shape, which
is outlined in Table 4. We use the same setup as von Kügelgen et al. (2021),
where the hare is trying to blend into the background and spotlight, while the
dragon tries to stand out. The colors of the other objects are spread out
across the color ring.

Finally, for the object shape, we use a noisy identity map over time. With a
probability of 5%, we change the current object shape with a newly sampled one
from a uniform distribution. This introduces additional noise to the object
shape besides the interventions.

Object shape | Object hue goal
---|---
Teapot | $0$
Armadillo | $\frac{1}{5}\cdot 2\pi$
Hare | $\text{avg}(\text{hue\\_spot},\text{hue\\_back})$
Cow | $\frac{2}{5}\cdot 2\pi$
Dragon | $\pi+\text{avg}(\text{hue\\_spot},\text{hue\\_back})$
Head | $\frac{3}{5}\cdot 2\pi$
Horse | $\frac{4}{5}\cdot 2\pi$
Table 4: The causal relation of the object shape, background hue, and
spotlight hue to the object hue goal $g$ with which we determine the next step
as
$\text{hue\\_o}^{t+1}=f\left(g,\text{hue\\_o}^{t},\epsilon^{t}_{ho}\right)$.
The angle mean is defined as
$\text{avg}(\alpha,\beta)=\text{atan2}\left(\frac{\sin(\alpha)+\sin(\beta)}{2},\frac{\cos(\alpha)+\cos(\beta)}{2}\right)$.

To showcase the dependency among causal variables, we plot the marginal
distribution of tuples of variables in Figure 15. The distributions are
plotted based on a histogram of a dataset with 250,000 samples. Despite the
occasional interventions, a clear correlation among variables with a
confounder and ancestor-descendant relations can be seen. Overall, this shows
that the chosen functions in Equation 32 to 39 introduce strong correlations
among variables, which makes disentangling the factors a difficult task.

(a) $p(\text{pos\\_y},\text{pos\\_z})$

(b) $p(\text{pos\\_y},\text{rot\\_}\alpha)$

(c) $p(\text{pos\\_z},\text{hue\\_b})$

Figure 15: Overview of the joint marginal distributions of selected variables
in the Temporal Causal3DIdent dataset, showcasing the correlations among
variables. The figures show histograms over the two variables, where yellow
indicates a high likelihood/frequency, while dark blue has close to zero
samples. (a) The two causal variables pos_y and pos_z share a common
confounder, $\text{rot\\_}\alpha$, which causes them to follow a circle with
radius 1.5. (b) $\text{rot\\_}\alpha$ is a parent of pos_y, and one can see
that the marginal closely follows its functional form (see Equation 33). (c)
The hue of the background is an ancestor of pos_z, with rot_$\alpha$ in
between the two variables. Yet, the marginal clearly follows the cosine
signal, showing that the correlation goes beyond parents.

###  C.2 Interventional Pong

|  |  |  |  |  |
---|---|---|---|---|---|---
(1) - | (2) - | (3) - | (4) - | (5) - | (6) ball_vel_dir | (7) ball_vel_dir
|  |  |  |  |  |
(8) pright_y | (9) - | (10) - | (11) pright_y | (12) - | (13) ball_x | (14) pright_y
|  |  |  |  |  |
(15) pleft_y | (16) - | (17) - | (18) ball_y | (19) pleft_y | (20) pleft_y | (21) ball_x
Figure 16: An example sequence with 21 frames in the Interventional Pong
dataset. The causal variables denoted below each image indicate the variables
which were intervened on at this time step. For instance, when transitioning
from the first to the second image, all variables were sampled according to
their temporal dependency except. The ball velocity direction is encoded in a
4th channel, which is not visualized here. It encodes the velocity by showing
the ball moved by the velocity at the current time step.

The Interventional Pong dataset models the dynamics of the popular Atari game
Pong (Bellemare et al., 2013). We show a sequence of 21 frames in Figure 16.
Thereby, each frame has a resolution of $32\times 32$.

In this environment, we have six underlying causal factors:

  * •

The ball x-position (ball_x) describes the x-position (left to right) of the
ball. The value range is limited to $[-1,1]$ where $-1$ corresponds to hitting
the left border, and $1$ the right. Under interventions, we sample the
x-position uniformly in the space between the two paddles, which corresponds
to a value range of approximately $[-0.7,0.7]$.

  * •

The ball y-position (ball_y) describes the y-position (top to bottom) of the
ball. The value range is limited to $[-1,1]$ where $-1$ corresponds to hitting
the upper border, and $1$ the lower border. Under interventions, we sample the
y-position uniformly from $U(-1,1)$.

  * •

The ball velocity direction (ball_vel_dir) describes the angle of the velocity
of the ball. The value range is $[0,2\pi)$ where $0$ corresponds to the ball
heading to the right. Under interventions, we sample the velocity direction
uniformly from $U(0,2\pi)$.

  * •

The paddle-left y-position (pleft_y) describes the y-position of the left
paddle (blue in the figures). It has the same value range as the causal factor
ball_y. Over time, the paddle takes a step towards the ball_y position of the
previous time step. The step size is sampled from a Gaussian mean of $0.05$
and standard deviation of $0.017$. If an intervention is performed on pleft_y,
it flips the sign of the step, _i.e._ going up instead of down, with a chance
of 50%. This replaces the policy of following the ball with a random one.

  * •

The paddle-right y-position (pright_y) describes the y-position of the right
paddle (green in the figures), and follows the same dynamics as pleft_y.

  * •

The score (score) describes the player’s score (gray background numbers). It
contains two scores (left player and right player), each with an integer value
between $0$ and $4$. When the ball hits one of the two border, the score for
the corresponding player is increased. Further, if a player scores their 5th
goal, the score is reset to $0,0$. No interventions are provided on this
causal factor.

|  |  |  |  |  |
---|---|---|---|---|---|---
Figure 17: An example sequence of 7 frames (left to right) in the
Interventional Pong dataset, where the environment is being reset at the
fourth time step. The ball is returned to the center position and gets a new
velocity direction, as well as the paddles being moved to new random
positions. The score of the respective player is increased as well.

Given a previous time step, the game engine performs the following steps in
order to determine the next step:

  1. 1.

Using the velocity direction at the previous time step, we determine the
velocity in $x$ and $y$ direction. We then move the x-position of the ball
accordingly.

  2. 2.

If this new ball-x position hits one of the two borders, we increase the score
of the corresponding player and reset the game. Resetting the game includes
setting the ball to the center of the field, randomly sampling a ball velocity
direction, and replacing the position of the paddles with a new random
position between $-0.66$ and $0.66$. The latter is done to prevent the paddles
of ending up in one of the two corners when the ball is reset. An example of
resetting the environment is shown in Figure 17.

  3. 3.

If the game was not reset, we do the following steps:

     * •

We first move the two paddles according to their dynamics, as described
before, moving towards the ball.

     * •

We move the ball y position by the velocity in y direction. If the new ball_y
position collides with the upper or lower boundary, we calculate the new
y-position it would have after the collision, and mirror the velocity angle on
the x-axis.

     * •

If the new ball x and y position collide with one of the two paddle positions
in the previous time step, we reflect the ball accordingly and calculate its
new x position, as well as mirroring the velocity angle on the y-axis.

     * •

Finally, we add Gaussian noise with standard deviation of $0.015$ to the ball
x and y position, as well as the velocity direction. The small standard
deviation was chosen to keep the dynamics similar to Pong and not divert too
much into random movements.

  4. 4.

With a chance of about 35%, we do not perform any intervention. Otherwise, we
randomly sample one out of the five causal factors (excluding the score), and
perform an intervention as described before.

We provide the full code of the dynamics and dataset generation for this Pong
environment in the supplementary. The causal graph implied by these dynamics
are shown in Figure 18. Similarly to the Temporal Causal3DIdent dataset, we
generate a dataset by sampling one long sequence with a dataset size of
$100$,$000$.

ball_xball_yball_vel_dirpleft_ypright_yscoreball_xball_yball_vel_dirpleft_ypright_yscore

Figure 18: Causal relations between all dimensions of the causal factors of
the Interventional Pong dataset. Note that although the graph is very dense,
not all relations have an effect on the causal variables at each time step.
For instance, the velocity only influences the paddles and the score when the
ball x position moved by the velocity ends up in one of the two boundaries.

###  C.3 Experimental Design

In this section, we give a more detailed description of elements used in the
experiments. We first discuss the evaluation via the correlation metrics and
triplet evaluation, and then describe the implementation details of the target
classifier.

####  C.3.1 Correlation Metrics

As described in Section 6.1, we evaluate the models by measuring the
correlation of the learned latent variables to the ground truth causal
factors. Typically, in the literature, the latent variables are one-
dimensional. In our setup, however, the latent variables span over multiple
dimensions and can describe a causal factor in a non-linear relation. Thus,
plotting the correlation per variable would not reveal the true correlation
between the learned, multidimensional latent variables and causal factors.

As a first step, we therefore need to learn a mapping between sets of latent
variables and true causal factors. We emphasize that during this whole
evaluation, the models are frozen and will not be updated. We only make use of
their encoders to map observations to latent vectors. For CITRIS, we therefore
apply one MLP per set of latent variables that, according to the learned
assignment function $\psi$, belong to the same causal factor. The MLP is then
trained to predict _all_ causal factors per set of latents, on which we
measure the correlation. To this end, we use a two-layer MLP with a latent
dimensionality of 128. The MLP size was chosen sufficiently large to uncover
any possible correlations between the latent and causal factors. Using larger
MLPs did not show any noticeable change in the correlation found, verifying
that the MLP is large enough to learn any potential non-linear mapping.

To train the MLP, we split the test dataset into two subsets: one for training
the MLP ($40\%$ of the dataset), and one for measuring the correlation metrics
on ($60\%$ of the dataset). The test datasets are generated by randomly
sampling sets of causal variables of the respective dataset. We use a dataset
size of $25$,$000$ images to prevent any form of overfitting, and ensure a
sufficiently large sample size for the correlation metrics. However, in our
datasets, the causal factors are correlated themselves. This means that the
optimal disentanglement function would yet measure a correlation from, _e.g._
, the latents describing the ball_y position to the paddles’ y positions. This
makes it difficult to spot spurious correlations between latents and causal
factors, since with sufficient latent dimensions, a model can describe the
same causal factors in multiple sets of latent dimensions. Since this would
not follow the goal of a good disentanglement, we prevent measuring such
correlations by perform this evaluation on test datasets without temporal
correlations among causal factors. In other words, we generate samples by
sampling each causal factor independently, and mapping those into the
observation space.

As a loss function for training the MLP, we use a mean-squared-error loss
(MSE) for all continuous values excluding the circular values (angles and
hues). For the circular values between 0 and $2\pi$, using the MSE loss would
be disadvantageous since predicting a value of $2\pi-\epsilon$, when the label
is $0$, causes a large error although on a unit circle, the two points would
be very close. Hence, as a better alternative, we train the MLP by predicting
a vector in two dimensions. We then use the cosine distance between the
predicted vector and the vector of the ground truth angle projected onto the
unit circle. This gives us a loss that respects the circular nature of the
causal factor. Since the loss is independent of the length of the predicted
vector, we apply a small regularizer to stabilize the vector length. Finally,
for the categorical causal factors, we use cross entropy as a training loss.

After training the MLPs, we can use them to create predictions for all causal
factors for a set of latents. Then, we determine the correlation metrics
between these predictions and the ground truth values for the individual
causal factors. Note that for multidimensional causal factors, we calculate
the metrics for each dimension individually, and take the average correlation
coefficient over dimensions afterwards.

The $R^{2}$ correlation coefficient (Spearman, 1904) is a metric which
compares the prediction of a variable from a conditional input set to the
average prediction. For continuous variables, this is:

$R^{2}=1-\frac{\sum_{i}(x_{i}-\hat{x}_{i})^{2}}{\sum_{i}(x_{i}-\bar{x})^{2}}$
(40)

where $x_{i}$ is the ground truth value of the data point $i$, $\hat{x}_{i}$
is the prediction, and $\bar{x}=\sum_{i}x_{i}/N$ is the average of the ground
truth values. For angles, we adjust the $R^{2}$ metric by replacing the
difference with the cosine distance, and take the angle mean to calculate
$\bar{x}$, since the standard mean of two angles, i.e.
$\frac{\alpha+\beta}{2}$, does not take into account the circular nature of
the angle values. The mean over more than 2 angles is defined as:

$\text{avg}(\alpha_{1},...,\alpha_{n})=\text{atan2}\left(\frac{1}{n}\sum_{i}\sin(\alpha_{i}),\frac{1}{n}\sum_{i}\cos(\alpha_{i})\right)$
(41)

with atan2 being the 2-argument $\arctan$. For the categorical variables, we
replace $\bar{x}$ with the most frequent category in the dataset, and distance
by a simple equal, _i.e._ , $\delta_{x_{i}=\hat{x_{i}}}$. While one could also
apply more suitable methods for finding correlations on categorical variables,
we found this simple adjustment sufficient for the two categorical variables
we consider in this paper (the object shape in Temporal Causal3DIdent and the
score in Pong).

The Spearman’s rank correlation (Wright, 1921) ranks both the prediction and
ground truth values, and checks the difference in ranking using the Pearson
correlation coefficient. We can apply it as is for both categorical and
continuous values, excluding angles. For the circular values (angles and
hues), we report the average Spearman’s rank of their sine and cosine
transform, since the original values do not have a unique ranking.

Both metrics result in a correlation matrix, where we have on one dimension
the sets of latent variables, and the ground truth causal variables on the
other. We show examples of these matrices in Figure 21 and Figure 22. The
optimal matrix would contain 1s where the sets of latents was assigned to its
respective causal factor by $\psi$ in CITRIS, and zero everywhere else. To
summarize such a matrix, we therefore determine the average of these matching
values (usually the diagonal). Further, to show the entanglement with other
causal factors, we determine for each set of latents the causal factor with
the highest correlation besides the matching one, and average these
correlations across latents. Since the baselines, SlowVAE and iVAE∗, do not
learn an assignment function like $\psi$, we first need to determine one
before calculating above’s metrics. We do this by first calculating all
metrics for all latent dimensions independently, and then assigning a latent
variable to the causal factor it has the highest correlation to according to
the $R^{2}$ correlation coefficient.

####  C.3.2 Triplet Evaluation

The triplet evaluation provides a parameteric-free evaluation, which can
reveal complex dependencies between latent variables in any disentanglement
model. To repeat the general approach, we create triplets of images: the first
two are randomly sampled test images, while the third one is created based on
a random combination of causal factors of the first two images. The test
images are taken from a test dataset which contains a sequence of $10$,$000$,
generated in the same fashion as the training dataset. Hence, it also contains
the common correlations between causal factors. The third image is generated
by, _e.g._ , taking the spotlight rotation and object shape from image 1, and
all other causal factors from image 2. We overall create a set of $10$,$000$
of such triplets.

For evaluation, we then encode the two test images independently with the
model that is supposed to be evaluated. We then perform the combination of
ground-truth causal factors as done for the third image in latent space. This
means that for each latent dimension $z_{i}$, we pick the value of image 1 if
the causal factor $C_{j}$, where $C_{j}$ is the factor $z_{i}$ is assigned to,
is taken from image 1 for the ground truth third image. Otherwise, we pick the
value of image 2. We repeat this procedure for every latent dimension, and
then use the decoder of the model to generate a new image. Note that for the
normalizing flow, this also includes inverting the flow before applying the
decoder.

Since the reconstruction error is not descriptive of the errors being made,
_e.g._ the rotation in Figure 7, we train an additional CNN in a supervised
manner that maps images to the causal factors. This CNN has the same
architecture as the VAE encoder (see Section C.4), and is trained on the same
training dataset while using the ground truth causal factors as labels. We use
the same training methods for angles as for the MLPs in the correlation
metrics. With this model, we can extract the causal factors from the generated
image, and report the average distance of these to the ground truth causal
factors. However, one common difference between the model’s prediction and the
ground truth is that the prediction is more blurry, and can cause surprising
mispredictions of the CNN when it focused on clear corners in the original
dataset. We found it to be beneficial for both the VAE and the AE+NF approach
to train the CausalEncoder on reconstructions from the pretrained autoencoder.
Since the autoencoder is able to reconstruct the images relatively well and
only has minor blurring effects compared to the VAE, the CNN still managed to
predict most factors very accurately and, for the Pong dataset and Teapot,
with no noticeable difference to the original ground truth data.

####  C.3.3 Target Classifier Implementation

Figure 19: Visualization of the target classifier. From the latents of the
first time step, $z^{t}$, and masked latents of the second time step,
$z_{\psi_{i}}^{t+1}$ (here shown in the mask setting), the network tries to
predict the intervention targets. If the intervention targets are independent
and the model has learned a good disentanglement, only $I_{i}^{t+1}$ should be
predicted well.

In this section, we give a detailed description of the target classifier. The
target classifier is an additional small network with the goal of predicting
the intervention targets from the latent variables over time. While one cannot
identify the intervention target 100% of the time, the distribution change
between time steps is yet often detectable. This way, the target classifier
ensures that information about all causal variables is encoded in the latent
space, while guiding the disentanglement in latent space and the assignment
function $\psi$.

In detail, the classifier takes as input the latents $z^{t}$ and a masked set
of $z^{t+1}$ that, according to a sample from $\psi$, are assigned to the
causal variable $C_{i}$. We denote this mask with
$T^{\psi_{i}}_{j}\in\\{0,1\\}^{M}$, where $T^{\psi_{i}}_{j}$ is sampled from
the Gumbel Softmax distribution, modeling the assignment of latent dimension
$j$ to the causal factor $i$. Then, the classifier is trained to recover all
intervention targets from those latents by optimizing:

$\mathcal{L}^{\varphi}_{\text{target}}=-\mathbb{E}_{T,z}\left[\sum_{i=0}^{K}\log
p_{\varphi}\left(I^{t+1}|z^{t},z^{t+1}\odot
T^{\psi}_{i},T^{\psi}_{i}\right)\right]$ (42)

where $\varphi$ are the parameters of the classifier. Intuitively, we train
the parameters $\varphi$ to predict as many intervention targets from any
input as possible, thus revealing any correlations between latents and
intervention targets.

When training the classifier for a causal factor $C_{i}$, we also take
gradients with respect to the latent variables $z^{t}$ and $z^{t+1}$ as well
as the assignment $\psi$ to optimize the target classification for
$I^{t+1}_{i}$. However, for all other target variables, $z^{t+1}\odot
T^{\psi}_{i}$ should be uninformative since the intervention target
$I_{i}^{t+1}$ only affects the causal variable $C_{i}^{t+1}$ while being
conditionally independent to all other causal variables. Hence, in this case,
the latents and $\psi$ are optimized to make the classifier’s prediction
follow the target’s marginal $p(I^{t+1}_{j}|I^{t+1}_{i})$ (or $p(I^{t+1}_{j})$
if $I_{i}^{t+1}\mathchoice{\mathrel{\hbox
to0.0pt{$\displaystyle\perp$\hss}\mkern
2.0mu{\displaystyle\perp}}}{\mathrel{\hbox
to0.0pt{$\textstyle\perp$\hss}\mkern 2.0mu{\textstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptstyle\perp$\hss}\mkern
2.0mu{\scriptstyle\perp}}}{\mathrel{\hbox
to0.0pt{$\scriptscriptstyle\perp$\hss}\mkern
2.0mu{\scriptscriptstyle\perp}}}I_{j}^{t+1}$). This way, the latents and their
assignment to causal variables are trained to only contain information from
the causal variables that they are assigned to by $\psi$, and not any others.
This training objective can be summarized as:

$\mathcal{L}^{z,\psi}_{\text{target}}=-\mathbb{E}_{T,z}\Bigg{[}\sum_{i=0}^{K}\log
p_{\varphi}\left(I_{i}^{t+1}|Tz_{i}\right)+\sum_{j\neq
i}\sum_{l\in\\{0,1\\}}p\left(I^{t+1}_{j}=l|I_{i}^{t+1}\right)\log
p_{\varphi}\left(I_{j}^{t+1}=l|I_{i}^{t+1},Tz_{i}\right)\Bigg{]}$ (43)

where $Tz_{i}=[z^{t},z^{t+1}\odot T^{\psi}_{i},T^{\psi}_{i}]$. The second part
of the objective is a cross-entropy between the marginal
$p(I^{t+1}_{j}|I^{t+1}_{i})$ and the classifier’s prediction. Hence, the full
loss for CITRIS is
$\mathcal{L}=\mathcal{L}_{\text{ELBO}}+\beta_{\text{class}}\cdot\left(\mathcal{L}^{\varphi}_{\text{target}}+\mathcal{L}^{z,\psi}_{\text{target}}\right)$,
where $\beta_{\text{class}}>0$ is a hyperparameter. In common Deep Learning
frameworks like PyTorch (Paszke et al., 2019), this two-folded loss can be
implemented by stopping the gradients of $z$ and $\psi$ for
$\mathcal{L}^{\phi}_{\text{target}}$, and $\varphi$ for
$\mathcal{L}^{z,\psi}_{\text{target}}$. We empirically verify that this simple
addition provides considerable benefits in disentanglement and overall
optimization of the VAE (see Section 6.2).

###  C.4 Hyperparameters

####  C.4.1 Network Architectures

All models are implemented in the Deep Learning framework PyTorch (Paszke et
al., 2019) and PyTorch Lightning (Falcon & The PyTorch Lightning team, 2019).
The trainings have been performed on a single NVIDIA TitanRTX GPU. Below, we
give details about the specific architectures used.

##### VAEs

For all VAE architectures, we make use of the same network architecture to
ensure a fair comparison. The encoder is a convolutional architecture,
applying several convolutions in sequence with stride 2 for decreasing the
resolution. Its output is two parameters per latent variable, which denote the
mean and standard deviation for each latent variable. Experiments with
increased complexity of the predicted distribution did not show any
improvements. For the decoder, we found it beneficial to replace transposed
convolutions with bilinear upsampling and using residual blocks. The full
architectures are outlined in Table 5. Further, we make use of Batch
Normalization (Ioffe & Szegedy, 2015) and the SiLU activation function
(Hendrycks & Gimpel, 2016; Ramachandran et al., 2017) for non-linearity.
Together, the encoder and decoder have slightly less than 1 million
parameters.

##### Autoencoders

For the standard autoencoder training, we found it to be beneficial to
increase the complexity of the decoder by applying two residual blocks per
resolution instead of one. We also experimented with this increased complexity
for the VAE approaches. However, it did not show any noticeable improvement
while requiring a smaller batch size and considerably longer training time.

Table 5: The architecture for both the encoder and decoder. Every residual
block consists of two convolutions, each 64 channels, kernel 3 and stride 1
with BatchNorm+SiLU non-linearity. The upsampling represent bilinear
upsampling of the features with a scaling factor of 2. The Tanh activation
function on the output scales the output between -1 and 1.

| Layer | Feature Dimension | Kernel | Stride | Activation Function
---|---|---|---|---|---
|  | (H $\times$ W $\times$ C) |  |  |
$f_{\textrm{enc}}$ | Conv |  32 $\times$ 32 $\times$ 64 | 3 | 2 | BatchNorm+SiLU
Conv |  32 $\times$ 32 $\times$ 64 | 3 | 1 | BatchNorm+SiLU
Conv |  16 $\times$ 16 $\times$ 64 | 3 | 2 | BatchNorm+SiLU
Conv |  16 $\times$ 16 $\times$ 64 | 3 | 1 | BatchNorm+SiLU
Conv |  8 $\times$ 8 $\times$ 64 | 3 | 2 | BatchNorm+SiLU
Conv |  8 $\times$ 8 $\times$ 64 | 3 | 1 | BatchNorm+SiLU
Conv |  4 $\times$ 4 $\times$ 64 | 3 | 2 | BatchNorm+SiLU
Conv |  4 $\times$ 4 $\times$ 64 | 3 | 1 | BatchNorm+SiLU
Reshape |  1 $\times$ 1 $\times$ 1024 | - | - | -
Linear |  1 $\times$ 1 $\times$ 256 | - | - | LayerNorm+SiLU
Linear |  1 $\times$ 1 $\times$ $2\cdot$num_latents | - | - | -
$f_{\textrm{dec}}$ | Linear |  1 $\times$ 1 $\times$ 256 | - | - | LayerNorm+SiLU
Linear |  1 $\times$ 1 $\times$ 1024 | - | - | -
Reshape |  4 $\times$ 4 $\times$ 64 | - | - | -
Upsample |  8 $\times$ 8 $\times$ 64 | - | - | -
ResidualBlock |  8 $\times$ 8 $\times$ 64 | 3 | 1 | -
Upsample |  16 $\times$ 16 $\times$ 64 | - | - | -
ResidualBlock |  16 $\times$ 16 $\times$ 64 | 3 | 1 | -
Upsample |  32 $\times$ 32 $\times$ 64 | - | - | -
ResidualBlock |  32 $\times$ 32 $\times$ 64 | 3 | 1 | -
Upsample |  64 $\times$ 64 $\times$ 64 | - | - | -
ResidualBlock |  64 $\times$ 64 $\times$ 64 | 3 | 1 | -
Pre-Activation |  64 $\times$ 64 $\times$ 64 | - | - | BatchNorm+SiLU
Conv |  64 $\times$ 64 $\times$ 64 | 1 | 1 | BatchNorm+SiLU
Conv |  64 $\times$ 64 $\times$ 3 | 1 | 1 | Tanh

##### Prior distribution

For CITRIS, the prior distribution
$p_{\phi}(z_{\psi_{i}}^{t+1}|z^{t},I_{i}^{t+1})$ is implemented using an
autoregressive model. We use a single autoregressive model that can model the
distributions for all latent sets $z_{\psi_{i}}$ by providing a mask on the
intervention vector $I^{t+1}$. The autoregressive model follow a MADE
architecture (Germain et al., 2015) where we assign 16 neurons per layer to
each latent variable, and the input to these neurons are the features of all
previous latents. Additionally, to the latents, we use the latents of the
previous layer as input. The prior is 2 layers deep, and uses the SiLU
activation function. Finally, the prior predicts a Gaussian distribution per
latent. Similarly to the encoder distribution, one can make the prior more
flexible by using normalizing flows. However, in experiments, this showed to
not provide any improvement compared to increased computational cost and
parameters. For the iVAE∗, we use a two-layer MLP with a hidden dimensionality
of 128.

##### Normalizing Flows

For the normalizing flow learned on top of the autoencoder, we found a simple
affine autoregressive flow to work well. More flexible coupling
transformations such as a MixtureCDF layer (Ho et al., 2019) did not show to
improve the results. For the autoregressive network, we use the same MADE
architecture as for the prior distribution. Further, in between the coupling
layer, we use common flow layers such as Activation Normalization and
Invertible 1x1 Convolutions (Kingma & Dhariwal, 2018).

##### Target classifier

The target classifier uses a single-layer MLP with hidden dimensionality of
128, Layer Normalization (Ba et al., 2016) and a SiLU activation function.

####  C.4.2 Training Hyperparameters

Table 6: An overview of the hyperparameter used for CITRIS-VAE, iVAE∗ and SlowVAE. This parameter configuration resulted in a training time of about 2 days on one NVIDIA TitanRTX GPU for the Temporal Causal3DIdent, and 12 hours for the Interventional Pong dataset. Hyperparameter | Value
---|---
Batch size | 512
Optimizer |  Adam (Kingma & Ba, 2015)
Learning rate | 1e-3
Learning rate scheduler | Cosine Warmup (100 steps)
KL divergence factor $\beta$ | 1.0
KL divergence factor $\psi_{0}$ ($\lambda$) | 0.01
Number of latents | 32 (Causal3DIdent), 16 (Pong)
Number of epochs | 600
Target classifier weight | 2.0
Gumbel Softmax temperature | 1.0
Table 7: An overview of the hyperparameter used for the standard autoencoder. This parameter configuration resulted in a training time of about 2 days on one NVIDIA TitanRTX GPU for the Temporal Causal3DIdent, and 12 hours for the Interventional Pong dataset. Hyperparameter | Value
---|---
Batch size | 512
Optimizer |  Adam (Kingma & Ba, 2015)
Learning rate | 1e-3
Learning rate scheduler | Cosine Warmup (100 steps)
Number of latents | 32 (Causal3DIdent), 16 (Pong)
Gaussian noise std | 0.05
Number of epochs | 1000
Table 8: An overview of the hyperparameter used for CITRIS-NF. This parameter configuration resulted in a training time of 6-10 hours on one NVIDIA TitanRTX GPU. Hyperparameter | Value
---|---
Batch size | 1024
Optimizer |  Adam (Kingma & Ba, 2015)
Learning rate | 1e-3
Learning rate scheduler | Cosine Warmup (100 steps)
KL divergence factor $\beta$ | 1.0
KL divergence factor $\psi_{0}$ ($\lambda$) | 0.01 (Causal3D), 0.1 (Pong)
Number of latents | 32 (Causal3DIdent), 16 (Pong)
Number of coupling layers | 4 (Causal3DIdent Teapot, Pong), 6 (Causal3DIdent 7-shapes)
Number of epochs | 1000
Target classifier weight | 2.0
Gumbel Softmax temperature | 1.0

We provide an overview of the used hyperparameters in Table 6 (VAE models),
Table 7 (Autoencoder training), and Table 8. In general, we use the Adam
optimizer (Kingma & Ba, 2015) with a learning rate of 1e-3. We use a Cosine
Warmup learning rate scheduler with 100 steps warmup. For the VAEs and
autoencoder, we use the MSE loss as reconstruction loss, which corresponds to
predicting a Gaussian mean per pixel with fixed standard deviation. We found
this to perform considerably better than discrete approaches such as a softmax
over the 256 pixel values. To balance the KL divergence and reconstruction
loss, we have experimented with scaling the KL divergence like in the
$\beta$-VAE (Higgins et al., 2017) in the range of
$\beta\in\\{0.25,0.5,1.0,2.0,4.0\\}$. However, the best results for all three
models were achieved when using a value of $\beta=1$. Further, we ran a grid
search over the $\gamma$ hyperparameter for SlowVAE for the values
$\gamma\in\\{1,2,4,8\\}$, with the model achieving the best disentanglement
results for $\gamma=1$. Finally, we also experimented with a KL warmup, which
slowly increases $\beta$ to its maximum value. However, no noticeable
improvement was found for all three models on the Temporal Causal3DIdent
dataset, such that we did not include it for the final experiments.

The autoencoder is trained via the MSE reconstruction loss. As mentioned in
Section 4.2, we add small Gaussian noise to the latents during training to
prevent the latent space from collapsing into Dirac deltas. We have used a
noise standard deviation of $0.05$, which did not show to negatively impact
the reconstruction loss.

In the normalizing flow, we use $4$ coupling layers for the Interventional
Pong dataset and the Temporal Causal3DIdent Teapot dataset. However, the more
challenging dataset, Temporal Causal3DIdent 7-shapes, showed to benefit from
using more coupling layers, hence using 6 here. Note that one considerable
benefit of the normalizing flow approach is that it trains much quicker than
the VAE once the autoencoder has been trained. Since one autoencoder can be
used for multiple different settings, it reduced the overall computational
cost and simplified the hyperparameter search.

##  Appendix D Additional Experiments and Results

###  D.1 Temporal Causal3DIdent

####  D.1.1 Details of the Main Results

Table 9: Experimental results for the Temporal-Causal3DIdent Teapot dataset,
including standard deviations over 3 seeds. See Table 1 for a detailed
discussion on the table and metrics.

|  Triplet evaluation distances $\downarrow$ | Correlation metrics
---|---|---
| pos_x | pos_y | pos_z | rot_$\alpha$ | rot_$\beta$ | rot_s | hue_s | hue_b | hue_o | Mean |  $R^{2}$ diag $\uparrow$ |  $R^{2}$ sep $\downarrow$ |  Spearman diag $\uparrow$ |  Spearman sep $\downarrow$
Oracle | 0.02 | 0.02 | 0.02 | 0.02 | 0.03 | 0.01 | 0.02 | 0.01 | 0.02 | 0.02 | - | - | - | -
SlowVAE | 0.13 | 0.10 | 0.12 | 0.50 | 0.59 | 0.22 | 0.64 | 0.21 | 0.17 | 0.30 | 0.65 | 0.20 | 0.62 | 0.27
(stds) |  $\pm$0.002 |  $\pm$0.004 |  $\pm$0.000 |  $\pm$0.009 |  $\pm$0.009 |  $\pm$0.005 |  $\pm$0.004 |  $\pm$0.006 |  $\pm$0.003 |  $\pm$0.003 |  $\pm$0.007 |  $\pm$0.007 |  $\pm$0.005 |  $\pm$0.012
iVAE∗ | 0.04 | 0.03 | 0.04 | 0.25 | 0.31 | 0.03 | 0.58 | 0.02 | 0.05 | 0.15 | 0.78 | 0.21 | 0.77 | 0.17
(stds) |  $\pm$0.002 |  $\pm$0.001 |  $\pm$0.002 |  $\pm$0.102 |  $\pm$0.243 |  $\pm$0.001 |  $\pm$0.178 |  $\pm$0.002 |  $\pm$0.010 |  $\pm$0.012 |  $\pm$0.036 |  $\pm$0.098 |  $\pm$0.045 |  $\pm$0.043
CITRIS-VAE | 0.05 | 0.04 | 0.05 | 0.10 | 0.20 | 0.03 | 0.08 | 0.02 | 0.05 | 0.07 | 0.96 | 0.02 | 0.95 | 0.04
(stds) |  $\pm$0.001 |  $\pm$0.001 |  $\pm$0.000 |  $\pm$0.038 |  $\pm$0.055 |  $\pm$0.000 |  $\pm$0.002 |  $\pm$0.001 |  $\pm$0.001 |  $\pm$0.010 |  $\pm$0.007 |  $\pm$0.002 |  $\pm$0.013 |  $\pm$0.002
CITRIS-NF | 0.04 | 0.03 | 0.04 | 0.06 | 0.10 | 0.03 | 0.04 | 0.01 | 0.04 | 0.04 | 0.98 | 0.01 | 0.97 | 0.05
(stds) |  $\pm$0.001 |  $\pm$0.001 |  $\pm$0.001 |  $\pm$0.002 |  $\pm$0.004 |  $\pm$0.003 |  $\pm$0.002 |  $\pm$0.004 |  $\pm$0.005 |  $\pm$0.002 |  $\pm$0.001 |  $\pm$0.004 |  $\pm$0.001 |  $\pm$0.007

Table 10: Experimental results for the Temporal-Causal3DIdent 7 shapes
dataset, including standard deviations over 3 seeds. See Table 1 for a
detailed discussion on the table and metrics.

|  Triplet evaluation distances $\downarrow$ | Correlation metrics
---|---|---
| pos_x | pos_y | pos_z | rot_$\alpha$ | rot_$\beta$ | rot_s | hue_s | hue_b | hue_o | obj_s | Mean |  $R^{2}$ diag $\uparrow$ |  $R^{2}$ sep $\downarrow$ |  Spearman diag $\uparrow$ |  Spearman sep $\downarrow$
Oracle | 0.08 | 0.06 | 0.08 | 0.06 | 0.09 | 0.04 | 0.04 | 0.01 | 0.04 | 0.00 | 0.05 | - | - | - | -
SlowVAE | 0.44 | 0.25 | 0.41 | 0.69 | 0.75 | 0.25 | 0.57 | 0.10 | 0.14 | 0.37 | 0.40 | 0.61 | 0.23 | 0.59 | 0.27
(stds) |  $\pm$0.035 |  $\pm$0.030 |  $\pm$0.021 |  $\pm$0.012 |  $\pm$0.003 |  $\pm$0.003 |  $\pm$0.064 |  $\pm$0.001 |  $\pm$0.004 |  $\pm$0.040 |  $\pm$0.018 |  $\pm$0.053 |  $\pm$0.005 |  $\pm$0.038 |  $\pm$0.006
iVAE∗ | 0.26 | 0.23 | 0.34 | 0.58 | 0.65 | 0.10 | 0.31 | 0.02 | 0.09 | 0.14 | 0.27 | 0.80 | 0.29 | 0.77 | 0.28
(stds) |  $\pm$0.013 |  $\pm$0.038 |  $\pm$0.015 |  $\pm$0.051 |  $\pm$0.043 |  $\pm$0.011 |  $\pm$0.003 |  $\pm$0.000 |  $\pm$0.001 |  $\pm$0.039 |  $\pm$0.002 |  $\pm$0.002 |  $\pm$0.013 |  $\pm$0.001 |  $\pm$0.019
CITRIS-VAE | 0.15 | 0.13 | 0.23 | 0.54 | 0.71 | 0.07 | 0.05 | 0.02 | 0.06 | 0.18 | 0.21 | 0.89 | 0.10 | 0.88 | 0.12
(stds) |  $\pm$0.002 |  $\pm$0.000 |  $\pm$0.001 |  $\pm$0.009 |  $\pm$0.009 |  $\pm$0.002 |  $\pm$0.002 |  $\pm$0.000 |  $\pm$0.002 |  $\pm$0.050 |  $\pm$0.004 |  $\pm$0.001 |  $\pm$0.007 |  $\pm$0.002 |  $\pm$0.006
CITRIS-NF | 0.12 | 0.08 | 0.11 | 0.09 | 0.14 | 0.05 | 0.05 | 0.02 | 0.06 | 0.00 | 0.07 | 0.98 | 0.04 | 0.97 | 0.08
(stds) |  $\pm$0.001 |  $\pm$0.000 |  $\pm$0.001 |  $\pm$0.000 |  $\pm$0.004 |  $\pm$0.001 |  $\pm$0.001 |  $\pm$0.004 |  $\pm$0.002 |  $\pm$0.000 |  $\pm$0.001 |  $\pm$0.000 |  $\pm$0.003 |  $\pm$0.000 |  $\pm$0.007

In this section, we provide additional results on the Temporal Causal3DIdent
dataset. Firstly, we report the standard deviations over 3 seeds for the
results in Table 1, which can be found in Table 9 and Table 10. Generally,
CITRIS-VAE and CITRIS-NF showed to be stable across seeds and reach similar
performance with different seeds.

For the teapot experiments, the greatest difference is found in the modeling
of the rotation angles, since the reconstruction error of a VAE can already be
optimized well when modeling the teapot as a sphere. Furthermore, the most
common failure mode in the rotations is having a difference of 180 degrees,
which increases the triplet distance considerably. The reason is yet again the
form of the teapot. Rotating the teapot by 180 degrees in the $\beta$ angle
switches the handle with the spout, which yet again gives a lower
reconstruction error than other rotations. CITRIS-NF suffered much less from
this problem since the disentanglement is performed independently of the
reconstruction error, and the pretrained autoencoder precisely modeled the
different rotations.

The iVAE∗, on the other hand, showed a more instable behavior. Especially, the
hue of the spotlight was very often entangled with different dimensions such
as the hue of the object or background. Similarly to CITRIS-VAE, it struggled
most with the rotation angles, which also had a high variance in the triplet
evaluation. Finally, the SlowVAE showed to be stable across seeds, but achieve
lower disentanglement than the other models due to assuming independent causal
factors, which was not the case here.

In the experiments on 7 shapes, the most common entanglement was the object
shape with the rotation angle. For modeling a change in the object shape
independent of the rotation angle, the models need to learn the default
rotation angles for each object. Otherwise, one cannot align the different
shapes correctly. Furthermore, since having multiple object shapes
considerably increased the modeling complexity, the VAE models often had
blurry predictions on which the rotation could not be determined. Thus,
CITRIS-VAE was able to disentangle most factors well, except the rotation and
the object shape. However, CITRIS-NF did not have this difficulty due to
relying on an autoencoder, which was trained independently of any latent space
prior regularization. This is why CITRIS-NF yet disentangled the different
causal factors well, even the rotation. We show some examples of the triplet
generation of CITRIS-NF on this dataset in Figure 20.

|  |  |
---|---|---|---
Image 1 | Image 2 | Ground Truth | Prediction
|  |  |
Image 1 | Image 2 | Ground Truth | Prediction
|  |  |
Image 1 | Image 2 | Ground Truth | Prediction
|  |  |
Image 1 | Image 2 | Ground Truth | Prediction

Row 1 | Image 1 | rot_s, obj_s
---|---|---
Image 2 | pos_o, rot_o, hue_o, hue_s, hue_b
Row 2 | Image 1 | rot_o, rot_s, hue_o, hue_s
Image 2 | pos_o, rot_o, hue_b
Row 3 | Image 1 | pos_o, rot_s, hue_b, obj_s
Image 2 | rot_o, hue_o, hue_s
Row 4 | Image 1 | rot_s, hue_o, obj_s
Image 2 | pos_o, rot_o, hue_b, hue_s

Figure 20: Visualizations of triplet generations. Each row represents one
example of the triplet evaluation. The table on the right shows which causal
factors were combines from Image 1 and 2 respectively. The predictions were
generated by CITRIS-NF trained on all 7 shapes. Despite the blurriness of some
predictions, the model still clearly identifies the correct causal factors.
Table 11: Experimental results for the Temporal-Causal3DIdent 5 shapes dataset
with additional testing on 2 unknown shapes (Cow and Head), including standard
deviations over 3 seeds for CITRIS-NF. See Table 2 for a detailed discussion
on the table and metrics. Further, we provide iVAE∗ as a baseline being
trained on the same setting.

|  Triplet evaluation distances $\downarrow$ | Correlation metrics
---|---|---
| pos_x | pos_y | pos_z | rot_$\alpha$ | rot_$\beta$ | rot_s | hue_s | hue_b | hue_o | obj_s | Mean |  $R^{2}$ diag $\uparrow$ |  $R^{2}$ sep $\downarrow$ |  Spearman diag $\uparrow$ |  Spearman sep $\downarrow$
Oracle | 0.08 | 0.06 | 0.08 | 0.06 | 0.09 | 0.04 | 0.04 | 0.01 | 0.04 | 0.00 | 0.05 | - | - | - | -
CITRIS-NF |  |  |  |  |  |  |  |  |  |  |  |  |  |  |
5 seen shapes | 0.14 | 0.10 | 0.14 | 0.12 | 0.16 | 0.06 | 0.06 | 0.02 | 0.07 | 0.01 | 0.09 | 0.98 | 0.05 | 0.97 | 0.10
(stds) |  $\pm$0.003 |  $\pm$0.004 |  $\pm$0.006 |  $\pm$0.010 |  $\pm$0.012 |  $\pm$0.002 |  $\pm$0.001 |  $\pm$0.002 |  $\pm$0.002 |  $\pm$0.000 |  $\pm$0.004 |  $\pm$0.001 |  $\pm$0.002 |  $\pm$0.000 |  $\pm$0.003
2 unseen shapes | 0.32 | 0.26 | 0.35 | 0.36 | 0.52 | 0.12 | 0.10 | 0.03 | 0.11 | 0.10 | 0.23 | 0.94 | 0.15 | 0.93 | 0.19
(stds) |  $\pm$0.026 |  $\pm$0.027 |  $\pm$0.021 |  $\pm$0.037 |  $\pm$0.048 |  $\pm$0.007 |  $\pm$0.009 |  $\pm$0.005 |  $\pm$0.008 |  $\pm$0.033 |  $\pm$0.018 |  $\pm$0.003 |  $\pm$0.012 |  $\pm$0.005 |  $\pm$0.012
iVAE∗ |  |  |  |  |  |  |  |  |  |  |  |  |  |  |
5 seen shapes | 0.26 | 0.18 | 0.29 | 0.66 | 0.68 | 0.13 | 0.25 | 0.03 | 0.08 | 0.15 | 0.27 | 0.80 | 0.21 | 0.81 | 0.25
2 unseen shapes | 0.37 | 0.28 | 0.70 | 0.95 | 0.94 | 0.15 | 0.24 | 0.03 | 0.09 | 0.86 | 0.46 | 0.69 | 0.19 | 0.68 | 0.25
CITRIS-NF |  |  |  |  |  |  |  |  |  |  |  |  |  |  |
\- Cow shape | 0.32 | 0.25 | 0.36 | 0.38 | 0.53 | 0.10 | 0.09 | 0.03 | 0.10 | 0.09 | 0.22 | 0.93 | 0.14 | 0.92 | 0.22
(stds) |  $\pm$0.030 |  $\pm$0.024 |  $\pm$0.017 |  $\pm$0.035 |  $\pm$0.049 |  $\pm$0.002 |  $\pm$0.006 |  $\pm$0.004 |  $\pm$0.004 |  $\pm$0.013 |  $\pm$0.015 |  $\pm$0.004 |  $\pm$0.010 |  $\pm$0.004 |  $\pm$0.009
\- Head shape | 0.32 | 0.27 | 0.35 | 0.33 | 0.51 | 0.13 | 0.11 | 0.04 | 0.12 | 0.12 | 0.23 | 0.94 | 0.15 | 0.93 | 0.16
(stds) |  $\pm$0.023 |  $\pm$0.029 |  $\pm$0.026 |  $\pm$0.038 |  $\pm$0.047 |  $\pm$0.012 |  $\pm$0.011 |  $\pm$0.005 |  $\pm$0.012 |  $\pm$0.052 |  $\pm$0.021 |  $\pm$0.001 |  $\pm$0.014 |  $\pm$0.006 |  $\pm$0.015
iVAE∗ |  |  |  |  |  |  |  |  |  |  |  |  |  |  |
\- Cow shape | 0.46 | 0.27 | 0.98 | 0.97 | 0.96 | 0.14 | 0.25 | 0.03 | 0.09 | 0.90 | 0.51 | 0.69 | 0.17 | 0.68 | 0.26
\- Head shape | 0.27 | 0.29 | 0.42 | 0.95 | 0.94 | 0.16 | 0.24 | 0.03 | 0.10 | 0.83 | 0.42 | 0.69 | 0.20 | 0.69 | 0.25

In correspondence to the generalization experiments on the Temporal-
Causal3DIdent dataset (Table 2), we show a detailed version of the results in
Table 11. Overall, the generalization performance shows to be stable across
the 2 unseen shapes, Cow (Crane, 2021) and Head (Rusinkiewicz et al., 2021).
The most difficult causal factors to generalize are the position and rotation,
since both of them are heavily dependent on the object shape and hence also
entangled in the autoencoder’s representation. Nonetheless, the performance
well above random holds promise for future work on exploring these
generalization capabilities. Furthermore, we provide a baseline in Table 2 by
training an iVAE∗ on the same setting. To reduce computational cost and
previous experiments showing stable performance across seeds, we run the iVAE∗
for a single seed. As one would expect, iVAE∗ heavily struggles with
generalizing to unseen shapes since it’s en- and decoder have not been trained
on these shapes and the new shapes require non-trivial extrapolation. Hence,
the model generates arbitrary shapes for the new objects, from which no
rotations can be inferred. While the factors independent of the object shape
obtain similar performance as for the known shapes, the model’s overall
disentanglement performance degrades considerably. This underlines the promise
of CITRIS-NF to generalize to unseen instances of causal factors.

####  D.1.2 Correlation Matrices

|
---|---
(a) CITRIS-NF, $R^{2}$ correlation matrix | (b) CITRIS-NF, Spearman correlation matrix
|
(c) CITRIS-VAE, $R^{2}$ correlation matrix | (d) CITRIS-VAE, Spearman correlation matrix
|
(e) iVAE∗, $R^{2}$ correlation matrix |  (f) iVAE∗, Spearman correlation matrix
|
(e) SlowVAE, $R^{2}$ correlation matrix | (f) SlowVAE, Spearman correlation matrix
Figure 21: Correlation matrices for the experiments on the Temporal-Causal3DIdent Teapot dataset. The y-axis shows the sets of latent dimensions that were assigned to a certain causal factor. The set $z_{\psi_{0}}$ is represented by ’no variable’ in the plots of CITRIS. The x-axis shows the ground truth causal factors with all dimensions, _i.e._ , pos_o represented by pos_x, pos_y, pos_z. The heatmap is the correlation matrix between those factors ($R^{2}$ left, Spearman right). |
---|---
(a) CITRIS-NF, $R^{2}$ correlation matrix | (b) CITRIS-NF, Spearman correlation matrix
|
(c) CITRIS-VAE, $R^{2}$ correlation matrix | (d) CITRIS-VAE, Spearman correlation matrix
|
(e) iVAE∗, $R^{2}$ correlation matrix |  (f) iVAE∗, Spearman correlation matrix
|
(e) SlowVAE, $R^{2}$ correlation matrix | (f) SlowVAE, Spearman correlation matrix
Figure 22: Correlation matrices for the experiments on the Temporal-
Causal3DIdent 7-shapes dataset. The y-axis shows the sets of latent dimensions
that were assigned to a certain causal factor. The set $z_{\psi_{0}}$ is
represented by ’no variable’ in the plots of CITRIS. The x-axis shows the
ground truth causal factors with all dimensions, _i.e._ , pos_o represented by
pos_x, pos_y, pos_z. The heatmap is the correlation matrix between those
factors ($R^{2}$ left, Spearman right).

The results of Table 1 summarized the correlation matrices by reporting the
mean on the diagonal and the maximum for any other causal factor. In this
section, we additionally show examples of full correlation matrices for all
models on the Temporal Causal3DIdent dataset. Figure 21 shows the results for
the Temporal Causal3dIdent Teapot dataset. The x-axis shows all dimensions of
all causal factors, which includes the $x$, $y$ and $z$ dimensions of the
position causal factor. The optimal model would have 1s between pos_o and
pos_x, pos_y, pos_z. The figure shows the difficulty with the rotation for all
models. Further, in the iVAE∗, one can see the difficulty the model had with
the hue of the spotlight.

The results for the 7-shapes datasets are shown in Figure 22. As one could
already see from the results in Table 1, only CITRIS-NF is able to disentangle
all causal factors well while the other models have especially difficulties
with the rotation.

####  D.1.3 Confounded Interventions

Table 12: Experimental results for the Temporal-Causal3DIdent Teapot dataset
with a limited intervention set, including standard deviations over 3 seeds.
CITRIS performs very similar as in the setting of independent interventions
(Table 9), showing that it can handle such intervention sets as well.

|  Triplet evaluation distances $\downarrow$ | Correlation metrics
---|---|---
| pos_x | pos_y | pos_z | rot_$\alpha$ | rot_$\beta$ | rot_s | hue_s | hue_b | hue_o | Mean |  $R^{2}$ diag $\uparrow$ |  $R^{2}$ sep $\downarrow$ |  Spearman diag $\uparrow$ |  Spearman sep $\downarrow$
Oracle | 0.02 | 0.02 | 0.02 | 0.02 | 0.03 | 0.01 | 0.02 | 0.01 | 0.02 | 0.02 | - | - | - | -
SlowVAE | 0.10 | 0.08 | 0.10 | 0.39 | 0.56 | 0.14 | 0.61 | 0.07 | 0.10 | 0.24 | 0.66 | 0.27 | 0.65 | 0.23
(stds) |  $\pm$0.007 |  $\pm$0.003 |  $\pm$0.001 |  $\pm$0.078 |  $\pm$0.068 |  $\pm$0.000 |  $\pm$0.002 |  $\pm$0.003 |  $\pm$0.003 |  $\pm$0.016 |  $\pm$0.016 |  $\pm$0.014 |  $\pm$0.012 |  $\pm$0.008
iVAE∗ | 0.10 | 0.07 | 0.09 | 0.23 | 0.37 | 0.06 | 0.31 | 0.02 | 0.06 | 0.15 | 0.81 | 0.18 | 0.80 | 0.18
(stds) |  $\pm$0.002 |  $\pm$0.001 |  $\pm$0.001 |  $\pm$0.124 |  $\pm$0.213 |  $\pm$0.001 |  $\pm$0.268 |  $\pm$0.001 |  $\pm$0.014 |  $\pm$0.006 |  $\pm$0.045 |  $\pm$0.111 |  $\pm$0.042 |  $\pm$0.068
CITRIS-VAE | 0.05 | 0.03 | 0.05 | 0.09 | 0.24 | 0.03 | 0.05 | 0.01 | 0.05 | 0.07 | 0.97 | 0.03 | 0.96 | 0.04
(stds) |  $\pm$0.002 |  $\pm$0.001 |  $\pm$0.002 |  $\pm$0.026 |  $\pm$0.091 |  $\pm$0.000 |  $\pm$0.005 |  $\pm$0.000 |  $\pm$0.002 |  $\pm$0.014 |  $\pm$0.001 |  $\pm$0.002 |  $\pm$0.006 |  $\pm$0.007
CITRIS-NF | 0.05 | 0.03 | 0.04 | 0.04 | 0.07 | 0.03 | 0.04 | 0.02 | 0.04 | 0.04 | 0.98 | 0.05 | 0.98 | 0.08
(stds) |  $\pm$0.002 |  $\pm$0.000 |  $\pm$0.001 |  $\pm$0.001 |  $\pm$0.003 |  $\pm$0.001 |  $\pm$0.001 |  $\pm$0.000 |  $\pm$0.001 |  $\pm$0.001 |  $\pm$0.002 |  $\pm$0.015 |  $\pm$0.002 |  $\pm$0.023

The previous experiments on the Temporal Causal3DIdent dataset were performed
on interventions that were independently sampled to show that CITRIS can
handle both single-target and joint interventions. In this section, we repeat
the experiments on the Teapot dataset, but with a different intervention
settings. Instead of sampling each intervention target independently, we
consider 6 possible intervention sets:

  1. 1.

The observational regime, _i.e._ , no causal variable being intervened upon

  2. 2.

A single target intervention on the object position (pos_o)

  3. 3.

A single target intervention on the object rotation (rot_o)

  4. 4.

A joint intervention on the rotation and hue of the spotlight (rot_s, hue_s)

  5. 5.

A joint intervention on the hue of the object and spotlight (hue_o, hue_s)

  6. 6.

A joint intervention on the hue of the object and background (hue_o, hue_b)

Note that this set of interventions fulfills the intervention conditions of
Theorem 3.3 for all causal variables. At each time step, we perform one of
these interventions, and use the same temporal dependencies and dynamics as
before.

The results over 3 seeds for these experiments are summarized in Table 12. All
models achieved very similar results as for the setting with independent
intervention targets. This verifies that CITRIS can also disentangle the
different causal factors in a setting with limited diversity in the observed
intervention settings.

####  D.1.4 Learning the Causal Graph

pos_orot_orot_shue_shue_bhue_opos_orot_orot_shue_shue_bhue_o

pos_orot_orot_shue_shue_bhue_oobj_spos_orot_orot_shue_shue_bhue_oobj_s

Figure 23: The learned graphs of CITRIS-NF from the Temporal-Causal3DIdent
dataset. False positive edges are colored red. Left: Temporal-Causal3DIdent
Teapot. The learned causal graph is identical to the ground truth graph. Since
we have a single object shape, we do not include obj_s as a causal variable in
the graph. Additionally, the parent set of hue_o is reduced due to the
constant object shape. Right: Temporal-Causal3DIdent 7-shapes. The edge
obj_s$\to$rot_o is predicted due to disentangling these two variables being
the hardest task of the dataset, and small correlations remain.

The transition prior $p_{\phi}$ of CITRIS takes as input all latent variables
from the previous time-step, $z^{t}$, in order to model a distribution over
$z^{t+1}$. Once the model has converged, we can try to learn a sparser prior,
which removes unnecessary edges between sets of latent variables that have
been assigned to different causal factors. More specifically, we are trying to
find the causal graph between
$z_{\Psi_{0}}^{t},z_{\Psi_{1}}^{t},...,z_{\Psi_{K}}^{t}$ and
$z_{\Psi_{0}}^{t+1},z_{\Psi_{1}}^{t+1},...,z_{\Psi_{K}}^{t+1}$, where the
directions of the directions of the edges are determined by $t\to t+1$, and no
edges within a time step exist. We find this graph by using ENCO (Lippe et
al., 2022), a continuous-optimization causal discovery method which supports
the usage of arbitrary neural networks. The causal graph is parameterized by a
weight matrix $\bm{\gamma}\in\mathbb{R}^{K+1\times K+1}$, where
$\sigma(\gamma_{ij})$ represents the probability of having the edge
$z_{\Psi_{i}}^{t}\to z_{\Psi_{j}}^{t+1}$ in the causal graph. To estimate the
conditional likelihoods under different causal graphs, we learn a new
autoregressive prior $p_{\varepsilon}(z^{t+1}_{\Psi_{i}}|z^{T}\cdot
M,I_{i}^{t+1})$ where $M$ is mask on the latents at time-step $t$ according to
the causal graph:
$M_{j}=\hat{M}_{\psi(j)},\hat{M}_{k}\sim\text{Bernoulli}(\gamma_{ki})$. In
other words, we sample a causal graph, and then mask a set of latents
$z^{t}_{\Psi_{j}}$ if the causal variable $C_{j}$ does not have an edge to
$C_{i}$: $z_{\Psi_{j}}^{t}\not\to z_{\Psi_{i}}^{t+1}$. With this prior, we can
then determine the gradients for $\bm{\gamma}$ on a held-out test set
following ENCO’s gradient updates. We use a sparsity regularizer of
$\lambda=0.05$ and train for 100 epochs.

We perform the experiments on two learned CITRIS-NF, one for Temporal-
Causal3DIdent Teapot, and one for Temporal-Causal3DIdent 7-shapes. The learned
graphs are shown in Figure 23. For the teapot experiments, the graph is
identical to the ground truth, showing that CITRIS has indeed learned the
causal variables. For the 7 shapes experiment, the model has one false
positive prediction, the edge obj_s$\to$rot_o. This fits our previous
discussion on the results of the 7 shape experiments (Table 10) since the
object shape and rotation was the most difficult to disentangle for all
models. Hence, there remains some correlation between the object shape and
rotation in the CITRIS-NF model.

####  D.1.5 Interventions on a Subset of Variables

|
---|---
(a) CITRIS-NF, $R^{2}$ correlation matrix | (b) CITRIS-NF, Spearman correlation matrix
Figure 24: Correlation matrices for the experiments on the Temporal-Causal3DIdent 7-shapes dataset where no intervention targets for the causal variable hue-background are shown (left: $R^{2}$ correlation, right: Spearman correlation). The set $z_{\psi_{0}}$ is represented by ’no variable’ in the plots of CITRIS. As intended, CITRIS learns to map the background hue to $z_{\psi_{0}}$, i.e. the latents that represent all information that does not belong to the causal variables for which we are given intervention targets. The remaining causal variables are disentangled as well as when interventions were provided for all variables (compare to Figure 22). |
---|---
(a) Excludes rot_s, $R^{2}$ correlation matrix | (b) Excludes rot_s, Spearman correlation matrix
|
(c) Excludes hue_s, $R^{2}$ correlation matrix | (d) Excludes hue_s, Spearman correlation matrix
Figure 25: Correlation matrices for the experiments on the Temporal-
Causal3DIdent Teapot dataset where no intervention targets for the causal
variable rotation-spotlight (a,b) and hue-spotlight (c,d) are shown (left:
$R^{2}$ correlation, right: Spearman correlation). The set $z_{\psi_{0}}$ is
represented by ’no variable’ in the plots of CITRIS. CITRIS assigns the
variables without intervention targets to $z_{\psi_{0}}$, and the remaining
causal variables only show slightly worse disentanglement than for the full
intervention set (compare to Figure 21).

One property of CITRIS is that when interventions are only provided for a
subset of variables, CITRIS is yet able to disentangle these variables, while
all remaining information is grouped into the latents $z_{\Psi_{0}}$. We have
used this to disentangle the score in the Interventional Pong experiments (see
Section 6.3), and show here that it can also be applied to the more complex
Temporal Causal3DIdent dataset. Specifically, we reuse the datasets that have
been used for the experiments in Section 6.2, and remove the intervention
targets for a subset of variables. This keeps all temporal relations intact,
while providing an additional challenge of modeling a multimodal distribution
for the excluded variables.

As a first setting, we train CITRIS-NF on the Temporal Causal3DIdent 7-shapes
dataset, where the interventions on the hue of the background have been
omitted. The correlation matrices in Figure 24 show that CITRIS learned to
disentangle all causal variables in a similar accuracy as in the full
experiment (see Figure 22), while the background hue is assigned to the
latents of $z_{\psi_{0}}$. Additionally, we show in Figure 21 results on the
Temporal Causal3DIdent Teapot dataset, in which we exclude intervention
targets for the spotlight rotation (a,b) and spotlight hue (c,d) respectively.
Both models assign the correct information to the ’no-variable’ slot. The
disentanglement for all other variables shows to still be very good, but
slightly worse than the results on the full intervention set, which comes from
additional entanglement in $z_{\psi_{0}}$. Still, the experiments show that
CITRIS can also handle variables without interventions in complex 3D rendered
scenes.

####  D.1.6 Visualizing the Learned Assignment Function $\psi$

Figure 26: Learned latent-to-causal variable assignment map $\psi$ for an
CITRIS-NF trained on the Temporal Causal3DIdent 7-shapes datasets with no
interventions on the hue of the background. The map is learned via a Gumbel-
Softmax, meaning that we have a softmax distribution over the possible causal
variables for each latent variable. Thus, the figure shows the probability
that latent dimension $i$ (x-axis) is assigned to a specific causal variable
(y-axis), where the distribution usually gets strongly peaked as training
progresses. The model has 32 dimensions, of which 5 are assigned to pos_o, 7
to rot_o, 4 to rot_s, 2 to hue_o, 3 to hue_s, 6 to obj_s, and 5 to ’no
variable’, i.e. $z_{\psi_{0}}$. In general, we find that the most complex
causal variable, here rot_o and obj_s, are assigned the most dimensions, while
circular values have always at least 2 dimensions (the model commonly encodes
them as two or more shifted sine-waves).

CITRIS learns an assignment between latent and causal variable, $\psi$. To get
an intuition of what a learned function of such looks like, we visualize the
parameters of $\psi$ of a trained CITRIS-NF in Figure 26. The model is trained
on the Temporal Causal3DIdent 7-shapes dataset, where no interventions have
been shown for the background hue (i.e., the same model as in Figure 24).
Overall, the probabilities of a latent variable belonging to a certain causal
variable are very peaked after training, with the second-highest value per
latent variable to be between $1e$-$6$ and $1e$-$9$. This shows that the soft
relaxation of $\psi$ becomes discrete as training progresses. Further, we find
that the most latent variables are assigned to the intuitively most complex
causal variables. In the Temporal Causal3DIdent 7-shapes dataset, this
corresponds to the rotation of the object and the shape of the object,
assigned 7 and 6 latent respectively. The property of learning a flexible
latent space size per causal variable is thereby important, since when the
causal variables are not known, it would be difficult to manually estimate the
number of latent variables per causal variable. Moreover, an interesting
property we observe across all models is that circular values like the hues
are always encoded in at least two latent dimensions. On closer inspection,
the model encodes the angles as two or more shifted sine-waves, verifying our
motivation of using multidimensional latent representation even for single-
dimensional causal variables.

###  D.2 Interventional Pong

Table 13: Experimental results for the Interventional Pong dataset, extension
of Table 3. The six causal factors are the ball x and y position, the velocity
direction, the left paddle y position, the right paddle y position, and the
player’s score. In the correlation metrics, diag refers to the average score
on its diagonal (optimal 1), and sep for the average of the maximum
correlation per causal variable besides itself (optimal 0).

|  Triplet evaluation distances $\downarrow$ | Correlation metrics
---|---|---
| ball_x | ball_y | ball_vel_dir | pleft_y | pright_y | score | Mean |  $R^{2}$ diag $\uparrow$ |  $R^{2}$ sep $\downarrow$ |  Spearman diag $\uparrow$ |  Spearman sep $\downarrow$
Oracle | 0.01 | 0.01 | 0.01 | 0.01 | 0.01 | 0.00 | 0.01 | - | - | - | -
SlowVAE | 0.06 | 0.21 | 0.99 | 0.48 | 0.07 | 0.23 | 0.34 | 0.61 | 0.17 | 0.66 | 0.23
(stds) |  $\pm$0.006 |  $\pm$0.012 |  $\pm$0.002 |  $\pm$0.121 |  $\pm$0.089 |  $\pm$0.001 |  $\pm$0.001 |  $\pm$0.004 |  $\pm$0.012 |  $\pm$0.007 |  $\pm$0.013
iVAE∗ | 0.03 | 0.21 | 0.05 | 0.03 | 0.03 | 0.20 | 0.09 | 0.91 | 0.04 | 0.92 | 0.06
(stds) |  $\pm$0.008 |  $\pm$0.255 |  $\pm$0.005 |  $\pm$0.011 |  $\pm$0.009 |  $\pm$0.015 |  $\pm$0.039 |  $\pm$0.135 |  $\pm$0.056 |  $\pm$0.131 |  $\pm$0.058
CITRIS-VAE | 0.02 | 0.03 | 0.05 | 0.04 | 0.03 | 0.00 | 0.03 | 0.99 | 0.01 | 0.99 | 0.05
(stds) |  $\pm$0.001 |  $\pm$0.004 |  $\pm$0.009 |  $\pm$0.006 |  $\pm$0.002 |  $\pm$0.001 |  $\pm$0.001 |  $\pm$0.001 |  $\pm$0.000 |  $\pm$0.001 |  $\pm$0.010
CITRIS-NF | 0.02 | 0.02 | 0.02 | 0.03 | 0.03 | 0.00 | 0.02 | 1.00 | 0.04 | 1.00 | 0.10
(stds) |  $\pm$0.002 |  $\pm$0.002 |  $\pm$0.001 |  $\pm$0.002 |  $\pm$0.001 |  $\pm$0.001 |  $\pm$0.001 |  $\pm$0.001 |  $\pm$0.011 |  $\pm$0.001 |  $\pm$0.023

|
---|---
(a) CITRIS-NF, $R^{2}$ correlation matrix | (b) CITRIS-NF, Spearman correlation matrix
|
(c) CITRIS-VAE, $R^{2}$ correlation matrix | (d) CITRIS-VAE, Spearman correlation matrix
|
(e) iVAE∗, $R^{2}$ correlation matrix |  (f) iVAE∗, Spearman correlation matrix
|
(e) SlowVAE, $R^{2}$ correlation matrix | (f) SlowVAE, Spearman correlation matrix
Figure 27: Correlation matrices for the experiments on the Interventional Pong
dataset. The y-axis shows the sets of latent dimensions that were assigned to
a certain causal factor. The set $z_{\psi_{0}}$ is represented by ’no
variable’ in the plots of CITRIS, where no interventions on score_left and
score_right have been provided. The x-axis shows the ground truth causal
factors. The heatmap is the correlation matrix between those factors ($R^{2}$
left, Spearman right).

The full experimental results for the interventional Pong dataset can be found
in Table 13, and we plot the correlation matrices for all models in Figure 27.
The triplet evaluation distances show that SlowVAE had difficulties with
disentangling all the factors. This is since they are highly correlated, while
SlowVAE assumes independence in the factors. Both the iVAE∗ and SlowVAE had
troubles with disentangling the y position of the ball (ball_y), which was
often entangled together with the ball velocity direction (ball_vel_dir) in
the models. Further, the player’s score was only consistently recovered by
CITRIS. Note that for the triplet evaluation, since there are no interventions
given on the player’s score, we use the variables $z_{\psi_{0}}$ of CITRIS of
that image of which the score it used for the third image. CITRIS-NF slightly
improves upon CITRIS-VAE, mostly due to better modeling of ball_vel_dir. This
is because it only affects a single channel and hence has a lower effect in
the reconstruction error. Nonetheless, the experiments show that for simpler
datasets, CITRIS-VAE can disentangle the causal factors as well as the
normalizing flow.

###  D.3 Ball-in-Boxes Example

|  |  |  |  |  |
---|---|---|---|---|---|---
(1) - | (2) - | (3) ball_y | (4) ball_x | (5) ball_y | (6) - | (7) -
|  |  |  |  |  |
(8) - | (9) - | (10) ball_x | (11) ball_x | (12) - | (13) - | (14) -
Figure 28: An example sequence with 14 frames in the Ball-in-Boxes dataset.
The causal variables denoted below each image indicate the variables which
were intervened on at this time step, where an intervention on ball_x changes
the box identifier (left or right), and ball_y the y position of the ball.

Throughout the paper, we have referred to the example of a ball in a 2D space,
which can be in one of two boxes. It can move freely within a box, but cannot
switch between them by itself. Only under interventions on its $x$ position,
we switch the box. However, this intervention does not influence the internal
position of the ball in the box. Under our theoretical setup of CITRIS, the
box identifier represents the minimal causal variable of $x$, while the
internal position should be modeled in $z_{\psi_{0}}$.

In this section, we give empirical evidence that this is indeed also the case
in practice. For simplicity, we implement it by generating frames of $32\times
32$ resolution similar to our Pong environment (see Figure 28). However, the
dynamics are replaced by the ball-in-boxes example: we consider two causal
variables, ball_x and ball_y. Both ball_x and ball_y follow a truncated
Gaussian distribution, _i.e._ move randomly with Gaussian noise over time in
the boundaries of the box the ball is in. Under interventions on ball_y, we
uniformly resample the position of the ball. Under interventions on ball_x, we
switch the box without influencing the internal ball position. We generate the
dataset by using a sequence of $100,000$ frames, and sample the intervention
targets from $I^{t}_{i}\sim\text{Bernoulli}(0.2)$. On this dataset, we apply
CITRIS-VAE with the same experimental setup as for the Pong dataset, including
the small regularizer $\lambda=0.01$, to show that little change in the KL
divergence already encourages the disentanglement to $z_{\psi_{0}}$.

Table 14: Experimental results for the Ball-in-Boxes dataset. The three
evaluation causal factors are the box identifier ball_b, the internal ball x
position ball_xin, and the ball y position. In the correlation metrics, diag
refers to the average score on its diagonal (optimal 1), and sep for the
average of the maximum correlation per causal variable besides itself (optimal
0). For the correlation metrics, we consider $z_{\psi_{0}}$ to be the latent
variables in which ball_xin should be modeled, _i.e._ have a high correlation
with. From the results, one can clearly see that CITRIS has indeed learned to
disentangle ball_b and ball_xin, showing that CITRIS learns the minimal causal
variables.

|  Triplet evaluation distances $\downarrow$ | Correlation metrics
---|---|---
| ball_b | ball_y | ball_xin | Mean |  $R^{2}$ diag $\uparrow$ |  $R^{2}$ sep $\downarrow$ |  Spearman diag $\uparrow$ |  Spearman sep $\downarrow$
Oracle | 0.00 | 0.01 | 0.01 | 0.01 | - | - | - | -
CITRIS-VAE | 0.00 | 0.01 | 0.02 | 0.01 | 0.99 | 0.00 | 1.00 | 0.02
(stds) |  $\pm$0.001 |  $\pm$0.003 |  $\pm$0.005 |  $\pm$0.002 |  $\pm$0.001 |  $\pm$0.001 |  $\pm$0.001 |  $\pm$0.002

|
---|---
(a) CITRIS-VAE, $R^{2}$ correlation matrix | (b) CITRIS-VAE, Spearman correlation matrix
Figure 29: The learned correlation matrices of CITRIS-VAE on the Ball-in-Boxes
dataset. The y-axis shows the sets of latent dimensions which are assigned to
causal variable 1 (ball_b/ball_x), causal variable 2 (ball_y), and to none of
them ($z_{\psi_{0}}$). The x-axis shows the true causal factors. Each element
of the matrix therefore describes the correlation between a set of latent
factors and a ground truth causal factor. CITRIS disentangles the causal
factors as intended, and assigning the internal box position ball_xin to
$z_{\psi_{0}}$ since it is independent of the provided interventions.

For evaluation purposes, we perform the correlation and triplet evaluation by
predicting both the box identifier ball_b and the box internal position
ball_xin instead of just ball_x. Note that this choice does not influence the
training of CITRIS in any way, but simply allows us to identify whether the
split between ball_b and ball_xin has been learned or not. The results are
shown in Table 14. One can clearly see that CITRIS learned to disentangle the
internal ball position ball_xin and ball_b, since $z_{\psi_{0}}$ has a $R^{2}$
correlation score of almost one to ball_xin, and close to zero for ball_b. We
show a visualization of the $R^{2}$ and Spearman correlation matrix in Figure
29. In conclusion, this experiment validates that CITRIS indeed learns the
minimal causal variables, _i.e._ the information that strictly depends on the
individual interventions.