HIDDEN PARAMETER RECURRENT STATE SPACE MODELS FOR CHANGING DYNAMICS SCENARIOS
Year: 2021

Abstract:
Recurrent State-space models (RSSMs) are highly expressive models for learning patterns in time series data and system identification. However, these models assume that the dynamics are fixed and unchanging, which is rarely the case in realworld scenarios. Many control applications often exhibit tasks with similar but not identical dynamics which can be modeled as a latent variable. We introduce the Hidden Parameter Recurrent State Space Models (HiP-RSSMs), a framework that parametrizes a family of related dynamical systems with a low-dimensional set of latent factors. We present a simple and effective way of learning and performing inference over this Gaussian graphical model that avoids approximations like variational inference. We show that HiP-RSSMs outperforms RSSMs and competing multi-task models on several challenging robotic benchmarks both on real-world systems and simulations.

heading: 1 INTRODUCTION
text: 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.

heading: 2 PRELIMINARIES
text: In this section, we introduce the fundamental concepts that underlie the design of our proposed model and inference procedure.

heading: 2.1 GAUSSIAN STATE SPACE MODELS
text: 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.

heading: 2.2 BAYESIAN CONTEXT AGGREGATION
text: 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.

heading: 3 HIDDEN PARAMETER RECURRENT STATE SPACE MODELS
text: 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.

heading: 3.2 LEARNING AND INFERENCE
text: 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:

heading: 3.2.1 INFERRING THE LATENT TASK VARIABLE (CONTEXT UPDATE)
text: In this stage, we infer the posterior over the Gaussian latent task variable l based on an observed context set Cl. For any HiP-RSSM instance defined over a target sequence T = (ot,at,ot+1, ...,ot+N ,at+N ), over which we intend to perform state estimation/prediction, we maintain a fixed context set Cl. Cl in our HiP-SSM formalism can be obtained as per the algorithm designer’s choice. We choose a Cl consisting of tuples of the current state/observation, current action and next state/observation for the previous N time steps, i.e. Cl = {olt−n,alt−n,olt−n+1}Nn=1. We map these tuples to a set of intermediate representation {rln, ( σln )2}Nn=1 using a context encoder,
i.e., p ( rln ) = N ( rln|l, diag ( σ2l,n )) , with
rln = encr(o l t−n,a l t−n,o l t−n+1), ( σln )2 = encσ(olt−n,a l t−n,o l t−n+1). (5)
We then impose a factorized Gaussian prior p0(l) ∼ N (l|µ0, diag(σ20)) and arrive at a Gaussian aggregation model which allows to derive the parameters of the posterior distribution p(l|Cl) = N (µl, diag(σ2l )) based on Bayesian principles via Gaussian conditioning as detailed in Section 2.2.

heading: 3.2.2 INFERRING PRIOR LATENT STATES OVER THE NEXT TIME STEP (TIME UPDATE)
text: The goal of this step is to compute the prior marginal p(zt|w1:t−1,a1:t, Cl) = ∫∫ p(zt|zt−1,at, l)p(zt−1|w1:t−1,a1:t−1, Cl)p(l|Cl)dzt−1dl. (6)
We assume a dynamical model of the following form: p(zt|zt−1,at, l) = N (At−1zt−1 +Bat +Cl,Σtrans) and denote the posterior from the previous time-step by p(zt−1|w1:t−1,a1:t−1, Cl) = N ( z+t−1,Σ + t−1 ) . Following Shaj et al. (2020) we assume the action at is known and not subject to noise.
At any time t, the posterior over the belief state (zt−1|w1:t−1,a1:t−1,Cl), posterior over the latent task variable l|Cl and the action at are independent of each other since they form a “common-effect” / v-structure (Koller et al., 2007) with the unobserved variable zt. With these independencies and Gaussian assumptions, according to Gaussian Identity 1, it can be shown that calculating the integral
in equation 6 has a closed form solution as follows, p(zt|w1:t−1,a1:t, Cl) = N (z−t ,Σ−t ), where
z−t = At−1z + t−1 +Bat +Cµl, (7)
Σ−t = At−1Σ + t−1A T t−1 +CΣlC T +Σtrans. (8)
Gaussian Identity 1. If u ∼ N (µu + b,Σu) and v ∼ N (µv,Σv) are normally distributed independent random variables and if the conditional distribution p (y|u,v) = N (Au+b+Bv,Σ), then marginal p(y) = ∫ p(y|u,v)p(u)p(v))dudv = N (Aµu+b+Bµv,AΣuAT +BΣvBT + Σ)
A derivation for Identity 1 is provided in the Appendix. Equation 7 is similar to the Kalman Time Update / Predict step, with an additional task variable. We call this the task conditional Kalman time update step.
We use a locally linear transition model At−1 as in Karl et al. (2016); Fraccaro et al. (2017); Becker et al. (2019) and a non-linear control model as in Shaj et al. (2020). For the latent task transformation we can either use a linear, locally-linear or non-linear transformation. More details regarding the latent task transformation model can be found in the appendix F.5. Our ablations (Figure 3c) show that a non-linear feedforward neural network f(.) that outputs mean and variances and interact additively gave the best performance in practice. f(.) transforms the latent task moments µl and σl directly into the latent space of the state space model via additive interactions. The corresponding time update equations are given below:
z−t = At−1z + t−1 + b(at) + f(µl),
Σ−t = At−1Σ + t−1A T t−1 + f(σl) +Σtrans.

