Unsupervised Separation of Dynamics from Pixels

TL;DR

This paper introduces an unsupervised probabilistic model that learns object dynamics from image sequences by combining linear state-space models with non-linear rendering, enabling efficient inference and separation of object motions.

Contribution

It presents a novel probabilistic framework that separates object dynamics from pixel data using linear models and non-linear rendering, allowing flexible inference without retraining.

Findings

Effective separation of object dynamics from images

Efficient inference enabled by linear state-space models

Applicable to multiple objects in image sequences

Abstract

We present an approach to learn the dynamics of multiple objects from image sequences in an unsupervised way. We introduce a probabilistic model that first generate noisy positions for each object through a separate linear state-space model, and then renders the positions of all objects in the same image through a highly non-linear process. Such a linear representation of the dynamics enables us to propose an inference method that uses exact and efficient inference tools and that can be deployed to query the model in different ways without retraining.

Equations

$${\alpha }^n$$

$${\hat{x}}^n$$

$$x^n$$

$$v_t\sim p_{\theta }(v_t|a_t^{1:N})=\mathrm{Bernoulli}(\mathrm{sigm}(W^v{x}^N+b^v)).$$

$$h_t^n$$

$$a_t^n$$

$$h_1^n\sim p_{\theta }(h_1^n)=\sum _{k=1}^K{p}_{\theta }(z^n=k)p_{\theta }(h_1^n|z^n=k)=\sum _{k=1}^K{\pi }_k\mathcal{N}(h_1^n;{\mu }_k,{\Sigma }_k).$$

$$\displaystyle p_{\theta}(a^{1:N}_{1:T},h^{1:N}_{1:T},z^{1:N})=\prod_{n=1}^{N}\bigg{\{}\prod_{t=1}^{T}p_{\theta}(a^{n}_{t}|h^{n}_{t})\bigg{\}}p_{\theta}(h^{n}_{1}|z^{n})\,p_{\theta}(z^{n})\prod_{t=2}^{T}p_{\theta}(h^{n}_{t}|h^{n}_{t-1})\,,$$

$$p_{\theta}(v_{1:T},a^{1:N}_{1:T},h^{1:N}_{1:T},z^{1:N})=\bigg{\{}\prod_{t=1}^{T}p_{\theta}(v_{t}|a^{1:N}_{t})\bigg{\}}p_{\theta}(a^{1:N}_{1:T},h^{1:N}_{1:T},z^{1:N})\,.$$

$$\displaystyle\textrm{KL}\Big{(}q_{\phi}(a^{1:N}_{1:T},h^{1:N}_{1:T},z^{1:N}|v_{1:T})\,\Big{\|}\,p_{\theta}(a^{1:N}_{1:T},h^{1:N}_{1:T},z^{1:N}|v_{1:T})\Big{)}\,.$$

$$\displaystyle\textrm{KL}(q(a,h,z|v)\,\big{\|}\,p(a,h,z|v))$$

$$\displaystyle\hskip-48.36958pt=\bigl{<}\log q(a,h,z|v)-\log p(a,h,z,v)\bigr{>}_{q(a,h,z|v)}+\log p(v)\geq 0\,,$$

$$\displaystyle{\cal F}_{\theta,\phi}=-\bigl{<}\log q(a,h,z|v)\bigr{>}_{q(a,h,z|v)}+\bigl{<}\log p(a,h,z,v)\bigr{>}_{q(a,h,z|v)}\,.$$

$$\bigl{<}\nabla_{\phi}f_{\theta,\phi}(\epsilon_{1:T})\bigr{>}_{q_{\epsilon}(\epsilon_{1:T})}\approx\hskip 0.0pt\frac{1}{M}\sum_{m=1}^{M}\nabla_{\phi}f_{\theta,\phi}(\epsilon^{m}_{1:T}),\hskip 11.38092pt\epsilon^{m}_{1:T}\sim q_{\epsilon}(\epsilon^{m}_{1:T})\,.$$

$$p_{\theta }(a_{1:T}^{1:N},h_{1:T}^{1:N},z^{1:N}|v_{1:T})$$

$$\displaystyle\approx\bigg{\{}\prod_{n=1}^{N}p_{\theta}(h^{n}_{1:T},z^{n}|a^{n}_{1:T})\bigg{\}}q_{\phi}(a^{1:N}_{1:T}|v_{1:T})\,.$$

