figuret

# Relational Neural Expectation Maximization: Unsupervised Discovery of
Objects and their Interactions

Sjoerd van Steenkiste
Swiss AI Lab IDSIA, SUPSI, USI
Lugano, Switzerland
<EMAIL_ADDRESS>
&Michael Chang
UC Berkeley
Berkeley, United States AAAAA
<EMAIL_ADDRESS>
Klaus Greff
Swiss AI Lab IDSIA, SUPSI, USI
Lugano, SwitzerlandAAAAAAAAAAAAAAAAAAAAAA
<EMAIL_ADDRESS>
&Jürgen Schmidhuber
Swiss AI Lab IDSIA, SUPSI, USI
Lugano, Switzerland
<EMAIL_ADDRESS>Work performed while at IDSIA.

###### Abstract

Common-sense physical reasoning is an essential ingredient for any intelligent
agent operating in the real-world. For example, it can be used to simulate the
environment, or to infer the state of parts of the world that are currently
unobserved. In order to match real-world conditions this causal knowledge must
be learned without access to supervised data. To address this problem we
present a novel method that learns to discover objects and model their
physical interactions from raw visual images in a purely _unsupervised_
fashion. It incorporates prior knowledge about the compositional nature of
human perception to factor interactions between object-pairs and learn
efficiently. On videos of bouncing balls we show the superior modelling
capabilities of our method compared to other unsupervised neural approaches
that do not incorporate such prior knowledge. We demonstrate its ability to
handle occlusion and show that it can extrapolate learned knowledge to scenes
with different numbers of objects.

##  1 Introduction

Humans rely on common-sense physical reasoning to solve many everyday physics-
related tasks Lake et al., 2016. For example, it enables them to foresee the
consequences of their actions (simulation), or to infer the state of parts of
the world that are currently unobserved. This _causal_ understanding is an
essential ingredient for any intelligent agent that is to operate within the
world.

Common-sense physical reasoning is facilitated by the discovery and
representation of objects (a _core_ domain of human cognition Spelke &
Kinzler, 2007) that serve as primitives of a compositional system. They allow
humans to decompose a complex visual scene into distinct parts, describe
relations between them and reason about their dynamics as well as the
consequences of their interactions Ullman et al., 2017; Battaglia et al.,
2013; Lake et al., 2016.

The most successful machine learning approaches to common-sense physical
reasoning incorporate such prior knowledge in their design. They maintain
explicit object representations, which allow for general physical dynamics to
be learned between object pairs in a compositional manner Chang et al., 2016;
Battaglia et al., 2016; Watters et al., 2017. However, in these approaches
learning is _supervised_ , as it relies on object-representations from
external sources (e.g. a physics simulator) that are typically unavailable in
real-world scenarios.

Neural approaches that learn to directly model motion or physical interactions
in pixel space offer an alternative solution Sutskever et al., 2009;
Srivastava et al., 2015. However, while unsupervised, these methods suffer
from a lack compositionality at the representational level of objects. This
prevents such end-to-end neural approaches from efficiently learning functions
that operate on multiple entities and generalize in a human-like way (c.f.
Lake et al. 2016; Santoro et al. 2017; Battaglia et al. 2013, but see Perez et
al. 2017).

In this work we propose Relational N-EM (R-NEM), a novel approach to common-
sense physical reasoning that learns physical interactions between objects
from raw visual images in a purely _unsupervised_ fashion. At its core is
Neural Expectation Maximization (N-EM; Greff et al., 2017), a method that
allows for the discovery of compositional object-representations, yet is
unable to model interactions between objects. Therefore, we endow N-EM with a
relational mechanism inspired by previous work Santoro et al., 2017; Chang et
al., 2016; Battaglia et al., 2016, enabling it to factor interactions between
object-pairs, learn efficiently, and generalize to visual scenes with a
varying number of objects without re-training.

##  2 Method

Our goal is to learn common-sense physical reasoning in a purely unsupervised
fashion directly from visual observations. We have argued that in order to
solve this problem we need to exploit the compositional structure of a visual
scene. Conventional unsupervised representation learning approaches (eg. VAEs
Kingma & Welling 2013; GANs Goodfellow et al. 2014) learn a single distributed
representation that _superimposes_ information about the input, without
imposing any structure regarding objects or other low-level primitives. These
monolithic representations can not factorize physical interactions between
pairs of objects and therefore lack an essential inductive bias to learn these
efficiently. Hence, we require an alternative approach that can discover
objects representations as primitives of a visual scene in an unsupervised
fashion.

One such approach is _Neural Expectation Maximization_ (N-EM; Greff et al.
2017), which learns a separate distributed representation for each object
described in terms of the same features through an iterative process of
perceptual grouping and representation learning. The compositional nature of
these representations enable us to formulate _Relational N-EM_ (R-NEM): a
novel _unsupervised_ approach to common-sense physical reasoning that combines
N-EM (Section 2.1) with an interaction function that models relations between
objects efficiently (Section 2.2).

###  2.1 Neural Expectation Maximization

_Neural Expectation Maximization_ (N-EM; Greff et al. 2017) is a
differentiable clustering method that learns a representation of a visual
scene composed of primitive object representations. These representations
adhere to many useful properties of a symbolic representation of objects, and
can therefore be used as primitives of a compositional system Hummel et al.,
2004. They are described in the same format and each contain only information
about the object in the visual scene that they correspond to. Together, they
form a representation of a visual scene composed of objects that is learned in
an unsupervised way, which therefore serves as a starting point for our
approach.

The goal of N-EM is to group pixels in the input that belong to the same
object (perceptual grouping) and capture this information efficiently in a
distributed representation $\bm{\theta}_{k}$ for each object. At a high-level,
the idea is that if we were to have access to the family of distributions
$P(\bm{x}|\bm{\theta}_{k})$ (a statistical model of images given object
representations $\bm{\theta}_{k}$) then we can formalize our objective as
inference in a mixture of these distributions. By using Expectation
Maximization (EM; Dempster et al., 1977) to compute a Maximum Likelihood
Estimate (MLE) of the parameters of this mixture
($\bm{\theta}_{1},\ldots,\bm{\theta}_{K}$), we obtain a grouping (clustering)
of the pixels to each object (component) and their corresponding
representation. In reality we do not have access to
$P(\bm{x}|\bm{\theta}_{k})$, which N-EM learns instead by parameterizing the
mixture with a neural network and back-propagating through the iterations of
the unrolled generalized EM procedure.

Following Greff et al. 2017, we model each image $\bm{x}\in\mathbb{R}^{D}$ as
a spatial mixture of $K$ components parameterized by vectors
$\bm{\theta}_{1},\dots,\bm{\theta}_{K}\in\mathbb{R}^{M}$. A neural network
$f_{\phi}$ is used to transform these representations $\bm{\theta}_{k}$ into
parameters $\psi_{i,k}=f_{\phi}(\bm{\theta}_{k})_{i}$ for separate pixel-wise
distributions. A set of binary latent variables
$\bm{\mathcal{Z}}\in[0,1]^{D\times K}$ encodes the unknown true pixel
assignments, such that $z_{i,k}=1$ iff pixel $i$ was generated by component
$k$. The full likelihood for $\bm{x}$ given
$\bm{\theta}=(\bm{\theta}_{1},\dots,\bm{\theta}_{K})$ is given by:

$P(\bm{x}|\bm{\theta})=\prod_{i=1}^{D}\sum_{\bm{z}_{i}}P(x_{i},\bm{z}_{i}|\bm{\psi}_{i})=\prod_{i=1}^{D}\sum_{k=1}^{K}P(z_{i,k}=1)P(x_{i}|\psi_{i,k},z_{i,k}=1).$
(1)

Figure 1: Illustration of the different computational aspects of R-NEM when
applied to a sequence of images of bouncing balls. Note that
$\bm{\gamma},\bm{\psi}$ at the _Representations_ level correspond to the
$\bm{\gamma}$ (_E-step_), $\bm{\psi}$ (_Group Reconstructions_) from the
previous time-step. Different colors correspond to different cluster
components (object representations).The right side shows a computational
overview of $\Upsilon^{\text{R-NEM}}$, a function that computes the pair-wise
interactions between the object representations.

If $f_{\phi}$ has learned a statistical model of images given object
representations $\bm{\theta}_{k}$, then we can compute the object
representations for a given image $\bm{x}$ by maximizing
$P(\bm{x}|\bm{\theta})$. Marginalization over $\bm{z}$ complicates this
process, thus we use generalized EM to maximize the following lowerbound
instead:

$\mathcal{Q}(\bm{\theta},\bm{\theta}^{\text{old}})=\sum_{\mathbf{z}}P(\mathbf{z}|\bm{x},\bm{\psi}^{\text{old}})\log
P(\bm{x},\mathbf{z}|\bm{\psi}).$ (2)

Each iteration of generalized EM consists of two steps: the _E-step_ computes
a new estimate of the posterior probability distribution over the latent
variables $\gamma_{i,k}:=P(z_{i,k}=1|x_{i},\psi_{i}^{\text{old}})$ given
$\bm{\theta}^{\text{old}}$ from the previous iteration. It yields a new soft-
assignment of the pixels to the components (clusters), based on how accurately
they model $\bm{x}$. The generalized _M-step_ updates
$\bm{\theta}^{\text{old}}$ by taking a gradient ascent step on (2), using the
previously computed soft-assignments:
$\bm{\theta}_{k}^{\text{new}}=\bm{\theta}_{k}^{\text{old}}+\eta\cdot\partial\mathcal{Q}/\partial\bm{\theta}_{k}$.111We
can not compute
$\operatorname*{\arg\\!\max}_{\theta}\mathcal{Q}(\bm{\theta},\bm{\theta}^{\text{old}})$
analytically, due to non-linearity of $f_{\phi}$.

The unrolled computational graph of the generalized EM steps is
differentiable, which provides a means to train $f_{\phi}$ to implement a
statistical model of images given object representations. Using back-
propagation through time (eg. Werbos 1988; Williams 1989) we train $f_{\phi}$
to minimize the following loss:

$L(\bm{x})=-\sum_{i=1}^{D}\sum_{k=1}^{K}\underbrace{\gamma_{i,k}\log
P(x_{i},z_{i,k}|\psi_{i,k})}_{\text{intra-cluster
loss}}-\underbrace{(1-\gamma_{i,k})D_{KL}[P(x_{i})||P(x_{i}|\psi_{i,k},z_{i,k})]}_{\text{inter-
cluster loss}}.$ (3)

The intra-cluster term is identical to (2), which credits each component for
accurately representing pixels that have been assigned to it. The inter-
cluster term ensures that each representation only captures the information
about the pixels that have been assigned to it.

A more powerful variant of N-EM can be obtained (RNN-EM) by substituting the
generalized M-step with a recurrent neural network having hidden state
$\bm{\theta}_{k}$. In this case, the entirety of $f_{\phi}$ consists of a
recurrent encoder-decoder architecture that receives
$\bm{\gamma}_{k}(\bm{x}-\bm{\psi}_{k})$ as input at each step.

The learning objective in (3) is prone to trivial solutions in case of
overcapacity, which could prevent the network from modelling the statistical
regularities in the data that correspond to objects. By adding noise to the
input image or reducing $\bm{\theta}$ in dimensionality we can guide learning
to avert this. Moreover, in the case of RNN-EM one can evaluate (3) at the
following time-step (_predictive coding_) to encourage learning of object
representations and their corresponding dynamics. One intuitive interpretation
of using denoising or next-step prediction as part of the training objective
is to guide the network to learn about essential properties of objects, in
this case those that correspond to the Gestalt Principles of _prägnanz_ and
_common fate_ Hatfield & Epstein, 1985.

###  2.2 Relational Neural Expectation Maximization

RNN-EM (unlike N-EM) is able to capture the dynamics of individual objects
through a parametrized recurrent connection that operates on the object
representation $\bm{\theta_{k}}$ across consecutive time-steps. However, the
relations and interactions that take place _between_ objects can not be
captured in this way. In order to overcome this shortcoming we propose
_Relational N-EM_ (R-NEM), which adds relational structure to the recurrence
to model interactions between objects without violating key properties of the
learned object representations.

Consider a generalized form of the standard RNN-EM dynamics equation, which
computes the object representation $\bm{\theta}_{k}$ at time $t$ as a function
of all object representations
$\bm{\theta}:=[\bm{\theta}_{1},\ldots,\bm{\theta}_{K}]$ at the previous time-
step through an _interaction function_ $\Upsilon$:

$\bm{\theta}_{k}^{(t)}=\text{RNN}(\tilde{\bm{x}}^{(t)},\Upsilon_{k}(\bm{\theta}^{(t-1)})):=\sigma(\bm{W}\cdot\tilde{\bm{x}}^{(t)}+\bm{R}\cdot\Upsilon_{k}(\bm{\theta}^{(t-1)})).$
(4)

Here $\bm{W},\bm{R}$ are weight matrices, $\sigma$ is the sigmoid activation
function, and $\tilde{\bm{x}}^{(t)}$ is the input to the recurrent model at
time $t$ (possibly transformed by an encoder). When $\Upsilon_{k}^{\text{RNN-
EM}}(\bm{\theta}):=\bm{\theta}_{k}$, this dynamics model coincides with a
standard RNN update rule, thereby recovering the original RNN-EM formulation.

The inductive bias incorporated in $\Upsilon$ reflects the modeling
assumptions about the interactions between objects in the environment, and
therefore the nature of $\bm{\theta}_{k}$’s interdependence. If $\Upsilon$
incorporates the assumption that no interaction takes place between objects,
then the $\bm{\theta}_{k}$’s are fully independent and we recover
$\Upsilon^{\text{RNN-EM}}$. On the other hand, if we do assume that
interactions among objects take place, but assume very little about the
structure of the interdependence between the $\bm{\theta}_{k}$’s, then we
forfeit useful properties of $\bm{\theta}_{k}$ such as compositionality. For
example, if $\Upsilon:=\text{MLP}(\bm{\theta})$ we can no longer extrapolate
learned knowledge to environments with more or fewer than $K$ objects and lose
overall data efficiency Santoro et al., 2017. Instead, we can make efficient
use of compositionality among the learned object representations
$\bm{\theta}_{k}$ to incorporate general but guiding constraints on how these
may influence one another Battaglia et al., 2016; Chang et al., 2016. In doing
so we constrain $\Upsilon$ to capture interdependence between
$\bm{\theta}_{k}$’s in a compositional manner that enables physical dynamics
to be learned efficiently, and allow for learned dynamics to be extrapolated
to a variable number of objects.

