This repository contains the code for the model introduced in the paper

A Disentangled Recognition and Nonlinear Dynamics Model for Unsupervised Learning.

Marco Fraccaro, Simon Kamronn, Ulrich Paquet, Ole Winther

NIPS 2017, arXiv:1710.05741

_{* Equal contribution.}

The Kalman variational auto-encoder is a framework for unsupervised learning of sequential data that disentangles two latent representations: an objectâ€™s representation, coming from a recognition model, and a latent state describing its dynamics. The recognition model is represented by a convolutional variational auto-encoder and the latent dynamics model as a linear Gaussian state space model (LGSSM).

As shown in the paper, the KVAE can be trained end-to-end, and is able learn a recognition and dynamics model from the videos. The model can be used to generate new sequences, as well as to do missing data imputation without the need to generate high-dimensional frames at each time step.

Videos of simulations are available for many of the experiments here: link.

The Kalman Variational Auto-Encoder (KVAE) package can be installed running

```
# Install requirements
pip install tensorflow-gpu numpy pandas matplotlib seaborn
# Clone the kvae repository and install it
git clone https://github.com/simonkamronn/kvae
cd kvae
pip install -e .
```

- Python >= 2.7 or 3.5
- Tensorflow >= 1.1

The data used in the experiments is available for download here: link. Once downloaded, it needs to be placed in the `data`

folder in the root directory.
In Linux you can download the data by running from the root directory:

```
wget https://archive.compute.dtu.dk/archives/public/users/marco/kvae/data.zip
unzip data.zip
```

Alternatively, you can also find the data here: link
The scripts we used to generate the video sequences can be found in `kvae/datasets`

, and depend on pygame and the pymunk physics engine (and their dependencies).

In the `examples`

folder execute the command

`python run_kvae.py --gpu=0`

to train a model using the default bouncing ball dataset and parameters. To train a model with different parameters run

`python run_kvae.py --help`

to view the configuration options or look in `kvae/utils/config.py`

. Further details on the experimental setup can be found in the paper and in the supplementary material.
Plots and videos generated during training can be found in `examples/logs`

.

```
@article{Fraccaro2017,
Author = {Marco Fraccaro, Simon Kamronn, Ulrich Paquet, Ole Winther},
Title = {A Disentangled Recognition and Nonlinear Dynamics Model for Unsupervised Learning},
Journal = {Advances in Neural Information Processing Systems 30, NIPS},
Year = {2017}
}
```

- Free software: MIT license