$${\mathcal{F}}_{\theta ,\varphi }$$

$$\displaystyle=-\biggl{<}\log\frac{q(a|v)}{p(v|a)}\biggr{>}_{q(a|v)}+\Bigl{<}\underbrace{-\log p(h,z|a)+\log p(a,h,z)}_{\log p(a)}\Bigr{>}_{q(a|v)p(h,z|a)}$$

$$\displaystyle=-\biggl{<}\log\frac{q(a|v)}{p(v|a)}\biggr{>}_{q(a|v)}+\bigl{<}\log p(a)\bigr{>}_{q(a|v)}\,.$$

$${\cal F}_{\theta,\phi}=\bigl{<}\log p_{\theta}(v_{1:T}|a^{1:N}_{1:T})\bigr{>}_{q_{\phi}(a^{1:N}_{1:T}|v_{1:T})}-\textrm{KL}\Big{(}q_{\phi}(a^{1:N}_{1:T}|v_{1:T})\,\Big{\|}\,p_{\theta}(a^{1:N}_{1:T})\Big{)}\,.$$

$${\beta }_t^n$$

$${\hat{s}}_t^n$$

$$s_t^n$$

$$a_t^n={\mu }_{\varphi }(s_t^n)+{\sigma }_{\varphi }(s_t^n)\circ {ϵ}_t^n,$$

$$\log p_{\theta}(a^{1:N}_{1:T})=\sum_{n=1}^{N}\log\Bigg{(}\sum_{k=1}^{K}p_{\theta}(z^{n}=k)p_{\theta}(a^{n}_{1:T}|z^{n}=k)\Bigg{)}.$$

Full text

1 Introduction

The dynamics of moving objects can be described on a very low-dimensional manifold of positions and velocities, e.g. by using a state-space model representation of Newtonian laws. A wealth of literature in probabilistic tracking barshalom93estimation ; blackman99design allows us to ask questions such as where objects are going to be in the future or where they have been in the past from noisy observations of their current positions, if the latent dynamics are known or can be learned.

In this paper, we consider the more challenging task of learning the presence of multiple objects and their dynamics in an unsupervised way from image observations. The problem of learning dynamics from pixels has been studied for the simpler case of one object fraccaro17disentangled ; pearce2018comparing . The multiple-object scenario is much harder as an object-to-identity assignment problem is introduced. The unsupervised learning method should be able to reliably disentangle objects from a sequence of images—this requires a recurrent attention mechanism which should learn to track each object separately over time.

Deep recurrent neural networks have demonstrated impressive results on pixel-prediction and related tasks babaeizadeh18stochastic ; chiappa17recurrent ; denton17unsupervised ; Finn2016 ; Oh2015 ; srivastava15unsupervised ; Sun2016 . Whilst powerful and easy to design, in such methods the hidden states do not generally correspond to interpretable dynamics. Enforcing a desired hidden-state representation is challenging—attempts to get positions from pixels have so far succeeded only in the supervised scenario watters17visual . These methods also have shortcomings when asked to infer intermediate images from observed preceding and following images in a sequence.

Probabilistic extensions, such as modern stochastic variational approaches to hidden Markov models Fraccaro2016 ; Gao2016 ; Krishnan2017 , can achieve interpretable hidden-state representations and perform rich probabilistic reasoning. The capabilities of these methods can be further enhanced with the use of traditional probabilistic inference routines fraccaro17disentangled ; johnson16composing ; lin18variational . This paper contributes an approach in this direction. In Sec. 2.1, we introduce a model for generating images containing multiple objects in which positions are explicitly represented using auxiliary variables. In Sec. 2.2, we leverage this representation to introduce a method for performing inference and learning that makes use of exact and efficient inference techniques. Finally, in Sec. 3 we show how our approach performs on inferring latent positions from image sequences, and on image generation and interpolation, using an artificial dataset representing moving cannonballs.

2 A model for rendering and inferring multiple objects dynamics

We wish to learn the dynamics of $N$ objects from sequences of images $v_{1:T}\equiv v_{1},\ldots,v_{T}$ in an unsupervised way. We restrict ourselves to the case in which the objects move independently in the two-dimensional plane, and $N$ is known; and assume that each image is formed by white pixels representing objects positions and a black background.