heading: 3.2.3 INFERRING POSTERIOR LATENT STATES / OBSERVATION UPDATE
text: The goal of this step is to compute the posterior belief p(zt|o1:t,a1:t,Cl). We first map the observations at each time step ot to a latent space using an observation encoder (Haarnoja et al., 2016; Becker et al., 2019) which emits latent features wt along with an uncertainty in those features via a variance vector σtobs. We then compute the posterior belief p(zt|w1:t,a1:t−1,Cl), based on p(zt|w1:t−1,a1:t−1,Cl) obtained from the time update, the latent observation (wt,σobst ) and the observation model H . This is exactly the Kalman update step, which has a closed form solution as shown below for a time instant t,
Kalman Gain: Qt = Σ−t H T ( HΣ−t H T + I · σobst )−1 , Posterior Mean: z+t = z − t +Qt ( wt −Hz−t ) , Posterior Covariance: Σ+t = (I −QtH)Σ − t ,
where I denotes the identity matrix. This update is added as a layer in the computation graph (Haarnoja et al. (2016); Becker et al. (2019)). However, the Kalman update involves computationally expensive matrix inversions of the order of O(L3), where L is the dimension of the latent state zt. Thus, in order to make the approach scalable we follow the same factorization assumptions as in Becker et al. (2019). This provides a simple way of reducing the observation update equation to a set of scalar operations, bringing down the computational complexity from O(L3) toO(L). More mathematical details on the simplified update equation can be found in appendix F.7.

heading: 3.3 TRAINING HIP-RSSM FOR ADAPTATION TO CHANGING DYNAMICS
text: We collect a set of trajectories over which the dynamics changes overtime. This can be trajectories where a robot picks up objects of different weights or a robot traversing terrain of different slopes. Now we introduce a multi-task setting with a rather flexible definition of task, where each temporal segment of a trajectory can be considered to be a different “task” and the observed context set based on interaction histories from the past N time steps provides information about the current task setting. This definition allows us to have potentially infinite number of tasks/local
dynamical systems and the distribution over these tasks/systems is modelled using a hierarchical latent task variable l. The formalism is based on the assumption that over these local temporal segments, the dynamics is unchanging. This local consistency in dynamics holds for most real world scenarios (Nagabandi et al., 2018b;a). Yet our formalism can model the global changes in dynamics since we obtain a different instance of the HiP-RSSM for each temporal segment based on the observed context set. We also provide a detailed description of the multi-task dataset creation process in in the Appendix C and a pictorial illustration in Appendix B.
Batch Training. Let T ∈ RB×N×D be the batch of local temporal segments with different dynamics which we intend to model with the HiP-RSSM formalism. Given a target batch T , HiP-RSSM can be trained in a supervised fashion similar to popular recurrent architectures like LSTMs or GRUs. However for each local temporal sequence t ∈ T , over which the dynamics is being modelled, we also input a set of the N previous interactions, which forms the context set C ∈ RB×N×D for inferring the latent task as explained in Section 3.2.1. Processing the context set C adds minimal additional computational/memory constraints as we use a permutation invariant set encoder. The set encoder allows parallelization in processing the context set as opposed to the recurrent processing of target set.
The parameters that are learnable in the computation graph includes the locally linear transition model At, the non-linear control factor b, the linear/non-linear latent task transformation model C, the transition noise Σtrans, along with the observation encoder, context encoder and the output decoder.

heading: Loss Function.
text: We optimize the root mean square error (RMSE) between the predicted decoder output decµ(z−t+1) and the ground truth outputs. Here decµ() denotes the part of the output decoder that are responsible for decoding the latent mean. In principle, we could train on the Gaussian log-likelihood instead of the RMSE and hence model the variances using an additional variance decoder. Training on RMSE yields slightly better predictions and allows for a fair comparison with deterministic baselines like feed-forward neural networks, LSTMs and meta-learnng algorithms such as MAML. Thus, we report results with the RMSE loss. A probabilistic formulation for this loss function is provided in appendix D.
Gradients are computed using (truncated) backpropagation through time (BPTT) (Werbos (1990)) and clipped. We optimize the objective using the Adam (Kingma & Ba (2014)) stochastic gradient descent optimizer with default parameters.

heading: 4 RELATED WORK
text: Deep State Space Models. Recently, there have been several works on deep SSMs that offer tractable inference and scalability to high dimensional and large datasets. Haarnoja et al. (2016); Becker et al. (2019); Shaj et al. (2020) use neural network architectures based on closed-form solutions of the forward inference algorithm on SSMs and achieve the state of the art results in state estimation and dynamics prediction tasks. Krishnan et al. (2017); Karl et al. (2016); Fraccaro et al. (2017); Hafner et al. (2019) perform learning and inference on SSMs using variational approximations. However, most of these recurrent state-space models assume that the dynamics is fixed, which is a significant drawback since this is rarely the case in real-world applications such as robotics.
Recurrent Switching Dynamical Systems. Linderman et al. (2017); Becker-Ehmck et al. (2019); Dong et al. (2020) tries to address the problem of changing/multi-modal dynamics by incorporating additional discrete switching latent variables. However, these discrete states make learning and inference more involved. Linderman et al. (2017) uses auxiliary variable methods for approximate inference in a multi-stage training process, while Becker-Ehmck et al. (2019); Dong et al. (2020) uses variational approximations and relies on additional regularization/annealing to encourage discrete state transitions. We take a rather different formalism for modelling changing dynamics by viewing it as a multi-task learning problem with a continuous hierarchical hidden parameter that model the distribution over these tasks.
Hidden Parameter MDPs. Hidden Parameter Markov Decision Process (HiP-MDP) (Doshi-Velez & Konidaris, 2016; Killian et al., 2016) address the setting of learning to solve tasks with similar, but not identical, dynamics. In HiP-MDP formalism, the states are assumed to be fully observed. However, we formalize the partially observable setting where we are interested in modelling the dynamics in a latent space under changing scenarios. The formalism is critical for learning from high dimensional observations and dealing with partial observability and missing data. The formalism can be easily extended to HiP-POMDPs by including rewards into the graphical model 1, for planning in the latent space (Hafner et al., 2019; Sekar et al., 2020). However this is left as a future work.
Meta Learning For Changing Dynamics. There also exists a family of approaches that attempt online model adaptation to changing dynamics scenarios via meta-learning (Finn et al., 2017; Nagabandi et al., 2018a;b). These methods perform online adaptation on the parameters of the dynamics models through gradient descent from interaction histories. Our method fundamentally differs from these approaches in the sense that we do not perform a gradient descent step at every time step during test time, which is computationally impractical in modern robotics, where we often deal with high frequency data. Also, we focus on dynamics modelling in the latent space which have several advantages as discussed previously. We also empirically show that our approach adapts better especially in scenarios with non-markovian dynamics, a property that is often encountered in real world robots due to stiction, slip, friction, pneumatic/hydraulic/cable-driven actuators etc.

