R ELATIONAL N EURAL E XPECTATION M AXIMIZATION : U NSUPERVISED D ISCOVERY OF O BJECTS AND THEIR I NTERACTIONS
Anonymous authors
Paper under double-blind review
A BSTRACT
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 solve this problem, we present a novel method that incorporates prior knowledge about the compositional nature of human perception to factor interactions between object-pairs and to learn them efﬁciently.
It learns to discover objects and to model physical interactions between them from raw visual images in a purely unsupervised fashion.
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 show its ability to handle occlusion and that it can extrapolate learned knowledge to environments with different numbers of objects.
I NTRODUCTION
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 (Battaglia et al., 2016) .The
most successful machine learning approaches to commonsense 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 (Battaglia et al., 2016;
Chang et al., 2016; Watters et al., 2017).
Learning is supervised , however, since it relies on object-representations from external sources (e.g. a physics simulator), which are unavailable in most real-world scenarios.
Neural approaches that learn to directly model motion or physical interactions in pixel space offer an alternative solution (Srivastava et al., 2015; Sutskever et al., 2009).
While unsupervised, these methods suffer from a lack compositionality at the representational level of objects.
This prevents such end-to-end neural approaches from efﬁciently learning functions that operate on the entity level and generalize in a human way (c.f. Battaglia et al.
; Lake et al.
; Santoro et al.
(2017), but see Perez et al.
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 (Battaglia et al., 2016;
Chang et al., 2016; Santoro et al., 2017), enabling it to factor interactions between object-pairs, learn efﬁciently, and generalize to environments with a different number of objects without re-training.
M ETHOD
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 its compositional structure.
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 efﬁciently.
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 efﬁciently (Section 2.2).
2. N EURAL E XPECTATION M AXIMIZATION
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 efﬁciently in a distributed representation θ k for each object.
At a high-level, the idea is that if we were to have access to the family of distributions P ( x | θ k ) (a statistical model of images given object representations θ 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 ( θ 1 , . .
. , θ 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 ( x | θ k ) , which N-EM learns instead by parametrizing the mixture with a neural network, to then back-propagate through the iterations of the unrolled generalized EM procedure.
Following Greff et al.
(2017), we model each image x ∈ R D as a spatial mixture of K components parametrized by vectors θ 1 , . .
, θ K ∈ R M .
A neural network f φ is used to transform these representations θ k into parameters ψ i,k = f φ ( θ k ) i for separate pixel-wise distributions.
A set of binary latent variables Z ∈ [0 , 1] D × 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 x given θ = ( θ 1 , . .
. , θ K ) is given by:
[Equation]
If f φ has learned a statistical model of images given object representations θ k , then we can compute the object representations for a given image x by maximizing P ( x | θ ) .
Marginalization over z complicates this process and we use generalized EM to maximize the following lowerbound instead:
[Equation]
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 γ i,k :=
P ( z i,k = 1 | x i , ψ old i ) given θ old from the previous iteration.
It yields a new soft-assignment of the pixels to the components (clusters), based on how accurate they model x .
The generalized M-step updates θ old by taking a gradient ascent step on (2), using the previously computed soft-assignments: θ new k
[Equation]
The unrolled computational graph of the generalized EM steps is differentiable, which provides a means to train f φ to implement a statistical model of images given object representations.
Using back-propagation through time (eg.
Werbos (1988); Williams (1989))
we optimize φ in the direction that minimizes the following loss:
[Equation]
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 θ k .
In this case the entirety of f φ consists of a recurrent encoder-decoder architecture that receives γ k ( x − ψ k ) as input at each step.
The learning objective in (3) is prone to trivial solutions in case of overcapacity, which prevents the network from modelling the statistical regularities in the data that correspond to objects.
By adding noise to the input image, or reducing θ in dimension we can guide the network to consider statistical structure that corresponds to objects.
Additionally, RNN-EM can be applied to sequential data and guided to learn object-representations that capture dynamics by evaluating (3) at the next time-step ( predictive coding ).
Another interpretation of the denoising and next-step prediction 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 (Hatﬁeld & Epstein, 1985).
2. R ELATIONAL N EURAL E XPECTATION M AXIMIZATION
RNN-EM (unlike N-EM) is able to capture the dynamics of objects through a parametrized recurrent connection that operates on the object representation θ 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 short-coming 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 θ k at time t as a function of all object representations θ :=
[ θ 1 , . .
, θ K ] at the previous time-step through an interaction function Υ :
[Equation]
Here, W , R are weight matrices, σ is the sigmoid activation function, and ˜ x ( t ) is the input to the recurrent model at time t (possibly transformed by an encoder).
This dynamics model coincides with a standard RNN update rule when Υ RNN-EM k ( θ ) :
= θ k , such that the original RNN-EM is recovered.
The inductive bias incorporated in Υ reﬂects the modeling assumptions about the interactions between objects in the environment, and therefore the nature of θ k ’s interdependence.
If Υ incorporates the assumption that no interaction takes place between objects, then the θ k ’s are fully independent and we recover Υ 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 θ k ’s, then we forfeit useful properties of θ k , such as compositionality.
For example if Υ :=
MLP ( θ ) we can no longer extrapolate learned knowledge to environments with more or fewer than K objects and lose overall data efﬁciency (Santoro et al., 2017).
Instead, we can make efﬁcient use of compositionality among the learned object-representations θ k to incorporate general but guiding constraints on how these may inﬂuence one another (Battaglia et al., 2016; Chang et al., 2016).
In doing so we constrain Υ to capture interdependence between θ k ’s in a compositional manner that enables physical dynamics to be learned efﬁciently, and allow for learned dynamics to be extrapolated to a variable number of objects.
We propose a parametrized interaction function Υ R-NEM that incorporates these modeling assumptions and updates θ k based on the pairwise effects of the objects i (cid:54)
=
k on k : Υ R-NEM k [Equation] [Equation] E k ] with ˆ θ k
=
MLP enc
[Equation]
where [ · ; · ] is the concatenation operator and MLP ( · ) corresponds to a multi-layer perceptron.
First, each θ i is transformed using MLP enc to obtain ˆ θ i , which enables information that is relevant for the object dynamics to be made more explicit in the representation.
Next, each pair ( ˆ θ k , ˆ θ i ) is concatenated and processed by MLP emb , which computes a shared embedding ξ k,i that encodes the interaction between object k and object i .
Unlike other work (Battaglia et al., 2016)
this function is not commutative and 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 ξ k,i we compute e k,i : the effect of object i on object k ; and an attention coefﬁcient α k,i that encodes whether interaction between object i and object k takes place.
These attention coefﬁcients (Bahdanau et al., 2014; Xu et al., 2015) help to select relevant context objects, and can be seen as a more ﬂexible unsupervised replacement of the distance based heuristic that was used in previous work (Chang et al., 2016).
Finally we compute the total effect of θ i (cid:54) =
k on θ k as a weighted sum of the effects multiplied by their attention coefﬁcient.
A visual overview of Υ R-NEM can be seen in Figure 1.
R ELATED W ORK
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 (Battaglia et al., 2013; Ullman et al., 2017; Wu et al., 2015), which limits the to synthetic environments.
The latter group employs
G r ound T r u t h
R e c on s t r u c t i on
[Figure]
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; Chang et al., 2016; Fragkiadaki et al., 2015; Grzeszczuk et al., 1998) or for training purposes (Watters et al., 2017). By incorporating information (eg. position, velocity) about objects they have achieved excellent generalization and simulation capabilities.
Purely unsupervised approaches (Agrawal et al., 2016; Lerer et al., 2016; Sutskever et al., 2009) 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 and takes a ﬁrst step in that direction.
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 and 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 difﬁcult.
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).
They suffer from a similar problem in that the sequential nature of these models prevents a coherent representation to take shape.
Other work, has also taken steps toward combining the learnability of neural networks with the compositionality of symbolic programs in modeling physics (Battaglia et al., 2016; Chang et al., 2016), playing games (Denil et al., 2017; Kansky et al., 2017), learning algorithms (Bošnjak et al., 2017;
Cai et al., 2017; Li et al., 2016; Reed & De Freitas, 2015), visual understanding (Ellis et al., 2017; Johnson et al., 2017), and natural language processing (Andreas et al., 2016; Hu et al., 2017).
E XPERIMENTS
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
BCE
Relational BCE 0.
T e s t
Lo ss R e l a t
i v e t
o B a s e li ne 4 Bouncing Balls (Trained with 4 Balls)
RNNLSTMR-NEMR-NEM no att RNN-EM
BCE
Relational BCE 0.
T e s t
Lo ss R e l a t
i v e t
o B a s e li ne 6-8 Bouncing Balls
(Trained with 4 Balls) 4 Balls 6-8 Balls 0.
A R I S c
o r e
Compositionality Measure other unsupervised All experiments use ADAM (Kingma & Ba, 2014) with default parameters, on 50 K train + 10 K validation + 10 K test sequences and early stopping with a patience of 10 epochs. For MLP enc,emb,eff we used a single layer neural network with 250 rectiﬁed linear units.
For MLP att 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.
We train R-NEM on sequences of 64 × 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 θ k from which each component produces the group reconstruction ψ k does indeed only contain information about a unique object, such that together the θ 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 learn to model the dynamics of bouncing balls.
Comparison We compare the modelling capabilities of R-NEM to an RNN, LSTM (Hochreiter &
Schmidhuber, 1997) 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, as well as the relational BCE that only takes into objects that currently take part in collision.
2 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 γ and the ground-truth assignment of pixels to objects.
In the left column of the ARI plot (right side in Figure 3)
we ﬁnd 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 beneﬁcial (comparing to RNN, LSTM) and that (when comparing to RNN-EM in terms of the relational BCE) Υ R-NEM enables interactions to be modelled accurately.
Extrapolating learned knowledge We use a test-set with sequences containing 6-8 balls summarize the ability of each method to extrapolate their learned knowledge about physical interactions between four balls to environments with more balls.
As can be seen from the middle plot in Figurethat R-NEM again greatly outperforms all other methods.
Notice, that since we report performance 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.
For the
LSTM we observe far worse performance (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 very sophisticated solution for sequences with four balls that does not generalize to dataset with 6-8 balls.
R-NEM and RNN-EM scale far better to this dataset than LSTM.
Notice that although the RNN also suffers to a lesser extend from this type of “overﬁtting”, this is most likely due its inability to produce a reasonable solution in the ﬁrst place.
Therefore we conclude that the compositional representations that RNN-EM and R-NEM frequently produce (see right side of the right plot in Figure 3) yields superior extrapolation capabilities
Attention Further insight in the role of the attention mechanism can be gained by visualizing the attention coefﬁcients, as is done in Figure 2.
For each component k we draw α k,i ∗ ψ i on top of the reconstruction ψ 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 coefﬁcient α 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 it should 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 fails to extrapolate the learned knowledge (6-8 balls).
A likely reason for this behavior is that the range of the sum changes with K .
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 θ k .
In Figure 4 we compare the simulation capabilities of R-NEM to an RNN on the bouncing balls environment.
On the left it shows for each method a sequence with ﬁve 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 = x ( t ) . G r ound T r u t h
R e c on s t r u c t i on [Figure] accurately simulate 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 ﬁnd 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 an RNN.
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.
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.
Space Invaders To test the performance of R-NEM in a more challenging environment, we train it on sequences of 84 × 84 binarized images over 25 time-steps of game-play on Space Invaders from the Arcade Learning Environment (Bellemare et al., 2013)
We use K = 4 and also feed the action of the agent to the interaction function.
Figure 6 conﬁrms 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 ﬁnd that R-NEM groups together different columns of aliens together with the spaceship and three large “shields”.
These seem to be grouped based on their movement, which to some degree coincides with their semantic roles of the environment.
In other examples (not shown) we found that R-NEM frequently groups every other column of the aliens and the three large “shields” separately.
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.
D ISCUSSION AND C ONCLUSION
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 (Ullman et al., 2017) and developmental psychology that intuitive physics and reasoning capabilities are built upon the ability to perceive objects and their interactions Spelke (1988).
The fact that young infants already exhibit this ability, may even suggest an innate bias towards compositionality (Lake et al., 2016;
Munakata et al., 1997; Spelke & Kinzler, 2007).
Inspired by these observations we have proposed R-NEM, a method that incorporates inductive biases about their grouping. G r ound T r u t h R e c [Figure]
t [Figure] u c t i
[Figure]
the existence of objects and interactions, implemented by its clustering objective and interaction function respectively.
The speciﬁc nature of the objects, and their dynamics and interactions can then be learned efﬁciently just from visual observation.
In our experiments ﬁnd 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-speciﬁc 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.
These capabilities could beneﬁt any intelligent agent that needs to operate in a physical environment, with provides exciting future research opportunities.
For example, a reinforcement learner could beneﬁt from the compositional representations produced by R-NEM to learn a modular policy that easily generalizes to novel combinations of known objects.
Another example is to utilize the causal knowledge of the environment captured by R-NEM to foresee the consequences of ones actions.
R-NEM itself could also beneﬁt from top-down feedback to produce groupings, that are more targeted to the task at hand and may deviate from the built-in inductive biases
A E XPERIMENT D ETAILS
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.
We use early stopping when the validation loss has not improved for 10 epochs.
A. B OUNCING B ALLS
The bouncing balls data is similar to previous work Sutskever et al.
(2009) with a few modiﬁcations.
The data consists of sequences of 64 × 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.
we use a convolutional encoder-decoder architecture with a recurrent neural network as bottleneck, that is updated according to (4):
4 × 4 conv.
16 ELU.
stride 2.
layer norm 2. 4 × 4 conv.
32 ELU.
stride 2.
layer norm 3. 4 × 4 conv.
64 ELU.
stride 2.
layer norm 4. fully connected.
512 ELU.
layer norm 5.
recurrent.
250 Sigmoid.
layer norm on the output 6.
fully connected.
512 RELU.
layer norm 7. fully connected.
8 × 8 × 64 RELU.
layer norm 8.
4 × 4 reshape 2 nearest-neighbour, conv.
32 RELU.
layer norm 9.
4 × 4 reshape 2 nearest-neighbour, conv.
16 RELU.
layer norm 10.
4 × 4 reshape 2 nearest-neighbour, conv.
Sigmoid
Instead of using transposed convolutions (to implement the "de-convolution") we ﬁrst reshape the image using the default nearest-neighbour interpolation followed by a normal convolution in order to avoid frequency artifacts Odena et al.
Note that we do not add layer norm on the recurrent connection.
At each timestep t we feed γ : ,k ( ψ ( t − 1) : ,k − ˆ x ( t ) ) as input to the network, where ˜ x is the input with added bitﬂip 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 conﬂicting gradient updates by not back-propagating any gradients through γ .
The Interaction Function Υ R-NEM network is structured as follows:
• MLP enc : fully connected.
250 RELU.
layer norm • MLP emb : fully connected.
250 RELU.
layer norm • MLP eff : fully connected.
250 RELU.
layer norm • MLP att : fully connected.
100 Tanh.
layer norm - fully connected.
1 Sigmoid.
We experimented with deeper architectures, but were unable to observe signiﬁcant improvement.
Comparison and Extrapolation In the comparison experiment both R-NEM and RNN-EM are trained with K = 5 , following insights from Greff et al.
On the extrapolation task we adjust the number of components at test time to K =
When comparing to RNN-EM we use Υ = Υ RNN-EM .
For comparing to RNN we take K = 1 , and use Υ = Υ RNN-EM , which yields 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 MLP att , such that α : , : =
Simulation Since the E-step relies on the ground-truth, which was not available for simulation, we used a thresholded version of max k ψ 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 a single mass.
The curtain was spawned at a random location for each sequence.
A. S PACE I NVADERS
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 ﬁrst frame of this stack.
These frames are ﬁrst pre-processed according to Mnih et al.
and then we threshold the frames at 0 .
to obtain binary images.
Since the images are 84 × 84 we used a different encoder and decoder, given by:
4 × 4 conv.
16 ELU.
stride 2.
layer norm 2. 4 × 4 conv.
32 ELU.
stride 2.
layer norm 3. 4 × 4 conv.
32 ELU.
stride 2.
layer norm 4.
4 × 4 conv.
32 ELU.
stride 2.
layer norm 5.
fully connected.
512 ELU.
layer norm 6. recurrent.
250 Sigmoid.
layer norm on the output 7.
fully connected.
512 RELU.
layer norm 8.
fully connected.
8 × 8 × 64 RELU.
layer norm 9.
4 × 4 reshape 2 nearest-neighbour, conv.
32 RELU.
layer norm 10.
4 × 4 reshape 2 nearest-neighbour, conv.
32 RELU.
layer norm 11.
4 × 4 reshape 2 nearest-neighbour, conv.
16 RELU.
layer norm 12.
4 × 4 reshape 2 nearest-neighbour, conv.
Sigmoid
We used the same architecture for Υ 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.
We used a single layer MLP with 10 units and a ReLU activation function.
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).