Title: STRUCTURED OBJECT-AWARE PHYSICS PREDICTION FOR VIDEO MODELING AND PLANNING

Abstract: When humans observe a physical system, they can easily locate objects, understand their interactions, and anticipate future behavior. 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 thousands 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.

1. 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 to this is the idea of vision as inverse graphics (Grenander, 1976) , which postulates a data generating graphics process and phrases vision as 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, a number of probabilistic models for unsupervised scene understanding in single images have recently been proposed. The structured nature of approaches such as AIR (Eslami et al., 2016) , MONet (Burgess et al., 2019) , or IODINE (Greff et al., 2019) provides two key advantages over unstructured image models such as variational autoencoders (Kingma & Welling, 2014) or generative adversarial networks (Goodfellow et al., 2014) . First, it allows for the specification 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 as object positions, which may be used for downstream reasoning tasks.

Building such a structured model for videos instead of individual images is the natural next challenge. Not only could such a model be used in more complex domains, such as reinforcement learning, but the additional redundancy in the data can even simplify and regularize the object detection problem (Kosiorek et al., 2018) . To this end, the notion of temporal consistency may be leveraged Figure 1 : Overview of STOVE's architecture. (Center left) At time t, the input image x t is processed by an LSTM in order to obtain a proposal distribution over object states q(z t | x t ). (Top) A separate proposal q(z t | z t-1 ) is obtained by propagating the previous state z t-1 using the dynamics model. (Center) The multiplication of both proposal distributions yields the final variational distribution q(z t | z t-1 , x t ). (Right) We sample z t from this distribution to evaluate the generative distribution p(z t | z t-1 )p(x t | z t ), where p(z t | z t-1 ) shares means -but not variances -with q(z t | z t-1 ), and p(x t | z t ) can be obtained by direct evaluation of x t in the sum-product networks. Not shown is the dependence on x t-1 in the inference routine which allows for the inference of velocities. (Best viewed in color.) as an additional inductive bias, guiding the model to desirable behavior. In situations where interactions between objects are prevalent, understanding and explicitly modeling these interactions in an object-centric state-space is valuable for obtaining good predictive models (Watters et al., 2017) . Existing works in this area, such as SQAIR (Kosiorek et al., 2018) , DDPAE (Hsieh et al., 2018) , R-NEM (Van Steenkiste et al., 2018) , and COBRA (Watters et al., 2019) have explored these concepts, but have not 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 structured, object-aware video model. With STOVE, we combine image and physics modeling into a single state-space modelwhich explicitly reasons about object positions and velocities. It is trained 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 in this model, we provide a novel inference architecture, which reuses the learned generative 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.

Moving beyond unsupervised learning, we also demonstrate how STOVE can be employed for model-based reinforcement learning (RL). Model-based approaches to RL have long been viewed

z t-1 z t x t-1 x t a t-1 r t z o t,velo z o t,size z o t,pos z o t,latent q(z o t,pos | z t-1 ) • q(z o t,pos | x t ) q(z o t,velo | z t-1 ) • q(z o t,velo | x t , x t-1 ) q(z o t,size | x t ) q(z o t,latent | z t-1 )

Figure 2 : (Left) Depiction of the graphical model underlying STOVE. Black arrows denote the generative mechanism and red arrows the inference procedure. The variational distribution q(z t | z t-1 , x t , x t-1 ) is formed by combining predictions from the dynamics model p(z t | z t-1 ) and the object detection network q(z t | x t ). For the RL domain, our approach is extended by action conditioning and reward prediction. (Right) Components of z o t and corresponding variational distributions. Note that the velocities are estimated based on the change in positions between timesteps, inducing a dependency on x t-1 . as a 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 conditioning 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 evidence shows that an actor based on Monte-Carlo tree search (MCTS) (Coulom, 2007) on top of our model is competitive to model-free approaches such as Proximal Policy Optimization (PPO) (Schulman et al., 2017) , while only requiring a fraction of the samples.