heading: 5 EXPERIMENTS
text: This section evaluates our approach on a diverse set of dynamical systems from the robotics domain in simulations and real systems. We show that HiP-RSSM outperforms contemporary recurrent state-space models (RSSMs) and recurrent neural networks (RNNs) by a significant margin under changing dynamics scenarios. Further, we show that HiP-RSSM outperforms these models even under situations with partial observability/ missing values. We also baseline our HiP-RSSM with contemporary multi-task models and improve performance, particularly in modelling non-Markovian dynamics and under partial observability. Finally, the visualizations of the Gaussian latent task variables in HiP-RSSM demonstrates that they learn meaningful representations of the causal factors of variations in dynamics in an unsupervised fashion.
We consider the following baselines:
• RNNs - We compare our method to two widely used recurrent neural network architectures, LSTMs (Hochreiter & Schmidhuber, 1997) and GRUs (Cho et al., 2014).
• RSSMs - Among several RSSMs from the literature, we chose RKN (Becker et al., 2019) as they have shown excellent performance for dynamics learning (Shaj et al., 2020) and relies on deep architectures based on the forward inference algorithm as in our case.
• Multi Task Models - We also compare with state of the art multi-task learning models like Neural Processes (NPs) and MAML (Nagabandi et al., 2018a). Both models receive the same context information as in HiP-RSSM.
In case of recurrent models we replace the HiP-RSSM cell with a properly tuned LSTM, GRU and RKN Cell respectively, while fixing the encoder and decoder architectures. For the NP baseline we use the same context encoder and aggregation mechanism as in HiP-RSSM to ensure a fair comparison. We create partially observable settings by imputing 50% of the observations during inference. More details regarding the baselines and hyperparameters can be found in the appendix.

heading: 5.1 SOFT ROBOT PLAYING TABLE TENNIS
text: We first evaluate our model on learning the dynamics of a pneumatically actuated muscular robot. This four Degree of Freedom (DoF) robotic arm is actuated by Pneumatic Artificial Muscles (PAMs) (Büchler et al., 2016). The data consists of trajectories of hitting movements with varying speeds while playing table tennis (Büchler et al., 2020). The fast motions with high accelerations of PAMs are difficult to model due to hysteresis and hence require recurrent models (Shaj et al., 2020).
We show the prediction accuracy in RMSE in Table 1a. We observe that the HiP-RSSM can outperform the previous state of the art predictions obtained by recurrent models. Based on the domain knowledge, we hypothesise that the latent task variable captures multiple unobserved causal factors of variation that affect the dynamics in the latent space, which are not modelled in contemporary recurrent models. These causal factors could be, in particular, the temperature changes or the friction due to a different path that the Bowden trains take within the robot. Disentangling and interpreting these causal factors can be exciting, but it is out of scope for the current work.
Further, we find that the multi-task models like NPs and MAML fail to model these dynamics accurately compared to all the recurrent baselines because of the non-markovian dynamics resulting from the high accelerations in this pneumatically actuated robot.

heading: 5.2 ROBOT MANIPULATION WITH VARYING LOADS
text: We collected the data from a 7 DoF Franka Emika Panda manipulator carrying six different loads at its end-effector. It involved a mix of movements of different velocities from slow to swift motions covering the entire workspace of the robot. We chose trajectories with four different loads as training set and evaluated the performance on two unseen weights, which results in a scenario where the dynamics change over time.
We show the prediction errors in RMSE in Table 1c. HiP-RSSMs outperforms all recurrent statespace models, including RKN and deterministic RNNs, in modelling these dynamics under fully observable and partially observable conditions. The multi-task baselines of NPs and MAML perform equally well under full observability for this task because of the near Markovian dynamics of Franka Manipulator, which often does not need recurrent models. However, HiP-RSSMs have an additional advantage in that these are naturally suited for partially observable scenarios and can predict ahead in a compact latent space, a critical component for recent success in model-based control and planning (Hafner et al., 2019).

heading: 5.3 ROBOT LOCOMOTION IN TERRAINS OF DIFFERENT SLOPES
text: Wheeled mobile robots are the most common types of robots being used in exploration of unknown terrains where they may face uneven and challenging terrain. We set up an environment using a Pybullet simulator (Coumans & Bai, 2016) where a four-wheeled mobile robot traverses an uneven terrain of varying steepness generated by sinusoidal functions (Sonker & Dutta, 2020) as shown
in 4a. This problem is challenging due to the highly non-linear dynamics involving wheel-terrain interactions. In addition, the varying steepness levels of the terrain results in a changing dynamics scenario, which further increases the complexity of the task.
We show the prediction errors in RMSE in Table 1c. When most recurrent models, including RSSMs and deterministic RNNs, fail to model these dynamics, HiP-RSSMs are by far the most accurate in modelling these challenging dynamics in the latent space. Further, the HiP-RSSMs perform much better than state of the art multi-task models like NPs and MAML.
We finally visualize the latent task representations using TSNE in Figure 3a. As seen from the plot, those instances of the HiP-SSMs under similar terrain slopes get clustered together in the latent task space, indicating that the model correctly identifies the causal effects in the dynamics in an unsupervised fashion.

