D ISENTANGLED S TATE S PACE M ODELS : U NSUPERVISED L EARNING OF D YNAMICS ACROSS H ETEROGENEOUS E NVIRONMENTS
A BSTRACT
Sequential data often originates from diverse environments .
Across them exist both shared regularities and environment speciﬁcs.
To learn robust crossenvironment descriptions of sequences we introduce disentangled state space models (DSSM) .
In the latent space of DSSM environment-invariant state dynamics is explicitly disentangled from environment-speciﬁc information governing that dynamics.
We empirically show that such separation enables robust prediction, sequence manipulation and environment characterization.
We also propose an unsupervised VAE-based training procedure to learn DSSM as Bayesian ﬁlters.
In our experiments, we demonstrate state-of-the-art performance in controlled generation and prediction of bouncing ball video sequences across varying gravitational inﬂuences
I NTRODUCTION
Learning dynamics and models from sequential data is a central task in various domains of science (Durbin & Koopman, 2012).
This includes managing input of diverse complexity e.g. natural language (Graves, 2013), videos (Srivastava et al., 2015) or ﬁnancial time-series (Øksendal, 2003).
It is also crucial for building interactive agents which use reinforcement and control algorithms on top (Finn & Levine, 2017).
Traditional choice in engineering are state space models (SSM) (Koller et al., 2009), typically found in form of Kalman ﬁlters (Gelb, 1974) where well-crafted, relatively simple state representations and (normally linear) functional forms are used.
To improve ﬂexibility, new solutions rather learn model-free SSM "from scratch".
Due to their non-autoregressive architecture they make an attractive alternative to recurrent neural networks.
Several recent works have already recognized the beneﬁts of introducing additional structure into
SSM: the requirement of separating confounders from actions, observations and rewards (Lu et al., 2018) or content from dynamics (Yingzhen & Mandt, 2018; Fraccaro et al., 2017), especially for transfer learning and extrapolation (Kansky et al., 2017).
Complementary to these approaches, we focus on learning structured SSM to decouple system dynamics into its generic (enviromentinvariant) and environment-speciﬁc components.
Some examples of sequential data which naturally admit this structure are given in ﬁgure 1.
Dynamics of these are deﬁned by some constant external factors which we jointly refer to as environment.
More concretely, we explore a panel data setting in which we are given multiple sequences describing the same time-evolving phenomena, one or more per environment e .
We would like to learn a robust non-parametric SSM to represent the dynamics of that phenomena across these environments, and robustly extrapolate to the unseen ones.
To do so, we explicitly model e as a learnable static element of the latent space.
Our idea is based on the assumption that one can decouple sequence dynamics to: (i) the generic part which is invariant across environments; and (ii) the environmentspeciﬁc part.
In other words, true e integrates all unobserved environment-speciﬁc inﬂuences which bias generic system dynamics.
Our hypothesis is that considering disentangled , implicitly causal structure of SSM enhances predictive robustness, domain adaptation, and allows for environment characterization and reasoning under interventions e.g. counterfactual inference.
O UR KEY CONTRIBUTIONS
DSSM.
We introduce a class of non-parametric SSM tailored to exploit invariance from sequential data originating from heterogeneous environments.
Disentangled state space models (DSSM) (see ﬁgure 2d) form a joint environment model while explicitly decoupling what is generic in sequence dynamics from what is environment-speciﬁc.
This enhances robustness and the ability to extrapolate knowledge to unseen environments.
Bayesian ﬁltering.
We extend on recent advances in amortized variational inference to design an unsupervised training procedure and implement DSSM in form of Bayesian ﬁlters.
In the spirit of (Karl et al., 2016), well-established reparameterization trick is applied such that the gradient propagates through time.
While VAE heuristic provides no convergence guarantees, it is fast, robust and allows end-to-end training.
Video prediction and manipulation.
We analyze video sequences of a bouncing ball, inﬂuenced by varying gravity (environment).
We outperform state-of-the-art K-VAE (Fraccaro et al., 2017) in predictions, and also do interventions by "swapping environments" i.e. we enforce a speciﬁc dynamic behaviour by using an environment from another sequence which exhibits the desired behaviour.
Example videos are available at: https://sites.google.com/view/dssm .
R ELATED W ORK
Closely related to our proposal are approaches which consider structured and disentangled representation of videos, separating the pose from the content (Denton et al., 2017; Tulyakov et al., 2017; [Figure]
Villegas et al., 2017; Yingzhen & Mandt, 2018) or the object appearance from its dynamics (Fraccaro et al., 2017).
Proposed models were shown to improve the prediction (Villegas et al., 2017; Denton et al., 2017) and enable controlled generation with "feature swapping" (Tulyakov et al., 2017;
Yingzhen & Mandt, 2018).
This so called content-based disentanglement was also performed in speech analysis where the structure imposed the explicit separation of sequence- and segment-level attributes (Hsu et al., 2017).
VAE frameworks have already been extended to sequence modeling (Marino et al., 2018), and applied to speech (Bayer & Osendorfer, 2014; Chung et al., 2015; Fraccaro et al., 2016; Goyal et al., 2017), videos (Yingzhen & Mandt, 2018) and text (Bowman et al., 2015).
However, these (mainly) recurrent neural network-based approaches are autoregressive and hence not always suitable e.g. for planning and control from raw pixel space (Watter et al., 2015; Hafner et al., 2018).
To "image the world" (Ha & Schmidhuber, 2018) from the latent space directly and circumvent the autoregressive feedback, alternative methods learn SSM instead (Karl et al., 2016;
Fraccaro et al., 2017; Krishnan et al., 2017).
In DVBF (Karl et al., 2016)
SSM is trained using VAE-based learning procedure which allows gradient to propagate through time during training.
K-VAE by Fraccaro et al.
(2017) is a two-layered model which decomposes object’s representation from its dynamics.
DKF (Krishnan et al., 2015), and very closely related DMM (Krishnan et al., 2017), admit SSM structure but the state inference is conditioned on both past and future observations, so the structure of a ﬁlter is not preserved.
This is problematic as noted by Karl et al.
Similar issues can be found in (Yingzhen & Mandt, 2018).
As opposed to content-based methods which focus on the observation model, our work is focused on dynamics-based disentanglement.
This makes our approach complementary to existing (see also ﬁgure 2).
For example, while DSA can represent and manipulate the shape or color of a bouncing ball, our method can manipulate its trajectory.
To implement our Bayesian ﬁlter, we blended some recent ideas in amortized variational inference (Karl et al., 2016; Marino et al., 2018) and adapted them to ﬁt our novel DSSM architecture.
V ARIATIONAL B AYESIAN FILTERING FOR DSSM
In this work, we assume that the underlying system is deterministic i.e. the latent process noise β and observation noise ω are both uncorrelated in time.
We consider the following DSSM description:
[Equation]
where f and g represent arbitrary ﬂexible functions.
X i ∈
R D represents the observation in time step i and S i ∈ R N is the corresponding latent state.
Σ β and Σ ω are noise covariances which we for simplicity assume are isotropic Gaussian.
Our goal is to jointly learn the generative model which consists of the transition function f and observation function g , together with the corresponding recognition networks φ encβ , φ encE and φ encS which infer the process noise residual β
i , the environment
E , and the initial state S 0 respectively.
The framework overview is given in ﬁgure 3.
Generative model.
Given an observed sequence (cid:126)X of length T , the joint distribution is:
[Equation]
This follows from ﬁgure 2d and the assumption that the process noise is serially uncorrelated.
We set the prior probabilities of the initial state p 0 ( S 0 ) , environment p 0 ( E ) and process noise p 0 ( β i ) =
p 0 ( β ) to be zero-mean unit-variance Gaussian.
Conditioned on β i and E , state transition is deterministic and the probability p
( S i | S i − 1 , E, β i ) is a Dirac function with the peak deﬁned by equation (2).
The emission probability p ( X i | S i ) is deﬁned by equation (1). [Figure]
Inference.
Joint variational distribution over the unobserved random variables E , (cid:126)S and (cid:126)β , for a sequence of observation (cid:126)X of length T factorizes as:
[Equation]
Here, the conditionals S − i | S i − 1 , E and S i | β i , S − i are deterministic and deﬁned by equation (2).
The remaining factors are given as follows:
[Equation]
[Equation]
Learning.
To match the posterior distributions of E , S 0 and (cid:126)β to the assigned prior probabilities p 0 ( E ) , p 0 ( S 0 ) and p 0 ( β ) , we utilize reparametrization trick (Kingma & Welling, 2013; Rezende et al., 2014).
This enables end-to-end training.
To deﬁne the objective function we derive the variational lower bound L , which we consequently attempt to maximize during the training.
We start from the well-known equality (Kingma & Welling, 2013):
[Equation]
Due to the conditional independence of the observations given the latent states, we can decompose the ﬁrst term as:
[Equation]
The KL term can be shown to simplify into a sum of the following KL terms:
[Equation]
where we dropped the conditional dependency β
i | S i − 1 , X i , E in q to ease the notation.
Full L derivation is given in Appendix.
Algorithm 1 shows the details of the training procedure for one iteration, for a batch of size 1.
The extension to the batch training is trivial.
Implementation.
We model [ f , g , φ encS , φ encE , φ encβ ] as neural networks.
In our experiments, g and φ encβ are convolutional/deconvolutional networks.
φ encS and φ encE are given as bi-directional LSTM followed by a multilayer perceptron to convert LSTM-based sequence embedding into S 0 and E , and f as an LSTM cell as elaborated in ﬁgure 3d. LSTM equations are taken from (Graves, 2013).
Optimization Challenges.
Some performance improvements were observed with an additional heuristic regularization term, which ensures the consistency during the inference of environment E .
Namely, we penalize the step-wise change in time embeddings produced by bi-directional LSTM used to model φ encE , in order to enforce E to remain time-invariant.
To that end, we add an additional term to our objective function, the moment matching regularization term deﬁned as:
[Equation]
where h i is the hidden state of the φ encE LSTM cell in step i .
This idea is related to the approaches based on the maximum mean discrepancy (Gretton et al., 2012).
Namely, enforcing equality of consecutive cell states corresponds to matching of their ﬁrst moments.
Furthermore, similarly to (Bowman et al., 2015; Karl et al., 2016) we used a KL annealing scheme.
This was helpful for circumventing local minimum and preventing the KL term to converge to zero too early during the training.
The exact details are given in our experiments.
P REDICTION , MANIPULATION AND GENERATION OF VIDEO SEQUENCES
Bouncing ball in varying gravity settings.
We test our framework on a 2D bouncing ball problem where the ball kinematics is affected by a varying gravity vector.
The idea is to evaluate model robustness across environments – gravitational settings.
See also visualizations at: https:// sites.google.com/view/dssm .
Using the physics engine code from (Fraccaro et al., 2017), we simulate video sequences of a bouncing ball.
During generation, we randomly change the gravity such that it remains constant within a sequence, but may vary across sequences.
The gravity vector takes 4 values, depending on whether the gravity points up, down, left or right.
The magnitude is kept ﬁxed.
Each video frame is a 32x32 binary image.
We generate 16’000 trajectories across 40 time steps for training, and another 2000 trajectories across 70 time steps for testing.
We ﬁrst perform long-term forecasting analysis comparing our method against one of the state-of-the-art approaches, the K-VAE from (Fraccaro et al., 2017).
Next, we demonstrate controlled generation, by manipulating video sequences.
Effectively we perform interventions by "swapping environments" between video sequences, and similarly we swap initial states.
Furthermore, we show the ability to perform uncontrolled generation where both initial state and gravity vector are sampled from a prior.
Finally we visualize the environment embeddings to provide further intuiton.
Forecasting ball trajectory.
We use OpenCV inbuilt functions to detect ground truth ball position p t in each time frame.
The exact algorithm for the position extraction is provided in the Appendix
B.2.
Following (Hsieh et al., 2018;
Chang et al., 2016)
we deﬁne the ball velocity as v t =
p t +1 − p t and compute the relative error in the predicted velocities of the balls for forecasting.
Evaluated models observe the ﬁrst 25 frames of a test sequence and then forecast the next 45.
The results for relative error in magnitude and cosine similarity of the velocities are then averaged across all test sequences.
This is shown in ﬁgure 4.
We observe an increase in prediction quality with respect to both metrics in comparison to the benchmarking model K-VAE.
Video manipulation for controlled generation.
Firstly, the initial state which consists of the velocity vector and the ball position, is extracted from the baseline video sequence and then "injected" into a series of other test sequences.
Similarly, we performed the gravity environment replacement.
We then enrolled the sequence effectively performing controlled generation (see ﬁgure 5).
Environment identiﬁcation.
We trained an auxiliary multi layer perceptron classiﬁer to map E to true gravity value.
The cross-validation results performed on the training set rendered accuracy of 99 . 15% .
Visualized embeddings (for E ∈ R 3 ) are given in ﬁgure 6a.
Well-deﬁned clusters can be observed, indicating that E indeed represents the true gravity.
Uncontrolled Generation.
We demonstrate the uncontrolled generation of the sequences where s 0 and E are sampled from the priors p 0 ( E ) and p 0 ( S 0 ) respectively.
In ﬁgure 6b we observe how the generated sequences preserve natural bouncing ball dynamics.
[Figure]
C ONCLUSION
This work proposes a novel view on data-driven learning of dynamics from diverse environments.
We proposed a new class of state space models particularly crafted to exploit this kind of a setting.
In disentangled state space models one separates generic system dynamics which is assumed to be invariant across environments and environment-speciﬁc information which governs this dynamics.
We showed that such separation is beneﬁcial and allows us to learn robust cross-environment models which hold promise to generalize on unseen environments.
Our particular application was learning of the video dynamics of a bouncing ball affected by varying gravitational inﬂuences where we achieved state-of-the-art results.
Our future work will include other types of data.
A L OWER BOUND DERIVATION ( SECTION 3)
[Table]
log-likelihood term in L
[Equation]
[Table]
the conditional independence follows from the state space model formulation
[Table]
[Equation]
[Table]
[Equation] [Table]
[Equation]
[Table]
[Equation]
where
[Table]
since (cid:126)s 0 | s 0 , E, (cid:126)β is deterministic)
[Equation]
p 0 ( E )) + KL ( q ( S 0 | (cid:126)X [Equation] |
[Equation]
p
[Equation]
( S 0 )) + T (cid:88) [Equation]
[Equation]
(where we have p 0 ( β i | E, s i ) = p 0 ( β ) by design)
B E XPERIMENTS (S ECTION 4) B. D ETAILS
To get compressed representation of each frame, the images are ﬁrst passed through a shallow convolutional network.
Kernel size was set to 3x3, while the network depth was 64.
The step size was 1 in both directions.
We used ReLU activation units.
All of the hidden latent states were equal to 64.
To parameterize g we used a deconvolutional network with transposed convolutions.
The kernel size was set to 5.
Following the insights from (Bowman et al., 2015), we tried different settings for KL annealing in the model.
Since we have three KL terms in our model which have different roles, we do not penalize KL terms of time-invariant components i.e. KL ( q ( S 0 | (cid:126)X ) ||
p 0 ( S )) and KL ( q ( E | (cid:126)X ) ||
p 0 ( E )) as forcefully as KL ( q ( β i | S i , X i ) ||
p 0 ( β )) during training.
This makes it relatively easier for the model to learn the time-invariant components.
Similarly to (Fraccaro et al., 2017), we also found that down-weighing the reconstruction term helps in faster convergence.
In particular we applied scaling coefﬁcients of [0.1,0.2,0.3,1.0] for terms
E q ( (cid:126)S | (cid:126)X ) [ log p ( (cid:126)X | (cid:126)S )] , KL ( q ( E | (cid:126)X ) ||
p 0 ( E )) ,
KL ( q ( S 0 | (cid:126)X ) || p 0 ( S )) and KL ( q ( β i | S i , X i ) || p 0 ( β ))] respectively.
We use ADAM as the optimizer with 0.0008 as the initial learning rate, and weight decay of 0.
applied every 20 epochs.
B. A LGORITHM FOR D ETECTING B ALL P OSITIONS
We use OpenCV’s (Itseez, 2015) inbuilt functions to detect the pixel level positions of the ball in the images. import
[Table]
imutils def find_positions(
[Table]
):ret,
[Table]
[Table]
=
[Table]
.threshold(
[Table]
[Table]
[Table]
[Table]
.THRESH_BINARY)
binary_mask =
[Table]
.erode(
[Table]
[Table]
[Table]
[Table]
=
[Table]
) binary_mask =
[Table]
.dilate( [Table] [Table] [Table]
[Table] [Table] [Table]
[Table] [Table] [Table] [Table]
[Table] [Table] [Table] [Table]
[Table]
C T RAINING ALGORITHM (S ECTION 3)
Algorithm
Input: sequence (cid:126)x of length T [ µ E , Σ E ] =
φ encE ( (cid:126)x ) ,
E ∼ N ( µ E , Σ E ) [ µ S , Σ S ] =
φ encS ( (cid:126)x ) , s 0 ∼ N ( µ S , Σ S ) for i
= 1 to T do
Predict step: S − i =
f ( S i − 1 , E )
Estimate residual: [ µ βi , Σ βi ] =
φ encβ ( S − i , X i ) , β i ∼ N ( µ βi , Σ βi )
Update step: S i =
S − i
+ β i
Predict observation:
ˆ X i =
g ( S i ) end for ll_loss = LL ( (cid:126)X, (cid:126)
ˆ X ) (see Eq (9))
kl_loss
=
KL ( E, S 0 , (cid:126)β ) ; (see Eq (10))
mm_loss = MM ( φ encE ( (cid:126)X )) ; (see Eq (11))
Backpropagate(ll_loss, kl_loss, mm_loss)