We propose a parametrized interaction function $\Upsilon^{\text{R-NEM}}$ that
incorporates these modeling assumptions and updates $\bm{\theta}_{k}$ based on
the pairwise effects of the objects $i\neq k$ on $k$:

$\begin{split}&\Upsilon^{\text{R-NEM}}_{k}(\bm{\theta})=[\bm{\hat{\theta}}_{k};\bm{E}_{k}]\
\ \text{with}\ \ \bm{\hat{\theta}}_{k}=\text{MLP}^{\textit{\hskip
1.0ptenc}}(\bm{\theta}_{k})\ \ ,\ \ \bm{E}_{k}=\sum_{i\neq
k}\alpha_{k,i}\cdot\bm{e}_{k,i}\\\ &\ \
\alpha_{k,i}=\text{MLP}^{\textit{\hskip 1.0ptatt}}(\bm{\xi}_{k,i})\ \ ,\ \
\bm{e}_{k,i}=\text{MLP}^{\textit{\hskip 1.0pteff}}(\bm{\xi}_{k,i})\ \ ,\ \
\bm{\xi}_{k,i}=\text{MLP}^{\textit{\hskip
1.0ptemb}}([\bm{\hat{\theta}}_{k};\bm{\hat{\theta}}_{i}])\end{split}$ (5)

where $[\cdot;\cdot]$ is the concatenation operator and $\text{MLP}^{(\cdot)}$
corresponds to a multi-layer perceptron. First, each $\bm{\theta}_{i}$ is
transformed using $\text{MLP}^{\textit{\hskip 1.0ptenc}}$ to obtain
$\bm{\hat{\theta}}_{i}$, which enables information that is relevant for the
object dynamics to be made more explicit in the representation. Next, each
pair $(\bm{\hat{\theta}}_{k},\bm{\hat{\theta}}_{i})$ is concatenated and
processed by $\text{MLP}^{\textit{\hskip 1.0ptemb}}$, which computes a shared
embedding $\bm{\xi}_{k,i}$ that encodes the interaction between object $k$ and
object $i$. Notice that we opt for a clear separation between the _focus_
object $k$ and the _context_ object i as in previous work Chang et al., 2016.
From $\bm{\xi}_{k,i}$ we compute $\bm{e}_{k,i}$: the effect of object $i$ on
object $k$; and an attention coefficient $\alpha_{k,i}$ that encodes whether
interaction between object $i$ and object $k$ takes place. These attention
coefficients Bahdanau et al., 2014; Xu et al., 2015 help to select relevant
context objects, and can be seen as a more flexible unsupervised replacement
of the distance based heuristic that was used in previous work Chang et al.,
2016. Finally, we compute the total effect of $\bm{\theta}_{i\neq k}$ on
$\bm{\theta}_{k}$ as a weighted sum of the effects multiplied by their
attention coefficient. A visual overview of $\Upsilon^{\text{R-NEM}}$ can be
seen on the right side of Figure 1.

##  3 Related Work

Machine learning approaches to common-sense physical reasoning can roughly be
divided in two groups: symbolic approaches and approaches that perform state-
to-state prediction. The former group performs inference over the parameters
of a symbolic physics engine Wu et al., 2015; Ullman et al., 2017; Battaglia
et al., 2013, which restricts them to synthetic environments. The latter group
employs machine learning methods to make state-to-state predictions, often
describing the state of a system as a set of compact object-descriptions that
are either used as an input to the system Battaglia et al., 2016; Fragkiadaki
et al., 2015; Chang et al., 2016; Grzeszczuk et al., 1998 or for training
purposes Watters et al., 2017. By incorporating information (eg. position,
velocity) about objects these methods have achieved excellent generalization
and simulation capabilities. Purely unsupervised approaches for state-to-state
prediction Sutskever et al., 2009; Michalski et al., 2014; Agrawal et al.,
2016; Lerer et al., 2016 that use raw visual inputs as state-descriptions have
yet to rival these capabilities. Our method is a purely unsupervised state-to-
state prediction method that operates in pixel space, taking a first step
towards unsupervised learning of common-sense reasoning in real-world
environments.

The proposed interaction function $\Upsilon^{\text{R-NEM}}$ can be seen as a
type of Message Passing Neural Network (MPNN; Gilmer et al. 2017) that
incorporates a variant of neighborhood attention Duan et al., 2017. In light
of other recent work Zaheer et al., 2017 it can be seen as a permutation
equivariant set function.

R-NEM relies on N-EM Greff et al., 2017 to discover a compositional object
representation from raw visual inputs. A closely related approach to N-EM is
the TAG framework Greff et al., 2016, which utilizes a similar mechanism to
perform inference over group representations, but in addition performs
inference over the group assignments. In recent work TAG was combined with a
recurrent ladder network Ilin et al., 2017 to obtain a powerful model
(RTagger) that can be applied to sequential data. However, the lack of a
single compact representation that captures all information about a group
(object) makes a compositional treatment of physical interactions more
difficult. Other unsupervised approaches rely on attention to group together
parts of the visual scene corresponding to objects Eslami et al., 2016; Gregor
et al., 2015. These approaches suffer from a similar problem in that their
sequential nature prevents a coherent object representation to take shape.

Other related work have also taken steps towards combining the learnability of
neural networks with the compositionality of symbolic programs in modeling
physics Chang et al., 2016; Battaglia et al., 2016, playing games Kansky et
al., 2017; Denil et al., 2017, learning algorithms Cai et al., 2017; Bošnjak
et al., 2017; Reed & De Freitas, 2015; Li et al., 2016, visual understanding
Johnson et al., 2017; Ellis et al., 2017, and natural language processing
Andreas et al., 2016; Hu et al., 2017.

Figure 2: R-NEM applied to a sequence of $4$ bouncing balls. Each column
corresponds to a time-step, which coincides with an EM step. At each time-
step, R-NEM computes $K=5$ new representations $\bm{\theta}_{k}$ according to
(4) (see also _Representations_ in Figure 1) from the input $\bm{x}$ with
added noise (bottom row). From each new $\bm{\theta}_{k}$ a group
reconstruction $\bm{\psi}_{k}$ is produced (rows 2-6 from bottom) that
predicts the state of the environment at the _next_ time-step. Attention
coefficients are visualized by overlaying a colored reconstruction of a
context object on the white reconstruction of the focus object (see
_Attention_ in Section 4). Based on the prediction accuracy of $\bm{\psi}$,
the _E-step_ (see Figure 1) computes new soft-assignments $\bm{\gamma}$ (row 7
from bottom), visualized by coloring each pixel $i$ according to their
distribution over components $\bm{\gamma}_{i}$. Row 8 visualizes the total
prediction by the network ($\sum_{k}\bm{\psi}_{k}\cdot\bm{\gamma}_{k}$) and
row 9 the ground-truth sequence at the next time-step.

##  4 Experiments

Figure 3: Performance of each method on the _bouncing balls_ task. Each method
was trained on a dataset with 4 balls, evaluated on a test set with $4$ balls
(left), and on a test-set with 6-8 balls (middle). The losses are reported
relative to the loss of a baseline for each dataset that always predicts the
current frame. The ARI score (right) is used to evaluate the degree of
compositionality that is achieved.