heading: 5.4 VISUALIZING CHANGING HIDDEN PARAMETERS AT TEST TIME OVER TRAJECTORIES
text: WITH VARYING DYNAMICS
We finally perform inference using the trained HiP-RSSM in a multi-task / changing dynamics scenario where the dynamics continuously changes over time. We use the inference procedure described in Appendix B based on a fluid definition for “task” as the local dynamics in a temporal segment. We plot the global variation in the latent task variable, captured by each instance of the HiP-RSSM defined over these local temporal segments, using the dimensionality reduction technique UMAP (McInnes et al., 2018). As seen in figures 4b and 4c, the latent task variable captures these causal factors of variations in an interpretable manner.

heading: 6 CONCLUSION
text: We proposed HiP-RSSM, a probabilistically principled recurrent neural network architecture for modelling changing dynamics scenarios. We start by formalizing a new framework, HiP-SSM, to address the multi-task state-space modelling setting. HiP-SSM assumes a shared latent state and action space across tasks but additionally assumes latent structure in the dynamics. We exploit the structure of the resulting Bayesian network to learn a universal dynamics model with latent parameter l via forward inference algorithm and backpropagation through time. The resulting recurrent neural network, namely HiP-RSSM, learns to cluster SSM instances with similar dynamics together in an unsupervised fashion. Our experimental results on various robotic benchmarks show that HiPRSSMs significantly outperform state of the art recurrent neural network architectures on dynamics modelling tasks. We believe that modelling the dynamics in the latent space which disentangles the state, action and task representations can benefit multiple future applications including planning/control in the latent space and causal factor identification.

heading: A PROOF FOR GAUSSIAN IDENTITY 1
text: This section provides a proof for Gaussian Identity 1. First we derive an expression for the joint distribution p(u,v,y).
Gaussian Identity 2. If u ∼ N (µu + b,Σu) and v ∼ N (µv,Σv) are normally distributed independent random variables and if conditional distribution p(y|u,v) = N (Au + b + Bv,Σ), the joint distribution has an expression as follows:(
u v y
) ∼ N ( µuµv Aµu + b+Bµv ) ,  Σu 0 ΣuA⊤0 Σv ΣvB⊤ AΣ⊤u BΣ ⊤ v AΣuA ⊤ +BΣvB ⊤ +Σ 

heading: Proof for Identity 2
text: Let displacement of a variable u be denoted by ∆u = u− ⟨u⟩. Since u and v are independent, covariances 〈 ∆u∆v⊤ 〉 = 0. We can write y = Au+b+Bv+ϵ, where ϵ ∼ N (0,Σ) and b is a constant. Then we have covariance〈 ∆u∆y⊤ 〉 = 〈 ∆u(A∆u+B∆v +∆ϵ)⊤ 〉 = 〈 ∆u∆u⊤ 〉 A⊤ + 〈 ∆u∆v⊤ 〉 B⊤ + 〈 ∆u∆ϵ⊤ 〉 .
Since 〈 ∆u∆v⊤ 〉 = 〈 ∆u∆ϵ⊤ 〉 = 0 we therefore have 〈 ∆u∆y⊤ 〉 = ΣuA ⊤.
The derivations for covariances 〈 ∆v∆y⊤ 〉 follows similarly and the corresponding covariance has
the expression, 〈 ∆v∆y⊤ 〉 = ΣvB ⊤.
Similarly, 〈 ∆y∆y⊤ 〉 = 〈 (A∆u+B∆v +∆ϵ)(A∆u+B∆v +∆ϵ)⊤ 〉 = A 〈 ∆u∆u⊤ 〉 A⊤ +
B 〈 ∆v∆v⊤ 〉 B⊤ + 〈 ∆ϵ∆ϵ⊤ 〉 = AΣuA ⊤ +BΣvB ⊤ +Σ. The result follows.
Gaussian Identity 3 (Gaussian Marginalization). If( u v y ) ∼ N ( µuµv µz ) ,  Σuu Σuv ΣuyΣ⊤uv Σvv Σvy Σ⊤uy Σ ⊤ vy Σyy  then marginal over y is given as p(y) = ∫ u,v p(y|u,v)p(u)p(v)dudv = N (µy,Σyy)

heading: Proof For Identity 3
text: We refer to Bishop (2006) for the derivation, which requires calculation of the Schur complement as well as completing the square of the Gaussian p.d.f. to integrate out the variable. The given derivation (Bishop, 2006) for two variable multivariate gaussians can be extended to 3 variable case WLOG.
Proof For Identity 1 is immediate from Identity 2 and Identity 3.

heading: B HIP-RSSM DURING TEST TIME / INFERENCE
text: We perform inference using HiP-RSSM at test time on a trajectory with varying dynamics using algorithm 1. A pictorial representation of the same is given in the figure 5. We use this inference scheme to visualize how the latent variable l, that describe different instances of a HiP-RSSM over different temporal segments, evolve at a global level. The visualizations are reported in figures 4b and 4c in the main text.
Algorithm 1: HiP-RSSM Test Time Inference Required: Trained HiP-RSSM Model Required: A time series τ of length K >> N Divide the time series τ into non-overlapping windows Tl of length N. Let T = {T1, T2, T3, ., ., .} be the ordered list of all temporal segments, sorted in the ascending order of time of occurance. foreach each time window Tl ∈ T do 1. maintain a context set Cl consisting of N previous interactions; 2. infer posterior latent task variable p(l|Cl) using context update stage as in section 3.2.1; 3. using the posterior over latent task variable l|Cl and observations in sequence Tl to
perform sequential bayesian inference over the state space model using Kalman observation update (3.2.3) and task conditional Kalman time update; (3.2.2)
end

