Hidden Parameter Recurrent State Space Models For Changing Dynamics Scenarios

1 INTRODUCTION . System identification , i.e. , learning models of dynamical systems from observed data ( Ljung ( 1998 ) ; Gevers ( 2005 ) ) , is a key ingredient of model-predictive control ( Camacho & Alba ( 2013 ) ) and model-based reinforcement learning ( RL ) . State space models ( Hamilton ( 1994 ) ; Jordan ( 2004 ) ; Schön et al . ( 2011 ) ) ( SSMs ) provide a principled framework for modelling dynamics . Recently there have been several works that fused SSMs with deep ( recurrent ) neural networks achieving superior results in time series modelling ( Haarnoja et al. , 2016 ; Karl et al. , 2016 ) and system identification tasks ( Becker et al. , 2019 ; Hafner et al. , 2019 ; Shaj et al. , 2020 ) . Learning the dynamics in an encoded latent space allows us to work with high dimensional observations like images and was proven to be successful for planning from high dimensional sensory inputs . However , most of these tasks assume a fixed , unchanging dynamics , which is quite restrictive in real-world scenarios . Several control applications involve situations where an agent must learn to solve tasks with similar , but not identical , dynamics . A robot playing table tennis may encounter several bats with different weights or lengths , while an agent manipulating a bottle may face bottles with varying amounts of liquid . Humanoids learning to walk may face different terrains with varying slopes or friction coefficients . Any real-world dynamical system might change over time due to multiple reasons , some of which might not be directly observable or understood . For example , in soft robotics small variations in temperature or changes in friction coefficients of the cable drives of a robot can significantly change the dynamics . Similarly , a robot may undergo wear and tear over time which can change its dynamics . Thus , assuming a global model that is accurate throughout the entire state space or duration of use is a limiting factor for using such models in real-world applications . We found that existing literature on recurrent models fails to model the dynamics accurately in these situations . Thus , we introduce hidden parameter state-space models ( HiP-SSM ) , which allow capturing the variation in the dynamics of different instances through a set of hidden task parameters . We formalize the HiP-SSM and show how to perform inference in this graphical model . Under Gaussian assumptions , we obtain a probabilistically principled yet scalable approach . We name the resulting probabilistic recurrent neural network as Hidden Parameter Recurrent State Space Model ( HiP-RSSM ) . HiP-RSSM achieves state-of-the-art performance for several systems whose dynamics change over time . Interestingly , we also observe that HiP-RSSM often exceeds traditional RSSMs in performance for dynamical systems previously assumed to have unchanging global dynamics due to the identification of unobserved causal effects in the dynamics . 2 PRELIMINARIES . In this section , we introduce the fundamental concepts that underlie the design of our proposed model and inference procedure . 2.1 GAUSSIAN STATE SPACE MODELS . State space models are Bayesian probabilistic graphical models ( Koller & Friedman , 2009 ; Jordan , 2004 ) that are popular for learning patterns and predicting behavior in sequential data and dynamical systems . Let f ( · ) be any arbitrary dynamic model , and let h ( · ) be any arbitrary observation model . Then , the dynamic systems of a Gaussian state space model can be represented using the following equations zt = f ( zt−1 , at ) + ut , ut ∼ N ( 0 , Q ) ( 1 ) ot = h ( zt ) + vt , vt ∼ N ( 0 , R ) . ( 2 ) Here zt , at and ot are the latent states , control inputs/actions and observations at time t. The vectors ut ∼ N ( 0 , Q ) and vt ∼ N ( 0 , R ) denote zero-mean Gaussian noise . Further , based on the assumptions made for f ( · ) and h ( · ) , we can have different variants of Gaussian state space models ( Ribeiro , 2004 ; Wan & Van Der Merwe , 2000 ) . In the linear case , inference can be performed efficiently using Kalman Filters . 2.2 BAYESIAN CONTEXT AGGREGATION . Inferring a latent task variable l ∈ L using observed context data Cl = { xln } Nn=1 has been explored previously by different neural process architectures ( Gordon et al. , 2018 ; Garnelo et al. , 2018 ) . Neural processes are a meta-learning technique that rely on deep set functions ( Zaheer et al. , 2017 ) which output a latent representation out of varying number of context points in a permutation invariant manner . Volpp et al . ( 2020 ) formulated context data aggregation as a Bayesian inference problem , where the information contained in Cl is directly aggregated into the statistical description of l based on a factorized Gaussian observation model of the form p ( rln|l ) , where p ( rln|l ) = N ( rln|l , diag ( ( σln ) 2 ) ) , rln = encr ( xln ) , ( σln ) 2 = encσ ( xln ) . ( 3 ) Here n is the index of an element from context set Cl . Given a prior p0 ( l ) = N ( l|µ0 , diag ( σ20 ) ) we can compute the posterior p ( l|Cl ) using Bayes rule . The Gaussian assumption allows us to get a closed form solution for the posterior estimate of the latent task variable , p ( l|Cl ) based on Gaussian conditioning . The factorization assumption further simplifies this update rule by avoiding computationally expensive matrix inversions into a simpler update rule as σ2l = ( ( σ20 ) ⊖ + N∑ n=1 ( ( σln ) 2 ) ⊖ ) ⊖ , µl = µ0 + σ 2 l ⊙ N∑ n=1 ( rln − µ0 ) 2 ⊘ ( σln ) 2 , ( 4 ) where ⊖ , ⊙ and ⊘ denote element-wise inversion , product , and division , respectively . Intuitively the mean of the latent task variable µl is a weighted sum of the individual latent observations rln , while the variance of the latent task variable σ2l gives the uncertainty of this task representation . 3 HIDDEN PARAMETER RECURRENT STATE SPACE MODELS . Our goal is to learn a state space model that can model the dynamics of partially observable robotic systems under scenarios where the dynamics changes overtime . Often , the dynamics of real systems may differ in significant ways from the system our models were trained on . However , it is unreasonable to train a model across all possible conditions an agent may encounter . Instead , we propose a state space model that learns to account for the causal factors of variation observed across tasks at training time , and then infer at test time the model that best describe the system . Thus , we introduce a new formulation namely Hidden Parameter State Space Models ( HiP-SSMs ) , a framework for modeling families of SSMs with different but related dynamics using low-dimensional latent task embeddings . In this section , we formally define the HiP-SSM family in Section 3.1 , propose a probabilistically principled inference scheme for HiP-SSMs based on the forward algorithm ( Jordan , 2004 ) in Section 3.2 , and finally provide training scheme for the setting in Section 3.3 . 3.1 HIDDEN PARAMETER STATE SPACE MODELS ( HIP-SSMS ) We denote a set of SSMs with transition dynamics fl that are fully described by hidden parameters l and observation model h as a Hidden Parameter SSM ( HiP-SSM ) . In this definition we assume the observation model h to be independent of the hidden parameter l as we only focus on cases where the dynamics changes . HiP-SSMs allows us to extend SSMs to multi-task settings where dynamics can vary across tasks . Defining the changes in dynamics by a latent variable unifies dynamics across tasks as a single global function . In dynamical systems , for example , parameters can be physical quantities like gravity , friction of a surface , or the strength of a robot actuator . Their effects influence the dynamics but not directly observed ; and hence l is not part of the observation space and is treated as latent task parameter vector . Formally , a HiP-SSM is defined by a tuple { Z , A , W , L , f , h } , where Z , A and W are the sets of hidden states z , actions a , and observations w respectively . L is the set of all possible latent task parameters and let p0 ( l ) be the prior over these parameters . Thus , a HiP-RSSM describes a class of dynamics and a particular instance of that class is obtained by fixing the parameter vector l ∈ L. The dynamics f for each instance depend on the value of the hidden parameters l. Each instance of a HiP-SSM is an SSM conditioned on l. We also make the additional assumption that the parameter vector l is fixed for the duration of the task , and thus the latent task parameter has no dynamics . This assumption considerably simplifies the procedure for inferring the hidden parametrization and is reasonable since dynamics can be assumed to be locally-consistent over small trajectories in most applications ( Nagabandi et al. , 2018a ) . The definition is inspired from related literature on HiP-MDPs ( Doshi-Velez & Konidaris , 2016 ) , where the only unobserved variable is the latent task variable . One can connect HiP-SSMs with HiP-MDPs by including rewards to the definition and formalize HiP-POMDPs . However this is left for future work . 3.2 LEARNING AND INFERENCE . We perform inference and learning in the HiP-SSM borrowing concepts from both deep learning and graphical model communities following recent works on recurrent neural network models ( Haarnoja et al. , 2016 ; Becker et al. , 2019 ) , where the architecture of the network is informed by the structure of the probabilistic state estimator . Aside from the simplicity of the training procedure , one of the key advantages of this approach is the ability to incorporate arbitrary nonlinear components into the observation and transition functions . We denote the resulting probabilistic recurrent neural network architecture as Hidden Parameter Recurrent State Space Model ( HiP-RSSM ) . The structure of the Bayesian network shown in Figure 1 allows tractable inference of latent variables by the forward algorithm ( Jordan , 2004 ; Koller & Friedman , 2009 ) . Since we are dealing with continuous dynamical systems , we assume a Gaussian multivariate distribution over all variables ( both observed and hidden ) for the graph shown in Figure 1 . This assumption has several advantages . Firstly , it makes the inference very similar to the well studied Kalman Filtering approaches . Secondly , the Gaussian assumptions and conditional in-dependencies allows us to have a closed form solution to each of these update procedures which are fully differentiable and can be backpropagated to the deep encoders . The belief update over the hidden variables zt and l happens in three stages . Similar to Kalman filtering approaches , we have two recursive belief state update stages , the time update and observation update which calculate the prior and posterior belief over the latent states respectively . However , we have an additional hierarchical latent variable l which models the ( unobserved ) causal factors of variation in dynamics in order to achieve efficient gener- alization . Hence , we have a third belief update stage to calculate the posterior over the latent task variable based on the observed context set . Each of these three stages are detailed in the sections below :

Summary: The paper proposes a method to learn a probabilistic recurrent state-space model for time-varying dynamics. The proposed method combines the Kalman filtering-based update rule with deep network-based encoder and decoder model. Effectively, the method can be used to replace RNN cells in a recurrent model, and is shown to outperform baseline models in modeling various robotic tasks.