We proceed by introducing the two main components of STOVE: a structured image model and a dynamics model. We show how to perform joint inference and training, as well as how to 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 a video with frames x 1 , . . . , x T from a probabilistic perspective, assuming a sequence of Markovian latent states z 1 , . . . , z T , which decompose into the properties of a fixed number O of objects, i.e. z t = (z 1 t , . . . , z O t ). In the spirit of compositionality, we propose to specify and train such a model by explicitly combining a dynamics prediction model p(z t+1 | z t ) and a scene model p(x t | z t ). This yields a state-space model, which can be trained on pure video data, using variational inference and an approximate posterior distribution q(z | x). Our model differs from previous work that also follows this methodology, most notably SQAIR and DDPAE, in three major ways:

• We propose a more compact architecture for the variational distribution q(z | x), which reuses the dynamics model p(z t+1 | z t ), and avoids the costly double recurrence across time and objects which was present in previous work.

• We parameterize the dynamics model using a graph neural network, taking advantage of the decomposed nature of the latent state z.

• Instead of treating each z o t as an arbitrary latent code, we explicitly reserve the first six slots of this vector for the object's position, size, and velocity, each in x, y direction, and use this information for the dynamics prediction task. We write

z o t = (z o t,pos , z o t,size , z o t,velo , z o t,latent ).

We begin by briefly introducing the individual components before discussing how they are combined to form our state-space model. Fig. 1 visualises the computational flow of STOVE's inference and generative routines, Fig. 2 (left) specifies the underlying graphical model.

2.1. OBJECT-BASED MODELING OF IMAGES USING SUM-PRODUCT ATTEND-INFER-REPEAT:

A variety of object-centric image models have recently been proposed, many of which are derivatives of attend-infer-repeat (AIR) (Eslami et al., 2016) what with a neural network and rendering the resulting image patches in the prescribed location, a generative model p(x | z) is obtained. Inference is accomplished using a recurrent neural network, which outputs distributions over the latent objects q(z o | x), 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-Infer-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 computation of p(x | z where ), effectively integrating out the appearance parameters z what analytically. This has been shown to drastically accelerate learning, as the reduced inference workload significantly lowers the variance of the variational objective. Since the focus of SuPAIR on interpretable object parameters fits our goal of building a structured video model, we apply it as our image model p(x t | z t ). Similarly, we use a recurrent inference network as in SuPAIR to model q(z t,where | x t ). For details on SuPAIR, we refer to Stelzner et al. (2019) .

2.2. MODELING PHYSICAL INTERACTIONS USING GRAPH NEURAL NETWORKS:

In order to successfully capture complex dynamics, the state transition distribution p(

z t+1 | z t ) = p(z 1 t+1 , . . . , z O t+1 | z 1 t , . . . , z O t )

needs to be parameterized using a flexible, non-linear estimator. A critical property that should be maintained in the process is permutation invariance, i.e., the output should not depend on the order in which objects appear in the vector z t . This type of function is well captured by graph neural networks, cf. (Santoro et al., 2017) , which posit that the output should 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; Sanchez-Gonzalez et al., 2018; Zhou et al., 2018) .

Following this line of work, we build a dynamics model of the basic form

ẑo t+1,pos , ẑo t+1,velo , ẑo t+1,latent = f   g(z o t ) + o =o α(z o t , z o t )h(z o t , z o t )   (1)

where f, g, h, α represent functions parameterized by dense neural networks. α is an attention mechanism outputting a scalar which allows the network to focus on specific object pairs. We assume a constant prior over the object sizes, i. e., ẑo t+1,size = z o t,size . The full state transition distribution is then given by the Gaussian p(z o t+1 | z o t ) = N (ẑ o t+1 , σ), using a fixed σ.

2.3. JOINT STATE-SPACE MODEL:

Next, we assemble a 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(x, z

) = p(z 0 )p(x 0 | z 0 ) t p(z t | z t-1 )p(x t | z t ).