heading: C MULTI-TASK DATASET CREATION
text: Algorithm 2: Multi Task Dataset Creation For Training HiP-RSSM Required:A set S of trajectories of changing dynamics D ← ϕ foreach each trajectory τ ∈ S do
1. divide the trajectory τ into non-overlapping windows Tl of length N. Let T = {T1, T2, T3, ., ., .} be the list of all temporal segments/time-series. foreach each time window Tl ∈ T do
(a) maintain a context set Cl consisting of N previous interactions; (b) update D ← D ∪ {Cl, Tl}
end
end Output: Output D consisting of batch of context and target sets.

heading: D PROBABILISTIC FORMULATION OF HIP-RSSM TRAINING OBJECTIVE
text: The training objective for the HiP-RSSM involves maximizing the posterior predictive loglikelihood given below:
T∑ t=1 N∑ b=1 log p(ôbt+1|Cl,ob1:t) = T∑ t=1 N∑ b=1 log ∫ p(ôbt+1|zbt+1)p(zbt+1|om1:t,Cl)dzbt+1. (9)
Here, for any time t, ôt+1 are the ground truth observations for the next time step and the latent state prior p(zbt+1|om1:t,Cl) as discussed in the Section 3.2.2, have closed form solutions. We approximate the objective in Eq.(9) using a deterministic variational approximation, similar to Becker et al. (2019).
We employ a Gaussian approximation of the posterior predictive log likelihood of the form p(ôbt+1|Cl,ob1:t) ≈ N (µobt+1 , diag((σobt+1) 2)). Here µobt+1 = decµ(z b t+1) and σobt = decΣ(Σbt+1)), where decµ() and decΣ() denote the parts of the output decoder that are responsible for decoding the latent means and latent variances respectively. This decoder can be interpreted as a “moment matching network”, computing the moments of obt given the moments of z b t (Volpp et al., 2020; Becker et al., 2019). This allows to evaluate an approximation to the objective in a deterministic manner.
Alternatively, we could ignore the variances and only focus on the mean predictions by training our model with a RMSE loss. Training on RMSE yields slightly better predictions and allows for a fair
comparison with deterministic baselines (Feed Forward NN, LSTM, GRU, MAML etc) and hence we report results with the RMSE loss.

heading: E ADDITIONAL EXPERIMENTS

heading: E.1 ABLATION ON CONTEXT ENCODER
text: In table 2, we report the details of evaluating our bayesian aggregation based context set encoder (discussed in 3.2.1) against the a causal / recurrent encoder that takes into account the temporal structure of the context data. We used a probabilistic recurrent encoder (Becker et al., 2019), whose mean and variance from the last time step is used to infer the posterior latent task distribution p(l|Cl) = N (µl, diag(σ2l )). The dimensions of latent task parameters obtained from both the set and recurrent encoders are kept the same (60).
The reported experiments are conducted on data from wheeled mobile robot discussed in section 5.3. As reported in Table 2 the permutation invariant set encoder outperforms the recurrent encoder by a good margin in terms of prediction accuracy for both fully and partially observable scenarios. Additionally the set encoder is far more efficient in terms of computational time required for training
as seen from the time taken per epoch for each of these cases.

heading: E.2 MULTI-STEP AHEAD PREDICTIONS
text: In figure 6, we compare the results of multistep ahead predictions for upto 50 steps of HiP-RSSM with recurrent baselines like RKN(Becker et al., 2019) and LSTM. Inorder to train the recurrent models for multi step ahead predictions, we removed three-quarters of the observations from the temporal sequence and tasked the models with imputing those missing observations, only based on the knowledge of available actions/control commands, i.e., we train the models to perform action conditional future predictions to impute missing observations. The imputation employs the model for multi-step ahead predictions in a convenient way (Shaj et al., 2020). One could instead also go for a dedicated multi-step loss function as in approaches like Finn et al. (2016).
As seen in figure 6, HiP-RSSM clearly outperforms contemporary recurrent models for multi-step ahead prediction tasks since it takes into account additional causal factors of variation (slopes of the terrain for this robot) in the latent dynamics in an unsupervised manner.
F IMPLEMENTATION DETAILS

heading: F.1 CONTEXT SET ENCODER AND LATENT TASK REPRESENTATION
text: The HiP-RSSM maps a set of previous interaction histories, {olt,n,alt,n,olt+1,n}Nn=1, to a set of latent features and an estimate of uncertainty in those features, {rln, (σln)2}Nn=1, using a context encoder. We use a feed-forward neural network as the encoder in all of our experiements since we
deal with high dimensional vectors as our observations. However, depending upon the nature of observations, we could use different encoder architectures. The set of latent features and uncertainties, {rln, ( σln )2}Nn=1, are further aggregated in a probabilistically principled manner using bayesian aggregation operator discussed in 2.2 and 3.2.1 to get a gaussian latent task variable, with a mean (µl) and diagonal covariance (σl). Intuitively the context encoder learns to weight the contribution from each observation in the context set based on bayesian priniciples and emits a probabilistic representation of the latent task.

