Structured Object-Aware Physics Prediction for Video Modeling and Planning
Authors: Jannik Kossen
Karl Stelzner
Marcel Hussing
Claas Voelcker
Kristian Kersting
Venue:  (2020)
Keywords: self-supervised learning
probabilistic deep learning
structured models
video prediction
physics prediction
planning
variational auteoncoders
model-based reinforcement learning
VAEs
unsupervised
variational
graph neural networks
tractable probabilistic models
attend-infer-repeat
relational learning
AIR
sum-product networks
object-oriented
object-centric
object-aware
MCTS

Abstract: When humans observe a physical system, they can easily locate components, understand their interactions, and anticipate future behavior, even in settings with complicated and previously unseen interactions. For computers, however, learning such models from videos in an unsupervised fashion is an unsolved research problem.  In this paper, we present STOVE, a novel state-space model for  videos, which explicitly reasons about objects and their positions, velocities, and interactions. It is constructed by combining an image model and a dynamics model in compositional manner and improves on previous work by reusing the dynamics model for inference, accelerating and regularizing training. STOVE predicts videos with convincing physical behavior over hundreds of timesteps, outperforms previous unsupervised models, and even approaches the performance of supervised baselines. We further demonstrate the strength of our model as a simulator for sample efficient model-based control, in a task with heavily interacting objects.