Two examples of sequences, each containing two cannonballs moving in opposite directions, are given in Fig. 1—the images are overlaid in time such that lighter shades correspond to more recent images.

The observed dynamics in the pixel space are high-dimensional and non-linear. However, the intrinsic dynamics can be described on the low-dimensional manifold of positions and velocities by simple dynamical systems. In the sections that follow, we show that the explicit representation of such latent dynamics enables us to infer positions using exact and efficient techniques. We also show that our approach enables us to answer different questions using the learned dynamics without the need to retrain or modify the model.

2.1 Generative model

We assume that the generative process underlying the observed images consists of two main parts: A part that describes the objects dynamics in the low-dimensional manifold of positions and velocities through a linear state-space model, and a part that renders the latent positions of the $N$ objects into the images through a highly non-linear process.

2.1.1 Rendering latent object positions into images

We use a set of auxiliary variables $a^{n}_{1:T}$ to explicitly represent the latent positions of object $n$ in the two-dimensional plane. Given the positions of all objects at time $t$, $a^{1:N}_{t}\equiv a^{1}_{t},\ldots,a^{N}_{t}$, an image $v_{t}$ is generated by recurrently rendering each $a^{n}_{t}$ on $v_{t}$. We start with $x^{0}=\tanh(\theta_{x^{0}})$, which represents a latent state vector from which an empty image canvas is generated, with $\theta_{x^{0}}$ denoting an unknown parameter vector. For $n=1,\ldots,N$, we iterate

[TABLE]

and finally generate the image as

[TABLE]

The symbol $\circ$ indicates element-wise vector multiplication, and $\mathrm{sigm}(a)\equiv 1/(1+\mathrm{e}^{-a})$. The $W$’s and $b$’s indicate weight matrices and biases, and are all included in the unknown parameters $\theta$ of the generative model.

This rendering process is illustrated in Fig. 2. The final state $x^{N}$, which is transformed in Eq. (2) so that $v_{t}$ can be sampled, should contain information from all $N$ objects. To achieve that, the state is iteratively updated through Eq. (1) to incorporate the contribution from object $n$, $\hat{x}^{n}$. This is obtained through an attention mask vector $\alpha^{n}$ with elements in the interval $[0,1]$, which specifies what information from object $n$ should be included and what information from $x^{n-1}$ should be retained.

2.1.2 Latent dynamics

We model the latent positions of each object, $a^{n}_{1:T}$, using a hidden Markov model with linear Gaussian hidden-state and output, also known as linear Gaussian state-space model (LGSSM) barber11inferenceA ; chiappa06phd , i.e.

[TABLE]

where ${\cal N}(x;\mu,\Sigma)$ denotes the density of a Gaussian random variable $x$ with mean $\mu$ and covariance $\Sigma$. We use the constraints $A=[I,\ \ \delta I;\ \ 0,\ \ I]\in\mathbb{R}^{4\times 4}$ (where $[\cdot,\cdot;\cdot,\cdot]$ indicates horizontal and vertical matrix concatenation, $I\in\mathbb{R}^{2\times 2}$ denotes the identity matrix, and $\delta$ denotes the sampling period) and $B=[I,\ \ 0]\in\mathbb{R}^{2\times 4}$ to obtain a description of Newtonian laws, such that the vector $h_{t}$ represents positions and velocities, and $u_{t}$ the force (which is assumed to be a constant $u$ over time). We include $\delta$, $u$, $\Sigma_{H}$ and $\Sigma_{A}$ in the generative model parameters $\theta$.

Objects may start moving from disjoint sets of initial positions and velocities; for instance from the left, the right but not the center of $v_{1}$. To allow the model to not put probability mass on initial positions and velocities that might never be helpful in explaining $a_{1:T}^{n}$, the initial positions and velocities $h^{n}_{1}$ are drawn from a $K$ component Gaussian mixture, i.e.

[TABLE]

We additionally include $\pi_{1:K}$, $\mu_{1:K}$ and $\Sigma_{1:K}$ in the generative model parameters $\theta$. The joint density of all random variables factorizes as

[TABLE]