heading: F.2 OBSERVATION ENCODER AND LATENT STATE REPRESENTATION
text: HiP-RSSM transforms the observations at each time step ot to a latent space using an encoder network which emits latent features wt and an estimate of the uncertainty in those features via a variance vector σobst .
The probabilistic recurrent module uses a latent state vector zt and corresponding covariance Σt whose transitions are governed by the update rules in the HiP-RSSM cell. The latent state vector zt has been designed to contain two conceptual parts, a vector pt for holding information that can directly be extracted from the observations and a vector mt to store information inferred over time, e.g., velocities. The former is referred to as the observation or upper part and the latter as the memory or lower part of the latent state as in Becker et al. (2019). For an ordinary dynamical system and images as observations the former may correspond to positions while the latter corresponds to velocities. The corresponding posterior and prior covariance matrices Σ+t and Σ − t have a chosen factorized representation to yield simple Kalman updates, i.e.,
Σt =
[ Σut Σ s t
Σst Σ l t
] ,
where each of Σut ,Σ s t,Σ l t ∈ Rm×m is a diagonal matrix. The vectors σut ,σlt and σst denote the vectors containing the diagonal values of those matrices. This structure with Σst ensures that the correlation between the memory and the observation parts are not neglected as opposed to the case of designing Σt as a diagonal covariance matrix. This representation was exploited to avoid the expensive and numerically problematic matrix inversions involved in the KF equations as shown below.

heading: F.3 LOCALLY LINEAR TRANSITION MODEL
text: The state transitions in the time update stage (section 3.2.2) of the HiP-RSSM Cell is governed by a locally linear transition model as in Becker et al. (2019). To obtain a locally linear transition model, the HiP-RSSM gloablly learns K constant transition matrices A(k) and interpolate between them using state dependent coefficients α(k)(z+t ), i.e., At = ∑K k=0 α (k)(z+t )A (k), where z+t is the mean of the posterior distribution at time step t. So we locally linearize At around the posterior mean, which can be interpreted as equivalent to what an EKF do with a Jacobian. A small neural network with softmax output is used to learn α(k).

heading: F.4 CONTROL MODEL
text: To achieve action conditioning within the recurrent cell, we include a control model b(at) in addition to the locally linear transition model At. b(at) = f(at), where f(.) can be any non-linear function approximator. We use a multi-layer neural network regressor with ReLU activations (Shaj et al., 2020).

heading: F.5 LATENT TASK TRANSFORMATION MODEL
text: To achieve latent task conditioning within the recurrent cell, we include a task transformation model (c), in addition to the locally linear transition model At and control model b in the time update stage (section 3.2.2). Though in section 3.2.2, we used the notation for a linear task transformation
matrix, C, to motivate the additive interaction of latent task variables, µl and σl in the latent space, the task transformation function can be designed in several ways, i.e.:
(i) Linear: c = C, where C is a linear transformation matrix. The corresponding time update equations are given below:
z−t = At−1z + t−1 + b(at) +Cµl,
Σ−t = At−1Σ + t−1A T t−1 +C(I · σl)CT +Σtrans.
(ii) Locally-Linear: c = Ct, where Ct = ∑K k=0 β (k)(zt)C
(k) is a linear combination of k linear control models C(k). A small neural network with softmax output is used to learn β(k). The corresponding time update equations are given below:
z−t = At−1z + t−1 + b(at) +Ctµl,
Σ−t = At−1Σ + t−1A T t−1 +Ct(I · σl)CTt +Σtrans.
(iii) Non-Linear: c = f , where f(.) can be any non-linear function approximator. We use a multi-layer neural network regressor with ReLU activations, which transforms the latent task moments µl and σl directly into the latent space of the state space model via additive interactions. The corresponding time update equations are given below:
z−t = At−1z + t−1 + b(at) + f(µl),
Σ−t = At−1Σ + t−1A T t−1 + f(σl) +Σtrans.
In our ablation study (Figure 3c), for the linear and locally linear task transformation models, we assume that the dimension of the latent context variable l and the latent state space zt are equal. This allows us to work with square matrices which are more convenient. For the non-linear transformation we are free to choose the size of the latent context variable. However for a fairer comparison we keep the dimension of latent task variable to be similar in all three cases. We choose the non-linear task transformation model in HiP-RSSM architecture as this gave the best performance in practice.

heading: F.6 OBSERVATION MODEL
text: The latent state space Z = Rn of the HiP-RSSM is related to the observation spaceW by the linear latent observation model H = [ Im 0m×(n−m) ] , i.e., w|z ∼ N (Hz,Σobs) with w ∈ W and z ∈ Z , where Im denotes the m×m identity matrix and 0m×(n−m) denotes a m× (n−m) matrix filled with zeros. Typically, m is set to n/2. This corresponds to the assumption that the first half of the state can be directly observed while the second half is unobserved and contains information inferred over time (Becker et al. (2019)).

heading: F.7 OBSERVATION UPDATE STEP
text: The observation update discussed in section 3.2.3 involves computing the Kalman gain matrix Qt, which requires computationally expensive matrix inversions that are difficult to backpropagate, at least for high dimensional latent state representations. However, the choice of a locally linear transition model, the factorized covariance Σt, and the special observation model simplify the Kalman update to scalar operations as shown below. As the network is free to choose its own state representation, it finds a representation where such assumptions works well in practice Becker et al. (2019).
Similar to the state, the Kalman gain matrix Qt is split into an upper Qut and a lower part Qlt. Both Q u t and Q l t are squared matrices. Due to the simple latent observation model H =[ Im 0m×(n−m) ] and the factorized covariances, all off-diagonal entries of Qut and Q l t are zero and one can again work with vectors representing the diagonals, i.e., qut and q l t. Those are obtained by
qut = σ u,− t ⊘ ( σu,−t + σ obs t ) qlt = σ s,− t ⊘ ( σu,−t + σ obs t ) ,
where ⊘ denotes an elementwise vector division. The update equation for the mean therefore simplifies to
z+t = z − t + [ qut qlt ] ⊙ [ wt − zu,−t wt − zu,−t ] ,
where ⊙ denotes the elementwise vector product. The update equations for the individual parts of covariance are given by
σu,+t = (1m − qut )⊙ σ u,− t , σs,+t = (1m − qut )⊙ σ s,− t ,
σl,+t = σ l,− t − qlt ⊙ σ s,− t ,
where 1m denotes the m dimensional vector consisting of ones.