In this section we evaluate R-NEM on three different physical reasoning tasks
that each vary in their dynamical and visual complexity: bouncing balls with
variable mass, bouncing balls with an invisible curtain and the Arcade
Learning Environment Bellemare et al., 2013. We compare R-NEM to other
_unsupervised_ neural methods that do not incorporate any inductive biases
reflecting real-world dynamics and show that these are indeed
beneficial.222Code is available at
https://github.com/sjoerdvansteenkiste/Relational-NEM.

All experiments use ADAM Kingma & Ba, 2014 with default parameters, on 50K
train + 10K validation + 10K test sequences and early stopping with a patience
of 10 epochs. For each of $\text{MLP}^{\textit{\hskip 1.0ptenc,emb,eff}}$ we
used a unique single layer neural network with 250 _rectified linear_ units.
For $\text{MLP}^{\textit{\hskip 1.0ptatt}}$ we used a two-layer neural
network: 100 _tanh_ units followed by a single _sigmoid_ unit. A detailed
overview of the experimental setup can be found in Appendix A.

#### Bouncing Balls

We study the physical reasoning capabilities of R-NEM on the _bouncing balls_
task, a standard environment to evaluate physical reasoning capabilities that
exhibits low visual complexity and complex non-linear physical
dynamics.333Videos are available at https://sites.google.com/view/r-nem-gifs/.
We train R-NEM on sequences of $64\times 64$ binary images over 30 time-steps
that contain four bouncing balls with different masses corresponding to their
radii. The balls are initialized with random initial positions, masses and
velocities. Balls bounce elastically against each other and the image window.

#### Qualitative Evaluation

Figure 1 presents a qualitative evaluation of R-NEM on the bouncing balls
task. After 10 time-steps it can be observed that the pixels that belong to
each of the balls are grouped together and assigned to a unique component
(with a saturated color); and that the background (colored grey) has been
divided among all components (resulting in a grey coloring). This indicates
that the representation $\bm{\theta}_{k}$ from which each component produces
the group reconstruction $\bm{\psi}_{k}$ does indeed only contain information
about a unique object, such that together the $\bm{\theta}_{k}$’s yield a
compositional object representation of the scene. The total reconstruction
(that combines the group reconstructions and the soft-assignments) displays an
accurate reconstruction of the input sequence at the next time-step,
indicating that R-NEM has learned to model the dynamics of bouncing balls.

#### Comparison

We compare the modelling capabilities of R-NEM to an RNN, LSTM Hochreiter &
Schmidhuber, 1997; Gers et al., 1999 and RNN-EM in terms of the Binomial
Cross-Entropy (BCE) loss between the predicted image and the ground-truth
image of the last frame,444Since the E-step in R-NEM and RNN-EM utilizes the
ground-truth for reconstruction, we substitute it with a simple $\max$
operator. The resulting loss serves as an _upperbound_ to the true BCE loss.
as well as the _relational_ BCE that only takes into account objects that
currently take part in collision. Unless specified we use $K=5$.

On a test-set with sequences containing four balls we observe that R-NEM
produces markedly lower losses when compared to all other methods (left plot
in Figure 3). Moreover, in order to validate that each component captures only
a single ball (and thus compositionality is achieved), we report the Adjusted
Rand Index (ARI; Hubert & Arabie 1985) score between the soft-assignments
$\bm{\gamma}$ and the ground-truth assignment of pixels to objects. In the
left column of the ARI plot (right side in Figure 3) we find that R-NEM
achieves an ARI score of 0.8, meaning that in roughly $80\%$ of the cases each
ball is modeled by a single component. This suggests that a compositional
object representation is achieved for most of the sequences. Together these
observations are in line with our qualitative evaluation and validate that
incorporating real world priors is greatly beneficial (comparing to RNN, LSTM)
and that $\Upsilon^{\text{R-NEM}}$ enables interactions to be modelled more
accurately compared to RNN-EM in terms of the relational BCE.

Similar to Greff et al. 2017 we find that further increasing the number of
components during training (leaving additional groups empty) increases the
quality of the grouping, see R-NEM $K=8$ in Figure 3. In addition we observe
that the loss (in particular the relational BCE) is reduced further, which
matches our hypothesis that compositional object representations are greatly
beneficial for modelling physical interactions.

Figure 4: Left: Three sequences of 15 time-steps ground-truth (top), R-NEM
(middle), RNN (bottom). The last ten time-steps of the sequences produced by
R-NEM and RNN are simulated. Right: The BCE loss on the entire test-set for
these same time-steps.

#### Extrapolating learned knowledge

We use a test-set with sequences containing 6-8 balls to evaluate the ability
of each method to _extrapolate_ their learned knowledge about physical
interactions between four balls to environments with more balls. We use $K=8$
when evaluating R-NEM and RNN-EM on this test-set in order to accommodate the
increased number of objects. As can be seen from the middle plot in Figure 3,
R-NEM again greatly outperforms all other methods. Notice that, since we
report the loss relative to a baseline, we roughly factor out the increased
complexity of the task. Perfect extrapolation of the learned knowledge would
therefore amount to _no change_ in relative performance. In contrast, we
observe far worse performance for the LSTM (relative to the baseline) when
evaluated on this dataset with extra balls. It suggests that the gating
mechanism of the LSTM has allowed it to learn a sophisticated and overly
specialized solution for sequences with four balls that does not generalize to
a dataset with 6-8 balls.

R-NEM and RNN-EM scale markedly better to this dataset than LSTM. Although the
RNN similarly suffers to a lesser extend from this type of “overfitting”, this
is most likely due its inability to learn a reasonable solution on sequences
of four balls to begin with. Hence, we conclude that the superior
extrapolation capabilities of RNN-EM and R-NEM are inherent to their ability
to factor a scene in terms of permutation invariant object representations
(see right side of the right plot in Figure 3).

#### Attention

Further insight in the role of the attention mechanism can be gained by
visualizing the attention coefficients, as is done in Figure 2. For each
component $k$ we draw $\alpha_{k,i}*\bm{\psi}_{i}$ on top of the
reconstruction $\bm{\psi}_{k}$, colored according to the color of component
$i$. These correspond to the colored balls (that are for example seen in time-
steps 13, 14), which indicate whether component $k$ took information about
component $i$ into account when computing the new state (recall (5)). It can
be observed that the attention coefficient $\alpha_{k,i}$ becomes non-zero
whenever collision takes place, such that a colored ball lights up in the
_following_ time-steps. The attention mechanism learned by R-NEM thus assumes
the role of the distance-based heuristic in previous work Chang et al., 2016,
matching our own intuitions of how this mechanism would best be utilized.

A quantitative evaluation of the attention mechanism is obtained by comparing
R-NEM to a variant of itself that does not incorporate attention (_R-NEM no
att_). Figure 3 shows that both methods perform equally well on the regular
test set (4 balls), but that _R-NEM no att_ performs worse at extrapolating
from its learned knowledge (6-8 balls). A likely reason for this behavior is
that the range of the sum in (5) changes with $K$. Thus, when extrapolating to
an environment with more balls the total sum may exceed previous boundaries
and impede learned dynamics.

#### Simulation