where $p_{\theta}(a^{n}_{t}|h^{n}_{t})={\cal N}(a^{n}_{t};Bh^{n}_{t},\Sigma_{A})$ and $p_{\theta}(h^{n}_{t}|h^{n}_{t-1})={\cal N}(h^{n}_{t};Ah^{n}_{t-1}+u,\Sigma_{H})$. The full generative model combines Eq. (2) with Eq. (5) to yield

[TABLE]

The backbone of the model is illustrated in Fig. 2.

The advantage of this formulation for the latent dynamics is that quantities such as the smoothed distribution $p_{\theta}(h^{n}_{t}|a^{n}_{1:T},z^{n})$, the likelihood $p_{\theta}(a^{n}_{1:T}|z^{n})$, or the most likely mixture component $\operatorname{argmax}_{k}p_{\theta}(z^{n}=k|a^{n}_{1:T})$, can be computed exactly in ${\cal O}(T)$ operations, using message passing algorithms such as the Kalman filtering and Rauch-Tung-Striebel smoothing barber11inferenceA ; chiappa06phd .

2.2 Inference and Learning

The non-linearity of the rendering process makes the computation of $p_{\theta}(v_{1:T})$, $p_{\theta}(a^{1:N}_{1:T},h^{1:N}_{1:T},z^{1:N}|v_{1:T})$, and of quantities like $p_{\theta}(a^{n}_{t}|v_{1:T})$ needed for estimating the positions of object $n$, intractable. We address this problem using a recent approach to variational methods known as variational auto-encoding (VAE) kingma14autoencoding ; rezende14stochastic .

The basic principle of variational methods is to introduce a tractable approximating distribution111Whilst in practice we need to consider all observed sequences in the KL, to simplify the notation we focus the exposition on one sequence only. $q_{\phi}(a^{1:N}_{1:T},h^{1:N}_{1:T},z^{1:N}|v_{1:T})$ to the intractable distribution $p_{\theta}(a^{1:N}_{1:T},h^{1:N}_{1:T},z^{1:N}|v_{1:T})$ via the Kullback-Leibler divergence

[TABLE]

Given that

[TABLE]

where we omitted super and subscript indices and used the notation \bigl{<}\cdot\bigr{>}_{q(\cdot)} to indicate averaging wrt $q(\cdot)$, we obtain a lower bound ${\cal F}_{\theta,\phi}$ on $\log p_{\theta}(v)$, i.e. $\log p_{\theta}(v)\geq{\cal F}_{\theta,\phi}$ with

[TABLE]

If ${\cal F}_{\theta,\phi}$ were tractable and we were able to perform marginalization on $q_{\phi}(a,h,z|v)$, we could find the optimal $q_{\phi}(a,h,z|v)$ (parameters $\phi$) and $\theta$ by maximizing the bound; see chiappa08bayesian ; chiappa14explicit for traditional approaches to variational methods in the temporal setting. However, this is not the case for our generative model choice, and thus we instead use the more recent VAE approach to variational methods, where a Monte-Carlo approximation of the intractable ${\cal F}_{\theta,\phi}$ is deployed.

The VAE approach consists in rewriting the bound in the form {\cal F}_{\theta,\phi}=\bigl{<}f_{\theta,\phi}(\epsilon_{1:T})\bigr{>}_{q_{\epsilon}(\epsilon_{1:T})} for a parameter free distribution $q_{\epsilon}(\epsilon_{1:T})$, such that the gradient of ${\cal F_{\theta,\phi}}$ with respect to $\phi$ is given by \nabla_{\phi}{\cal F_{\theta,\phi}}=\bigl{<}\nabla_{\phi}f_{\theta,\phi}(\epsilon_{1:T})\bigr{>}_{q_{\epsilon}(\epsilon_{1:T})}—this is often called reparemetrization trick. We can then approximate the gradient with the Monte-Carlo estimate

[TABLE]

In our case, the formulation of the latent dynamics described above enables us to avoid employing a full approximation of $p_{\theta}(a^{1:N}_{1:T},h^{1:N}_{1:T},z^{1:N}|v_{1:T})$, and instead to decompose this distribution as a product of the exact tractable distribution $p_{\theta}(h^{1:N}_{1:T},z^{1:N}|a^{1:N}_{1:T})$ and an approximation $q_{\phi}(a^{1:N}_{1:T}|v_{1:T})$ of $p_{\theta}(a^{1:N}_{1:T}|v_{1:T})$, i.e.