heading: G ROBOTS AND DATA
text: The experiments are performed on data from two different real robots. The details of robots, data and data preprocessing is explained below:

heading: G.1 MUSCULOSKELETAL ROBOT ARM
text: Observation and Data Set: For this soft robot we have 4 dimensional observation inputs(joint angles) and 8 dimensional action inputs(pressures). We collected the data of a four DoF robot actuated by Pneumatic Artificial Muscles (PAMs). The robot arm has eight PAMs in total with each DoF actuated by an antagonistic pair. The robot arm reaches high joint angle accelerations of up to 28, 000deg/s2 while avoiding dangerous joint limits thanks to the antagonistic actuation and limits on the air pressure ranges. The data consists of trajectories collected while training with a model-free reinforcement learning algorithm to hit balls while playing table tennis. We sampled the data at 100Hz. The hysteresis associated with the pneumatic actuators used in this robot is challenging to model and is relevant to the soft robotics in general.
Training Procedure: For the fully observable case, we trained HiP-RSSM for one-step ahead prediction using an RMSE loss. For partially observable case, during training, we randomly removed half of the observations from the sequences and tasked the models with imputing those missing states, only based on the knowledge of available actions/control commands, i.e., we train the models to perform action conditional future predictions to impute missing states. The imputation employs the model for multi-step ahead predictions in a convenient way. Other recurrent baselines (RKN, LSTM, GRU) are trained in a similar fashion except that, we don’t maintain a context set of interaction histories during training/inference.

heading: G.2 FRANKA EMIKA PANDA ROBOT ARM
text: Observation and Data Set: We collected the data from a 7 DoF Franka Emika Panda manipulator during free motion and while manipulating loads with weights 0kg (free motion), 0.5 kg, 1 kg, 1.5 kg, 2 kg and 2.5 kg. Data is sampled at high frequencies (1kHz). The training trajectories were motions with loads 0kg(free motion), 1kg, 1.5kg, 2.5 kgs, while the testing trajectories contained motions with loads of 0.5kg and 2 kgs. The observations for the forward model consist of the seven joint angles in radians, and the corresponding actions were joint Torques in Nm. We divide the data into sequences of length 600 while training the recurrent models for forward dynamics, with 300 time-steps (corresponding to 300 milli-seconds) used as context set and rest 300 is used for the recurrent time-series modelling.
Training Procedure: For the fully observable case, we trained HiP-RSSM for one-step ahead prediction using an RMSE loss. Similar to the training procedure for partially observable case as in G.1, during training we randomly removed half of the observations(joint angles) from the sequences and tasked the models with imputing those missing observations, only based on the knowledge of available actions/control commands. Other recurrent baselines (RKN, LSTM, GRU) are trained in a similar fashion except that, we dont maintain a context set of interaction histories during training/inference.

heading: G.3 WHEELED MOBILE ROBOT
text: Observation and Data Set: We collected 50 random trajectories from a Pybullet simulator a wheeled mobile robot traversing terrain with sinusoidal slopes. Data is sampled at high frequencies (500Hz). 40 out of the 50 trajectories were used for training and the rest 10 for testing. The observations consists of parameters which completely describe its location and orientation of the robot. The observation of the robot at any time instance t consists of the following features:.
ot = [x, y, z, cos(α), sin(α), cos(β) sin(β), cos(γ), sin(γ)]
where, x, y, z - denote the global position of the Center of Mass of he robot, α, β, γ− Roll, pitch and yaw angles of the robot respectively, in the global frame of reference (Sonker & Dutta, 2020). We divide the data into sequences of length 300 while training the recurrent models for forward dynamics, with 150 time-steps (corresponding to 300 milli-seconds) used as context set and rest 150 is used for the recurrent time-series modelling.
Training Procedure: For the fully observable case, we trained HiP-RSSM for one-step ahead prediction using an RMSE loss. Similar to the training procedure for partially observable case as in G.1, during training we randomly removed half of the observations from the sequences and tasked the models with imputing those missing observations, only based on the knowledge of available actions/control commands. Other recurrent baselines (RKN, LSTM, GRU) are trained in a similar fashion except that, we dont maintain a context set of interaction histories during training/inference.

heading: H HYPERPARAMETERS

heading: H.1 PNEUMATIC MUSCULOSKELTAL ROBOT ARM

heading: Recurrent Models
text: Hyperparameter HiP-RSSM RKN LSTM GRU Learning Rate 8e-4 8e-4 1e-3 1e-3 Latent Observation Dimension 15 15 15 15 Latent State Dimension 30 30 75 75 Latent Task Dimension 30 - - -
Context Encoder (HiP-RSSM): 1 fully connected + linear output (elu + 1)
• Fully Connected 1: 240, ReLU
Observation Encoder (HiP-RSSM,RKN,LSTM,GRU): 1 fully connected + linear output (elu + 1)
• Fully Connected 1: 120, ReLU
Observation Decoder (HiP-RSSM,RKN,LSTM): 1 fully connected + linear output:
• Fully Connected 1: 120, ReLU
Transition Model (HiP-RSSM,RKN): number of basis: 15
• α(zt): No hidden layers - softmax output
Control Model (HiP-RSSM,RKN): 3 fully connected + linear output
• Fully Connected 1: 120, ReLU
• Fully Connected 2: 120, ReLU
• Fully Connected 3: 120, ReLU