Once a scene has been accurately modelled, R-NEM can approximately simulate
its dynamics through recursive application of (4) for each
$\bm{\theta}_{k}$.555Note that in this case the input to the neural network
encoder in component $k$ corresponds to
$\bm{\gamma}_{k}(\bm{x}^{(t)}-\bm{\psi}^{(t-1)})$, such that the output of the
encoder $\tilde{\bm{x}}^{(t)}\approx\bm{0}$ when
$\bm{\psi}_{k}^{(t-1)}=\bm{x}^{(t)}$. In Figure 4 we compare the simulation
capabilities of R-NEM to RNN-EM and an RNN on the bouncing balls
environment.footnote 3 On the left it shows for R-NEM and an RNN a sequence
with five normal steps followed by 10 simulation steps, as well as the ground-
truth sequence. From the last frame in the sequence it can clearly be observed
that R-NEM has managed to accurately simulate the environment. Each ball is
approximately in the correct place, and the shape of each ball is preserved.
The balls simulated by the RNN, on the other hand, deviate substantially from
their ground-truth position and their size has increased. In general we find
that R-NEM produces mostly very accurate simulations, whereas the RNN
consistently fails. Interestingly we found that the cases in which R-NEM
frequently fails are those for which a single component models more than one
ball. The right side of Figure 4 summarizes the BCE loss for these same time-
steps across the _entire_ test-set. Although this is a crude measure of
simulation performance (since it does not take into account the identity of
the balls), we still observe that R-NEM consistently outperforms RNN-EM and an
RNN.

Figure 5: R-NEM applied to a sequence of bouncing balls with an invisible
curtain. The ground truth sequence is displayed in the top row, followed by
the prediction of R-NEM (middle) and the soft-assignments of pixels to
components (bottom). R-NEM models objects, as well as its interactions, even
when the object is completely occluded (step 36). Only a subset of the steps
is shown.

#### Hidden Factors

Occlusion is abundant in the real world, and the ability to handle hidden
factors is crucial for any physical reasoning system. We therefore evaluate
the capability of R-NEM to handle occlusion using a variant of bouncing balls
that contain an invisible “curtain.” Figure 5 shows that R-NEM accurately
models the sequence and can maintain object states, even when confronted with
occlusion.footnote 3 For example, note that in step 36 the “blue” ball, is
completely occluded and is about to collide with the “orange” ball. In step 38
the ball is accurately predicted to re-appear at the bottom of the curtain
(since collision took place) as opposed to the left side of the curtain. This
demonstrates that R-NEM has a notion of _object permanence_ and implies that
it understands a scene on a level beyond pixels: it assigns persistence and
identity to the objects.

In terms of test-set performance we find that R-NEM (_BCE:_ $46.22$,
_relational BCE:_ $2.33$) outperforms an RNN (_BCE:_ $94.64$, _relational
BCE:_ $4.14$) and an LSTM (_BCE:_ $59.32$, _relational BCE:_ $2.72$).

#### Space Invaders

To test the performance of R-NEM in a visually more challenging environment,
we train it on sequences of $84\times 84$ binarized images over 25 time-steps
of game-play on Space Invaders from the Arcade Learning Environment Bellemare
et al., 2013.666Binarization ensures that the color group of the entities on
the screen does not give away their grouping. We use $K=4$ and also feed the
action of the agent to the interaction function. Figure 6 confirms that R-NEM
is able to accurately model the environment, even though the visual complexity
has increased. Notice that these visual scenes comprise a large numbers of
(small) primitive objects that behave similarly. Since we trained R-NEM with
four components it is unable to group pixels according to individual objects
and is forced to consider a different grouping. We find that R-NEM assigns
different groups to every other column of aliens together with the spaceship,
and to the three large “shields.” These groupings seem to be based on
movement, which to some degree coincides with their semantic roles of the
environment. In other examples (not shown) we also found that R-NEM frequently
assigns different groups to every other column of the aliens, and to the three
large “shields.” Individual bullets and the space ship are less frequently
grouped separately, which may have to do with the action-noise of the
environment (that controls the movement of the space-ship) and the small size
of the bullets at the current resolution that makes them less predictable.

Figure 6: R-NEM accurately models a sequence of frames obtained by an agent
playing Space Invaders. A group no longer corresponds to an object, but
instead assumes the role of high-level entities that engage in similar
movement patterns.

##  5 Discussion and Conclusion

We have argued that the ability to discover and describe a scene in terms of
objects provides an essential ingredient for common-sense physical reasoning.
This is supported by converging evidence from cognitive science and
developmental psychology that intuitive physics and reasoning capabilities are
built upon the ability to perceive objects and their interactions Ullman et
al., 2017; Spelke, 1988. The fact that young infants already exhibit this
ability, may even suggest an innate bias towards compositionality Spelke &
Kinzler, 2007; Lake et al., 2016; Munakata et al., 1997. Inspired by these
observations we have proposed R-NEM, a method that incorporates inductive
biases about the existence of objects and interactions, implemented by its
clustering objective and interaction function respectively. The specific
nature of the objects, and their dynamics and interactions can then be learned
efficiently purely from visual observations.

In our experiments we find that R-NEM indeed captures the (physical) dynamics
of various environments more accurately than other methods, and that it
exhibits improved generalization to environments with different numbers of
objects. It can be used as an approximate simulator of the environment, and to
predict movement and collisions of objects, even when they are completely
occluded. This demonstrates a notion of object permanence and aligns with
evidence that young infants seem to infer that occluded objects move in
connected paths and continue to maintain object-specific properties Spelke,
1990. Moreover, young infants also appear to expect that objects only interact
when they come into contact Spelke, 1990, which is analogous to the behaviour
of R-NEM to only attend to other objects when a collision is imminent. In
summary, we believe that our method presents an important step towards
learning a more human-like model of the world in a completely unsupervised
fashion.

Current limitations of our approach revolve around grouping and prediction.
What aspects of a scene humans group together typically varies as a function
of the task in mind. One may perceive a stack of chairs as a whole if the goal
is to move them to another room, or as individual chairs if the goal is to
count the number of chairs in the stack. In order to facilitate this _dynamic_
grouping one would need to incorporate top-down feedback from an agent into
the grouping procedure to deviate from the built-in inductive biases. Another
limitation of our approach is the need to incentivize R-NEM to produce useful
groupings by injecting noise, or reducing capacity. The former may prevent
very small regularities in the input from being detected. Finally the
interaction in the E-step among the groups makes it difficult to increase the
number of components above ten without causing harmful training instabilities.
Due to the multitude of interactions and objectives in R-NEM (and RNN-EM) we
find that they are sometimes challenging to train.

In terms of prediction we have implicitly assumed that objects in the
environment behave according to rules that can be inferred. This poses a
challenge when objects deform in a manner that is difficult to predict (as is
the case for objects in Space Invaders due to downsampling). However in
practice we find that (once pixels have been grouped together) the masking of
the input helps each component in quickly adapting its representation to any
unforeseen behaviour across consecutive time steps. Perhaps a more severe
limitation of R-NEM (and of RNN-EM in general) is that the second loss term of
the outer training objective hinders in modelling more complex varying
backgrounds, as the background group would have to predict the “pixel prior”
for every other group.

We argue that the ability to engage in common-sense physical reasoning
benefits any intelligent agent that needs to operate in a physical
environment, which provides exciting future research opportunities. In future
work we intend to investigate how top-down feedback from an agent could be
incorporated in R-NEM to facilitate dynamic groupings, but also how the
compositional representations produced by R-NEM can benefit a reinforcement
learner, for example to learn a modular policy that easily generalizes to
novel combinations of known objects. Other interactions between a controller C
and a model of the world M (implemented by R-NEM) as posed in Schmidhuber 2015
constitute further research directions.

## Acknowledgements