To initialize the state, we model p(z 0 , z 1 ) using simple uniform and Gaussian distributions. Details are given in Appendix C.3.

Our model is trained on given video sequences x by maximizing the evidence lower bound (ELBO) E q(z|x) [log p(x, z) -log q(z | x)]. This requires formulating a variational distribution q(z | x) to approximate the true posterior p(z | x). A natural approach is to factorize this distribution over time, i.e. q(z | x) = q(z 0 | x 0 ) t q(z t | z t-1 , x t ), resembling a Bayesian filter. The distribution q(z 0 | x 0 ) is then readily available using the inference network provided by SuPAIR.

The formulation of q(z t | z t-1 , x t ), however, is an important design decision. Previous work, including SQAIR and DDPAE, have chosen to unroll this distribution over objects, introducing a 2019), we avoid this cost by reusing the dynamics model for the variational distribution. First, we construct the variational distribution q(z o t,pos | z o t-1 ) by slightly adjusting the dynamics prediction p(z o t,pos | z o t-1 ), using the same mean values but separately predicted standard deviations. Together with an estimate for the same object by the object detection network q(z o t,pos | x t ), we construct a joint estimate by multiplying the two Gaussians and renormalizing, yielding another Gaussian:

EQUATION

Intuitively, this results in a distribution which reconciles the two proposals. A double recurrence is avoided since q(z t | x t ) does not depend on previous timesteps and may thus be computed in parallel for all frames. Similarly, q(z t | z t-1 ) may be computed in parallel for all objects, leading to 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 q(z t | x t , z t-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 t | z t ) is backpropagated through the dynamics model, which has been shown to be beneficial for efficient training (Karl et al., 2017; Becker-Ehmck et al., 2019) . The same procedure is applied to reconcile velocity estimates from the two networks, where for the image model, velocities z o t,velo are estimated from position differences between two consecutive timesteps. The object scales z o t,scale are inferred solely from the image model. The latent states z o t,latent increase the modelling capacity of the dynamics network, are initialised to zero-mean Gaussians, and do not interact with the image model. This then gives the inference procedure for the full latent state

z o t = (z o t,pos , z o t,size , z o t,velo , z o t,latent

), as illustrated in Fig. 2 (right). Despite its benefits, this technique has thus far only been used in environments with a single object or with known state information. A challenge when applying it in a multi-object video setting is to match up the proposals of the two networks. Since the object detection RNN outputs proposals for object locations in an 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 z t , we assign the object order that results in the minimal difference of ||z t,pos -z t-1,pos ||, where || • || is the Euclidean norm. The resulting Euclidean bipartite matching problem can be solved in cubic time using the classic Hungarian algorithm (Kuhn, 1955) .

2.4. CONDITIONING ON ACTIONS:

In reinforcement learning, an agent interacts with the environment sequentially through actions a t to optimize a cumulative reward r. To extend STOVE to operate in this setting, we make two changes, yielding a distribution p(z t , r t | z t-1 , a t-1 ).

First, we condition the dynamics model on actions a t , enabling a conditional prediction based on both state and action. To keep the model invariant to the order of the input objects, the action information is concatenated to each object state z o t-1 before they are fed into the dynamics model. The model has to learn on its own which of the objects in the scene are influenced by the actions. To facilitate this, we have found it helpful to also concatenate appearance information from the extracted object patches to the object state. While this patch-wise code could, in general, be obtained using some neural feature extractor, we achieved satisfactory performance by simply using the mean values per color channel when given colored input.

The second change to the model is the addition of reward prediction. In many RL environments, rewards depend on the interactions between objects. Therefore, the dynamics prediction architecture, presented in Eq. 1, is well suited to also predict rewards. We choose to share the same encoding of object interactions between reward and dynamics prediction and simply apply two different output networks (f in Eq. 1) to obtain the dynamics and reward predictions. The total model is again optimized using the ELBO, this time including the reward likelihood p(r t | z t-1 , a t-1 ).