Published conference paper at ICLR 2020
STRUCTURED OBJECT-AWARE PHYSICS PREDICTION FOR VIDEO MODELING AND PLANNING
Jannik Kossen"), Karl Stelzner',, Marcel Hussing', Claas Voelcker Kristian Kersting*
Department of Physics and Astronomy, Heidelberg University
kossen @ @uuud wuu-unnieee mri-heiice ber.........
2.3 Department of Computer Science, TU Darmstadt
2 (sstel zner, kersting]@cs cu-da rmstaadt.cc. de
marccll. hasssigg, e. voelcker feaerrr@suuuu @@stud. tu-darmstadt. de
ABSTRACT
When humans observe a physical system, they can easily locate objects, under- stand their interactions, and anticipate future behavior. For computers, however, learning such models from videos in an unsupervised fashion is an unsolved re- search problem. In this paper, we present STOVE, a novel state-space model for videos, which explicitly reasons about objects and their positions, velocities, and interactions. It is constructed by combining an image model and a dynam- ics model in compositional manner and improves on previous work by reusing the dynamics model for inference, accelerating and regularizing training. STOVE pre- dicts videos with convincing physical behavior over thousands of timesteps, out- performs previous unsupervised models, and even approaches the performance of supervised baselines. We further demonstrate the strength of our model as a sim- ulator for sample efficient model-based control in task with heavily interacting objects.
INTRODUCTION
Obtaining structured knowledge about the world from unstructured, noisy sensory input is a key challenge in artificial intelligence. Of particular interest is the problem of identifying objects from visual input and understanding their interactions. One longstanding approach this is the idea of vision as inverse graphics (Grenander) 1976). which postulates a data generating graphics process and phrases vision posterior inference in the induced distribution. Despite its intuitive appeal, vision as inference has remained largely intractable in practice due to the high-dimensional and multimodal nature of the inference problem. Recently, however, probabilistic models based on deep neural networks have made promising advances in this area. By composing conditional distributions parameterized by neural networks, highly expressive yet structured models have been built. At the same time, advances in general approximate inference, particularly variational techniques, have put the inference problem for these models within reach (Zhang et al. 2019).
Based on these advances, number of probabilistic models for unsupervised scene understanding in single images have recently been proposed. The structured nature of approaches such as AIR (Es- lami et al., (2016), MONet (Burgess et al. 2019). or IODINE (Greff et al.. (2019) provides two key advantages over unstructured image models such variational autoencoders (Kingma & Welling 2014) or generative adversarial networks (Goodfellow al. 2014). First. allows for specifica- tion of inductive biases, such as spatial consistency of objects, which constrain the model and act as regularization. Second, it enables the use of semantically meaningful latent variables, such object positions, which may be used for downstream reasoning tasks.
Building such a structured model for videos instead individual images the natural next chal- lenge. Not only could such model be used in more complex domains, such as reinforcement learn- ing. but the additional redundancy in the data problem (Kosiorek et al. 2018). this end, the notion of temporal consistency may be leveraged Both authors contrihuted equally this work.
simplify and regularize the object detection

---

Published conference paper at ICLR 2020
Legend
Axes
from
previous timestep
Computation
pos velo
size
Intentt
Dynamics Model GNN
predict
fox
Inference
Generative
R2-1)
plas (e-1)
Input image
LSTM
sample
ST
PL
Object SPN
pl
( m))
AL-1: xt)
BC
SPN
to next timestep
p(ay [x)
Image Model
SuPAIR
Inference
Generative
Figure 1: Overview of STOVE's architecture. (Center left) At time the input image x, is processed by an LSTM order to obtain proposal distribution over object states q( 22)). (Top) A separate proposal g(z (2-1) is obtained by propagating the previous state 2 using the dynamics model. (Center) The multiplication of both proposal distributions yields the final variational distribution g(z Zt)) (Right) We sample 2; from this distribution to evaluate the generative distribution p(= x<-1)((x; =,), where p(z= zą-1) shares means but not variances with g(=s =+-1,, and p(r; ze) can be obtained by direct evaluation of 2; in the sum-product networks. Not shown is the dependence on xt- inference routine which allows the inference of velocities. (Best viewed in color.)
as an additional inductive bias, guiding the model to desirable behavior. In situations where inter- actions between objects are prevalent, understanding and explicitly modeling these interactions in an object-centric state-space valuable for obtaining good predictive models (Watters al. 2017). Existing works this area, such as SQAIR (Kosioekk et al. (2018). DDPAE (Hsieh et al. 2018)) R- NEM (Van Steenkiste al., 2018). and COBRA (Watters et 2019) have explored these concepts, but have demonstrated realistic long term video predictions on par with supervised approaches to modeling physics.
To push the limits of unsupervised learning of physical interactions, we propose STOVE, a struc- tured, object-aware video model. With STOVE, we combine image and physics modeling into single state-space modelwhich explicitly reasons about object positions and velocities. end-to-end on pure video data in a self-supervised fashion and learns to detect objects, to model their interactions, and to predict future states and observations. To facilitate learning via variational inference this model, we provide a novel inference architecture, which reuses learned gen- erative physics model in the variational distribution. As we will demonstrate, our model generates convincing rollouts over hundreds of time steps, outperforms other video modeling approaches, and approaches the performance of the supervised baseline which has access to the ground truth object states.
trained
Moving beyond unsupervised learning, we also demonstrate how STOVE can be employed for model-hased reinforcement learning (RI.). Model-based approaches RI. have long been viewed

---

Published conference paper at ICLR 2020
0
(44 (=-1))
9(2 xx)
ppps
( z2-1)
g(=Eivto Ie 14-1
2 velo
F1-1
2 size
2 Ialeent
( 2--1)
Ft-
Zt
(Left) Depiction of the graphical model underlying STOVE. Black arrows denote the generative mechanism and red arrows the inference procedure. The variational distribution g(z: 2t-1, Wt, TT-1) 1) is formed by combining predictions from the dynamics model p(=, =t-1) and the object detection network g(z, 2x). For the RL domain, our approach is extended by action conditioning and reward prediction. (Right) Components of and corresponding variational distri- butions. Note that the velocities are estimated based on the change in positions between timesteps, inducing dependency on Ir-1.
Figure 2:
as potential remedy to the often prohibitive sample complexity of model-free RL, but obtaining learned models of sufficient quality has proven difficult in practice (Sutton & Barto 2011). By con- ditioning state predictions on actions and adding reward predictions to our dynamics predictor, we extend our model to the RL setting, allowing it to be used for search or planning, Our empirical evi- dence shows that actor based Monte-Carlo tree search (MCTS) (Coulom 2007) top of our model is competitive to model-free approaches such as Proximal Policy Optimization (PPO) (Schul- man et al. 2017), while only requiring fraction of the samples.
We proceed by introducing the two main components of STOVE: structured image model and a dynamics model. We show how perform joint inference and training, as well how extend the model to the RL setting. We then present our experimental evaluation, before touching on further related work and concluding.
2 STRUCTURED OBJECT-AWARE VIDEO MODELING
We approach the task of modeling video with frames 21. 2T from a probabilistic perspective. assuming sequence of Markovian latent states 21..... zr. which decompose into the properties a fixed number O of objects. i.e. 2t = (2),.... ,......,,,.. 2)) In the spirit of compositionality, we propose specify and train such model by explicitly combining dynamics prediction model p(#_11 2)) scene model p(are 2,). This yields a state-space model, which can trained on pure video data, using variational inference and an approximate posterior distribution q(2 | x). Our model differs from previous work that also follows this methodology, most notably SQAIR and DDPAE, three major ways:
We propose more compact architecture for variational distribution q(= | Ir), which reuses the dynamics model p(=++1 =z), and avoids the costly double recurrence across time and objects which present in previous work.
We parameterize the dynamics model using a graph neural network, taking advantage of the decomposed nature of the latent state
Instead of treating each 2f as an arbitrary latent code, we explicitly reserve the first six slots of this vector for the object's 's position, size, and velocity, each in x, y direction, and use this
20 information for the dynamics prediction task. We write of = (2f "t,size (,velo 2t,letem)) )
>D
We begin by briefly introducing the individual components before discussing how they are combined to form our state-space space model. Fig. [1 visualises the computational flow of STOVE's inference and generative routines, Fig. (left) specifies the underlying graphical model.

---

Published conference paper at ICLR 2020
2.1 OBJECT-BASH MODELING OF IMAGES USING SUM-PRODUCT ATTEND-INEEN INFER-REPEAT
A variety of object- image models have recently been proposed, many of which are derivatives of attendiinfer-repen (AIR) (Eslami et al. 2016). AIR postulates that each image consists of a set of objects, each of which occupies rectangular region in the image, specified by positional parameters Sheere = (A ). The visual content of each object described by a latent code Whaat By decoding what with a neural network and rendering the resulting image patches in the prescribed location, generative model p(2 =) obtained. Inference is accomplished using recurrent neural network, which outputs distributions over the latent objects q(=") æ), attending to one object at a time. AIR is also capable of handling varying numbers of objects, using an additional set of latent variables.
Sum-Product Attend -efedd--fee--eeeee Repeat (SuPAIR) (Stelzner et al. 2019) utilizes sum-product networks (SPNs) instead of a decoder network to directly model the distribution over object appearances. The tractable inference capabilities of the SPNs used in SuPAIR allow for the exact and efficient com- putation of p( | Zwhere), effectively integrating out the appearance parameters Zwbat analytically. This has been shown to drastically accelerate learning, as the reduced inference workload signifi- cantly lowers the variance of the variational objective. Since the focus of SuPAIR on interpretable object parameters fits our goal of building structured video model, we apply as our image model p(r; 22). Similarly, we use a recurrent inference network as in SuPAIR to model g( where zz). For details SuPAIR, we refer to Stelzner et al. (019))
2.2 MODELING PHYSICAL INTERACTIONS USING GRAPH NEURAL NETWORKS
In order to successfully capture complex dynamics, the state transition distribution p(z++1 | =))
1 = = . -.,2, = needs to be parameterized using a flexible, non-linear estimator. A critical property that should be maintained in the process is permutation invariance, ie.. the output should not depend on the order in which objects appear in the vector 2y This type of function is well captured by graph neural networks, cf. (Santoro depend on the sum of pairwise interactions between objects. Graph neural networks have been extensively used for modeling physical processes in supervised scenarios (Battaglia et al. 2016 2018 Sonzaezz al. 2018 Zhou al. 2018).
2017), which posit that the output should
Following this line of work, we build a dynamics model of the basic form
2l 1.pss ++ Ieelo 1t = =1 (((5)+ +
(1)
o'g0o
where J. ht, represent functions parameterized by dense neural networks. is an attention mech- anism outputting scalar which allows the network to focus on specific object pairs. We assume a constant prior over the object sizes, i.e., L = "L,zize- The full state transition distribution is then given by the Gaussian p(2 1 =() -8) using fixed o.
2.3 JOINT STATE-SPACE MODEL
Next, we assemble state-space model from the two separate models for image modeling and physics prediction. The interface between the two components are the latent positions and velocities. The scene model infers them from images and the physics model propagates them forward in time.
Combining the two yields the state-space model p(z, 2) p(zo )p(2 z0) I[ p(z, (x-1)p(2;
2t). To initialize the state, we model p(20, 21) using simple uniform and Gaussian distributions.
Details are given in Appendix
Our model is trained given video sequences by maximizing the evidence lower bound (ELBO) Ey(=t)) [log p(2, z) - log q(= =:).. This requires formulating a variational distribution g(= | x)) to
approximate the true posterior p(2 | 2).. A natural approach is to factorize this distribution over time, i.e. 9(= x) = a(20 (ra) 1l q(2; 24),, resembling a Bayesian filter. The distribution g(=o (ro) is then readily available using the inference network provided by SuPAIR.
The formulation of g(2, 2,)1, however, is an important design decision. Previous work, including SQAIR and DDPAE, have chosen to unroll this distribution objects, introducing a

---

Published conference paper at ICLR 2020
real
outs
sqaic
supervised
real
ours
sqair
supervised
Figure 3: Visualisation of object positions from the real environment and predictions made by our model, SQAIR, and the supervised baseline, for the billiards and gravity environment after the first 8 frames were given. Our model achieves realistic behaviour, outperforms the unsupervised baselines, and approaches the quality of the supervised baseline, despite being fully unsupervised. For full effect, the reader is encouraged to watch animated versions of the sequences in repository githb..oom/jkk/SSTVV (Best viewed in color.)
costly double recurrence over time and objects, requiring T' O sequential recurrence steps in total. This increases the variance of the gradient estimate, slows down training. and hampers scalability. Inspired by (Becker-Ehmee et al. (2019). we avoid this cost by reusing the dynamics model for the variational distribution. First, we construct the variational distribution q=22ppo xf-1) by slightly ad- justing the dynamics prediction p(={, 22-1), using the same mean values but separately predicted standard deviations. Together with an estimate for the same object by the object detection network g(=ff irt), we construct a joint estimate by multiplying the two Gaussians and renormalizing, yielding another Gaussian:
q= 21), q 24-1) {= xx).
(2)
Intuitively, this results in distribution which reconciles the two proposals. A double recurrence is avoided since g(2g xx) does not depend on previous timesteps and may thus be computed parallel for all frames. Similarly, g( 2=-1) may be computed in parallel for all objects, leading only T + O sequential recurrence steps total. An additional benefit of this approach is that the information learned by the dynamics network is reused for inference - if g( Zt; 2-1) were just another neural network, it would have to essentially relearn the environment's dynamics from scratch, resulting in a waste of parameters and training time. A further consequence is that the image likelihood p(x; xx) is backpropagated through the dynamics model, which has been shown to be beneficial for efficient training (Karl et al., 2017; [BeckerEEhmck et al., 2019). The same procedure applied reconcile velocity estimates from the two networks, where for the image model, velocities 2i veloo are estimated from position differences between two consecutive timesteps. The object scales =f.saale are inferred solely from the image model. The latent states &f,tatent increase the modelling capacity of the dynamics network, not interact with the image model. This then gives the inference procedure for the full latent state 2f (zf siiiz velol 2f, meee)) as illustrated in Fig, 2 (right).
initialised to zero-mean Gaussians, and do
Despite its benefits, this technique has thus far only been used in environments with single object or with known state information. challenge when applying match up the proposals of the two networks. Since the object detection RNN outputs proposals for object locations in indeterminate order. it is not immediately clear how to find the corresponding proposals from the dynamics network. We have, however, found that a simple matching procedure results in good performance: For each 25. we assign the object order that results in the minimal difference of || žt- 1.-ss Il. where ||-|| is the Fuclidean norm. The resulting Euclidean bipartite matching problem can be solved in cubic time using the classic Hungarian algorithm (Kuhn 1955).
a multi-object video setting to
2.4 CONDITIONING ON ACTIONS
In reinforcement learning, an agent interacts with the environment sequentially through actions a, optimize cumulative reward r. To extend STOVE to operate in this setting, we make two changes, yielding a distribution p(2,, 2t 1,))
First, we condition the dynamics model on actions as, enabling conditional prediction based on both state and action. To keep the model invariant to the order of the input objects, the action information concatenated to each object state =ي_ before they are fed into the dynamics model. The model has to leam on its own which of the objects in the scene are influenced by the actions. To facilitate this, have found it helpful to also concatenate appearance information from the

---

Published conference paper at ICLR 2020
Billiards
Cravity
Pixel error
0.20
0.05
0.10
liner
VRNN
0.00
0.00
Position error
0.50
0.50
Imear
SOAIR
DOAPI
0.25
0.25
sugertvised
0.00
0.00
0
20
40
60
80
100
0
20
40
60
80
100
Frame number
Frame number
Figure 4: Mean test set performance of our model compared to baselines. Our approach (STOVE) clearly outperforms all unsupervised baselines and is almost indistinguishable from the supervised dynamics model on the billiards task. (Top) Mean squared errors over pixels the video pre- diction setting (the lower, the better). (Bottom) Mean Euclidean distances between predicted and true positions (the lower, the better). All position and pixel values are in [0, 1]. In all experiments, the first eight frames are given, remaining frames are then conditionally generated. The shading indicates the max and min values over multiple training runs with identical hyperparameters. (Best viewed in color.)
extracted object patches to the object state. While this patch-wist wise code could. in general, be obtained using some neural feature extractor, achieved satisfactory performance by simply using the values per color channel when given colored input.
The second change to the model is the addition of reward prediction. many RL environments, re- wards depend on the interactions between objects. Therefore, the dynamics prediction architecture, presented in Eq. [1 is well suited also predict rewards. We choose share the same encoding of object interactions between reward and dynamics prediction and simply apply two different out- put networks (f in Rq. (T to obtain the dynamics and reward predictions. The total model is again optimized using the ELBO, this time including the reward likelihood p(ra 2e-1, 0e-1).
3 EXPERIMENTAL EVIDENCE
order to evaluate our model, we compare it to baselines three different settings: First, pure video prediction, where the goal is to predict future frames of video given previous ones. Second, the prediction of future object positions, which may be relevant for downstream tasks. Third, we extend one of the video datasets to a reinforcement learning task and investigate how our physics model may be utilized for sample-efficien, model-based reinforcement learning. With this paper, we also release PyTorch implementation of STOVE!!
3.1 VIDEO AND STATE MODELING
Inspired by Watters et al. 2017). we consider grayscale videos of objects moving according physical laws. In particular, we opt for the commonly used bouncing billiards balls dataset, as well
dataset of gravitationally interacting balls. For further details the datasets, Appendix [ When trained using single GTX 1080 Ti, STOVE converges after about 20 hours. baselines, we compare to VRNNs (Chung et al. 2015), SQAIR (Kosiorek et al. 2018). and DDPAE (Hsich et al. 2018
To allow for a fair comparison, we fix the number of objects predicted by SQAIR and DDPAE
The code can be found the GitHub repository crthub.comm/kko//STO. It also contains animated sions of the videos predicted model and the hasclines.

---

Published conference paper at ICLR 2020
0.04
STOVE
0.03
SQAIR
DDPAE
Energy
real
0.02
0.01
0.00
100
70000
Frame number
Frame number
Frame number
Figure 5: Comparison of the kinetic energies of the rollouts predicted by the models, computed based position differences between successive states. Only STOVE's predictions reflect the conserva- tion total kinetic energy in the billiards data set. This is a quantitive measure the convincing physical behavior in the rollout videos. (Left, center) Averages are over 300 trajectories from the test set. Shaded regions indicate one standard deviation. STOVE correctly predicts trajectories with constant energy, whereas SQAIR and DDPAE quickly diverge. (Right) Rolling average over a sin- gle, extremely long-term run. We conjecture that STOVE predicts physical behavior indefinitely. (Best viewed in color.)
the correct amount. Furthermore, compare to a supervised baseline: Here, we consider the ground truth positions and velocities to be fully observed, train our dynamics model on them, resembling the setting of Battaglia et al.] (206)) Since our model needs to infer object states from pixels, this baseline provides upper bound on the predictive performance with our model. In turn, the size of the performance gap between two is good indicator the quality of our state-space model. We also report the results obtained by combining our image model with simple linear physics model, which linearly extrapolates the objects' trajectories. Since VRNN does not reason about object positions, we only evaluate it on the video prediction task. Similarly, the supervised baseline does not reason about images and is considered for the position prediction task only. For more information the baselines, Appendix E
can hope to achieve
Fig. depicts the reconstruction and prediction errors of the various models: Each model is given eight frames of video from the test set as input, which it then reconstructs. Conditioned on this input, the models predict the object positions or resulting video frames for the following 92 timesteps. The predictions pixels and the Euclidean distance between positions based on the best available object matching. We outperform all baselines on both the state and the image prediction task by large margin. Additionally, we perform strikingly close the supervised model.
evaluated on ground truth data by computing the mean squared error between
For the gravitational data, the prediction task appears easier, as all models achieve lower errors than on the billiards task. However, in this regime of easy prediction, precise access to the object states becomes more important, which is likely the reason why the gap between approach and the supervised baseline is slightly more pronounced. Despite this, STOVE produces high-quality rollouts and outperforms the unsupervised baselines.
Table T underlines these results with concrete numbers. We also report results for three ablations of STOVE, which are obtained by (a) training a separate dynamics networks for inference with the same graph neural network architecture, instead of sharing weights with the generative model as argued for in section 2.3, (b) no longer explicitly modelling velocities Avelo in the state, and (c) removing the latent state variables Ptaterx. The ablation study shows that each of these components contributes positively curves for the ablations.
performance STOVE. See Appendix for comparison of training
Fig. 3 illustrates predictions on future object positions made by the models. after each of them was given eight consecutive frames from the datasets. Visually, we find that STOVE predicts physically plausible sequences over long timeframes. This desirable property is not captured by the rollout er- ror: Due to the chaotic nature of our environments, infinitesimally close initial states diverge quickly and model which perfectly follows the ground truth states cannot exist. After this divergence has

---

Published conference paper at ICLR 2020
Table 1: Predictive performance of our approach, the baselines, and ablations (lower better, best unsupervised values are bold). STOVE outperforms all unsupervised baselines and is almost indis- tinguishable from the supervised model on the billiards task. The values computed by summing the prediction errors presented in Fig. in the time interval & E [9. 18), i... the first ten predicted timesteps. In parentheses, standard deviations across multiple training runs are given.
Billiards (pixels)
Billiards (positions)
Gravity (pixels)
Gravity (positions)
STOVE (ours)
0.24
0.418(20)
0.040(3)
0.142(7)
VRNN
0.526(14) (4]]
0.055(12)
SQAIR
0.591
0.804
0.070
0.194
DDPAR
0.405
0.482
0.120
0.298
Linear
0.841(5)
1.348(15)
0.196(2)
0. (9)
Supervised
(3232(37)
0.013(2)
Abl: Double Dynamics
0.262
0.458
0.042
0.154
Abl: No Velocity
0.272
0.460
0,053
0.174
Abl: No Latent
0.338
0.050
0.089
0.235
occurred, the rollout eтror no longer provides any information on the quality of the learned phys- ical behavior. We therefore turn investigating the total kinetic energy of the predicted billiards trajectories. Since the collisions in the training set are fully elastic and frictional forces are not present, the initial energy should be conserved. Fig. shows the kinetic energies of trajectories pre- dicted by STOVE and its baselines, computed based on the position differences between consecutive timesteps. While the energies of SQAIR and DDPAE diverge quickly in less than 100 frames, the mean energies of STOVE's E's rollouts stay constant and are good estimates of the true energy. We have confirmed that STOVE predicts constant energies and therefore displays realistic looking behavior for at least 100 000 steps. This is in stark contrast to the baselines, which predict teleporting, stopping, or overlapping objects after less than 100 frames. In the billiards dataset used by us and the literature, the total energy is the same for all sequences in the training set. See Appendix B for a discussion of how STOVE handles diverse energies.
3.2 MODEL-BASER CONTROL
To explore the usefulness of STOVE for reinforcement learning, we extend the billiards dataset into a reinforcement learning task. Now, the agent controls one of the balls using nine actions, which correspond to moving in one of the eight (inter)cardinal directions and staying rest. The goal is to avoid collisions with the other halls, which elastically bounce off of each other, the walls, and the controlled ball. A negative reward of -1 is given whenever the controlled ball collides with one of the others. To allow the models to recognize the object controlled by the agents we now provide with RGB input in which the balls are colored differently. Starting with a random policy, we iteratively gather observations from the environment, i. sequences of images, actions, and rewards. Using these, train our model as described in Sec. [2.4] To obtain policy based on our world model, we use Monte-Carlo tree search (MCTS), leveraging our model as a simulator for planning. Using this policy, we gather more observations and apply them to refine the world model. As upper bound obtained by MCTS when the real environment used for planning. As a model-free baseline, we consider PPO (Schulman et al. (017). which is a state-of-the-art algorithm on comparable domains such as Atari games. To explore the effect of the availability of state information, we also PPO on a version the environment which, instead of images, the ground-trum object positions and velocities are observed directly.
the performance achievable in this manner, we report the results
Learning curves for each of the agents are given in Fig. 6 (left), reported at intervals of 10000 samples taken from the environment, up total of 000. For model, we collect the first 50 000 samples using random policy to provide an initial training set. After that, the described training loop is used, iterating between collecting 10 000 observations using an MCTS- bassd policy and refining the model using examples sampled from the pool of previously seen observations. After 030 samples, PPO has not yet seen enough samples to converge, whereas our model quickly leams to meaningfully model the environment and thus produces a better policy at this stage. Even when PPO is trained ground truth states, MCTS based on STOVE remains comparable.

---

Published conference paper at ICLR 2020
roward convergeasee
Camutaaiio zoward 100 stepss
-0.5
-4
-1.5
-6
-2.0
-8
-2.5
-3.0
(ours)
MCTS real eav
PPO on cav state
images
-14
-4.5
Random baseline
-5.0
10
30
100 110
120
130
enviromment samples (O0)
Figure Comparison of all models sample efficiency and final performance. (Left) Mean cumulative reward over steps on environment, averaged over 100 environments, using the specified policy. The shaded regions correspond to one-tenth of a standard deviation. In addition the training curves, two constant baselines shown, representing random policy and corresponding to the MCTS based policy when using the real environment as a simulator. (Right) Final performance of all approaches, after training each model to convergence. The shaded region corresponds to standard deviation. (Best viewed in color.)
After training each model to convergence, the final performance of all approaches is reported Fig. (right). In this case, PPO achieves slightly better results, however it only converges after training approximately 000 000 steps, while our approach only uses 130 000 samples. After around 5000000 steps. PPO does eventually surpass the performance of STOVE-based MCTS. Additionally, we find that MCTS on STOVE yields almost the same performance as the real environment, indicating that it can be used anticipate and avoid collisions accurately.
4 RELATED WORK
Multiple lines of work with the goal of video modeling or prediction have emerged recently. Promi- nently, the supervised modeling of physical interactions from vidcos has been investigated by Fragkiadaki et al. (2015). who train model play billiards with single ball. Similarly, graph neural networks have been trained a supervised fashion predict the dynamics objects from images (Watters et al. 2017; 2017: Sanchez- et al.] 2018 [Sun et al., 2018: 2019) or ground truth states (Kipf al. 2018 Wang al.) 2018 Chang et al. (207)). A number of works learn object interactions games in terms of rules instead continuous dynamics (Guzdial al. 2017 Ersen & Sariel 2014). Janner et al. (2019) show successful planning based on learned interactions. but assume access image segmentations. Several unsupervised approaches address the problem by fitting parameters a physies engine to data (Taques al. 2019 Wu et al.) 2016 2015). This necessitates specifying in advance which physical laws govern the observed interactions. In the fully unsupervised setting, mainly unstructured variational approaches have been explored (Babaeizadeh) al. 2017 Chung al.. 2015 Krishnan al., 2015). However, without the explicit notion ob- jects, their performance in scenarios with interacting objects remains limited. Nevertheless, unstruc- tured video models have recently been applied model-based RL and have been shown improve sample efficiency when used simulator for the real environment (Oh et al., 2015 Kaiser et al., 2020).
Only small number of works incorporate objects into unsupervised video models. Xu et al. (2019) and Ehrhardt et al. (2018) take con-probabillitic autoencoding approaches discovering objects real-world videos. COBRA (Watters et al., (2019) represents model--ased RL approach based on MONet, but is restricted environments with interacting objects and only uses one-step search build its policy. Closest to STOVE
small number of probabilistic models, namely SQAIR (Kosiorek et al. 2018), R-NEM (Van Steenkiste et al. 2018 Greff et al. (2017), and DDPAE (Hsich) et al.) 2018). R-NEM learns mixture model via expectation maximization unrolled through time and handles interactions between objects in factorized fashion. However, it lacks explicitly

---

Published conference paper at ICLR 2020
structured latent space, and requires noise the input data to avoid local minima. Both DDPAE and SQAIR extend the AIR approach to work on videos using standard recurrent architectures. As discussed, this introduces a double recurrence over objects and time, which is detrimental for performance. However, SQAIR capable of handling varying number of objects, which is not something we consider in this paper.
5 CONCLUSION
We introduced STOVE, structured, object-aware model for unsupervised video modeling and plan- ning. It combines recent advances in unsupervised image modeling and physics prediction into a single compositional state-space model. The resulting joint model explicitly reasons about object po- sitions and velocities, and is capable of generating highly accurate video predictions in domains fea- turing complicated non-linea interactions between objects. As our experimental evaluation shows, it outperforms previous unsupervised approaches and even approaches the performance and visual quality of supervised model.
Additionally, we presented an extension of the video learning framework the RL setting. Our periments demonstrate that our model may be utilized for sample-efficient model-based control a visual domain, making headway towards a long standing goal of the model-basee -ased RL community. In particular, STOVE yields good performance with more than order of magnitude fewer samples compared to the model-free baseline, even when paired with a relatively simple planning algorithm like MCTS.
At the same time, STOVE also makes several assumptions for the sake of simplicity. Relaxing them provides interesting avenues for future research. First, we assume fixed number objects, which may be avoided by performing dynamic object propagation and discovery like in SQAIR. Second, we have inherited the assumption of rectangular object masks from AIR. Applying flexible model such as MONet (Burgess et al. 2019) GENESIS (Engeleke et al.) (020)) may alleviate this, but also poses additional challenges, especially regarding the explicit modeling of movement. Finally, the availability of high-quality learned state-space models enables the use of more sophisticated planning algorithms in visual domains (Chua et al. 2018). In particular, by com- bining planning with policy and value networks, model-free and model- based RL may be integrated into comprchensive system (Buckman al. 2018).
The authors thank Adam Kosiorek for assistance with the SQAIR experi- ments and Emilien Dupont for helpful discussions about conservation laws in dynamics models. KK acknowledges the support of the thine-Main universities" network for "Deep Continuous-Discrete Machine Learning" (DeCoDeML.).
Acknovledgmentt.
REFERENCES
Mohammad Babacizadeh, Chelsea Finn, Dumitru Erhan, Roy H Campbell, and Sergey Levine. Stochastic variational video prediction. In Proceedings of ICLR, 2017.
Peter Battaglia, Razvan Pascanu, Matthew Lai, Danilo Jimenez Rezende, et al. Interaction networks for learning about objects, relations and physics. In Proceedings of NewrIPS, pp. 4502-4510, 2016.
Peter W Battaglia, Jessica B Hamrick. Victor Bapst, Alvaro Sanchez-Gonzalez Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, al. Relational inductive biases, deep learning, and graph networks. arXiv preprint 1806.v0800000 2018.
Philip Becker-Ehmed Jan Peters, and Patrick Van Der Smagt. Switching linear dynamics for varia- tional bayes filtering. Proceedings of ICML, 2019.
Jacob Buckman, Danijar Hafner, George Tucker, Eugene Brevdo, and Honglak Lec. Sample- efficient reinforcement learning with stochastic ensemble value expansion. In Proceedings of NeurIPS, pp. 8234- 2018.
10

---

Published conference paper at ICLR 2020
Christopher P Burgess, Loic Matthey, Nicholas Watters, Rishabh Kabra, Irina Higgins, Matt Botvinick, and Alexander Lerchner. Monet: Unsupervised scene decomposition and represen- tation. arXiv preprint arXi:: 1901.1139, 2019.
Michael Chang, Tomer Ullman, Antonio Torralba, and Joshua B Tenenbaum. A compositional object-based approach learning physical dynamics. In Proceedings of ICLR, 2017.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learning in handful of trials using probahilistic dynamics models. In Proceedings NeurIPS, pp. 4754-4765, 2018.
Junyoung Chung, Kyle Kastner, Laurent Dinh, Kratarth Goel, Aaron Courville, and Yoshua Ben- gio. A recurrent latent variable model for sequential data. In Proceedings of NeurlPS, pp. 2980- 2988, 2015.
Rémi Coulom. Efficient selectivity and backup operators in monte-carlo tree search. In Computers and Games, pp. 72-83. Springer Berlin Heidelberg, 2007.
Sebastien Ehrhardt, Aron Monszpart, Niloy Mitra, and Andrea Vedaldi. Unsupervised intuitive physics from visual observations. In Proceedings of ACCV, pp. 700-716, 2018.
Martin Engelcke, Adam R. Kosiorek, Oiwi Parker Jones, and Ingmar Posner. Genesis: Generative scene inference and sampling with object-centric latent representations. Proceedings ICLR. 2020.
Mustafa Ersen and Sanem Sariel. Learning behaviors of and interactions among objects through spatio-temporal reasoning. IEEE TCIAIG, 7(1):75-87, 2014.
SM Eslami, Nicolas Heess, Theophane Weber, Yuval Tassa, David Szepesvari, Geoffrey Hin- ton, et al. Attend, infer, repeat: Fast scene understanding with generative models. In Proceedings of NeurtPS, pp. 3225 3225-3233, 2016.
Katerina Fragkiadaki, Pulkit Agrawal, Sergey Levine, and Jitendra Malik. Learning visual predictive models of physies for playing billiards. arXiv preprint arXiv: 15111.0444 2015.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Proceedings of NeurIPS, pp. 2672-2680, 2014.
Klaus Greff, Sjoerd van Steenkiste, and Jürgen Schmidhuher. Neural expectation maximization. Advances in Neural Information Processing Systems, pp. 6691-6701, 2017.
Klaus Greff, Raphaël Lopez Kaufmann, Rishab Kabra, Nick Watters, Chris Burgess, Daniel Zoran, Loic Matthey, Matthew Botvinick, and Alexander Lerchner. Multi-object representation learning with iterative variational inference. In Proceedings of ICML, 2019.
U. Grenander. Lectures in Pattern Theory: Vol. 2 Pattern Analysis. Springer-Verlag, 1976.
Matthew Guzdial, Boyang Li, and Mark O Riedl. Game engine learning from video. In Proceedings of ICAI, 2017.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8): 1735-1780, 1997.
Jun-Ting Hsieh, Bingbin Liu, De-An Huang, Li F Fei-Fei, and Juan Carlos Niebles. Learning decompose and disentangle representations for video prediction. In Proceedings of NeurIPS, pp. 517-526, 2018.
Michael Janner, Sergey Levine, William Freeman, Joshua B. Tenenbaum, Chelsea Finn, and Jiajun Wu. Reasoning about physical interactions with object-centric models. In Proceedings of ICLR, 2019.
Miguel Jaques, Michael Burke, and Timothy Hospedales. pervised learning of objects and physics from video. arXiv preprint arXiv: 1166,,, 2019.
Joint unsu-

---

Published conference paper at ICLR 2020
Lukasz Kaiser, Mohammad Babaeizadeh. Piotr Mitos, Blažej Osiński, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, Afroz Mohiuddin, Ryan Sepassi, George Tucker, and Henryk Michalewski. Model based reinforcement learning for atari. In Proceedings of ICLR, 2020.
Maximilian Karl, Maximilian Soelch, Justin Bayer, and Patrick van der Smagt. Deep Variational Bayes Filters: Unsupervised Learning of State Space Models from Raw Data. Proceedings of ICLR, 2017.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. Proceedings of ICLR, 2015.
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. In Proceedings of ICLR, 2014.
Thomas Kipf. Ethan Fetaya, Kuan-Chich Wang, Max Welling, and Richard Zemel. Neural relational inference for interacting systems. In Proceedings of ICML, 2018.
Adam Kosiorek, Hyunjik Kim, Yee Whye Teh, and Ingmar Posner. Sequential attend, infer, repeat: Generative modelling of moving objects. In Proceedings of NeurlPS, pp. 8606-8616, 2018.
Rahul G. Krishnan, Uri Shalit, and David Sontag. arXiv:18/2.08434, 2015.
Deep kalman filters.
arXiv preprint
Harold w Kuhn. The hungarian method for the assignment problem. Naval research logistics quarterly, 2(1-2):83 1955.
Junhyuk Oh, Xiaoxiao Guo, Honglak Lee, Richard I. Lewis, and Satinder Singh. Action-conditional video prediction using deep networks atari games. In Proceedings of NeurlPS, pp. 2845-2853, 2015.
Alvaro Sanchez-Gonzate Nicolas Heess, Jost Tobias Springenberg, Josh Merel, Martin Riedmiller, Raia Hadsell, and Peter Battaglia. Graph networks as learnable physics engines for inference and control. In Proceedings of ICMI,, 2018.
Adam Santoro, David Raposo, David G Barrett, Mateusz Malinowski, Razvan Pascanu, Peter Battaglia, and Timothy Lillicrap. simple neural network module for relational reasoning. Proceedings of NeurlPS pp. 4967-4976, 2017.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alee Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:/707.066 2017.
Karl Stelzner, Robert Peharz, and Kristian Kersting. Faster -npeea- with tractable proba- bilistic models. In Proceedings of ICML. pp. 5966-5975, 2019.
Chen Sun, Abhinav Shrivastava, Carl Vondrick, Kevin Murphy, Rahul Sukthankar, and Cordelia Schmid. Actor-centric relation network. In Proceedings of ECCV, 2018.
Chen Sun, Abhinav Shrivastava. Carl Vondrick, Rahul Sukthankar, Kevin Murphy, and Cordelia Schmid. Relational action forecasting. In Proceedings CVPR, 2019.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT Press, Cambridge, 2011.
Sjoerd Van Steenkiste, Michael Chang, Klaus Greff, and Jürgen Schmidhuber. Relational neural ex- pectation maximization: Unsupervised discovery objects and their interactions. In Proceedings ICLR, 2018
Antonio Vergari, Robert Peharz, Nicola Mauro, Alejandro Molina, Kristian Kersting, and Flo- riana Esposito. Sum-product autoencoding: Encoding and decoding representations using sum- product networks. In Proceedings of AAAI, 2018.
Tingwu Wang, Renjie Liao, Jimmy Ba, and Sanja Fidler. Nervenet: Learning structured policy with graph neural networks. In Proceedings of ICIR, 2018
12

---

Published conference paper at ICLR 2020
Nicholas Watters, Daniel Zoran, Theophane Weber, Peter Battaglia, Razvan Pascanu, and Andrea Tacchetti. Visual interaction networks: Learning a physics simulator from video. In Proceedings of NeurlPS, pp. 4539-4547, 2017.
Nicholas Watters, Loic Matthey, Matko Bosnjak, Christopher Burgess, and Alexander Lerchner. Cobra: Data-efficient model-based rl through unsupervised object discovery and curiosity-driven exploration. arXiv preprint arXiv: 9905099755 2019.
Jiajun Wu, liker Yildirim, Joseph Lim, Bill Freeman, and Josh Tenenbaum. Galileo: Perceiving physical object properties by integrating a physics engine with deep learning. Proceedings of NeurIPS, pp. 127-135, 2015.
Jiajun Wu, Joseph Lim, Hongyi Zhang, Joshua B Tenenbaum, and William Freeman. Physics 101: Learning physical object properties from unlabeled videos. In Proceedings of BMVC, 2016.
Zhenjia Xu, Zhijian Liu, Chen Sun, Kevin Murphy, William T. Freeman, Joshua B. Tenenbaum, and Jiajun Wu. Modeling parts, structure, and system dynamics via predictive learning. In Proceed- ings of ICLR, 2019.
Cheng Zhang, Judith Bütepage, Hedvig Kjellström, and Stephan Mandt. inference. IEEE TPAMI, 2019.
Advances in variational
Jie Zhou, Ganqu Cui, Zhengyan Zhang, Cheng Yang, Zhiyuan Liu, and Maosong Sun. Graph neural networks: review methods applications. arXiv preprint arXiv: 1812.08434, 2018.
13

---

Published conference paper at ICLR 2020
A RECONSTRUCTIONS: SPRITES DATA
SuPAIR does not need latent description of the objects' appearances. Nevertheless, object recon- structions can obtained by using variant approximate MPE (most probable explanation) the sum-product networks as proposed by Vergari et al.] (2018). We follow the AIR approach and reconstruct each object separately and paste into the canvas using spatial transformers. Unlike AIR, SuPAIR explicitly models the background using separate background SPN. reconstruction of the background is also obtained using MPE.
To demonstrate the capabilities of our image model, we also trained our model on a variant of the gravity data which the round balls were replaced by random selection of four different sprites of the same size. Fig. shows the reconstructions obtained from SuPAIR when trained on these more complex object shapes.
B STUDY OF ENERGIES
As discussed in Sec. 3.1 the energies of the ground truth data were constant for all sequences during the training of STOVE. However, initial velocities drawn from random normal distribution. This is the standard procedure of generating the bouncing balls data set as used by previous pub- lications. Under these cireumstances. STOVE does indeed learn to discover and replicate the total energies of the system, while SQAIR and DDPAE do not. Even if trained constant energy data, STOVE does to some extent generalise to unseen energies. Observed velocities and therefore to- tal energies are highly correlated with the true total kinetic energies of the sequences. However as prediction starts, STOVE quickly regresses to the energy of the training set, see Fig. X (left). If trained on a dataset diverse total energies, the performance modelling sequences of different energies increases, see Fig. (right). Rollouts now initially represent the true energy of the observed sequence, although this estimate of the true energy diverges to a constant but wrong energy value. This is an improvement over the model trained on constant energy data, where the regression to the training data energy happens much quicker within around 10 frames. Note that this does not drastically decrease the visual quality of the rollouts the change of total energy over 500 frames is gradual enough. We leave the reliable prediction of rollouts with physically valid constant energy for sequences of varying energies for future work.
time span of around 500 frames
C MODEL DETAILS
Here, present additional details the architecture and hyperparameters of STOVE.
C.I INFERENCE ARCHITECTURE
The object detection network @ (2,) is realised by LSTM (Hochreiter & Schmidhuber
1997) with 256 hidden units, which outputs the mean and standard deviation of the objects' two- dimensional position and size distributions, i.c. g( size 22) with 2 2 = 8 parameters per
object. Given such position distributions for two consecutive timesteps q popos xx- 1), g 2t
2t). with parameters Peg_..m O2f 1,00 Man Jap "pppa
the following velocity estimate based the
difference in position is constructed:
( 1(, = NN Me 22
1.p0s
As described in Sec. positions and velocities are also inferred from the dynamics model g(2p 1t and q( [0 velo a.-1). A joint estimate, including information from both image model and dynamics prediction, is obtained by multiplying the respective distributions and renormalizing. Since both g-distributions Gaussian, the normalized product again Gaussian, with mean and
14

---

Published conference paper at ICLR 2020
Bal
Figure 7: Reconstructions obtained from our image model when using more varied shapes.
standard deviation are given by
2(<4 ae. 4-11) a(26 22) ((2 2--1) = -N(ZG Mei d.)) NNZii o.,a
-N(GG of)
14
Ri ot.
+
0
itd
where we relax our notation for readability E [=] 24 velo and the indices i and refer to the parameters obtained from the image and dynamics model. This procedure is applied independently for the positions and velocities each object.
For 2 latent> we choose dimension 12, such that a full state =f = (=f,pos =F,izer =i Ff,atten) is 18-dimensional.
C.2 GRAPH NEURAL NETWORK
The dynamics prediction is given by the following series of transformations applied to each input state of shape (batch size, number of objecta,/), where / 16, since currently, size information is not propagated through the dynamics prediction.
S Encode input state with linear layer 21.
S2: Apply linear layer [21, 21] to S₁ followed by ReLU non-linearity.
Sy: Apply linear layer [21, 21] S2 and add result S2. This gives the dynamics prediction without relational effects, corresponding to ((²)) in Eq
C₁:
The following steps obtain the relational aspects of dynamics pre- Concatenate the encoded
diction, corresponding to =f²= in Eq. state Si pairwise with all state encoding. yielding tensor of shape (batch size, number of objects, numiber of jects,,
Cz: Apply linear layer [41, 4I] C followed by ReLUU.
Ca: Apply linear layer [41, 20] to C2 followed by ReLU.
Cs: Apply linear layer [21, 21] to C3 and add to C3.
A₁ To obtain attention coefficients a<=;; =f'), apply linear layer [41,41] to C, followed by ReLU.
A2: Apply linear layer [41, 2/] to A₁ followed by ReLU.
Ay: Apply linear layer [21, 1] Ag and apply exponential function.
R1: Multiply C4 with As. where diagonal elements of A3 are masked out to ensure that R? only covers cases where o + o'.
Ry: Sum over R1 for all d, to obtain tensor of shape (batch size, number of objects, s. 2(). This is the relational dynamics prediction.
15

---

Published conference paper at ICLR 2020
Mean Predicted Energy
Trained on Constant Energy
Trained on Diverse Energies
Coestact Energy Short Praduction Prodiction Gbwwroed
0.04
0.02
0.00
0.00
0.01
0.02
0.03
0.04
0.00
0.01
0.02
0.03
0.04
True Energy
True Energy
Figure 8: Mean kinetic energy observed/predicted by STOVE over true energy of the sequences. (left) STOVE is trained sequences of constant kinetic energy. As can be seen from the blue scatter points, STOVE manages to predict sequences of arbitrary lengths which, on average, preserve the constant energy of the test set. When STOVE is applied to sequences of different energies, it manages to infer these energies from observed frames fairly well, with inaccuracies compounding larger energies (red). quickly to the energy value of the training set (orange and green). (right) STOVE is now trained on sequences of varying energies. Compared to the constant energy training. energies from observed as well predicted energies improve drastically. The predictions no longer immediately regress towards specific value (orange). However after 100 frames, the quality of the predicted energies still regresses to wrong value (green). (all) The observed values refers energies obtained as the mean energy value over six initially observed frames. The short (long) time frame refers to an energy obtained as the mean energy over the first 10 (100) frames of prediction. (Best viewed in color.)
following prediction, however,
mean predicted energies diverge
Dj: Sum relational dynamics R2 and self-dynamics Sa, obtaining the input to f in Eq.
Dy: Apply linear layer [21, 20] to D1 followed by tanh non-lincarity.
Dy: Apply linear layer [2l, 2[] D2 followed by tanh non-linearity and add result to Dg.
D4: Concatenate D3 and S1, and apply linear layer [41, 2/] followed by tanh.
Dy: Apply linear layer [21, 21] to D2 and add result to D2 to obtain final dynamics predic- tion.
The output D5 has shape (batch size, number of objects, 21), twice the size of means and standard deviations over the next predicted state.
For the model-based control scenario, the one-hot encoded actions (batch size, action space) are transformed with a linear layer [action space, number of objects encoding size] and reshaped to (action space, number of objects, encoding size). The action embedding and the object appearances (batch size, number of objects, 3) are then concatenated to the input state. The rest of the dynamics prediction follows above. The reward prediction consists of the following steps:
Hi: Apply linear layer [21, 20 to Dj followed by ReL.U.
Hy: Apply linear layer [2], 20] H1.
Hy: Sum over object dimension to obtain tensor shape (batch size,().
H4: Apply linear layer [1, to Hg followed by ReLU.
Hs: Apply linear layer [1/2, 1/4] to H, followed by ReLU.
Hş: Apply linear layer [t/4,0 to H4 followed by a sigmoid non-linearity.
Hs then gives the final reward prediction.
16

---

the gravity Published conference paper at ICLR 2020
C.3 STATE INITIALIZATION
In the first two timesteps, cannot yet apply STOVE's main inference step g( 2/e, "0, described above. order to initialize the latent state over the first two frames, we apply a simplified architecture and only use a partial state at f 0.
t = 0, 2o ~ (( size) 20) is given purely by the object detection network, since no previous states, which could be propagated, exist. 20 is incomplete insofar it does not contain velocity information or latents. At = 1. g(21,pos sine ;1) is still given purely based on the object detec- tion network. Note that for a dynamics prediction of #1, velocity information at = 0 would need be available. However, the previously inferred object positions. We sample Jilttt from the prior Gaussian distribution to assemble the first full initial state 21. At t 2 2, the full inference network can be run: States are inferred both from the object detection network g(zz 11, as well propagated using the dynamics model (( 21-1).
velocities can be constructed based the differences between
In the generative model, similar adjustments are made: p(20.pon size) is given by a uniform prior, velocities and latents are omitted. At t I, velocities are sampled from uniform distribution in planar coordinates p(21.elo) and positions are given by a simple linear dynamics model p(21,pu pus 21, volo) = (20.pos + =1.vllo o). Latents 2all are sampled from a Gaussian prior. Starting t 2, the full dynamics model is used.
C.4 TRAINING PROCEDURE
Our model was trained using the Adam optimizer (Kingma & Ba 2015), with a learning rate of 2 x 10-3 exp(-40 x 10 step) for a total of 83 000 steps with batch size of 256.
D DATA DETAILS
For the billiards and gravitational data, 1000 sequences of length 100 were generated for training. From these, subsequences of lengths 8 were sampled and used optimize ELBO. test dataset of 300 sequences of length 100 was also generated and used for all evaluations. The pixel resolution of the dataset was 32 32 for the billiards data and 50 x 50 for the gravity data. All models for video prediction were learned on grayscale data, with objects of identical appearance. The 3 balls were initialised with uniformly random positions and velocities, rejecting configuration with overlap. They rendered using amti-aliasing. The billiards data models the balls circular objects, which perform elastic collision with each other or the walls data, the halls are modeled as point masses, where, following Watters et al.] (2017). we clip the gravitational force to avoid slingshot effects. Also, add additional basin of attraction towards the center of the canvas and model the balls in their center off mass system avoid drift. Velocities here are initialised orthogonal to the center of the canvas for a stabilising effect. For full details we refer to the file envs py in the provided code.
environment.
E BASELINES FOR VIDEO MODELING
Following Kosiorek et al. 2018), we experimented with different hyperparameter configurations for VRNNs. We varied the sizes the hidden and latent states [h. z), experimenting with the val- ues [256, 16). [512, 32). [1024, 64), and [2048, 32). We found that increasing the model capacity beyond| [112, 32] did not yield large increases in performance, which is why we chose the configura- tion [5 1,, 32] for our experiments. Our VRNN implementation is written Py PyToch and based on httpp::/giitae httpub.ccmmeee emi ttd/Veriiaaa i RecurrentNeurrreeee INetwrr'
SQAIR can handle variable number of objects in each sequence. However, to allow for fairer comparison STOVE, we fixed the number of objects the correct number. This means that in the first timestep, exactly three objects are discovered, which are then propagated in all following timesteps, without further discoveries. Our implementation is based on the original implementation provided by the authors at athttps: //github..tttt. com/ wtt./oosiooee///eooo.ooooo/ kkosiorek/sqaa....
17

---

Published conference paper at ICLR 2020
Billiards
Gravity
0.5
0.4
Position Error
All: Double Dynamks
0.4
0.3
ABL: No velecly
Abl:
0.2
0.3
0.1
0.2
0
20000
40000
60000
20000
40000
60000
Step
Step
Displayed is the mean predicted position error over rollout length of 8 frames as training progresses for the billiards (left) and gravity (right) scenario for STOVE and its ablations. (Best viewed in color.)
Figure 9:
performed using the implementation available https: / github. con. DttthsseeeeDDA---deeee----e
The DDPAE experiments
Default parameters for training DDPAE with billiards datasets are provided with the code. However, the resolution of our bil- liards (32 pixels) and gravity (64 pixels) datasets is different the resolution DDPAE expects (64 pixels). While we experimented with adjusting DDPAE parameters such
latent space dimen- sion to fit our different resolution, best results were obtained when bilinearly scaling our data to the resolution DDPAE expects. DDPAE was trained for 400 000 steps, which sufficed for convergence of models' test set error.
The linear baseline was obtained as follows: For the first frames, we infer the full model state using STOVE. We then take the last inferred positions and velocities of each object and predict future positions by assuming constant, uniform motions for each object. We do not allow objects leave the frame, e. when objects reach the canvas boundary after some timesteps, they stick to it.
Since our dynamics model requires only object positions and velocities as input, it is trivial to construct supervised baseline for our physics prediction by replacing the SuPAIR-inferm states with real, ground-truth states. On these, the model can then be trained in supervised fashion.
F TRAINING CURVES OF ABLATIONS
Fig. 9 we display learning curves STOVE and presented ablations. As mentioned in the main text, the ablations demonstrate the value of the reuse of the dynamics model, the explicit inclusion of a velocity value, and the presence of unstructured latent space in the dynamics model. (Best viewed color.)
G DETAILS ON THE REINFORCEMENT LEARNING MODELS
Our MCTS implementation uses standard UCT formulation for exploration/exploittioo. The e parameter is set to I. in all our experiments. Since the environment does not provide natural endpoint, off all rollouts depth of 20 timesteps. We found this to be good trade-off between runtime accuracy.
When expanding a node on the true environment, we compute the result of the most promising action, and then start rollout using random policy from the resulting state. For the final evaluation, total of 200 nodes expanded. To hetter utilize the GPU, slightly different approach is used for STOVE. When we expand node in this setting, we predict the results of all actions simultaneously, and compute rollout from each resulting position. In turn, only 50 nodes are expanded. To estimate the node value function, the average reward node's visit counter is increased by 1. Furthermore, we discount the reward predicted STOVE with a factor of 0.95 per timestep to account for the higher uncertainty of longer rollouts. This is not done the baseline running the real environment, since it hehaves deterministically
all rollouts propagated back to the root and each

---

Published conference paper at ICLR 2020
For PPO, we employ a standard convolutional neural network an actor-critic for the evaluation on images and a MLP for the evaluation on states. The image network consists of two convolutional layers, each using 32 output filters with a kernel size of 4 and 3 respectively and a stride of 2. The MLP consists of two fully connected layers with 128 and 64 hidden units, In both cases, an addi- tional fully connected layer links the outputs of the respective base to an actor and a critic head. For the convolutional base, this linking layer employs 512 hidden units, for the MLP 64. All previously mentioned layers use rectified linear activations. The actor head predicts a probability distribution over next actions using softmax activation function while the critic head outputs a value estimation for the current state using linear prediction. We tested several hyperparameter configurations but found the following be the most efficient one. To update the actor-critic architecture, sample 32 trajectories of length 16 from separate environments in every batch. The training uses an Adam optimizer with learning rate of 2 x 10-1 and and value of x 10-5, The clipping parameter of PPO is of the sampled data. The value loss is weighted at 5 x 10" and the entropy coefficient is set to I x 10 10-
1 10-1 We update the network for epochs in each batch using 32 mini-batche
19


Reviews:
[{"review_id": "rylOjQdDuS", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper introduces a structured deep generative model for video frame prediction, with an object recognition model based on the Attend, Infer, Repeat (AIR) model by Eslami et al. (2016) and a graph neural network as a latent dynamics model. The model is evaluated on two synthetic physics simulation datasets (N-body gravitational systems and bouncing billiard balls) for next frame prediction and on a control task in the billiard domain. The model can produce accurate predictions for several time steps into the future and beats a variational RNN and SQAIR (sequential AIR variant) baseline, and is more sample-efficient than a model-free PPO agent in the control task.\n\nOverall, the paper is well-structured, nicely written and addresses an interesting and challenging problem. The experiments use simple domains/problems, but give good insights into how the model performs.\n\nRelated work is covered to a satisfactory degree, but a discussion of some of the following closely related papers could improve the paper:\n* Chang et al., A Compositional Object-Based Approach To Learning Physical Dynamics, ICLR 2017\n* Greff et al., Neural Expectation Maximization, NeurIPS 2017\n* Kipf et al., Neural Relational Inference for Interacting Systems, ICML 2018\n* Greff et al., Multi-object representation learning with iterative variational inference, ICML 2019\n* Sun et al., Actor-centric relation network, ECCV 2018\n* Sun et al., Relational Action Forecasting, CVPR 2019\n* Wang et al., NerveNet: Learning structured policy with graph neural networks, ICLR 2018\n* Xu et al., Unsupervised discovery of parts, structure and dynamics, ICLR 2019\n* Erhardt et al., Unsupervised intuitive physics from visual observations, ACCV 2018\n\nIn terms of clarity, the paper could be improved by making the used model architecture more explicit, e.g., by adding a model figure, and by providing an introduction to the SuPAIR model (Stelzner et al., 2019) — the authors assume that the reader is more or less familiar with this particular model. It is further unclear how exactly the input data is provided to the model; Figure 2 makes it seem that inputs are colored frames, section 3.1 mentions that inputs are grayscale videos (do all objects have the same appearance or different shades of gray?), which is in conflict with the statement on page 5 that the model is provided with mean values of input color channels. Please clarify.\n\nIn terms of novelty, the proposed modification of SQAIR (separating object detection and latent dynamics prediction) is novel and likely leads to a speed-up in training and evaluation. Using a Graph Neural Network for modeling latent physics is reasonable and has been shown to work on related problems before (see referenced work above and related work mentioned in the paper). Similarly, using such a model for planning/control is interesting and adds to the value of the paper, but has in related settings been explored before (e.g. Wang et al. (ICLR 2018) and Sanchez-Gonzalez (ICML 2018)).\n\nExperimentally, it would be good to provide ablation studies (e.g. a different object detection module like AIR instead of SuPAIR, not splitting the latent variables into position, velocity, size etc.) and run-time comparisons (wall-clock time), as one of the main contributions of the paper is that the proposed model is claimed to be faster than SQAIR. The overall model predictions are (to my surprise) somewhat inaccurate, when looking at e.g. the billiard ball example in Figure 2. In Steenkiste et al. (ICLR 2018), roll-outs appear to be more accurate. Maybe a quantitative experimental comparison could help?\n\nWhy does the proposed model perform worse than a model-free PPO baseline when trained to convergence on the control task? What is missing to close this gap?\n\nDo all objects have the same appearance (color/greyscale values) or are they unique in appearance? In the second case, a simpler encoder architecture could be used such as in Jaques et al. (2019) or Xu et al. (ICLR 2019).\n\nOverall, I think that this paper addresses an important issue and is potentially of high interest to the community. Nonetheless I think that this paper needs a bit more work and at this point I recommend a weak reject.\n\nOther comments:\n* This sentence is unclear to me: “An additional benefit of this approach is that the information learned by the dynamics model is reused for inference — […]”\n* What are the failure modes of the model? Where does it break down? \n* How does the model deal with partial occlusion?\n\n---------------------\nUPDATE (after reading the author response and the revised manuscript): My questions and comments are addressed and the additional ablation studies and experimental results on energy conservation are convincing and insightful. I think the revised version of the paper meets the bar for acceptance at ICLR.\n", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I did not assess the derivations or theory."}, "extra_text": {}}, {"review_id": "S1gNkPJnYr", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "In this paper the authors present a graph neural network for modeling the dynamics of objects in simple environments from video. The intuition of the presented system is that it first identifies the different objects from the image using Sum-Product Attend-Infer-Repeat (SuPAIR), which gives the objects positions and sizes. The system uses a “simple matching procedure” to map objects between frames, which allows for the system to extra the object’s velocities. Then a graph neural network is employed to model the dynamics of the particular environment (whether objects bounce, whether there are other forces at play like gravity, etc.). The authors present two environments (Billiards and Gravity) and two evaluations, one focused on predicting future states, and the second focused on using these predictions to play the game. \n\nI think that this paper presents an interesting approach and I agree with the authors of the importance of developing approaches that allow AI to make good predictions of future environments. However, I’m not convinced of many of the technical details in the paper. \n\nI am not certain whether I would classify this work as unsupervised learning. While it’s certainly true that there are no labels in the raw video, the object-finding can be understood as a preprocessing step after which the data is in fact in a fairly standard supervised learning framework. The authors use the term “self-supervised” in the first section, which I believe describes the work more clearly. \n\nThe primary technical contributions of the work appear to be the graph network, the experiments, and their results. While I would have preferred more detail on the graph network in an appendix, it’s acceptable to instead have access to the code. However, the experiments seem set up primarily to evaluate the system as a whole. For example, the inclusion of a supervised learning version of the system where the object’s positions are given exactly sheds light on the quality of SuPAIR. However, SuPAIR is taken from prior work. I would have thought that an entirely different approach, like that used by Ha and Schmidhuber in their World Models paper would have been more appropriate as a comparison as it represents an alternate approach entirely. \n\nThere is a repeated claim made in the paper that the system presents output that is “convincing” and “realistic” over hundreds of time steps. There is no clear definition given for what this means. Figure 1 only presents pixel and positional error for 80 frames, and the error appears to go pretty large (~15%) after only forty frames. The results presented in Figure 4 suggests a much larger timescale, but it’s unclear the quality of the output predictions from it. Some clarity on this or scaling back the claims would improve the paper. \n\nIn terms of related work Guzdial and Riedl’s 2017 “Game Engine Learning from Gameplay Video” appear to use a very similar approach (but with OpenCV instead of SuPAIR and search instead of a graph network) as does Ersen and Sariel’s 2015 “Learning behaviors of and interactions among objects through spatio–temporal reasoning”. These approaches also function over much more complex environments with variable numbers of objects. It would be helpful for the authors to continue adding some discussion of this and related papers. \n\n---\n\nEdit: In response to the author's changes I have increased my rating to a weak accept. This is in large part due to Figure 4, which provides a great deal of additional support to the author's claims and clarity on the technical value of the results. ", "extra_scores": {"experience_assessment": "I have published in this field for several years.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_derivations_and_theory": "N/A"}, "extra_text": {}}, {"review_id": "BJgAmQbAFB", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper presents STOVE, an object-centric structured model for predicting the dynamics of interacting objects. It extends SuPAIR, a probabilistic deep model based on Sum-Product Networks, towards modeling multi-object interactions in video sequences. Compared to prior work, the model uses graph neural networks for learning the transition dynamics and reuses the dynamics model for the state-space inference model, further regularising the learning process. The approach has been tested on simple multi-body physics tasks and performs well compared to other unsupervised and supervised baselines. Additionally, an action-conditional version of STOVE was tested on a visual MPC task (using MCTS for planning) and was shown to learn significantly faster compared to model-free baselines.\n\nThe paper is well written and clearly motivated but comes across as an incremental improvement on top of prior work. Here are a few comments:\n1. The idea of reusing the dynamics model for inference is neat as it helps to regularise the learning process and remove the costly double recurrence, potentially speeding up learning. It would be great if this could be evaluated experimentally via an ablation study — this can be done by using two separate instances of the transition model with separate weights. \n2. A keys step that allows to reconcile the transition model and the object detection network is the matching process. Currently, this is done via choosing the pair with the least position and velocity difference between subsequent time steps. This could give erroneous results in the case of object interactions when objects are fairly close to each other (or colliding). A potentially better way could be to additionally use the content/latent codes for this matching process — as long as the object’s appearance stays similar these can provide good signal that disambiguates different objects.\n3. The experiments presented in the paper are quite simplistic visually — it is not clear if this approach can generalise to more complicated visual settings. Additionally, it would be good to see further comparisons and ablations that quantifies the effect of the different components — e.g. comparing to a combination of image model + black-box MLP dynamics model can quantify the effect of the graph neural network. These results can add further strength to the paper. \n\nOverall, the approach presented in the paper is a bit incremental and the experiments are somewhat simplistic. Further comparisons and ablation experiments can significantly\tstrengthen the paper. I would suggest a borderline accept.", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper presents a method for modeling videos with object-centric structured representations. The paper is well written and clearly motivated. Using a Graph Neural Network for modeling latent physics is a sensible idea and can be beneficial for planning/control. Experimental results show improved performance over the baselines. After the rebuttal, many questions/concerns from the reviewers were addressed, and all reviewers recommend weak acceptance.
Decision: Accept (Poster)