The authors wish to thank Tom Griffiths and the anonymous reviewers for
helpful comments and constructive feedback. This research was supported by the
Swiss National Science Foundation grant 200021_165675/1, the EU project
“INPUT” (H2020-ICT-2015 grant no. 687795), and the Zeno Karl Schindler
Foundation Summerschool Grant. Chang would like to thank Christiane Born,
Sarah Craver, Cinzia Daldini, and the MIT MISTI Program for supporting his
stay in Switzerland. We are grateful to NVIDIA Corporation for donating us a
DGX-1 as part of the Pioneers of AI Research award, and to IBM for donating a
“Minsky” machine.

## References

  * Agrawal et al. 2016 Pulkit Agrawal, Ashvin V Nair, Pieter Abbeel, Jitendra Malik, and Sergey Levine.  Learning to poke by poking: Experiential learning of intuitive physics.  In _Advances in Neural Information Processing Systems_ , pp. 5074–5082, 2016.
  * Andreas et al. 2016 Jacob Andreas, Marcus Rohrbach, Trevor Darrell, and Dan Klein.  Neural module networks.  In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_ , pp. 39–48, 2016.
  * Bahdanau et al. 2014 Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio.  Neural machine translation by jointly learning to align and translate.  _arXiv preprint arXiv:1409.0473_ , 2014.
  * Battaglia et al. 2016 Peter Battaglia, Razvan Pascanu, Matthew Lai, Danilo Jimenez Rezende, and others.  Interaction networks for learning about objects, relations and physics.  In _Advances in Neural Information Processing Systems_ , pp. 4502–4510, 2016.
  * Battaglia et al. 2013 Peter W Battaglia, Jessica B Hamrick, and Joshua B Tenenbaum.  Simulation as an engine of physical scene understanding.  _Proceedings of the National Academy of Sciences_ , 110(45):18327–18332, 2013.
  * Bellemare et al. 2013 Marc G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling.  The arcade learning environment: An evaluation platform for general agents.  _J. Artif. Intell. Res.(JAIR)_ , 47:253–279, 2013.
  * Bošnjak et al. 2017 Matko Bošnjak, Tim Rocktäschel, Jason Naradowsky, and Sebastian Riedel.  Programming with a differentiable forth interpreter.  In _Proceedings of the 34th International Conference on Machine Learning_ , pp. 547–556, 2017.
  * Cai et al. 2017 Jonathon Cai, Richard Shin, and Dawn Song.  Making neural programming architectures generalize via recursion.  _arXiv preprint arXiv:1704.06611_ , 2017.
  * Chang et al. 2016 Michael B. Chang, Tomer Ullman, Antonio Torralba, and Joshua B. Tenenbaum.  A compositional object-based approach to learning physical dynamics.  _arXiv preprint arXiv:1612.00341_ , 2016.
  * Dempster et al. 1977 A. P. Dempster, N. M. Laird, and D. B. Rubin.  Maximum likelihood from incomplete data via the EM algorithm.  _Journal of the royal statistical society._ , pp. 1–38, 1977.
  * Denil et al. 2017 Misha Denil, Sergio Gómez Colmenarejo, Serkan Cabi, David Saxton, and Nando de Freitas.  Programmable agents.  _arXiv preprint arXiv:1706.06383_ , 2017.
  * Duan et al. 2017 Yan Duan, Marcin Andrychowicz, Bradly Stadie, Jonathan Ho, Jonas Schneider, Ilya Sutskever, Pieter Abbeel, and Wojciech Zaremba.  One-shot imitation learning.  In _Advances in Neural Information Processing Systems_ , pp. 1087–1098, 2017.
  * Ellis et al. 2017 Kevin Ellis, Daniel Ritchie, Armando Solar-Lezama, and Joshua B. Tenenbaum.  Learning to infer graphics programs from hand-drawn images.  _CoRR_ , abs/1707.09627, 2017.
  * Eslami et al. 2016 SM Ali Eslami, Nicolas Heess, Theophane Weber, Yuval Tassa, David Szepesvari, Geoffrey E Hinton, et al.  Attend, infer, repeat: Fast scene understanding with generative models.  In _Advances in Neural Information Processing Systems_ , pp. 3225–3233, 2016.
  * Fragkiadaki et al. 2015 Katerina Fragkiadaki, Pulkit Agrawal, Sergey Levine, and Jitendra Malik.  Learning visual predictive models of physics for playing billiards.  _arXiv preprint arXiv:1511.07404_ , 2015.
  * Gers et al. 1999 Felix A Gers, Jürgen Schmidhuber, and Fred Cummins.  Learning to forget: Continual prediction with LSTM.  In _Artificial Neural Networks, 1999. ICANN 99. Ninth International Conference on (Conf. Publ. No. 470)_ , volume 2, pp. 850–855, 1999.
  * Gilmer et al. 2017 Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl.  Neural message passing for quantum chemistry.  _arXiv preprint arXiv:1704.01212_ , 2017.
  * Goodfellow et al. 2014 Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio.  Generative adversarial nets.  In _Advances in neural information processing systems_ , pp. 2672–2680, 2014.
  * Greff et al. 2016 Klaus Greff, Antti Rasmus, Mathias Berglund, Tele Hao, Harri Valpola, and Juergen Schmidhuber.  Tagger: Deep unsupervised perceptual grouping.  In _Advances in Neural Information Processing Systems_ , pp. 4484–4492, 2016.
  * Greff et al. 2017 Klaus Greff, Sjoerd van Steenkiste, and Jürgen Schmidhuber.  Neural expectation maximization.  In _Advances in Neural Information Processing Systems_ , pp. 6694–6704, 2017.
  * Gregor et al. 2015 Karol Gregor, Ivo Danihelka, Alex Graves, Danilo Jimenez Rezende, and Daan Wierstra.  Draw: A recurrent neural network for image generation.  _arXiv preprint arXiv:1502.04623_ , 2015.
  * Grzeszczuk et al. 1998 Radek Grzeszczuk, Demetri Terzopoulos, and Geoffrey Hinton.  Neuroanimator: Fast neural network emulation and control of physics-based models.  In _Proceedings of the 25th annual conference on Computer graphics and interactive techniques_ , pp. 9–20. ACM, 1998.
  * Hatfield & Epstein 1985 Gary Hatfield and William Epstein.  The status of the minimum principle in the theoretical analysis of visual perception.  _Psychological Bulletin_ , 97(2):155, 1985.
  * Hochreiter & Schmidhuber 1997 Sepp Hochreiter and Jürgen Schmidhuber.  Long short-term memory.  _Neural computation_ , 9(8):1735–1780, 1997.
  * Hu et al. 2017 Ronghang Hu, Jacob Andreas, Marcus Rohrbach, Trevor Darrell, and Kate Saenko.  Learning to reason: End-to-end module networks for visual question answering.  _arXiv preprint arXiv:1704.05526_ , 2017.
  * Hubert & Arabie 1985 Lawrence Hubert and Phipps Arabie.  Comparing partitions.  _Journal of classification_ , 2(1):193–218, 1985\.
  * Hummel et al. 2004 John E. Hummel, Keith J. Holyoak, Collin Green, Leonidas AA Doumas, Derek Devnich, Aniket Kittur, and Donald J. Kalar.  A solution to the binding problem for compositional connectionism.  In _Compositional Connectionism in Cognitive Science: Papers from the AAAI Fall Symposium, Ed. SD Levy & R. Gayler_, pp. 31–34, 2004.
  * Ilin et al. 2017 Alexander Ilin, Isabeau Prémont-Schwarz, Tele Hotloo Hao, Antti Rasmus, Rinu Boney, and Harri Valpola.  Recurrent Ladder Networks.  _arXiv:1707.09219 [cs, stat]_ , July 2017.
  * Johnson et al. 2017 Justin Johnson, Bharath Hariharan, Laurens van der Maaten, Judy Hoffman, Li Fei-Fei, C Lawrence Zitnick, and Ross Girshick.  Inferring and executing programs for visual reasoning.  _arXiv preprint arXiv:1705.03633_ , 2017.
  * Kansky et al. 2017 Ken Kansky, Tom Silver, David A Mély, Mohamed Eldawy, Miguel Lázaro-Gredilla, Xinghua Lou, Nimrod Dorfman, Szymon Sidor, Scott Phoenix, and Dileep George.  Schema networks: Zero-shot transfer with a generative causal model of intuitive physics.  _arXiv preprint arXiv:1706.04317_ , 2017.
  * Kingma & Ba 2014 Diederik Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Kingma & Welling 2013 Diederik P Kingma and Max Welling.  Auto-encoding variational bayes.  _arXiv preprint arXiv:1312.6114_ , 2013.
  * Lake et al. 2016 Brenden M. Lake, Tomer D. Ullman, Joshua B. Tenenbaum, and Samuel J. Gershman.  Building machines that learn and think like people.  _Behavioral and Brain Sciences_ , pp. 1–101, 2016.
  * Lerer et al. 2016 Adam Lerer, Sam Gross, and Rob Fergus.  Learning physical intuition of block towers by example.  _arXiv preprint arXiv:1603.01312_ , 2016.
  * Li et al. 2016 Chengtao Li, Daniel Tarlow, Alexander L Gaunt, Marc Brockschmidt, and Nate Kushman.  Neural program lattices.  2016\.
  * Michalski et al. 2014 Vincent Michalski, Roland Memisevic, and Kishore Konda.  Modeling deep temporal dependencies with recurrent grammar cells"".  In _Advances in neural information processing systems_ , pp. 1925–1933, 2014.
  * Mnih et al. 2013 Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller.  Playing atari with deep reinforcement learning.  _arXiv preprint arXiv:1312.5602_ , 2013.
  * Munakata et al. 1997 Yuko Munakata, James L McClelland, Mark H Johnson, and Robert S Siegler.  Rethinking infant knowledge: toward an adaptive process account of successes and failures in object permanence tasks.  _Psychological review_ , 104(4):686, 1997.
  * Odena et al. 2016 Augustus Odena, Vincent Dumoulin, and Chris Olah.  Deconvolution and Checkerboard Artifacts.  _Distill_ , 2016.  doi: 10.23915/distill.00003.
  * Perez et al. 2017 Ethan Perez, Florian Strub, Harm de Vries, Vincent Dumoulin, and Aaron Courville.  FiLM: Visual Reasoning with a General Conditioning Layer.  _arXiv:1709.07871 [cs, stat]_ , September 2017.
  * Reed & De Freitas 2015 Scott Reed and Nando De Freitas.  Neural programmer-interpreters.  _arXiv preprint arXiv:1511.06279_ , 2015.
  * Santoro et al. 2017 Adam Santoro, David Raposo, David GT Barrett, Mateusz Malinowski, Razvan Pascanu, Peter Battaglia, and Timothy Lillicrap.  A simple neural network module for relational reasoning.  _arXiv preprint arXiv:1706.01427_ , 2017.
  * Schmidhuber 2015 Jürgen Schmidhuber.  On learning to think: Algorithmic information theory for novel combinations of reinforcement learning controllers and recurrent neural world models.  _arXiv preprint arXiv:1511.09249_ , 2015.
  * Spelke 1988 Elizabeth S Spelke.  Where perceiving ends and thinking begins: The apprehension of objects in infancy.  1988\.
  * Spelke 1990 Elizabeth S Spelke.  Principles of object perception.  _Cognitive science_ , 14(1):29–56, 1990.
  * Spelke & Kinzler 2007 Elizabeth S. Spelke and Katherine D. Kinzler.  Core knowledge.  _Developmental science_ , 10(1):89–96, 2007.
  * Srivastava et al. 2015 Nitish Srivastava, Elman Mansimov, and Ruslan Salakhudinov.  Unsupervised learning of video representations using lstms.  In _International Conference on Machine Learning_ , pp. 843–852, 2015.
  * Sutskever et al. 2009 Ilya Sutskever, Geoffrey E. Hinton, and Graham W. Taylor.  The recurrent temporal restricted boltzmann machine.  In _Advances in Neural Information Processing Systems_ , pp. 1601–1608, 2009.
  * Ullman et al. 2017 Tomer D Ullman, Elizabeth Spelke, Peter Battaglia, and Joshua B Tenenbaum.  Mind games: Game engines as an architecture for intuitive physics.  _Trends in Cognitive Sciences_ , 21(9):649–665, 2017.
  * Watters et al. 2017 Nicholas Watters, Andrea Tacchetti, Theophane Weber, Razvan Pascanu, Peter Battaglia, and Daniel Zoran.  Visual Interaction Networks.  _arXiv:1706.01433 [cs]_ , June 2017.
  * Werbos 1988 Paul J Werbos.  Generalization of backpropagation with application to a recurrent gas market model.  _Neural networks_ , 1(4):339–356, 1988.
  * Williams 1989 Ronald J Williams.  Complexity of exact gradient computation algorithms for recurrent neural networks.  Technical report, Technical Report Technical Report NU-CCS-89-27, Boston: Northeastern University, College of Computer Science, 1989.
  * Wu et al. 2015 Jiajun Wu, Ilker Yildirim, Joseph J Lim, Bill Freeman, and Josh Tenenbaum.  Galileo: Perceiving physical object properties by integrating a physics engine with deep learning.  In _Advances in neural information processing systems_ , pp. 127–135, 2015.
  * Xu et al. 2015 Kelvin Xu, Jimmy Ba, Ryan Kiros, Kyunghyun Cho, Aaron Courville, Ruslan Salakhudinov, Rich Zemel, and Yoshua Bengio.  Show, attend and tell: Neural image caption generation with visual attention.  In _International Conference on Machine Learning_ , pp. 2048–2057, 2015.
  * Zaheer et al. 2017 Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Ruslan R Salakhutdinov, and Alexander J Smola.  Deep sets.  In _Advances in Neural Information Processing Systems_ , pp. 3394–3404, 2017.