[TABLE]

Thanks to this representation, the bound can be expressed as

[TABLE]

This gives

[TABLE]

We model $q_{\phi}(a^{1:N}_{1:T}|v_{1:T})$ using a recurrent neural network with states $s^{1:N}_{1:T}$ as in Fig. 2. More specifically we assume $a^{n}_{t}\sim{\cal N}(a^{n}_{t}\,;\,\mu_{\phi}(s^{n}_{t}),\sigma^{2}_{\phi}(s^{n}_{t}))$ and use the reparametrization $a^{n}_{t}=\mu_{\phi}(s^{n}_{t})+\sigma_{\phi}(s^{n}_{t})\circ\epsilon^{n}_{t}$, under the assumption $q_{\epsilon}(\epsilon^{1:N}_{1:T})=\prod_{n,t}q_{\epsilon}(\epsilon^{n}_{t})$ with $\epsilon^{n}_{t}\sim{\cal N}(\epsilon^{n}_{t};0,I)$. We describe this inference network in more details in the next section.

2.2.1 Inference network

As shown in Fig. 2, the inference network iterates a latent state vector $s_{t}^{n}$ over $t=1,\ldots,T$ and $n=1,\ldots,N$. Starting with $s^{0}=s_{t}^{0}=0$ at time-step $t$, we recurrently iterate

[TABLE]

to compute a vector $s_{t}^{1:N}$ for each of the objects at time step $t$. Similar to the generative model in Sec. 2.1.1, there is an attention mask $\beta_{t}^{n}$ that specifies how much of each component of $s_{t}^{n-1}$ we should keep. The mask is a function of the visible image $v_{t}$, as well as the recurrently computed values for both the previous object at this time-step, and this object and the previous time-step. There is also a contribution $\hat{s}^{n}_{t}$ coming from image $v_{t}$ for object $n$. The combination of $s^{n-1}_{t}$ and $\hat{s}^{n}_{t}$ is used to update $s^{n}_{t}$. Samples from $q_{\phi}(a^{1:N}_{1:T}|v_{1:T})$ are generated as

[TABLE]

where $\epsilon_{t}^{n}\sim{\cal N}(\epsilon_{t}^{n};0,I)$. In this computation, external Gaussian noise is inserted in a computation graph, and transformed—this ensures that the Monte-Carlo estimate in Eq. (6) is fully differentiable.

The $W$’s and $b$’s denote weight matrices and biases, and are included in the inference network parameters $\phi$. The recurrent process at time-step $t=1$ depends on $s_{0}^{n}=\tanh(\phi_{s_{0}^{n}})$, and an initial state for each object $n$ is learned through parameters $\phi_{s_{0}^{n}}$222In practice, as the state $s_{0}^{n}$ encodes which way we can interrogate $v_{1}$ to infer $a_{1}^{n}$, we have obtained better results by learning separate $\phi_{s_{0}^{n}}$ that depend on the number of objects $N$ in the image..

Note that $v_{t}$ appears in Eqs. (8) and (9) as input to every step $n$ at time-step $t$. This is important: To infer the position $a_{t}^{n}$, we need to consider the latent representation $s_{t-1}^{n}$ of object $n$ (actually 1 to $n$) in the previous image, as well as $s_{t}^{n-1}$, which contains a rolled-up representation of objects 1 to $n-1$ in this image. Both these representations need to act on $v_{t}$ to infer $a_{t}^{n}$.

2.2.2 Learning

In the Kullback-Leibler divergence term in Eq. (7), \bigl{<}\log q_{\phi}(a^{1:N}_{1:T}||v_{1:T})\bigr{>}_{q} can be expressed in analytic form. Both the first term in Eq. (7) and \bigl{<}\log p_{\theta}(a^{1:N}_{1:T})\bigr{>}_{q} in the KL divergence term can be stochastically estimated using a sample $a^{1:N}_{1:T}\sim q_{\phi}(a^{1:N}_{1:T}|v_{1:T})$. For such a sample,

[TABLE]

To compute Eq. (10), we take $a^{n}_{1:T}$ as “observations”, and for each mixture component $z^{n}=k$ we run a Kalman filter to obtain the log-likelihood $\log p_{\theta}(a^{n}_{1:T}|z^{n}=k)$. The objective function is the sum of the bounds in Eq. (7) over all sequences in the training dataset. The negative of this objective is minimized via a stochastic gradient descent algorithm, using mini-batches from the dataset.

