Latent Matters: Learning Deep State-Space Models

1 Introduction . Many dynamical systems can only be ( partially ) observed , with the exact dynamics unknown . Yet , precise models are needed for prediction and control [ e.g . 19 , 21 , 25 , 3 ] . Learning such accurate models is subject of current research , especially in image-based domains [ e.g . 9 ] . Deep state-space models ( DSSMs ) [ e.g . 28 , 17 , 13 ] describe sequence data by a ( typically Markovian ) nonlinear transition model and a nonlinear observation model . The transition model is assumed to capture the dynamics underlying the observed data , and the observation model maps the latent variables to the domain of observable data and accounts for the measurement noise . In this paper , we show that DSSMs , however , often do not learn the correct system dynamics , which is suboptimal for accurate predictions or performing downstream tasks , such as model-based reinforcement learning . We identify in our experiments three main reasons causing this problem : ( i ) DSSMs are often trained by maximising the sequential evidence lower bound ( ELBO ) . High ELBO values , however , do not imply the model has learned the correct system dynamics . ( ii ) The prior/initial distribution is usually just a Gaussian . This often leads to an over-regularisation of the approximate posterior or even to a broken generative model , where the transition model is not optimised to process samples from the prior . ( iii ) Most DSSMs use RNNs to approximate or support Bayesian filtering/smoothing . Yet , RNNs often prove to be a limiting factor for learning accurate models of the system dynamics ; moreover , RNN-based transition models as in [ 8 , 9 ] can lead to a non-Markovian state-space , where the latent variables do not capture the entire information about the system ’ s state . To address these problems , we propose the following solutions : ( i ) we introduce a constrained optimisation ( CO ) framework as a general approach for learning DSSMs . It ensures a good reconstruction ∗Correspondence to a.klushyn @ tum.de . †Work done while at Machine Learning Research Lab , Volkswagen Group , Munich . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . quality and thus provides a necessary basis for learning the underlying system dynamics . To this end , we extend a recent method [ 16 ] presented in the context of variational autoencoders ( VAEs ) to DSSMs . We do this by formulating the sequential ELBO as the Lagrangian of a CO problem and introducing the associated optimisation algorithm . ( ii ) We complement the proposed CO framework with a powerful empirical Bayes prior . ( iii ) To obtain more accurate predictions of observed dynamical systems , we introduce the extended Kalman VAE ( EKVAE ) , where we dispense with RNNs by combining extended Kalman filtering/smoothing with amortised variational inference and a neural linearisation approach . Furthermore , we show that the EKVAE is capable of learning state-space representations where static and dynamic features are disentangled . We use this to validate the learned model in the context of model-based reinforcement learning . Our evaluation includes experiments on the image data of a moving pendulum [ 13 ] and on the reacher environment [ 26 ] , where we use angle as well as high-dimensional RGB image data as observations . We show that each of our proposed approaches significantly helps in learning accurate models of observed dynamical systems—and that applying our CO framework to established DSSMs leads to a substantial increase in their prediction accuracy . 2 Background : A Rate–Distortion Perspective on Deep State-Space Models . DSSMs [ e.g . 28 , 17 , 13 ] model an unknown distribution of observed sequence data x1 : T = ( x1 , x2 , . . . , xT ) by means of typically lower-dimensional latent variables z1 : T that represent the underlying state of the system . To achieve this , the Markov assumption is imposed . It states that the future state zt+1 as well as the current observation xt solely depend on zt : p ( x1 : T , z1 : T |u1 : T ) = pθ ( x1 : T |z1 : T ) pψ ( z1 : T |u1 : T ) = p ( z1 ) pθ ( x1| z1 ) T∏ t=2 pψ ( zt| zt−1 , ut−1 ) pθ ( xt| zt ) , ( 1 ) where u1 : T are optional control signals ( actions ) , and the use of different parameters ( θ , ψ ) will become important in the course of this paper . The model parameters in Eq . ( 1 ) are often learned through amortised variational inference [ e.g . 17 , 13 ] . This requires introducing a recognition model qφ ( z1 : T |x1 : T , u1 : T ) that learns—in combination with the transition model pψ ( zt| zt−1 , ut−1 ) —the dynamics underlying the observed data . The resulting objective function is known as sequential ELBO : log p ( x1 : T |u1 : T ) ≥ FELBO ( θ , ψ , φ ) = Eqφ [ log pθ ( x1 : T | z1 : T ) pψ ( z1 : T |u1 : T ) qφ ( z1 : T |x1 : T , u1 : T ) ] . ( 2 ) In the context of generative models , the ELBO can be divided into a reconstruction term ( distortion ) and a compression term ( rate ) [ 1 ] . We extend the theory in [ 1 ] to DSSMs , where the distortion D ( θ , φ ) = −Eqφ [ log pθ ( x1 : T | z1 : T ) ] optimises the model ’ s ability for reconstructing observations , whereas the rateR ( φ , ψ ) = KL ( qφ ( z1 : T |x1 : T , u1 : T ) ‖ pψ ( z1 : T |u1 : T ) ) enables learning the underlying dynamics . These definitions lead to the following general formulation of the ELBO : FELBO ( θ , ψ , φ ) = −D ( θ , φ ) −R ( ψ , φ ) . ( 3 ) Balancing the ratio between distortion and rate during optimisation can be an effective approach to improve the learning of DSSMs , as we discuss in the following section . 3 Constrained Optimisation Framework for Improved System Identification . High ELBO values do not necessarily imply that the model has learned the underlying system dynamics of the observed data , as we verify in Sec . 6.2 . This is because different combinations of rate and distortion can result in the same ELBO value . Previous work addresses this issue by introducing weighting schedules for either D ( θ , φ ) orR ( ψ , φ ) [ e.g . 4 ] since a different ratio favours either better reconstruction or compression [ 1 ] . However , we demonstrate in Sec . 6 that balancing reconstruction and compression with predefined annealing schedules often does not achieve the desired result . A more recent approach originates from the framework of VAEs : Rezende and Viola [ 23 ] and Klushyn et al . [ 16 ] define the VAE as a CO problem allowing for controlling the model ’ s reconstruction quality . We transfer this approach to DSSMs to ensure a good reconstruction—i.e . a low D ( θ , φ ) —and thus provide a sufficient basis for learning the underlying system dynamics . To this end , we formulate the sequential ELBO as the Lagrangian of a CO problem ( i ) by specifying the rate R ( ψ , φ ) in Eq . ( 3 ) as optimisation objective ; and ( ii ) by imposing the inequality constraint D ( θ , φ ) ≤ D0 . Here , D0 is a hyperparameter that defines the baseline for our desired reconstruction quality—we provide a heuristic for the simple determination of D0 in App . A.1 . The resulting Lagrangian is L ( θ , ψ , φ ; λ ) = R ( ψ , φ ) + λ ( D ( θ , φ ) −D0 ) , ( 4 ) where the Lagrange multiplier λ can be viewed as a weighting term for the distortion . The original EM algorithm [ e.g . 20 ] for optimising the ELBO , minθ , ψ minφ −FELBO ( θ , ψ , φ ) , provides the following connection to the CO problem : M-step︷︸︸︷ min θ , ψ E-step︷ ︸︸ ︷ max λ min φ L ( θ , ψ , φ ; λ ) s.t . λ ≥ 0 , ( 5 ) where , unlike in the original EM algorithm , we want qφ ( z1 : T |x1 : T , u1 : T ) to additionally satisfy the inequality constraint D ( θ , φ ) ≤ D0 in the E-step . However , as detailed in [ 16 ] , it can only be guaranteed that L ( θ , ψ , φ ; λ ) optimises a lower bound on log p ( x1 : T |u1 : T ) if and only if 1 ≥ λ ≥ 0 . 3.1 Learning the Initial Distribution . It is common practice to define the initial/prior distribution p ( z1 ) as a standard normal [ e.g . 17 , 8 ] . However , in this case , the prior KL in the ELBO can cause an over-regularisation ( cf . [ 16 ] ) of the approximate posterior and thus of the transition model . Furthermore , if the discrepancy between prior and posterior is too large , we may obtain a broken generative model , where p ( zt| zt−1 , ut−1 ) is not trained to process samples from p ( z1 ) . We provide empirical evidence in Sec . 6.1 . This issue often arises in the context of neural models trained by stochastic gradient methods : in order to alleviate possible vanishing-gradient problems , time-series data is typically cut into equally-sized short-length units . For sufficiently large datasets , the initial states can therefore be assumed to cover most possible states , which results in a nontrivial marginal approximate posterior EpD [ qφ ( z1|x , u ) ] . Since the optimal prior distribution is p∗ ( z1 ) = EpD [ qφ ( z1|x , u ) ] [ cf . 27 ] , an empirical Bayes prior pψ0 ( z1 ) must have the complexity to approximate pψ0 ( z1 ) ≈ p∗ ( z1 ) . For this reason , we propose to learn a hierarchical prior pψ0 ( z1 ) = ∫ pψ0 ( z1| ζ ) p ( ζ ) dζ as part of the DSSM by applying the variational approach in [ 16 ] . Klushyn et al . [ 16 ] define , by means of an approximate distribution qφ0 ( ζ|z1 ) , a ( VAE-like ) lower bound on the optimal empirical Bayes prior : Ep∗ ( z1 ) [ log pψ0 ( z1 ) ] ≥ EpD ( x1 : T , u1 : T ) Eqφ ( z1 : T |x1 : T , u1 : T ) [ Eqφ0 ( ζ| z1 ) [ log pψ0 ( z1| ζ ) p ( ζ ) qφ0 ( ζ| z1 ) ] ︸ ︷︷ ︸ =FVHP ( ψ0 , φ0 ; z1 ) ] , ( 6 ) where p ( ζ ) is a standard normal distribution and pD ( x1 : T , u1 : T ) is the empirical distribution of our data . This method is referred to as variational hierarchical prior ( VHP ) . Learning the VHP as part of the model is consistent with the CO problem since Eq . ( 6 ) introduces an upper bound on the rate , R ( ψ , φ , ψ0 ) ≤ R ( ψ , φ , ψ0 , φ0 ) ( see App . A.2 ) , and thus a lower bound on the ELBO . As a result , we obtain the Lagrangian L ( θ , ψ , φ , ψ0 , φ0 ; λ ) = R ( ψ , φ , ψ0 , φ0 ) + λ ( D ( θ , φ ) −D0 ) and arrive at the following optimisation problem : Empirical Bayes︷︸︸︷ min ψ0 , φ0 M-step︷︸︸︷ min θ , ψ E-step︷ ︸︸ ︷ max λ min φ L ( θ , ψ , φ , ψ0 , φ0 ; λ ) s.t . λ ≥ 0 . ( 7 )

Summary: This paper proposes various improvements to deep sequential latent variable models (or deep state space models). Namely, 1) the constrained optimization method from the Taming VAEs paper [Rezende & Viola, 2018], 2) a method for improving empirical priors over the first time step’s latent variable (referred to as VHP), and 3) a new model that combines the locally linear mappings of Karl et al., 2017 and the KVAE variable separation from Fraccaro et al., 2017. These improvements are demonstrated on pendulum and reacher datasets (both with images and joint angles), showing improvements in system identification, prediction accuracy, and the ability to use the state representations from MBRL. Overall, the paper provides several helpful tips for constructing and training deep state space models.