3. EXPERIMENTAL EVIDENCE:

In order to evaluate our model, we compare it to baselines in three different settings: First, pure video prediction, where the goal is to predict future frames of a 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-efficient, model-based reinforcement learning. With this paper, we also release a PyTorch implementation of STOVE.1 

3.1. VIDEO AND STATE MODELING:

Inspired by Watters et al. (2017) , we consider grayscale videos of objects moving according to physical laws. In particular, we opt for the commonly used bouncing billiards balls dataset, as well as a dataset of gravitationally interacting balls. For further details on the datasets, see Appendix D. When trained using a single GTX 1080 Ti, STOVE converges after about 20 hours. As baselines, we compare to VRNNs (Chung et al., 2015) , SQAIR (Kosiorek et al., 2018) , and DDPAE (Hsieh et al., 2018) . To allow for a fair comparison, we fix the number of objects predicted by SQAIR and DDPAE to the correct amount. Furthermore, we compare to a supervised baseline: Here, we consider the ground truth positions and velocities to be fully observed, and train our dynamics model on them, resembling the setting of Battaglia et al. (2016) . Since our model needs to infer object states from pixels, this baseline provides an upper bound on the predictive performance we can hope to achieve with our model. In turn, the size of the performance gap between the two is a good indicator of the quality of our state-space model. We also report the results obtained by combining our image model with a 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 on the baselines, see Appendix E. Fig. 4 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 are evaluated on ground truth data by computing the mean squared error between 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 a large margin. Additionally, we perform strikingly close to the supervised model.

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 our approach and the supervised baseline is slightly more pronounced. Despite this, STOVE produces high-quality rollouts and outperforms the unsupervised baselines.

Table 1 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 z velo in the state, and (c) removing the latent state variables z latent . The ablation study shows that each of these components contributes positively to the performance of STOVE. See Appendix F for a comparison of training curves for the ablations.

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 error: Due to the chaotic nature of our environments, infinitesimally close initial states diverge quickly and a model which perfectly follows the ground truth states cannot exist. After this divergence has occurred, the rollout error no longer provides any information on the quality of the learned physical behavior. We therefore turn to 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. 5 shows the kinetic energies of trajectories predicted 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 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-BASED 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 at rest. The goal is to avoid collisions with the other balls, 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 it with RGB input in which the balls are colored differently. Starting with a random policy, we iteratively gather observations from the environment, i. e. sequences of images, actions, and rewards. Using these, we train our model as described in Sec. 2.4. To obtain a 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 an upper bound on the performance achievable in this manner, we report the results obtained by MCTS when the real environment is used for planning. As a model-free baseline, we consider PPO (Schulman et al., 2017) , 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 run PPO on a version of the environment in which, instead of images, the ground-truth object positions and velocities are observed directly.

Learning curves for each of the agents are given in Fig. 6 (left), reported at intervals of 10 000 samples taken from the environment, up to a total of 130 000. For our model, we collect the first 50 000 samples using a 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-based policy and refining the model using examples sampled from the pool of previously seen observations. After 130 000 samples, PPO has not yet seen enough samples to converge, whereas our model quickly learns to meaningfully model the environment and thus produces a better policy at this stage. Even when PPO is trained on ground truth states, MCTS based on STOVE remains comparable. After training each model to convergence, the final performance of all approaches is reported in Fig. 6 (right). In this case, PPO achieves slightly better results, however it only converges after training for approximately 4 000 000 steps, while our approach only uses 130 000 samples. After around 1 500 000 steps, PPO does eventually surpass the performance of STOVE-based MCTS. Additionally, we find that MCTS on STOVE yields almost the same performance as on the real environment, indicating that it can be used to anticipate and avoid collisions accurately.

4. RELATED WORK:

Multiple lines of work with the goal of video modeling or prediction have emerged recently. Prominently, the supervised modeling of physical interactions from videos has been investigated by Fragkiadaki et al. (2015) , who train a model to play billiards with a single ball. Similarly, graph neural networks have been trained in a supervised fashion to predict the dynamics of objects from images (Watters et al., 2017; Sanchez-Gonzalez et al., 2018; Sun et al., 2018; 2019) or ground truth states (Kipf et al., 2018; Wang et al., 2018; Chang et al., 2017) . A number of works learn object interactions in games in terms of rules instead of continuous dynamics (Guzdial et al., 2017; Ersen & Sariel, 2014) . Janner et al. (2019) show successful planning based on learned interactions, but assume access to image segmentations. Several unsupervised approaches address the problem by fitting the parameters of a physics engine to data (Jaques et 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 et al., 2017; Chung et al., 2015; Krishnan et al., 2015) . However, without the explicit notion of objects, their performance in scenarios with interacting objects remains limited. Nevertheless, unstructured video models have recently been applied to model-based RL and have been shown to improve sample efficiency when used as a simulator for the real environment (Oh et al., 2015; Kaiser et al., 2020) .

Only a small number of works incorporate objects into unsupervised video models. Xu et al. (2019) and Ehrhardt et al. (2018) take non-probabilistic autoencoding approaches to discovering objects in real-world videos. COBRA (Watters et al., 2019) represents a model-based RL approach based on MONet, but is restricted to environments with non-interacting objects and only uses one-step search to build its policy. Closest to STOVE are a small number of probabilistic models, namely SQAIR (Kosiorek et al., 2018) , R-NEM (Van Steenkiste et al., 2018; Greff et al., 2017) , and DDPAE (Hsieh et al., 2018) . R-NEM learns a mixture model via expectation-maximization unrolled through time and handles interactions between objects in a factorized fashion. However, it lacks an explicitly structured latent space, and requires noise in 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 is capable of handling a varying number of objects, which is not something we consider in this paper.

5. CONCLUSION:

We introduced STOVE, a structured, object-aware model for unsupervised video modeling and planning. 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 positions and velocities, and is capable of generating highly accurate video predictions in domains featuring complicated non-linear interactions between objects. As our experimental evaluation shows, it outperforms previous unsupervised approaches and even approaches the performance and visual quality of a supervised model.

Additionally, we presented an extension of the video learning framework to the RL setting. Our experiments demonstrate that our model may be utilized for sample-efficient model-based control in a visual domain, making headway towards a long standing goal of the model-based RL community. In particular, STOVE yields good performance with more than one 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 a fixed number of 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 a more flexible model such as MONet (Burgess et al., 2019) or GENESIS (Engelcke et al., 2020) 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 combining planning with policy and value networks, model-free and model-based RL may be integrated into a comprehensive system (Buckman et al., 2018) .

A RECONSTRUCTIONS: SPRITES DATA SuPAIR does not need a latent description of the objects' appearances. Nevertheless, object reconstructions can be obtained by using a variant of approximate MPE (most probable explanation) in the sum-product networks as proposed by Vergari et al. (2018) . We follow the AIR approach and reconstruct each object separately and paste it into the canvas using spatial transformers. Unlike AIR, SuPAIR explicitly models the background using a separate background SPN. A 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 in which the round balls were replaced by a random selection of four different sprites of the same size. Fig. 7 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 are drawn from a random normal distribution. This is the standard procedure of generating the bouncing balls data set as used by previous publications. Under these circumstances, STOVE does indeed learn to discover and replicate the total energies of the system, while SQAIR and DDPAE do not. Even if trained on constant energy data, STOVE does to some extent generalise to unseen energies. Observed velocities and therefore total 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. 8 (left) . If trained on a dataset of diverse total energies, the performance of modelling sequences of different energies increases, see Fig. 8 (right) . Rollouts now initially represent the true energy of the observed sequence, although this estimate of the true energy diverges over a time span of around 500 frames 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 as 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.

C MODEL DETAILS:

Here, we present additional details on the architecture and hyperparameters of STOVE.