2.3 Limitations

One of the main limitations of our approach is that, as $q_{\phi}(a^{1:N}_{1:t-1}|v_{1:t-1})$ does not explicitly incorporate the LGSSM dynamics, the objective function can have many sub-optimal local maxima. As common in the VAE literature, we address such a decoupling between the generative model and variational distribution by annealing the KL term in the bound ${\cal F}_{\theta,\phi}$ to ensure that the dynamics are correctly accounted for during training. More advanced methods in the literature consist in incorporating the unknown latent dynamics into $q_{\phi}(a^{1:N}_{1:t-1}|v_{1:t-1})$ by essentially letting the dynamics be a “regularizer” to the parameters $\phi$. “Structured inference networks” provide a framework for achieving this lin18variational . In pearce2018comparing we show that this approach leads to better inference and more stable results than annealing for the case of learning the dynamics of one object from pixels—it not obvious how this can be extended to the multiple-object scenario.

A more general limitation of our approach is its applicability to images that contain arbitrary backgrounds and objects, and nonlinear interactions between objects.

3 Results

In this section, we evaluate our approach on inferring latent positions, and on image generation and interpolation using artificially generated images describing the movement of cannonballs, see the examples shown in Fig. 1.

3.1 Dataset

We generated image sequences of length $T=30$ describing the movement of up to three cannonballs. Noisy positions $a_{1:T}^{n}$ were generated with an LGSSM formulation of Newtonian laws, as described in Sec. 2.1.2, with sampling period $\delta=0.015$; force $u=-g(0\ \ 0.5\delta^{2}\ \ 0\ \ \delta)$, where $g=9.81$ is the gravitational constant; $\Sigma_{H}=0$; and $\Sigma_{A}=0.001I$.

Each ball was shot with random shooting angle $\gamma$ in the interval $[40^{\circ},60^{\circ}]$, from either the left side of the $x$-axis in the interval $[-0.5,-0.1]$ or the right side of the $x$-axis in the interval $[-0.5,-0.1]+0.9\max_{x}$, where $\max_{x}$ indicates the maximum possible displacement at $T=30$ when starting in $[-0.5,-0.1]$. The initial position on the $y$-axis was sampled in the interval $[-0.5,0.5]$. The initial velocity $i_{v}$ was sampled in the interval $[2,3]$. The resulting velocity on the $x$-axis, $i_{v}\cos(\gamma\pi/180.0)$, was flipped in sign if the ball was shot from the right side of the image. Some examples of trajectories are shown in black in Fig. 3, with circles indicating initial positions.

To render the positions into white patches of radius $R=2$ in the image, the generated positions $a_{1:T}^{1:N}$ were re-scaled to lie in the interval $[R,H-1-R]\times[R,W-1-R]$, where $H=48,W=48$ indicate the height and width of the image. This re-scaling ensured that each ball was always fully contained in the image. We also experimented with similar datasets with $T=50$ and $H=W=32$, obtaining similar results. With $H=W=32$ the problem is easier in terms of dimensionality, but the latent positions are less identifiable, as close positions in the latent space might induce the same position in the image.

3.2 Initialization and training

The dataset consists of sequences with $N\in\{1,2,3\}$ balls. Importantly, as the networks in Fig. 2 can dynamically unroll, the model was trained on all such sequences jointly. We used $N$ to inform the networks of how many steps to unroll for each image sequence.

As the initial cannonballs are roughly separated into two main clusters, we assumed two mixture components, i.e. $K=2$. Although a higher $K$ would induce a more refined grouping of initial positions and velocities, our experiments indicate that $K=2$ was sufficient to obtain accurate results.

We experimented with different types of initialization for the LGSSM. Good results were obtained as long as smoothness in the dynamics was enforced so that we could guide the inference network toward smooth trajectories (as explained below). One example of such an initialization is to set $A$ and $B$ as above with $\delta=0.1$, $u=0$, $\Sigma_{H}=0.001I$, and $\Sigma_{A}=I$; sample the part of $\mu_{k}$ corresponding to positions from a standard Gaussian distribution and set the part corresponding to the velocities to zero; and set $\Sigma_{k}=I$: This ensures symmetry breaking without imposing any meaningful prior on the clusters.

