Switching Linear Dynamics for Variational Bayes Filtering

1 INTRODUCTION . Learning dynamics from raw data ( also known as system identification ) is a key component of model predictive control and model-based reinforcement learning . Problematically , environments of interest often give rise to very complex and highly nonlinear dynamics which are seemingly difficult to approximate . However , switching linear dynamical systems ( SLDS ) approaches claim that those environments can often be broken down into simpler units made up of areas of equal and linear dynamics ( Ackerson & Fu , 1970 ; Chang & Athans , 1978 ) . Not only are those approaches capable of good predictive performance , which often is the sole goal of learning a system ’ s dynamics , they also encode valuable information into so called switching variables which determine the dynamics of the next transition . For example , when looking at the movement of an arm , one is intuitively aware of certain restrictions of possible movements , e.g . constraints to the movement due to joint constraints or obstacles . The knowledge is present without the need to simulate ; it ’ s explicit . Exactly this kind of information will be encoded when successfully learning switching dynamics . Our goal in this work will therefore entail the search for richer representations in the form of latent state space models which encode knowledge about the underlying system dynamics . In turn , we expect this to improve the accuracy of our simulation as well . Such a representation alone could then be used in a reinforcement learning approach that possibly only takes advantage of the learned latent features but not necessarily its learned dynamics . To learn richer representations , we identify one common problem with prevalent recurrent Variational Autoencoder models ( Karl et al. , 2017a ; Krishnan et al. , 2015 ; Chung et al. , 2015 ; Fraccaro et al. , 2016 ) : the non-probabilistic treatment of the transition dynamics often modeled by a powerful nonlinear function approximator . From the history of the Autoencoder to the Variational Autoencoder , we know that in order to detect features in an unsupervised manner , probabilistic treatment of the latent space is paramount . As our starting point , we will build on previously proposed approaches by Krishnan et al . ( 2017 ) and Karl et al . ( 2017a ) . The latter already made use of locally linear dynamics , but only in a deterministic fashion . We extend their approaches by a stochastic switching LDS model and show that such treatment is vital for learning richer representations and simulation accuracy . 2 BACKGROUND . We consider discretized time-series data consisting of continuous observations xt ∈ X ⊂ Rnx and control inputs ut ∈ U ⊂ Rnu that we would like to model by corresponding latent states zt ∈ Z ⊂ Rnz . We ’ ll denote sequences of variables by x1 : T = ( x1 , x2 , ... , xT ) . 2.1 SWITCHING LINEAR DYNAMICAL SYSTEMS . Switching Linear Dynamical System models ( SLDS ) enable us to model nonlinear time series data by splitting it into sequences of linear dynamical models . At each time t = 1 , 2 , ... , T , a discrete switch variable st ∈ 1 , ... , M chooses of a set LDSs a system which is to be used to transform our continuous latent state zt to the next time step ( Barber , 2012 ) . zt = A ( st ) zt−1 +B ( st ) ut−1 + ( st ) ( st ) ∼ N ( 0 , Q ( st ) ) xt = H ( st ) zt + η ( st ) η ( st ) ∼ N ( 0 , R ( st ) ) ( 1 ) Here A ∈ Rnz×nz is the state matrix , B ∈ Rnz×nu control matrix , the transition noise with covariance matrix Q and η the emission/sensor noise with covariance matrix R. Finally , the observation matrix H ∈ Rnx×nz defines a linear mapping from latent to observation space which we will replace by a nonlinear transformation parameterized by a neural net . These equations imply the following joint distribution : p ( x1 : T , z1 : T , s1 : T | u1 : T ) = T∏ t=1 p ( xt | zt ) p ( zt | zt−1 , ut−1 , st ) p ( st | zt−1 , ut−1 , st−1 ) ( 2 ) with p ( z1 | z0 , u0 , s1 ) = p ( z1 ) being the initial state distribution . The corresponding graphical model is shown in figure 1a . 2.2 STOCHASTIC GRADIENT VARIATIONAL BAYES p ( x ) = ∫ p ( x , z ) dz = ∫ p ( x | z ) p ( z ) dz ( 3 ) Given the simple graphical model in equation ( 3 ) , Kingma & Welling ( 2014 ) and Rezende et al . ( 2014 ) introduced the Variational Autoencoder ( VAE ) which overcomes the intractability of posterior inference of q ( z | x ) by maximizing the evidence lower bound ( ELBO ) of the model log-likelihood . LELBO ( x ; θ , φ ) = Eqφ ( z|x ) [ ln pθ ( x | z ) ] − DKL ( qφ ( z | x ) || p ( z ) ) ≤ log p ( x ) ( 4 ) Their main innovation was to approximate the intractable posterior distribution by a recognition network qφ ( z|x ) from which they can sample via the reparameterization trick to allow for stochastic backpropagation through both the recognition and generative model at once . Assuming that the latent state is normally distributed , a simple transformation allows us to obtain a Monte Carlo gradient estimate of Eqφ ( z|x ) [ ln pθ ( x|z ) ] w.r.t . to φ . Given that z ∼ N ( µ , σ2 ) , we can generate samples by drawing from an auxiliary variable ∼ N ( 0 , 1 ) and applying the deterministic and differentiable transformation z = µ+ σ . 2.3 THE CONCRETE DISTRIBUTION . One simple and efficient way to obtain samples d from a k-dimensional categorical distribution with class probabilities α is the Gumbel-Max trick : d = one_hot ( argmax [ gi + logαi ] ) , with g1 , . . . , gk ∼ Gumbel ( 0 , 1 ) ( 5 ) However , since the derivative of the argmax is 0 everywhere except at the boundary of state changes , where it is undefined , we can ’ t learn a parameterization by backpropagation . The Gumbel-Softmax trick approximates the argmax by a softmax which gives us a probability vector ( Maddison et al. , 2017 ; Jang et al. , 2017 ) . We can then draw samples via dk = exp ( ( logαk + gk ) /λ ) ∑n i=1 exp ( ( logαi + gi ) /λ ) , with g1 , . . . , gk ∼ Gumbel ( 0 , 1 ) ( 6 ) This softmax computation approaches the discrete argmax as temperature λ → 0 , for λ → ∞ it approaches a uniform distribution . 3 RELATED WORK . Our model can be viewed as a Deep Kalman Filter ( Krishnan et al. , 2015 ) with structured inference ( Krishnan et al. , 2017 ) . In our case , structured inference entails another stochastic variable model with parameter sharing inspired by Karl et al . ( 2017b ) and Karl et al . ( 2017a ) which pointed out the importance of backpropagating the reconstruction error through the transition . We are different to a number of stochastic sequential models like Bayer & Osendorfer ( 2014 ) ; Chung et al . ( 2015 ) ; Shabanian et al . ( 2017 ) ; Goyal et al . ( 2017 ) by directly transitioning the stochastic latent variable over time instead of having an RNN augmented by stochastic inputs . Fraccaro et al . ( 2016 ) has a transition over both a deterministic and a stochastic latent state sequence , wanting to combine the best of both worlds . Previous models ( Watter et al. , 2015 ; Karl et al. , 2017a ; Fraccaro et al. , 2017 ) have already combined locally linear models with recurrent Variational Autoencoders , however they provide a weaker structural incentive for learning latent variables determining the transition function . Van Steenkiste et al . ( 2018 ) approach a similar multi bouncing ball problem ( see section 5.1 ) by first distributing the representation of different balls into their own entities without supervision and then structurally hardwiring a transition function with interactions based on an attention mechanism . Recurrent switching linear dynamical systems ( Linderman et al. , 2016 ) uses message passing for approximate inference , but has restricted itself to low-dimensional observations and a multi-stage training process . Johnson et al . ( 2016 ) propose a similar model to ours but combine message passing for discrete switching variables with a neural network encoder for observations learned by stochastic backpropagation . Tackling the problem of propagating state uncertainty over time , various combinations of neural networks for inference and Gaussian processes for transition dynamics have been proposed ( Eleftheriadis et al. , 2017 ; Doerr et al. , 2018 ) . However , these models have not been demonstrated to work with high-dimensional observation spaces like images . One feature a switching LDS model may learn are interactions which have recently been approached by employing Graph Neural Networks ( Battaglia et al. , 2016 ; Kipf et al. , 2018 ) . These methods are similar in that they predict edges which encode interactions between components of the state space ( nodes ) . 4 PROPOSED APPROACH . Our goal is to fit a series of continuous state z1 : T and switching variables s2 : T to a given sequence of observations x1 : T . We assume a nonlinear mapping between observations and latent space which we generally approximate by neural networks , apart from the transition which is modeled by a locally linear function . Our generative model is shown in figure 1b an our inference model in figure 2a . 4.1 GENERATIVE MODEL . Our generative model for a single xt is described by p ( xt ) = ∫ s≤t ∫ z≤t p ( xt | zt ) p ( zt | zt−1 , st , ut−1 ) p ( st | st−1 , zt−1 , ut−1 ) p ( zt−1 , st−1 ) ( 7 ) which is close to the one of the original SLDS model ( see figure 1a ) . Latent states zt are continuous and represent the state of the system while states st are the switching variables determining the transition . We approximate the discrete switching variables by a continuous relaxation , namely the Concrete distribution.1 Differently to the original model , we do not condition the likelihood of the current observation pθ ( xt | zt ) directly on the switching variables . This limits the influence of the switching variables to choosing a proper transition dynamic for the continuous latent space . The likelihood model is parameterized by a neural network with either a Gaussian or a Bernoulli distribution as output depending on the data . There is both a transition on the continuous states zt and discrete latent states st. For the continuous state transition p ( zt | zt−1 , st , ut−1 ) we follow ( 1 ) and maintain a set of M base matrices { ( A ( i ) , B ( i ) , Q ( i ) ) | ∀i . 0 < i < M } as our linear dynamical systems to choose from . For the transition on discrete latent states p ( st | st−1 , zt−1 , ut−1 ) , we usually require the learning of a Markov transition matrix . However , since we approximate our discrete switching variables by a continuous relaxation , we can parameterize this transition by a neural network . Therefore , our entire generative model can be learned end-to-end by ( stochastic ) backpropagation . Finally , the resulting dynamics matrices are computed through a linear combination of the base matrices : At ( st ) = M∑ i=1 s ( i ) t A ( i ) , B ( st ) = M∑ i=1 s ( i ) t B ( i ) , Q ( st ) = M∑ i=1 s ( i ) t Q ( i ) ( 8 ) Both transition models – the continuous state transition pθ ( zt | zt−1 , st , ut−1 ) and concrete switching variables transition pθ ( st | st−1 , zt−1 , ut−1 ) – are shared with the inference model which is key for good performance . pθ ( zt | zt−1 , st , ut−1 ) = N ( µ , σ2 ) where [ µ , σ2 ] = fθ ( zt−1 , st , ut−1 ) pθ ( st | st−1 , zt−1 , ut−1 ) = Concrete ( α , λprior ) where α = gθ ( zt−1 , st−1 , ut−1 ) ( 9 )

Summary: This paper proposes a new model for switching linear dynamical systems. The standard model and the proposed model are presented. Together with the inference procedure associated to the new model. This inference procedure is based on variational auto-encoders, which model the transition and measurement posterior distributions, which is exactly the methodological contribution of the manuscript. Experiments on three different tasks are reported, and qualitative and quantitative results (comparing with different state-of-the-art methods) are reported.