C.1 INFERENCE ARCHITECTURE:

The object detection network for q(z t,where | x t ) is realised by an LSTM (Hochreiter & Schmidhuber, 1997) with 256 hidden units, which outputs the mean and standard deviation of the objects' twodimensional position and size distributions, i.e. q(z o t,pos, size | x t ) with 2 • 2 • 2 = 8 parameters per object. Given such position distributions for two consecutive timesteps q(z t-1,pos | x t-1 ), q(z t,pos | x t ), with parameters µ z o t-1,pos , σ z o t-1,pos , µ z o t,pos , σ z o t,pos , the following velocity estimate based on the difference in position is constructed:

q(z o t,velo | x t , x t-1 ) = N (µ z o t,pos -µ z o t-1,pos , σ 2 z o t,pos + σ 2 z o t-1,pos ).

As described in Sec. 2.3, positions and velocities are also inferred from the dynamics model as q(z o t,pos | z t-1 ) and q(z o t,velo | z t-1 ). A joint estimate, including information from both image model and dynamics prediction, is obtained by multiplying the respective distributions and renormalizing. Since both q-distributions are Gaussian, the normalized product is again Gaussian, with mean and (left) STOVE is trained on 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 at larger energies (red). In the following prediction, however, the mean predicted energies diverge 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 as predicted energies improve drastically. The predictions no longer immediately regress towards a specific value (orange). However after 100 frames, the quality of the predicted energies still regresses to a wrong value (green). (all) The observed values refers to energies obtained as the mean energy value over the 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.)

• D 1 : Sum relational dynamics R 2 and self-dynamics S 3 , obtaining the input to f in Eq. 1.

• D 2 : Apply linear layer [2l, 2l] to D 1 followed by tanh non-linearity.

• D 3 : Apply linear layer [2l, 2l] to D 2 followed by tanh non-linearity and add result to D 2 .

• D 4 : Concatenate D 3 and S 1 , and apply linear layer [4l, 2l] followed by tanh.

• D 5 : Apply linear layer [2l, 2l] to D 4 and add result to D 4 to obtain final dynamics prediction.

The output D 5 has shape (batch size, number of objects, 2l), 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 as above. The reward prediction consists of the following steps:

• H 1 : Apply linear layer [2l, 2l] to D 1 followed by ReLU.

• H 2 : Apply linear layer [2l, 2l] to H 1 .

• H 3 : Sum over object dimension to obtain tensor of shape (batch size, l).

• H 4 : Apply linear layer [l, l/2] to H 3 followed by ReLU.

• H 5 : Apply linear layer [l/2, l/4] to H 4 followed by ReLU.

• H 5 : Apply linear layer [l/4, l] to H 4 followed by a sigmoid non-linearity.

H 5 then gives the final reward prediction.

C.3 STATE INITIALIZATION:

In the first two timesteps, we cannot yet apply STOVE's main inference step q(z t | z t-1 , x t , x t-1 ) as described above. In order to initialize the latent state over the first two frames, we apply a simplified architecture and only use a partial state at t = 0.

At t = 0, z 0 ∼ q(z 0,(pos, size) | x 0 ) is given purely by the object detection network, since no previous states, which could be propagated, exist. z 0 is incomplete insofar as it does not contain velocity information or latents. At t = 1, q(z 1,pos, size | x 1 , x 0 ) is still given purely based on the object detection network. Note that for a dynamics prediction of z 1 , velocity information at t = 0 would need to be available. However, at t = 1, velocities can be constructed based on the differences between the previously inferred object positions. We sample z 1,latent from the prior Gaussian distribution to assemble the first full initial state z 1 . At t ≥ 2, the full inference network can be run: States are inferred both from the object detection network q(z t | x t , x t-1 ) as well as propagated using the dynamics model q(z t | z t-1 ).