The absence of force gives rise to positions that form straights lines, as the dots in Fig. 4, which represent the first two dimensions of $h^{n}_{1:T}$ for $n=1,\ldots,3$. The high emission noise gives rise to highly non-smooth $a^{n}_{1:T}$ (crosses). This initialization induces very different dynamics from the ground-truth and does not assume any clustering, but encodes prior information that objects move smoothly in time.

For the rendering and inference networks in Secs. 2.1.1 and 2.2.1, the weight matrices entries were randomly initialized from $\mathcal{N}(\,\cdot\,;0,1/\sqrt{d})$, where $d$ is the number of matrix elements. All biases were initialized to zero. The dimension of the latent state $s^{n}_{t}$ was set to 1024. We used the Adam optimizer with learning rate 0.001, mini-batch size 20, and default values $\beta_{1}=0.9$, $\beta_{2}=0.999$, $\epsilon=10^{-8}$. Training was stopped after $2\times 10^{5}$ iterations.

To guide the inference network toward smooth trajectories we kept the LGSSM parameters fixed to their initial values for the first $10^{4}$ iterations, and only optimized for all other parameters. This initialization loosely gives temporal coherence between the inference network and the renderer, for different images. The model was then trained jointly, first by changing the objective function by multiplying the KL term in Eq. (7) with a weight $\beta$, starting at $\beta=100$, and annealing $\beta$ down to one. This avoids the LGSSM parameters from too quickly modelling the output of a (still very sub-optimal) $q_{\phi}(a_{1:T}^{1:N}|v_{1:T})$, and stagnating at a local maximum. This process initializes the model, after which end-to-end training proceeds.

3.3 Inferring latent positions from image sequences

In Fig. 3 we show estimates of positions from the inference network. The ground-truth trajectories from which each image sequence is generated are shown in black. The inference network means $\mu_{\phi}(s^{n}_{1}),\ldots,\mu_{\phi}(s^{n}_{T})$ are plotted in blue. Notice that each plot is scaled differently to aid qualitative evaluation. For that reason, we also selected a run with minimum rotation, but the latent positions can only be retrieved up to re-scaling and rotation. The larger scale of the inferred trajectories is highlighted by re-plotting the ground-truth trajectories (red lines) together with the inferred positions. These figures demonstrate that our model can accurately infer latent positions from the sequences of images in an unsupervised way.

In Fig. 5(a), we show the inference network means for a batch of image sequences at three stages of training (for a case in which $N=2$), illustrating how the model learns over time. To highlight the challenge in learning to disentangle the objects dynamics, and the importance of initially strongly regularizing it to be close to the LGSSM during training, Fig. 5(b) shows an example in which we fail to learn an accurate inference network despite having learned accurate LGSSM dynamics. The ground-truth trajectories (black; left) and the trajectories generated from the learned LGSSM dynamics (red; right) are very similar. On the other hand, the inference network means (blue; right) swap the ball when reaching the middle part of the image: Rather than learning a successful attention mechanism, the inference network has learned to attend to the left part of the image for one ball and to the right part of the image for the other ball.

3.4 Multi-step ahead generation of images

To benchmark against a standard deep-learning model for multi-step ahead generation of images, we compared our model to an encoder-decoder long-short term memory model (ED-LSTM) hochreiter97long on the task of generating the 25 images following five observed images $v_{1:5}$.

The structure of the ED-LSTM is represented in Fig. 7: At each time-step $t$, the hidden state ${\sigma}_{t}$ generates an image $\hat{v}_{t}$ through a decoding transformation. For the first $\tau=5$ time-steps, ${\sigma}_{t}$ receives an encoded version of the previous ground-truth image $v_{t-1}$, as well as ${\sigma}_{t-1}$, as input. From time-step $\tau+1$ onward, ${\sigma}_{t}$ receives an encoded version of the previous ground-truth image $v_{t-1}$ during training, or an encoded version of the previous generated image $\hat{v}_{t-1}$ when the model is used in a multi-step ahead generation mode. We experimented with both convolutional and fully connected encoding and decoding transformations. The best results were obtained with one or two fully connected layers and dimension 2048 for ${\sigma}_{t}$. (All layers except the last were followed by a ReLU activation. The decoder last layer was followed by a sigmoid activation.) We used the same weights and biases initialization and optimizer settings as for our model.

