Sequential data often originates from diverse environments. Across them exist both shared regularities and environment specifics. To learn robust cross-environment 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-specific 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 filters. In our experiments, we demonstrate state-of-the-art performance in controlled generation and prediction of bouncing ball video sequences across varying gravitational influences. Learning dynamics and models from sequential data is a central task in various domains of science BID5 . This includes managing input of diverse complexity e.g. natural language BID11 , videos BID28 or financial time-series (Øksendal, 2003) . It is also crucial for building interactive agents which use reinforcement and control algorithms on top BID6 . Traditional choice in engineering are state space models (SSM) BID21 , typically found in form of Kalman filters BID9 where well-crafted, relatively simple state representations and (normally linear) functional forms are used. To improve flexibility, 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 benefits of introducing additional structure into SSM: the requirement of separating confounders from actions, observations and rewards BID24 or content from dynamics BID32 BID8 , especially for transfer learning and extrapolation BID18 . Complementary to these approaches, we focus on learning structured SSM to decouple system dynamics into its generic (enviromentinvariant) and environment-specific components. Some examples of sequential data which naturally admit this structure are given in figure 1. Dynamics of these are defined 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 environmentspecific part. In other words, true e integrates all unobserved environment-specific influences 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.Figure 1: Sequential systems across environments. Examples include, from left to right: (i) Michaelis-Menten model for enzyme kinetics, governed by reaction rate constants k; (ii) bouncing ball kinematics, determined by ball weight and playground characteristics; (iii) ODE dynamics, governed by model parameters; (iv) bat swinging motion, influenced by the person performing it. In each example, environments are defined differently, depending on what governs sequence dynamics. 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-specific information which governs this dynamics.We showed that such separation is beneficial 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 influences where we achieved state-of-the-art results. Our future work will include other types of data.A LOWER BOUND DERIVATION (SECTION 3) DISPLAYFORM0 (where the conditional independence follows from the state space model formulation) DISPLAYFORM1 (where we used the factorization of the variational and the prior distribution. S0 is vector S without S0) DISPLAYFORM2 q(S0| x)q(E| x)q( β| X, S0, E)q( S0|S0, E, β) log p0( S0|S0, E, β) q( S0|S0, E, β) (where we dropped the integral sums for which the corresponding term does not depend on) DISPLAYFORM3 + KL(q( β| X, E, S0)||p0( β|E, S0)) (where the last term vanishes since s0|s0, E, β is deterministic) DISPLAYFORM4 (where we have p0(βi|E, si) = p0(β) by design) B EXPERIMENTS (SECTION 4) B.1 DETAILS To get compressed representation of each frame, the images are first 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 BID1 , 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 | X)||p 0 (S)) and KL(q(E| 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 BID8 , we also found that down-weighing the reconstruction term helps in faster convergence. In particular we applied scaling coefficients of [0.1,0.2,0.3,1.0] for terms E q( S| X) [log p( X| S)], KL(q(E| X)||p 0 (E)), KL(q(S 0 | 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.6 applied every 20 epochs.
DISENTANGLED STATE SPACE MODELS The paper presents a generative state space model using a global latent variable E to capture environment-specific information.