In the generative model, similar adjustments are made: p(z 0,pos, size ) is given by a uniform prior, velocities and latents are omitted. At t = 1, velocities are sampled from a uniform distribution in planar coordinates p(z 1,velo ) and positions are given by a simple linear dynamics model p(z 1,pos | z 0,pos , z 1,velo ) = N (z 0,pos + z 1,velo , σ). Latents z 1,latent are sampled from a Gaussian prior. Starting at 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 × 10 -3 exp(-40 × 10 -3 • step) for a total of 83 000 steps with a 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 to optimize the ELBO. A 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 × 50 for the gravity data. All models for video prediction were learned on grayscale data, with objects of identical appearance. The O = 3 balls were initialised with uniformly random positions and velocities, rejecting configurations with overlap. They are rendered using anti-aliasing. The billiards data models the balls as circular objects, which perform elastic collision with each other or the walls of the environment. For the gravity data, the balls are modeled as point masses, where, following Watters et al. (2017) , we clip the gravitational force to avoid slingshot effects. Also, we add an additional basin of attraction towards the center of the canvas and model the balls in their center off mass system to 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. The DDPAE experiments were performed using the implementation available at https:// github.com/jthsieh/DDPAE-video-prediction. Default parameters for training DDPAE with billiards datasets are provided with the code. However, the resolution of our billiards (32 pixels) and gravity (64 pixels) datasets is different to the resolution DDPAE expects (64 pixels). While we experimented with adjusting DDPAE parameters such as the latent space dimension 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 the models' test set error.

E BASELINES FOR VIDEO MODELING:

The linear baseline was obtained as follows: For the first 8 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 to leave the frame, i. 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 a supervised baseline for our physics prediction by replacing the SuPAIR-inferred states with real, ground-truth states. On these, the model can then be trained in supervised fashion.

F TRAINING CURVES OF ABLATIONS:

In Fig. 9 we display learning curves for 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 in color.)

G DETAILS ON THE REINFORCEMENT LEARNING MODELS:

Our MCTS implementation uses the standard UCT formulation for exploration/exploitation. The c parameter is set to 1. in all our experiments. Since the environment does not provide a natural endpoint, we cut off all rollouts at a depth of 20 timesteps. We found this to be a good trade-off between runtime and accuracy.

When expanding a node on the true environment, we compute the result of the most promising action, and then start a rollout using a random policy from the resulting state. For the final evaluation, a total of 200 nodes are expanded. To better utilize the GPU, a slightly different approach is used for STOVE. When we expand a node in this setting, we predict the results of all actions simultaneously, and compute a rollout from each resulting position. In turn, only 50 nodes are expanded. To estimate the node value function, the average reward over all rollouts is propagated back to the root and each 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 in the baseline running on the real environment, since it behaves deterministically.

For PPO, we employ a standard convolutional neural network as 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 additional 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 a softmax activation function while the critic head outputs a value estimation for the current state using a linear prediction. We tested several hyperparameter configurations but found the following to be the most efficient one. To update the actor-critic architecture, we sample 32 trajectories of length 16 from separate environments in every batch. The training uses an Adam optimizer with a learning rate of 2 × 10 -4 and and value of 1 × 10 -5 . The clipping parameter of PPO is set to 1 × 10 -1 . We update the network for 4 epochs in each batch using 32 mini-batches of the sampled data. The value loss is weighted at 5 × 10 -1 and the entropy coefficient is set to 1 × 10 -2 .

:

The code can be found in the GitHub repository github.com/jlko/STOVE. It also contains animated versions of the videos predicted by our model and the baselines.

Figures and Tables:

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 github.com/jlko/STOVE. (Best viewed in color.)

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 all pixels in the video prediction 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, all 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.)

Figure5: Comparison of the kinetic energies of the rollouts predicted by the models, computed based on position differences between successive states. Only STOVE's predictions reflect the conservation of total kinetic energy in the billiards data set. This is a quantitive measure of 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 single, extremely long-term run. We conjecture that STOVE predicts physical behavior indefinitely. (Best viewed in color.)

