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 a framework for learning neural network-based state-space models from raw data. In comparison to the standard RNN-based approach [1], there are several modifications: (i) a modified ELBO formulated in the context of a constrained optimization (CO) framework [2]. CO prioritizes reconstruction term at the beginning of the training; (ii) hierarchical prior [3] instead of a simple Gaussian; (iii) a modified Kalman-VAE framework [4] instead of RNN-based transition model. Experimental evidence on pendulum data and the reacher environment indicates that the resulting framework outperforms baselines in predictive modeling and ELBO values.   [1] A Recurrent Latent Variable Model for Sequential Data  [2] Taming VAEs  [3] Learning Hierarchical Priors in VAEs  [4]  A Disentangled Recognition and Nonlinear Dynamics Model for Unsupervised Learning