##  Appendix A Experiment Details

In all experiments we train the networks using ADAM Kingma & Ba, 2014 with
default parameters, a batch size of 64 and $50\,000$ train + $10\,000$
validation + $10\,000$ test inputs. The quality of the learned groupings is
evaluated by computing the Adjusted Rand Index (ARI; Hubert & Arabie 1985)
with respect to the ground truth, while ignoring the background and overlap
regions (as is consistent with earlier work Greff et al., 2017). We use early
stopping when the validation loss has not improved for 10 epochs.

###  A.1 Bouncing Balls

The bouncing balls data is similar to previous work Sutskever et al., 2009
with a few modifications. The data consists of sequences of $64\times 64$
binary images over 30 time-steps and balls are randomly sampled from two
types: one ball is six times heavier and 1.25 times larger in radius than the
other. The balls are initialized with random initial positions and velocities.
Balls bounce elastically against each other and the image window.

As in previous work Greff et al., 2017 we use a convolutional encoder-decoder
architecture with a recurrent neural network as bottleneck, that is updated
according to (4):

  1. 1.

$4\times 4$ conv. 16 ELU. stride 2. layer norm

  2. 2.

$4\times 4$ conv. 32 ELU. stride 2. layer norm

  3. 3.

$4\times 4$ conv. 64 ELU. stride 2. layer norm

  4. 4.