Figure 6: Comparison of all models on sample efficiency and final performance. (Left) Mean cumulative reward over 100 steps on the environment, averaged over 100 environments, using the specified policy. The shaded regions correspond to one-tenth of a standard deviation. In addition to the training curves, two constant baselines are shown, one representing a random policy and one 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 one standard deviation. (Best viewed in color.)

Figure 8: Mean kinetic energy observed/predicted by STOVE over true energy of the sequences.(left) STOVE is trained on 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 at larger energies (red). In the following prediction, however, the mean predicted energies diverge 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 as predicted energies improve drastically. The predictions no longer immediately regress towards a specific value (orange). However after 100 frames, the quality of the predicted energies still regresses to a wrong value (green). (all) The observed values refers to energies obtained as the mean energy value over the 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.)

Figure 9: Displayed is the mean predicted position error over a 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.)

Predictive performance of our approach, the baselines, and ablations (lower is better, best unsupervised values are bold). STOVE outperforms all unsupervised baselines and is almost indistinguishable from the supervised model on the billiards task. The values are computed by summing the prediction errors presented in Fig.4in the time interval t ∈ [9, 18], i.e., the first ten predicted timesteps. In parentheses, standard deviations across multiple training runs are given.

Table content: <table><tr><td/><td colspan="4">Billiards (pixels) Billiards (positions) Gravity (pixels) Gravity (positions)</td></tr><tr><td>STOVE (ours)</td><td>0.240(14)</td><td>0.418(20)</td><td>0.040(3)</td><td>0.142(7)</td></tr><tr><td>VRNN</td><td>0.526(14)</td><td>-</td><td>0.055(12)</td><td>-</td></tr><tr><td>SQAIR</td><td>0.591</td><td>0.804</td><td>0.070</td><td>0.194</td></tr><tr><td>DDPAE</td><td>0.405</td><td>0.482</td><td>0.120</td><td>0.298</td></tr><tr><td>Linear</td><td>0.844(5)</td><td>1.348(15)</td><td>0.196(2)</td><td>0.493(4)</td></tr><tr><td>Supervised</td><td>-</td><td>0.232(37)</td><td>-</td><td>0.013(2)</td></tr><tr><td>Abl: Double Dynamics</td><td>0.262</td><td>0.458</td><td>0.042</td><td>0.154</td></tr><tr><td>Abl: No Velocity</td><td>0.272</td><td>0.460</td><td>0.053</td><td>0.174</td></tr><tr><td>Abl: No Latent</td><td>0.338</td><td>0.050</td><td>0.089</td><td>0.235</td></tr></table>

FollowingKosiorek et al. (2018), we experimented with different hyperparameter configurations for VRNNs. We varied the sizes of the hidden and latent states [h, z], experimenting with the values[256, 16],[512, 32],[1024, 64], and [2048, 32]. We found that increasing the model capacity beyond[512, 32]  did not yield large increases in performance, which is why we chose the configuration[512, 32]  for our experiments. Our VRNN implementation is written in PyTorch and based on https://github.com/emited/VariationalRecurrentNeuralNetwork. SQAIR can handle a variable number of objects in each sequence. However, to allow for a fairer comparison to STOVE, we fixed the number of objects to 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 https://github.com/akosiorek/sqair.

Table content: <table><tr><td>Position Error</td><td>0.2 0.3 0.4 0.5</td><td>0</td><td>20000 40000 60000 Step Billiards</td><td>0.1 0.2 0.3 0.4</td><td>0</td><td>Step 20000 40000 60000 Gravity STOVE Abl: Double Dynamics Abl: No velocity Abl: No Latent</td></tr></table>

Peer review (iclr2020):
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. The paper is well written and clearly motivated but comes across as an incremental improvement on top of prior work. Here are a few comments: 1. 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. 2. 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. 3. 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. Overall, the approach presented in the paper is a bit incremental and the experiments are somewhat simplistic. Further comparisons and ablation experiments can significantly strengthen the paper. I would suggest a borderline accept.