To generate images with our model, we first inferred $a^{n}_{1:5}$ as the means $\mu_{\phi}(s^{n}_{1}),\ldots,\mu_{\phi}(s^{n}_{5})$ of the inference network, for each object $n$. Using the inferred $a^{n}_{1:5}$, we computed the most likely mixture component as $k^{*}=\operatorname{argmax}_{k}[p(z^{n}=k|a^{n}_{1:5})\propto p(a^{n}_{1:5}|z^{n}=k)\pi_{k}$] by running a Kalman filter. The filtered means of $p(h^{n}_{5}|a^{n}_{1:5},z^{n}=k^{*})$ were then used as initial conditions to generate $a^{n}_{6:30}$ with the learned LGSSM dynamics (Eq. (3)). Finally images $v_{6:30}$ were generated through rendering of $a^{1:N}_{6:30}$ (Eqs. (1) and (2)).

In Fig. 6, we show generated versus ground-truth images overlaid in time, for four sequences (more examples are given in the Appendix). For generation, our averaged test loss over the 25 time-steps was 0.691, using only eight as latent state dimension for the dynamics ($h_{t}^{1:2}$). In contrast, the ED-LSTM’s test loss was 0.693, with 2048 as latent state dimension for the dynamics. Some examples of trajectories generated by our model are shown in Fig. 8.

3.5 Inference using past and future observations

Our model can interpolate missing images (and positions) from past and future images, a task that cannot be solved by the ED-LSTM without model adjustment and retraining.

We evaluated how our model performs in inferring the latent positions and images in the intermediate time-steps $t=6,\ldots,25$, based on observing the first and last five images $v_{1:5}$ and $v_{26:30}$. The model should be able to use information from the future to produce more accurate latent positions than the ones that would be estimated by forward generation as in Fig. 8.

To solve the task, we first used $v_{1:5}$ to obtain $\mu_{\phi}(s^{n}_{1}),\ldots,\mu_{\phi}(s^{n}_{5})$ and the most likely mixture component $k^{*}$; images at time-steps $t=6,\ldots,25$ were then generated by the mechanism explained above. The inference network was run with the generated images to obtain a warmed-in state $s_{25}$, which was then used as initial state for another run of the inference network with observed images $v_{26:30}$ to infer $\mu_{\phi}(s^{n}_{26}),\ldots,\mu_{\phi}(s^{n}_{30})$. We finally used $\mu_{\phi}(s^{n}_{1}),\ldots,\mu_{\phi}(s^{n}_{5})$ and $\mu_{\phi}(s^{n}_{26}),\ldots,\mu_{\phi}(s^{n}_{30})$ as observations in a Rauch-Tung-Striebel smoother to interpolate the missing trajectories $a_{6:25}^{n}$, assuming that the observations at the intermediate time-steps were missing (integrated out from the model).

Some examples of obtained results are shown in Fig. 9 (more examples are given in the Appendix). We show the ground-truth trajectories in black, $\mu_{\phi}(s^{n}_{1}),\ldots,\mu_{\phi}(s^{n}_{30})$ in blue, the generated trajectories in red, and the interpolated trajectories in cyan. The interpolation corrects the generated trajectories by bringing them closer to $\mu_{\phi}(s^{n}_{1}),\ldots,\mu_{\phi}(s^{n}_{5})$ (obtained by observing the images at all time-steps) whilst maintaining the smoothness of dynamics constraints. To show the results of this correction mechanism in the pixel space, below the trajectories, from left to right, we show the generated versus ground-truth images and the interpolated versus ground-truth images, overlaid in time.

4 Conclusions

This paper describes an unsupervised approach to disentangle the dynamics of objects from pixels. We showed that it is possible for an inference network that is recurrent over both time and object number to sequentially parse each image to determine latent object positions. The model is regularized with a mixture of linear Gaussian state-space models, which encourages temporal coherence between latent positions extracted from images. Whilst the considered images of cannonballs are much simpler than those that would be encountered in most real-world applications, we nevertheless successfully demonstrated the usefulness of recovering interpretable latent structure in an unsupervised way and, more generally, of building structured generative models for high-dimensional visual stimuli.