fully connected. 512 ELU. layer norm

  5. 5.

recurrent. 250 Sigmoid. layer norm on the output

  6. 6.

fully connected. 512 RELU. layer norm

  7. 7.

fully connected. $8\times 8\times 64$ RELU. layer norm

  8. 8.

$4\times 4$ reshape 2 nearest-neighbour, conv. 32 RELU. layer norm

  9. 9.

$4\times 4$ reshape 2 nearest-neighbour, conv. 16 RELU. layer norm

  10. 10.

$4\times 4$ reshape 2 nearest-neighbour, conv. 1 Sigmoid

Instead of using transposed convolutions (to implement the "de-convolution")
we first reshape the image using the default nearest-neighbour interpolation
followed by a normal convolution in order to avoid frequency artifacts Odena
et al., 2016. Note that we do not add layer norm on the recurrent connection.

At each timestep $t$ we feed
$\gamma_{:,k}(\bm{\psi}_{:,k}^{(t-1)}-\bm{\hat{x}}^{(t)})$ as input to the
network, where $\bm{\tilde{x}}$ is the input with added bitflip noise
($p=0.2$). Consistent with earlier work Greff et al., 2017 R-NEM is trained
with a next-step prediction objective, the prior for each pixel in the data is
set to a Bernoulli distribution with $p=0$, and we prevent conflicting
gradient updates by not back-propagating any gradients through $\gamma$.

The Interaction Function $\bm{\Upsilon}^{\text{R-NEM}}$network is structured
as follows:

  * •

$\text{MLP}^{\hskip 1.0pt\textit{enc}}$: fully connected. 250 RELU. layer norm

  * •

$\text{MLP}^{\hskip 1.0pt\textit{emb}}$: fully connected. 250 RELU. layer norm

  * •

$\text{MLP}^{\hskip 1.0pt\textit{eff}}$: fully connected. 250 RELU. layer norm

  * •

$\text{MLP}^{\hskip 1.0pt\textit{att}}$: fully connected. 100 Tanh. layer norm
- fully connected. 1 Sigmoid.

We experimented with deeper architectures, but were unable to observe
significant improvement.

#### Comparison and Extrapolation

In the comparison experiment both R-NEM and RNN-EM are trained with $K=5$
(unless otherwise mentioned), following insights from Greff et al. 2017. On
the extrapolation task we adjusted the number of components at test time to
$K=8$.

When comparing to RNN-EM we used $\bm{\Upsilon}=\bm{\Upsilon}^{\text{RNN-
EM}}$. For comparing to RNN we set $K=1$, and used
$\bm{\Upsilon}=\bm{\Upsilon}^{\text{RNN-EM}}$, yielding a standard recurrent
autoencoder that receives at each time-step the difference between the
prediction and the noisy ground-truth as input. In case of LSTM, we
additionally replace the recurrent layer with an LSTM update. The R-NEM _no
att_ model is the same as R-NEM, without $\text{MLP}^{\hskip
1.0pt\textit{att}}$, such that $\alpha_{:,:}=1$

#### Simulation

Since the E-step relies on the ground-truth, which was not available for
simulation, we used a thresholded version of $\max_{k}\bm{\psi}$ at 0.1 (such
that everything below becomes 0 and everything above becomes 1) as a
replacement in stead.

#### Occlusion

On the occlusion dataset we used three balls with equal mass. The curtain was
spawned at a random location for each sequence. We trained R-NEM with $K=5$.

###  A.2 Space Invaders

We used a pre-trained DQN to produce a dataset with sequences of 25 time-
steps. The DQN receives a stack of four frames as input and we recorded every
first frame of this stack. These frames were first pre-processed as in Mnih et
al. 2013 and then thresholded at $0.0001$ to obtain binary images.

Since the images are $84\times 84$ we used a different encoder and decoder,
given by:

  1. 1.

$4\times 4$ conv. 16 ELU. stride 2. layer norm

  2. 2.

$4\times 4$ conv. 32 ELU. stride 2. layer norm

  3. 3.

$4\times 4$ conv. 32 ELU. stride 2. layer norm

  4. 4.

$4\times 4$ conv. 32 ELU. stride 2. layer norm

  5. 5.

fully connected. 512 ELU. layer norm

  6. 6.

recurrent. 250 Sigmoid. layer norm on the output

  7. 7.

fully connected. 512 RELU. layer norm

  8. 8.

fully connected. $8\times 8\times 64$ RELU. layer norm

  9. 9.

$4\times 4$ reshape 2 nearest-neighbour, conv. 32 RELU. layer norm

  10. 10.

$4\times 4$ reshape 2 nearest-neighbour, conv. 32 RELU. layer norm

  11. 11.

$4\times 4$ reshape 2 nearest-neighbour, conv. 16 RELU. layer norm

  12. 12.

$4\times 4$ reshape 2 nearest-neighbour, conv. 1 Sigmoid

We used the same architecture for $\Upsilon^{\text{R-NEM}}$, with the only
difference that at each time-step we concatenated an embedding of the action
produced by the agent to the hidden state. Here we used a single layer MLP
with 10 units and a _ReLU_ activation function to compute this embedding.

In the Atari experiment we trained with $K=4$ and reduced the input noise to
0.02, in order to preserve tiny elements such as bullets (that only occupy 1-2
pixels).


* [60] Z. Sun, J. A. Carrillo, and C.-W. Shu, An entropy stable high-order discontinuous Galerkin method for cross-diffusion gradient flow systems, arXiv preprint arXiv:1810.03221, (2018).
  * [61] E. Tadmor, The numerical viscosity of entropy stable schemes for systems of conservation laws. I, Math. Comput., 49 (1987), pp. 91–103.
  * [62] H. Terashima and G. Tryggvason, A front-tracking/ghost-fluid method for fluid interfaces in compressible flows, J. Comput. Phys., 228 (2009), pp. 4012–4037.
  * [63] S. Tokareva and E. F. Toro, HLLC-type Riemann solver for the Baer–Nunziato equations of compressible two-phase flow, J. Comput. Phys., 229 (2010), pp. 3573–3604.
  * [64] N. Wintermeyer, A. R. Winters, G. J. Gassner, and D. A. Kopriva, An entropy stable nodal discontinuous galerkin method for the two dimensional shallow water equations on unstructured curvilinear meshes with discontinuous bathymetry, J. Comput. Phys., 340 (2017), pp. 200–242.
  * [65] A. R. Winters and G. J. Gassner, Affordable, entropy conserving and entropy stable flux functions for the ideal MHD equations, J. Comput. Phys., 304 (2016), pp. 72–108.
  * [66] X. Zhang and C. Shu, On positivity-preserving high order discontinuous Galerkin schemes for compressible Euler equations on rectangular meshes, J. Comput. Phys., 229 (2010), pp. 8918–8934.
  * [67] X. Zhang and C.-W. Shu, On maximum-principle-satisfying high order schemes for scalar conservation laws, J. Comput. Phys., 229 (2010), pp. 3091–3120.