heading: Neural Process (NP) Baseline
text: Latent Task Dimension: 30 Learning Rate: 9e-4 Optimizer Used: Adam Optimizer
Context Encoder : 1 fully connected + linear output (elu + 1)
• Fully Connected 1: 240, ReLU
• Linear output: 30
Decoder 3 fully connected + linear output
• Fully Connected 1: 512, ReLU
• Fully Connected 2: 240, ReLU
• Fully Connected 2: 120, ReLU

heading: FFNN Baseline
text: Learning Rate: 1e-3 Optimizer Used: Adam Optimizer
2 fully connected + linear output
• Fully Connected 1: 6000, ReLU
• Fully Connected 2: 3000, ReLU

heading: MAML Baseline
text: Meta Optimizer Learning Rate: 3e-3 Inner Optimizer Learning Rate: 0.4 Number Of Gradient Steps: 1 Optimizer Used: Adam Optimizer (Meta Update) and SGD Optimizer (Inner Update) Model Feed Forward Neural Network with 3 fully connected + linear output
• Fully Connected 1: 512, ReLU
• Fully Connected 2: 240, ReLU
• Fully Connected 2: 120, ReLU

heading: H.2 FRANKA ROBOT ARM WITH VARYING LOADS

heading: Recurrent Models
text: Hyperparameter HiP-RSSM RKN LSTM GRU Learning Rate 1e-3 1e-3 3e-3 3e-3 Latent Observation Dimension 15 15 15 15 Latent State Dimension 30 30 75 75 Latent Task Dimension 30 - - -
Encoder (HiP-RSSM,RKN,LSTM,GRU): 1 fully connected + linear output (elu + 1)
• Fully Connected 1: 30, ReLU
Observation Decoder (HiP-RSSM,RKN,LSTM,GRU): 1 fully connected + linear output:
• Fully Connected 1: 30, ReLU
Transition Model (HiP-RSSM,RKN): number of basis: 32
• α(zt): No hidden layers - softmax output
Control Model (HiP-RSSM, RKN): 1 fully connected + linear output
• Fully Connected 1: 120, ReLU

heading: Neural Process (NP) Baseline
text: Latent Task Dimension: 30 Learning Rate: 1e-3 Optimizer Used: Adam Optimizer
Context Encoder : 1 fully connected + linear output (elu + 1)
• Fully Connected 1: 240, ReLU • Linear output: 30
Decoder 3 fully connected + linear output
• Fully Connected 1: 512, ReLU • Fully Connected 2: 240, ReLU • Fully Connected 2: 120, ReLU

heading: FFNN Baseline
text: Learning Rate: 1e-3 Optimizer Used: Adam Optimizer
2 fully connected + linear output
• Fully Connected 1: 6000, ReLU • Fully Connected 2: 3000, ReLU

heading: MAML Baseline
text: Meta Optimizer Learning Rate: 3e-4 Inner Optimizer Learning Rate: 0.04 Number Of Gradient Steps: 1 Optimizer Used: Adam Optimizer (Meta Update) and SGD Optimizer (Inner Update) Model Feed Forward Neural Network with 3 fully connected + linear output
• Fully Connected 1: 512, ReLU
• Fully Connected 2: 240, ReLU
• Fully Connected 2: 120, ReLU

heading: H.3 WHEELED ROBOT TRAVERSING SLOPES OF DIFFERENT HEIGHT

heading: Recurrent Models
text: Hyperparameter HiP-RSSM RKN LSTM GRU Learning Rate 9e-4 9e-4 1e-2 1e-2 Latent Observation Dimension 30 30 15 15 Latent State Dimension 60 60 75 75 Latent Task Dimension 60 - - -
Encoder (HiP-RSSM,RKN,LSTM,GRU): 1 fully connected + linear output (elu + 1)
• Fully Connected 1: 120, ReLU
Observation Decoder (HiP-RSSM,RKN,LSTM,GRU): 1 fully connected + linear output:
• Fully Connected 1: 240, ReLU
Transition Model (HiP-RSSM,RKN): number of basis: 15
• α(zt): No hidden layers - softmax output
Control Model (HiP-RSSM, RKN): 3 fully connected + linear output
• Fully Connected 1: 120, ReLU

heading: Neural Process (NP) Baseline
text: Latent Task Dimension: 30 Learning Rate: 1e-3 Optimizer Used: Adam Optimizer
Context Encoder : 1 fully connected + linear output (elu + 1)
• Fully Connected 1: 240, ReLU
• Linear output: 30
Decoder 3 fully connected + linear output
• Fully Connected 1: 512, ReLU
• Fully Connected 2: 240, ReLU
• Fully Connected 2: 120, ReLU
• Fully Connected 2: 60, ReLU

heading: FFNN Baseline
text: Learning Rate: 2e-3 Optimizer Used: Adam Optimizer
3 fully connected + linear output
• Fully Connected 1: 512, ReLU
• Fully Connected 2: 240, ReLU
• Fully Connected 2: 120, ReLU • Fully Connected 2: 60, ReLU

heading: MAML Baseline
text: Meta Optimizer Learning Rate: 3e-4 Inner Optimizer Learning Rate: 0.4 Number Of Gradient Steps: 1 Optimizer Used: Adam Optimizer (Meta Update) and SGD Optimizer (Inner Update). Model Feed Forward Neural Network with 3 fully connected + linear output
• Fully Connected 1: 512, ReLU • Fully Connected 2: 240, ReLU • Fully Connected 2: 120, ReLU • Fully Connected 3: 60, ReLU

Review summaries:
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.

This work proposes to extend latent state-space models (SSMs) with a latent variable that changes the dynamics. Update equations akin to Kalman filtering are provided, along with a training loss and method. Experiments on several robotics tasks appear to indicate that the method performs well relative to alternative methods that do not consider latent dynamics differences (up to the points below).