Hidden Parameter Recurrent State Space Models For Changing Dynamics Scenarios
Authors: Vaisakh Shaj
Dieter Büchler
Rohit Sonker
Philipp Becker
Gerhard Neumann
Venue: ICLR 2022 Poster (2022)
Keywords: State Space Models
Changing Dynamics
Recurrent Neural Networks
Multi Task Learning

Abstract: Recurrent State-space models (RSSMs) are highly expressive models for learning patterns in time series data and for system identification. However, these models are often based on the assumption that the dynamics are fixed and unchanging, which is rarely the case in real-world scenarios. Many control applications often exhibit tasks with similar, but not identical dynamics, that can be modelled as having a common latent structure. We introduce the Hidden Parameter Recurrent State Space Models (HiP-RSSMs), a framework that parametrizes a family of related state-space models with a low-dimensional set of latent factors. We present a simple and effective way of performing learning and 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 systems and simulations.

Published as a conference paper at ICLR 2022
HIDDEN PARAMETER RECURRENT STATE SPACE MODELS FOR CHANGING DYNAMICS SCENARIOS
Vaisakh Shaju2 Dieter Büchler", Rohit Sonker", Philipp Becker', and Gerhard Neumann'
Autonomous Learning Robots, KIT, Germany LLCAS, University Of Lincoln, UK
Max Planck Institute for Intelligent Systems, Tübingen, Germany IIdiaa Institute Of Technology, Kanpur
ABSTRACT
Recurrent State-space models (RSSMs) are highly expressive models for learning patterns in time series data and system identification. However, these models as- sume that the dynamics are fixed and unchanging. which rarely the case in real- world scenarios. Many control applications often exhibit tasks with similar but not identical dynamies which can be modeled as a latent variable. We introduce the Hidden Parameter Recurrent State Space Models (HiP- HHP--SSMMs. 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 vari- ational 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.
INTRODUCTION
System identification, i.e., learning models of dynamical systems from observed data ( (Ljung 1998);
Gevers (2005)) is a key ingredient of preeccireeiceive control (Camacho & Alba (2013) and masel-based reinforcement learning (RL). State space models (Hamilton) (1994): Jordan (2004);
Schön et al. (20TT) (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 (Haarooja al. 20166) Karl et al. (016)) and system identi- fication tasks (Becker et al., 2019 Hafner et al.) 2019 Shaj et al.) 2020). Learning the dynamics in an encoded latent space allows 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 fixed, unchanging dynamics, which is quite restrictive 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 bottle may face bottles with varying amounts of liquid. An agent driving car may encounter many different cars, each with unique handling char- acteristics. Humanoids learning to walk may face different terrains with varying slopes or friction coefficients. Any real-worid dynamical system might change over time due 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 robot can sig- nificantly change the dynamics. Similarly, robot may undergo wear and tear over time which can change its dynamics. Thus, assuming global model that is accurate throughout the entire state space or duration of use is 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, 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.
Correspondence Vaisakh Shaj c..shaj@kit.edae

---

Published as a conference paper at ICLR 2022
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 the identification of unobserved causal effects in the dynamics.
HIDDEN PARAMETER RECURRENT 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 (()) 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
2, = f(=-1, + ut
U ~ (O
O, h(z)) + Uss 2. (N(O R).
Here ze, ag and Og are the latent states, control inputs/actions and observations at time t. The vectors ~ (O, Q) and V, NN 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) [004] Wan & Van Der Merwe 2000). In the linear case, inference can be performed efficiently using Kalman Filters.
Our goal is to learn state space model that can model the dynamics of partially observable robotic systems under scenarios where the dynamics changes over time. Often, the dynamics of real systems may differ in significant ways from the system our models were trained However, it is unrea- sonable to train a model across all possible conditions an agent may encounter. Instead, we propose state space model that learns 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 intro- duce new formulation namely Hidden Parameter State Space Models (HiP-SSMss, 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 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 Section 3.3.
2.1 HIDDEN PARAMETER STATE SPACE MODELS (HIP-SSMs)
We denote set of SSMs with transition dynamics fi that are fully de- scribed by hidden parameters and observation model as a Hidden Parameter SSM (HiP-SSM) In this definition we assume the observa- tion model h to be independent of the hidden parameter i as we only focus on cases where the dynamics changes. HiP-SSMs allows to extend SSMs to multi-task settings where dynamics can vary across tasks. Defining the changes in dynamics by latent variable unifies dynamics across tasks as a single global function. In dynamical sys- tems, for example, parameters can be physical quantities like gravity, friction of a surface, or the strength of robot actuator. Their effects influence the dynamics but not directly observed: and hence is not part of the observation space and is treated as latent task parameter vector.
20
Figure 1: The graphical model for particular instance for the HiP-SSM. The transition dy- namics between latent states is a function of the previous la- tent state, and specific latent task parameter I which is in- ferred from a context set served past observations. Ac- tions are omitted for brevity.
Formally, HiP-SSM is defined by tuple Z, A and W are the sets of hidden states z, actions a, and observa- tions w respectively. L is the set of all possible latent task parameters and let (I) be the prior over these parameters. Thus, a HSP-RSSM describes a class of dynamics and particular instance of that class is obtained by fixing the parameter vector E L. The dynamics f for each instance depend on the value of the hidden parameters l.
J. where
2

---

Published as a conference paper at ICLR 2022
Dicent
[e).
20
Ulecmitaee
2
THee
Pxy Ni
Obocriaoe
Deordr
Op-1
(a)
(b)
Figure 2: (a) The Gaussian graphical model corresponding to particular instance I E L. for the HiP-RSSM. The nosterior of the latent task/context variable is inferred via Bayesian aggregation procedure described 2.2.1 based on a set N interaction histories, The prior over the latent states 2; inferred via task conditional Kalman time update described in 2.2.2 and the posterior over the latent states z," is inferred via Kalman measurement update described in 2.2.3] Here, the hollow arrows denote deterministic transformation leading to implicit distributions, using the context and observation encoders. (b) Depicts the schematic Of HiP-RSSM. The output of the task conditional "Time Update' stage, which forms the prior for the next time step (zz41: ) is decoded to get the prediction of the next observation.
Each instance of HiP-SSM is an SSM conditioned on L. We also make the additional assumption that the parameter vector / fixed for the duration of the task (i.e. a local segment of trajectory), 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.) (018a).
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 the definition and formalize HiP-POMDPs. However this is left for future work.
2.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. 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 vari- ables by the forward algorithm (Jordan| 2004 Koller & Friedman, 2009). Since we are dealing with continuous dynamical systems, we assume Gaussian multivariate distribution over all vari- ables (both observed and hidden) for the graph shown in Figure [ This assumption has several advantages. Firstly, makes the inference very similar to well studied Kalman Filtering ap- proaches. Secondly, the Gaussian assumptions and conditional in-dependencies allows 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 2; and 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 I which models the (unobserved) causal factors of variation in dynamics in order 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:
2.2.1 INFERRING THE LATENT TASK VARIABLE (CONTEXT UPDATE)
In this stage, we infer the posterior over the Gaussian latent task variable I based on an ob- served context set Cr. For any HiP-RSSM instance defined over target sequence T
3

---

Published as a conference paper at ICLR 2022
(or, ai, Ot+1...., a(+N), over which we intend to perform state estimation/prediction. maintain a fixed context set G. C in our HiP-SSM formalism can be obtained as per the algo- rithm designer's choice. We choose a Ct consisting a set of tuples C₁ = [o,-n: a a o Jn I==1- Here each set element is a tuple consisting of the current state/observation, current action and next state/observation for the previous N time steps.
Inferring a latent task variable 1 E L based on an observed context set C₁ explored previously by different neural process architectures (Gordon et al. 2018 Garnelo al., 2018). Neural processes are multi-task latent variable models that rely on deep set functions (Zaheer et at. 2017) to output a latent representation out of varying number of context points in permutation invariant manner. Similar to Volpp et al. (2020) we formulate context data aggregation as a Bayesian inference problem, where the information contained Cj directly aggregated into the statistical description of / based on a factorized Gaussian observation model of the form p(r n (2), where
has been
rn enc, ( di-w o(nn1),
(o%))² = enc,
(o/_.. cnc_(o)__,
ai-n; d--nt1).
Here n is the index of an element from context set C. Given a prior po(I) = po diag(of))
we can compute the posterior p([[Cr) using Bayes rule. The Gaussian assumption allows us to get a closed form solution for the posterior estimate of the latent task variable, p(([G) based on Gaussian conditioning. The factorization assumption further simplifies this update rule by avoiding computationally expensive matrix inversions into a simpler update rule as
N
of = (ot)
+
Σ
((o))))
((o1)
n=1
N
Σ
(r/ mo) (0%)
1 = Mo di
where e, C and Ø denote element-wise wise inversion, product, and division, respectively. Intuitively the mean of the latent task variable 14 is weighted sum of the individual latent observations r while the variance of the latent task variable of gives the uncertainty this task representation.
2.2.2 INFERRING PRIOR LATENT STATES OVER THE NEXT TIME STEP (TIME UPDATE)
The goal of this step to compute the prior marginal
p(z] a1., Ci) /ʃ p(zzlzz-1, at (na.
at,
We assume a dynamical model of the following form: p(z, -1 a,,
M (A,-12,-1- Ba, + Cl, Ema))) and denote the posterior from the previous time-step by p(z;-1 1if11/11--11 11:-1, a1.-1.1.... Ct) N (2t,, 1t Following Shaj et al. (200)) assume the action a, is known and not subject to noise.
At any time the posterior over the belief state 2 01:t-1, posterior over the latent task variable I|C; and the action a are independent of each other since they form a coommon-effere / v-structure (Koller al., (207)) with the unobserved variable z,. With these independencies and Gaussian assumptions, according to Gaussian Identity [ can be shown that calculating the integral
in equation has a closed form solution as follows, p(ze W1:1-1, a1.4, G) N Eg ). where
2,- = A1-12 + Ba: + Cl. E, A.-EE
(2)
CE +E Emruns
(3)
Gaussian Identity 1. If u ~ N(pn + b,) and v ~ N( En) are normally distributed independent random variables and the conditional distribution p (y v) Ni(u+b++ bBt E) then marginal p(y) = p(y u, p)p(()p(() ddudu = N(Auu (1AAuu+He +b+ Bu AE, + +BE, +
Σ)
We use a locally linear transition model A,- as in Becker et al. 2019) and a non-linear control model as in (Shaj et al. (2020). The local linearization around the posterior mean, be interpreted
4

---

Published as a conference paper at ICLR 2022
equivalent to an EKF. For the latent task transformation or non-linead transformation. More details regarding the latent task transformation model can be found in the appendix [S
either use linear, locally-
Our ablations (Figure Bcb show that 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 la and dr directly into the latent space of the state space model via additive interactions. The corresponding time update equations are given below:
2 = A + b(a+) + f(w)),
Σ, = A A + f(a)) + Irans
2.2.3 INFERRING POSTERIOR LATENT STATES / OBSERVATION UPDATE
The goal of this step is to compute the posterior belief p(z:)01: 01s, Ct). We first map the observations at each time step O, to a latent space using an observation en- coder (Haarnoja et al. 2016) Becker et al.) (019) which emits latent features w; along with an uncertainty in those features via a variance vector Tohe. We then compute the pos- p(za|W11t11, Ci) obtained from the time
terior belief p(zz W1.1, 01.,: Ct), based update, the latent observation (W, ofth)) and the observation model H. This is exactly the Kalman update step, which has closed form solution as shown below for time instant
Kalman Gain:
== H (HE, H' +I ozbe)
Posterior Mean:
_ = + Qu (e --, Σ = (I -))
Posterior Covariance:
where I denotes the identity matrix. This update is added as layer in the computation graph (Haarnoja et al. 2016); Becker et al. (019)). However, the Kalman update involves computation- ally 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 factorization provides a simple way of reducing the observation update equation to set of scalar operations, bringing down the computational complexity from O to OL). More mathematical details on the simplified update equation can be found in appendix F77
Discussion From a computational perspective, this a Gaussian conditioning layer, similar to sec- tion 2.2.1 Both outputs a posterior distribution over latent variables 2, given prior p(z) and an observation model p( (z), using Bayes rule: p(z|w) = p( p(w) x z)p(z) mx)pp)pww,))))))))))))))>))))( This has a closed form solution because of Gaussian assumptions, which is coded as a layer in the neural network. The observation model is assumed to have the following structure, P(w 2) = (HH2, Eohs).
2.3 TRAINING HIP- RSSM FOR EPISODIC ADAPTATION To CHANGING DYNAMICS
The latent task variable models a distribution over functions Garnelo et al. 2018 rather than a single function. In case these functions are latent dynamics functions. In order to train such a model we use a training procedure that reflects this objective, where we form datasets consisting of timeseries, each with different latent transition dynamics. Thus, collect set of trajectories over which the dynamics changes over time. 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 be 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 hierarchical latent task variable 1. The formalism is based on the assumption that over these local temporal segments, the dynamics unchanging. This local consistency in dynamics holds for most real world scenarios (Nagabandi et al. (018b Yet. our formalism can model the global changes in dynamics at test time, since we obtain a different instance of the HiP-RSSM for each temporal segment based the observed context set. We also provide a detailed description of the multi-task dataset creation process in the appendix and a pictorial illustration in appendix D
Batch Training. Let T E R BxNxD be the batch of local temporal segments with different dynam- ies which we intend to model with the HiP-RSSM formalism. Given target batch T, HiP-RSSM
5

---

Published as a conference paper at ICLR 2022
can be trained in supervised fashion similar to popular recurrent architectures like LSTMs or GRUs. However for each local temporal sequence E T. over which the dynamics is being mod- elled, we also imput set of the N previous interactions, which forms the context set C' E R Bx X for inferring the latent task as explained in Section Processing the context set C' adds minimal additional computational/memory constraints as we permutation invariant set encoder. The set encoder allows parallelization in processing the context set as opposed to the recurrent processing of target set.
The learnable parameters in the computation graph includes the locally linear transition model Ar. the non-linear control factor b, the linear/non-lineat latent task transformation model C, the transi- tion noise mmans along with the observation encoder, context encoder and the output decoder.
Loss Function. We optimize the root mean square error (RMSE) between the decoder output and the ground truth states. As in Shaj et al. (2020) we use the differences to next state as our ground truth states, as this results better performance for dynamics learning especially higher frequen- cies. In principle, we could train on the Gaussian log-likelihood instead of the RMSE and hence model the variances. Training on RMSE yields slightly better predictions and allows for a fair com- parison with deterministic baselines which use the same loss like feed-forward neural networks, LSTMs and meta-learung algorithms such as MAML. Thus, we report results with the RMSE loss. A probabilistic formulation for this loss function is provided in appendix B
Gradients are computed using (truncated) backpropagation through time (BPTT) (Werbos (990))
and clipped. We optimize the objective using the Adam (Kingma & Ba 2014 stochastic gradient descent optimizer with default parameters.
RELATED WORK
Deep State Space Models. Classical State-space models (SSMs) are popular due their tractable inference and interpretable predictions. However, inference and learning in SSMs with high dimen- sional and large datasets often not scalable. Recently, there have been several works deep SSMs that offer tractable inference and scalability to high dimensional and large datasets. Jaarnooj al. (2016); [Bekkr al. (2019); Shaj et al.] (2020) use neural network architectures based on closed- 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); 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. [inderman et al. (2017) Becker-Ehmek et al., (2019)
Dong et al. (2020) tries to address the problem of changing/multi-modde dynamics by incorporat- ing 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
et al. (2019): Dong et al.) (2020) uses variational approximations and relies additional regularization/anneallin crete state transitions. On the other hand |raccaro et al.] (2017) uses "soft" switches, which can be interpreted SLDS which interpolate linear regimes continuously rather than using truly discrete states. We take a rather different, simpler formalism for modelling changing dynamics by view- ing multi-task learning problem with continuous hierarchical hidden parameter that model the distribution over these tasks. Further our experiments in appendix C.I show that our model significantly outperfroms the soft switching baseline (Fraccaro et al., 2017).
encourage dis-
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 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 [ for planning and control in the latent space (Hafner et al. 2019; Sekar et al. 2020). However this is left as a future work.

---

Published as a conference paper at ICLR 2022
Inpuration
Iegetation
Ingerstion
50% Enguistiin
No Iegetaiion
SIFE Ingutiion
HJP RSSM
0.50 = 0.048
9.47 = 0.001
HIP
RSS-R
2,533 = 0.024
2.545 2-0.024
HIP RSSM
6.16+ 0:201
7:37 0.017
1.5MM
0.001
3630 0.007
LSTM
0.005
2.766 3,000
1STM
B:3600025
GRR
3.257 0.013
3.21 + 0.047
GRU
0:000
2.422 0.108
CRU
9.215-20000 2.00/3
50.11 0.32
FFNN
5.330 0.007
FINN
3.315 2 0.BL5
JENN
532 = 01022
2.789
00.04
4.57 -8013
(a) Pneumatic RMSE (10-))
(b) Franka RMSE (10-²)
(c) Mobile Robot RMSE (10-5)
Table 1: Prediction Error in RMSE for (a) pncumatic muscular arm 4.1) and (b) Franka Arm manipulating
varying loads 4.2 for both fully observable and partially observable scenarios.
Meta Learning For Changing Dynamics. There exists a family of approaches that attempt online model adaptation to changing dynamics scenarios via meta-learning (Nagabandi et al. 2018 b). They perform online adaptation on the parameters of the dynamics models through gradient de- scent (Finn et al. 2017) from interaction histories. Our method fundamentally differs from these approaches in the sense that we do not perform gradient descent step at every time step during test time, which is computationally impractical in modern robotics, where we often deal with high fre- quency data. 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, hawuumaatic/yddraauc/ccaeeeeeeed
actuators etc. SSmmundsson et al. (2018): : Accherroldd & Stueckler (202T) on the other hand learn context conditioned hierarchical dynamics model for control formalism similar to HiP-MDPs. The former meta-learn context conditioned gaussian process while the later use a context conditioned determisitic GRU. Our method on the other hand focuses on a principled probabilistic formulation and inference scheme for context conditioned re- current state space models, which is critical for learning under partial observabilly/igg dimensional observations and with noisy and missing data.
4 EXPERIMENTS
This section evaluates our approach on a diverse set of dynamical systems from the roboties 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, show that HiP-RSSM outperforms these models even un- der situations with partial observability/ missing values. We also baseline our HiP-RSSM with con- temporary multi-task models and improve performance, particularly in modelling non-Markovian dynamics and under partial observability. Finally, the visualizations of the Gaussian latent task vari- ables 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 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 these have shown excellent performance for dynamics learning (Shaj et al. 2020) and relies on exact inference 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 NNgabaand et al. 2018a). Both models receive the same context information as in HiP- HiP-RSSM.
In case of recurrent models we replace the HiP-RSSM cell with properly tuned LSTM, GRU and RKN Cell respectively, while fixing the encoder and decoder architectures. For the NP baseline use the same context encoder and aggregation mechanism comparison. We create partially observable settings by imputing 50% of the observations during inference. More details regarding the baselines and hyperparameters be found in the appendix.
in HiP-RSSM to ensure a fair
7

---

Published as a conference paper at ICLR 2022
6.55
6.0
Slat Lirie
0.4
0.55
400
Phouunte
Nobk:
Pianfa
(a)
(b)
(c)
Figure 3: (a) and (h) shows the iSNE (Van der Maaten & Hinton 2008) plots of the latent task embeddings
produced from randomly sampled instances HiP-RSSM for two different robots. (a) The wheeled robot
discussed section 4.3 traversing temains of varying slopes. The color coding indicates the average gradients of the terrain for each of these instances. These can have either positive negative values. (h) The Franka manipulator with loads of varying weights attached the end-effector. The color coding indicated weights ranging from to kgs. (c) An ablation on the performance different task transformation models discussed in section 2.2.2
4.1 SOFT ROBOT PLAYING TABLE TENNIS
We first evaluate our model on learning the dynamics of pneumatically actuated muscular robot. This four Degree of Freedom (DoF) robotic arm is actuated by Pneumatic Artificial Muscles (PAMs) (Büchler et al., 2016). The data consists of trajectories of hitting movements with vary- ing speeds while playing table tennis (Büchler et al., 2020). This robot's fast motions with high accelerations are complicated to model due to hysteresis and hence require recurrent models (Shaj
al. 2020
We show the prediction accuracy in RMSE in Table 5a We observe that the HiP-RSSM can out- perform the previous state of the predictions obtained by recurrent models. Based on the domain knowledge, we hypothesise that the latent context variable captures multiple unobserved causal fac- tors of variation that affect the dynamics in the latent space, which not modelled in contemporary recurrent models. These causal factors could be, in particular, the temperature changes or the friction due a different path that the Bowden trains take within the robot. Disentangling and interpreting these causal factors can be exciting and improve generalisation, 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- non-markovian dynamics resulting from the high accelerations this pneumatically actuated robot.
4.2 ROBOT MANIPULATION WITH VARYING LOADS
We collected the data from 7 DoF Franka Emika Panda manipulator carrying six different loads at its end-effector. It involved a mix movements of different velocities from slow to swift motions covering the entire workspace of the robot. We chose trajectories with four different loads training set and evaluated the performance on two unseen weights, which results in a scenario where the dynamics change over time. Here, the causal factor of variation in dynamics is the weights attached to the end-effector and assumed to be unobserved.
We show the prediction errors in RMSE in Table Ic HiP-RSSMs outperforms all recurrent state- space 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 per- form equally well under full observability for this task because of the near Markovian dynamics of Franka Manipulator, which often does need recurrent models. However, HiP-RSSMs have an additional advantage in that these are naturally suited for partially observable scenarios and can pre- dict ahead compact latent space, a critical component for recent success in model-based control and planning (Hafner et al. 2019).
4.3 ROBOT LOCOMOTION IN TERRAINS OF DIFFERENT SLOPES
Wheeled mobile robots are the most common types of robots being used exploration unknown terrains where they may face uneven and challenging terrain. We set up environment using
8

---

Published as a conference paper at ICLR 2022
Eine
(b)
(c)
Figure 4: (a) Simulated environment of Wheeled Robot traversing terrains of different slopes. (b) and (c) shows how the one dimensional UMAP (MeInnes et al 2018) embedding of the inferred latent task variable (top blue plots) changes at test time when an agent undergoes changes its dynamics mobile robot respectively. An indicator of the Ground Truth (GT) Task (bottom red plots) variables are also given. In case of the Franka Robot, the groundtruth (GT) tasks denotes the switching of dynamics between 0 kg (free motion) and 2.5 kg loads. In case of the mobile robot the groundtruth (GT) tasks denoted slopes of the terrain.
the franka robot and
Pybullet simulator (Coumans & Bai 2016) where four-wheeled mobile robot traverses an uneven terrain of varying steepness generated by sinusoidal functions (Sonker & Duttal 2020) 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 changing dynamics scenario, which further increases the complexity of the task.
We show the prediction errors in RMSE in Table 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 effeets in the dynamics in an unsupervised fashion.
4.4 VISUALIZING CHANGING HIDDEN PARAMETERS AT TEST TIME OVER TRAJECTORIES WITH VARYING DYNAMICS
We finally perform inference using the trained HiP-RSSM multi-task/ changing dynamics sce- nario where the dynamics continuously changes over time. We use the inference procedure described in appendix D based on fluid definition for "task" as the local dynamics in temporal segment. We plot the global variation in the latent task variable captured by each instance of the HiP-RSSM over these local temporal segments using the dimensionality reduction technique UMAP (McInnes) et al. 201 2018). As seen in figures 4b and 4c the latent task variable captures these causal factors of variations in an interpretable manner.
CONCLUSION
We proposed HiP-RSSM, 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 universal dynamics model with latent parame- ter 1 via exact inference and backpropagation through time. The resulting recurrent neural network, namely HiP-RSSM. learns to cluster SSM instances with similar dynamics together in an unsuper- vised fashion. Our experimental results on various robotic benchmarks show that HiP-RSSMs sig- nificantly 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, tion and task representations can benefit multiple future applications including planning/control in the latent space and causal factor identification.

---

Published as a conference paper at ICLR 2022
6 AKNOWLEDGEMENTS
We thank Michael Volpp, Onur Celik, Marc Hanheide and the anonymous reviewers for valuable remarks and discussions which greatly improved this paper. The authors acknowledge support by the state of Waalen-Württebberr through bwHPC and the Lichtenberg high performance computer of the TU Darmstadt. This work was supported by the EPSRC UK (project NCNR, National Centre for Nuclear Robotics.EP/RO25722
REFERENCES
Jan Achterhold and Joerg Stueckler. Explore the context: Optimal data collection for context- conditional dynamics models. In International Conference on Artificial Intelligence and Statisties, pp. 3529-3537, PMLR, 2021.
Philipp Becker, Harit Pandya, Gregor Gebhardt, Cheng Zhao, C James Taylor, and Gerhard Neu- mann. Recurrent kalman networks: Factorized inference high-dimensional deep feature spaces. In International Conference on Machine Learning, pp. 544-552. PMLR, 2019.
Philip Becker-Ehmck, Jan Peters, and Patrick Van Der Smagt. Switching linear dynamics for vari- ational bayes filtering. In International Conference on Machine Learning, pp. 553-562. PMLR, 2019.
Christopher M Bishop. Pattern recognition. Machine learning, 128(9), 2006.
Dieter Büchler, Heiko Ott, and Jan Peters. A lightweight robotic arm with pneumatic muscles for robot learning. In 2016 IEEE International Conference on Robotics and Automation (ICRA), pp- 4086-4092. IEEE, 2016.
Dieter Büchher, Simon Guist, Roberto Calandra, Vincent Berenz, Bernhard Schölkopf, and Jan Peters. Learning to play table tennis from scratch using muscular robots. arXiv preprint wwXiv:2006.05935. 2020.
Eduardo Camacho and Carlos Bordons Alba. Model predictive control. Springer science & business media, 2013.
Kyunghyun Cho, Bart van Merriënboer, Dzmitry Bahdanau, and Yoshua Bengio. On the properties of neural machine translation: Encoder-decoder approaches. In Proceedings of SSST-8, Eighth Workshop Syntax, Semantics and Structure in Statistical Translation, pp. 103-111, 111.1.......1.111 Association for Computational Linguistics, October 2014.
Erwin Coumans and Yunfei Bai. Pybullet, a python module for physics simulation for games, robotics and machine learning. 2016.
Zhe Dong, Bryan Seybold, Kevin Murphy, and Hung Bui. Collapsed amortized variational inference for switching nonlinear dynamical systems. International Conference on Machine Learning, pp. 2638-2647. PMLR, 2020.
Finale Doshi Velez and George Konidaris. Hidden parameter markov decision processes: A semi- parametric regression approach for discovering latent task parametrizations. In IJCAI: proceed- ings of the conference, volume 2016, pp. 1432. NIH Public Access, 2016.
Chelsea Finn, Ian Goodfellow, and Sergey Levine. Unsupervised learning for physical interaction through video prediction. Advances in neural information processing systems, 29:64-72, 2016.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. International Conference on Machine Learning, pp. 1126-1135. PMLR, 2017.
Marco Fraccaro, Simon Kamronn, UIrich Paquet, and Ole Winther. A disentangled recognition and nonlinear dynamics model for unsupervised learning. arXiv preprint arXiv: 170.00541, 2017.
Marta Gareel,, Jonathan Schwarz, Dan Rosenbaum, Fabio Viola, Danilo J Rezende, SM Eslami, and Yee Whye Teh. Neural processes. arXiv preprint arXiv:/ 2018
10

---

Published as a conference paper at ICLR 2022
Michel Gevers. Identification for control: From the early achievements to the revival of experiment design. European journal of control, 11(4-5):335-33 -52233355552,-- 2005.
Jonathan Gordon, John Bronskill, Matthias Bauer, Sebastian Nowozin, and Richard E Turner. Meta- learning probabilistic inference for prediction. arXiv preprint arXiv: 00991,, 2018.
Tuomas Haarnoja, Anurag Ajay, Sergey Levine, and Pieter Abbeel. Backprop kf: Learning discrim- inative deterministic state estimators. In Advances in neural information processing systems, pp- 4376-4384, 2016.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International Conference on Machine Learning, pp. 2555-2565. PMLR, 2019.
James D Hamilton. State-space models. Handbook of econometrics, 4:3039-3000 -0890, 1994.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Comput., 9(8): 1735-1780, November 1997. ISSN 0899-7667.
Michael Jordan. Graphical models. Statistical science, 19(1):140-1555 2004.
Maximilian Karl, Maximilian Soelch, Justin Bayer, and Patrick Van der Smagt. Deep varia- tional bayes filters: Unsupervised learning of state space models from raw data. arXiv preprint arXiv:1605.064432 2016.
Taylor Killian, George Konidaris, and Finale Doshi Transfer learning across patient varia- tions with hidden parameter markov decision processes. arXiv preprint arXiv:1612.00475 2016.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 1412.64412.6980 6980, 2014.
Daphne Koller and Nir Friedman. Probabilistic graphical models: principles and techniques. MIT press, 2009.
Daphne Koller, Nir Friedman, Lise Getoor, and Ben Taskar. Graphical models in a nutshell. Intro- duction statistical relational learning, 43, 2007.
Rahul Krishnan, Uri Shalit, and David Sontag. Structured inference networks for nonlincar state space models. In Proceedings of the AAAI Conference Artificial Intelligence, volume 31, 2017.
Scott Linderman, Matthew Johnson, Andrew Miller, Ryan Adams, David Blei, and Liam Paninski. Bayesian learning and inference in recurrent switching lincar dynamical systems. Artificial Intelligence and Statistics, pp. 914 922. PMLR, 2017.
Lennart Ljung. System identification. In Signal analysis and prediction, pp. 163-173. Springer, 1998.
Leland McInnes, John Healy, and James Melville. Umap:
Uniform manifold approximation and projection for dimension reduction. arXiv preprint arXiv: 1802.0342, 2018.
Anusha Nagabandi, Ignasi Clavera, Simin Liu, Ronald S Fearing, Pieter Abbeel, Sergey Levine, and Chelsea Finn. Learning to adapt in dynamic, real-world environments through meta- reinforcement learning. arXiv preprint / 1347, 2018a.
Anusha Nagabandi, Chelsea Finn, and Sergey Levine. Deep online learning via meta-learning: Continual adaptation for model-based arXiv preprint arXiv.1812.07677 2018b.
Maria Isabel Ribeiro. Kalman and extended kalman filters: Concept, derivation and properties. Institute for Systems and Robotics, 43:46, 2004.
Steindór Sæmundsson,, Katja Hofmann, and Mare Peter Deisenroth. Meta reinforcement learning with latent variable gaussian processes. arXiv preprint arXiv: 1803.0755), 2018.

---

Published as a conference paper at ICLR 2022
Thomas B Schön, Adrian Wills, and Brett Ninness. System identification of nonlinear state-space models. Automatica, 47(1):39 49, 2011.
Ramanan Sekar, Oleh Rybkin, Kostas Daniilidis, Pieter Abbeel, Danijar Hafner, and Deepak Pathak. Planning explore via self-supervised world models. In International Conference Machine Learning, pp. 8583-8592. PMLR, 2020.
Vaisakh Shaj, Philipp Becker, Dieter Buchler, Harit Pandya, Niels van Duijkeren, C James Tay- lor, Mare Hanheide, and Gerhard Neumann. Action-conditional recurrent kalman networks for forward and inverse dynamics learning, arXiv preprint arXiv:2010. 10201, 2020.
Rohit Sonker and Ashish Dutta. Adding terrain height to improve model learning for path tracking on uneven terrain by a four wheel robot. IEEE Robotics and Automation Letters, 6(1):239-246, 2020.
Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research, 9(11), 2008.
Michael Volpp,
Fabian Flürenbrock, Lukas Grossberger, Christian Daniel, and Gerhard Neumann. Bayesian context aggregation for neural processes. In International Conference on Learning Representations, 2020.
Eric A Wan and Rudolph Van Der Merwe. The unscented kalman filter for nonlinear estimation. In Proceedings of the IEEE 2000 Adaptive Systems for Signal Processing, Communications, and Control Symposium (Cat. No. 00EX373), pp. 153-158. leee, 2000.
Paul Werbos. Backpropagation through time: what does and how do it. Proceedings of the IEEE, 18(10):1500-1560 1990.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Ruslan Salakhutdinov, and Alexander Smola. Deep sets, arXiv preprint arXiv: 2017.
12

---

Published as a conference paper at ICLR 2022
A PROOF FOR GAUSSIAN IDENTITY 1
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 ~ ( + b. E.,) and v (u En) are normally distributed inde- pendent random variables and conditional distribution plylu, v) NiAu Bu, E), the joint distribution has an expression as follows:
AΣ, B + ))
(
u .V
-N
)
Σn
E,A
Hy Apu +b+ bBu
0 Σ,
B
AE
BE,
Proof for Identity 2
Let displacement of variable denoted by Au u (u).
Since and independent, covariances (AnAvT)
We can write = Auutb c where ~ /O Σ) and b is constant. Then we have covariance (AuAyT) > (Au(AAr BAv + Ac) = (Audu A + (AuAv") B (AuAcT) Since (AuAv ( (Aude )
0 we therefore have (AuAy E.A".
The derivations for covariances (AvAyT) follows similarly and the corresponding covariance has the expression, (AvAy E,B!
Similarly, (AyAyT) = ((AAu+ Bav (c)(AAu Ad) B&v + Ac) A (AnAnT) AT + B (Avâv) B (DedeT) = A + BE,B + E. The result follows.
Gaussian Identity (Gaussian Marginalization). If
Puu
Enn Σ
Enp
Euy
v N
Eve
Evy
y
Pzz
Σ
D
Σ Egg
ag
then marginal over is given as p(y)
Ju,v p(y ()p(u))p())duuls - (Hy: Eny)
Proof For Identity 3
We refer to [Bishop (2006) for the derivation, which requires calculation of the Schur complement well 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.
B PROBABILISTIC FORMULATION OF HIP-RSSM TRAINING OBJECTIVE
The training objective for the HiP-RSSM involves maximizing the posterior predictive log- likelihood given below:
T
T
N
Logp(Cs Ci, (o.
Σ log t=1
Ci)d2/11
(4)
1=1 1=1
1=1
Here, for any time â are the ground truth observations for the next time step and the latent state prior p(2f,11o7) Ci) as discussed in the Section 2.2.2 have closed form solutions. We approxi- mate the objective in Eq. using deterministic variational approximation, similar to Becker al. 2019).
We employ Gaussian approximation of the posterior predictive log likelihood of the form p(o), Cr z N(Ho, (aaug(());,,)"))) Here Mot,, = (ec, and a (ecy;(E):))) where (ec,,() and decx() denote the parts of the output decoder that are responsi- ble for decoding the latent means and latent variances respectively. This decoder be interpreted
13

---

Published as a conference paper at ICLR 2022
100
HIP-RSSM
HIP-RSSM
RKNN
RKN
Soft-Switching
RMSE(1e-5)
LSTM
Soft Switching
rmse.
Step Number Of Steps
Step
epochs
Figure 5: Comparison of different algorithms for Wheeled Mobile Robot in terms of (a) moving average of decoder error in normalized RMSE for the test set, plotted against training epochs (b) multi-step ahead prediction error in RMSE.
"moment matching network", computing the moments of of given the moments of 2 (Volpp et al. 2020: Becker et al. 2019). This allows to evaluate an approximation to the objective in deterministic manner.
Alternatively, we could ignore the variances and only focus on the mean predictions by training our model with RMSE loss. Training on RMSE yields slightly better predictions and allows for fair comparison with deterministic baselines (Feed Forward NN, LSTM, GRU, MAML etc) and hence we report results with the RMSE loss.
C ADDITIONAL EXPERIMENTS
C.I COMPARISON TO SOFT SWITCHING BASELINE
We also implement soft-switching baseline similar to Fraccaro et al. (2017), where soft mix- ture of dynamics is implemented the latent transition dynamics (Kalman time update). Similar to Fraccaro al. (007)), we now globally learn K constant transition matrices A (k) and control matri- ((k) An interpolation is done between these using "dynamics parameter network" (Fraccaro) et al. 2017) ag = as (Wort1)) The dynamics parameter network is implemented with a recurrent neural network with LSTM cells that takes each time step the mean of the encoded observation Wi as input and recurses de LSTM (w,-1, d and Q₂ softmax (dz). The output of the
dynamics parameter network is weights that sum to one, ΣK at (wo:--1 These weights choose and interpolate between K different operating modes:
A:
ai (k) A
B. = Σ (k) a
(wor-1) B
Authors interpret the weighted time-invariant matrices are combined using the time varying weights ax, practice, each the K sets {A (k) B(²)) } models different/changing dynamics, that will dominate when the corresponding (k) is high.
soft mixture of K different Linear Gaussian SSMs whose
ac
Figure compares HiP-RSSI RSSM with different recurrent architectures including the soft-switching baseline. As seen in figure 5 HiP-RSSM clearly outperforms the soit-switching baseline both in terms of convergence speed and also mult-step ahead predictions. More details regarding the multi- step ahead training procedure can be found in appendix C.33
C.2 ABLATION ON CONTEXT ENCODER
In table 2 report the details of evaluating bayesian aggregation based context set coder (discussed in [22.II] against the causal recurrent encoder that takes into account the temporal structure of the context data. We used probabilistic recurrent encoder (Becker et al.
14

---

Published as a conference paper at ICLR 2022
2019)
whose mean and variance from the last time step is used to infer the posterior latent task distribution p(1 ( = (i,, diag(of))) The dimensions of latent task parameters obtained from both the set and recurrent encoders are kept the same (60).
The reported experiments conducted on data from wheeled mobile robot dis- cussed in section 4.3 As reported in Ta- ble 2 the permutation invariant set en- coder outperforms the recurrent encoder by good margin in terms of prediction accuracy both fully and partially ob- servable scenarios. Additionally the set encoder far more efficient in terms of computational time required for training as seen from the time taken per epoch for each these cases.
Table 2: Comparison between the permutation invariant en- coder and recurrent encoder. The performance is measured in terms of prediction RMSE (10-5) and mean of the training time per epoch (in seconds) over runs.
No Impuration RMSE
50% Imputation RMSE
Training Time Per Epoch
Fnceder Recurent Envoder
0.212
0.327
5.10 0.041
30.12 0.112
14.13
C.3 MULTI-STEP AHEAD PREDICTIONS
In figure 5b we compare the results of multi-step ahead predictions for upto 50 steps of HiP-RSSM with recurrent baselines like RKN(Becker et al., 2019) and LSTM. Inorder 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 approaches like Finn
(2016).
As seen in figure 5b HiP-R HSP-RSSM clearly outperforms contemporary recurrent models for multi-step ahead prediction tasks since it takes into account additional causal factors of variation (slopes the terrain for this robot) in the latent dynamics in an unsupervised manner.
D HiP-RSSM DURING TEST TIME INFERENCE
We perform inference using HiP-RSSM at test time a trajectory with varying dynamics using algorithm[] [ A pictorial representation of the same given in the figure o We use this inference scheme to visualize how the latent variable I, that describe different instances of 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: time series r of length >> N
Divide the time series t into non- non-overlapee con-overlapping windows Tj of length N. Let
T (Ti, Ty.T...... .} be the ordered list of temporal segments, sorted in the ascending order of time of occurrence.
foreach each time window Ti do
1. maintain a context set C/ consisting of previous interactions;
2. infer posterior latent task variable p(1()) using context update stage in section 2.2.1
3. using the posterior over latent task variable ([c, and observations in sequence T to perform sequential Bayesian inference over the state space model using Kalman observation update (22.3) and task conditional Kalman time update; (22.22
end
15

---

Published as a conference paper at ICLR 2022
Time Series With Changing Dynamics
Context Set
Locally consistent dynamics modalled as a HIP-RSSM instance
Target Time Series
Figure 6: HiP-RSSM Inference Under Changing Dynamics Scenarios
E MULTI-TASK DATASET CREATION
Algorithm 2: Multi Task Dataset Creation For Training HiP-RSSM
Required:A set 8 of trajectories of changing dynamics
D á
foreach each trajectory
do 1. divide the trajectory into non-overlapping windows Ti of length N. Let T' = [ T2, .3.... be the list of all temporal segments/imme-srinss
foreach each time window Ti E T do
(a) maintain context set Ct consisting of N previous interactions;
(b) update D < DU (Cr, Ti}
end
end
Output: Output consisting of batch of context and target sets.
F IMPLEMENTATION DETAILS
F.1 CONTEXT SET ENCODER AND LATENT TASK REPRESENTATION
The HiP-RSSM maps a set of previous interaction histories, {o) a. 1,4 to a set of latent features and an estimate of uncertainty in those features, {ri, { 2 encoder. We use feed-forward neural network as the encoder in all of our experiements since we deal with high dimensional vectors our observations. However, depending upon the nature of observations, we could use different encoder architectures.
IN
Jn=1, using a context
The set of latent features and uncertainties, [r/, " are further aggregated in a probabilis- tically principled manner using bayesian aggregation operator discussed in 2.2.1 to get gaussian latent task variable, with mean (μt;) and diagonal covariance (ar)). Intuitively the context encoder learns to weight the contribution from each observation the context set based on bayesian prinici- ples and emits probabilistic representation of the latent task.
F.2 OBSERVATION ENCODER AND LATENT STATE REPRESENTATION
HiP-RSSM transforms the observations each time step Oe latent space using an encoder network which emits latent features and an estimate of the uncertainty in those features via variance vector ohhs
The probabilistic recurrent module uses latent state vector 2, and corresponding covariance Σ, whose transitions are governed by the update rules in the HiP-RSSM cell. The latent state vector ze has been designed to contain two conceptual parts, vector Pt for holding information that can directly be extracted from the observations and a vector m, to store information inferred over time, e.g., velocities. The former is referred to 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 observations the former may correspond to positions while the latter corresponds to velocities. The corresponding posterior and prior covariance matrices Σ and E, have a chosen factorized representation yield simple Kalman updates, i.e.,
16

---

Published as a conference paper at ICLR 2022
Σ
Σ Σ
Σ
where each of 2r, Ei, Ej E Ru is diagonal matrix. The vectors of o and o denote the vectors containing the diagonal values of those matrices. This structure with 2n ensures that the correlation between the memory and the observation parts are not neglected as opposed to the case of designing Σ as diagonal covariance matrix. This representation was exploited to avoid the expensive and numerically problematic matrix inversions involved in the KF equations shown below.
F.3 LOCALLY LINEAR TRANSITION MODEL
The state transitions in the time update stage (section 2.2.2 of the HiP-RSSM Cell is governed by locally linear transition model in Becker al. (019). To obtain locally linear transition model, the HiP-RSSM gloablly learns K constant transition matrices (k) and interpolate between them using state dependent coefficients a (2 i.e., A. = - a (x) ) A where 2 is the mean of the posterior distribution at time step So we locally linearize A, 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 used learn (R).
F.4 CONTROL MODEL
To achieve action conditioning within the recurrent cell, we include a control model b(a;) in addition the locally linear transition model A.. b(a,) = f(a.), where f(.) can be any non-linear function approximator. We use multi-layer neural network regressor with ReLU activations (Shaj et al.) 2020).
F.5 LATENT TASK TRANSFORMATION MODEL
To achieve latent task conditioning within the recurrent cell, we include task transformation model (c), in addition to the locally linear transition model A, and control model in the time update stage (section 2.2.2 Though section 2.2.2] we used the notation for linear task transformation matrix, C, to motivate the additive interaction of latent task variables, ly and of in the latent space, the task transformation function can be designed in several ways, i.e.:
ii Linear: C, where C is linear transformation matrix. The corresponding time update equations are given below:
z = A. 1 + b(a:) + CHu,
Σ, = A, Σ AT E C(I C.) + Σ trans
(ii) Locally-Linear: = Ct, where C, = Ek ((h) (2)) (z)) is a linear combination
linear control models ((k), small neural network with softmax output used to learn g(k) The corresponding time update equations are given below:
2 A,-12 + b(a,) + C.u,
Σ, - + Ce(I - + Exrass
(iii) Non-Linear: f. where f() can be any non-linear function approximator. We use
multi-layer neural network regressor with ReLU activations, which transforms the latent task moments 14 and or directly into the latent space of the state space model via additive interactions. The corresponding time update equations are given below:
z,- = A, x
b(a,) ) + f
Σ, A,
A f( Emns.
17

---

Published as a conference paper at ICLR 2022
In ablation study (Figure 3e for the linear and locally linear task transformation models, assume that the dimension of the latent context variable and the latent state space 2, are equal. This allows us to work with square matrices which 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 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.
F.6 OBSERVATION MODEL.
The latent state space Z R of the HiP-RSSM is related to the observation space W by the linear latent observation model H = I (x-m) ], i.e., w|z ~ W Eoba) with w E W and 2 E 2, where m. denotes the m x m identity matrix and 0m (n-m) denotes m x (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).
F.7 OBSERVATION UPDATE STEP
The observation update discussed in section 2.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 tran- sition model, the factorized covariance Σ,, and the special observation model simplify the Kalman update to scalar operations shown below. As the network is free to choose its own state rep- resentation, finds representation where such assumptions works well in practice Becker et al. 2019).
Similar to state, the Kalman gain matrix Q: is split into an upper Qi and lower part Q Both Qi and Q are squared matrices. Due to the simple latent observation model H I 0mx(n-m) and the factorized covariances, all off-diagonal entries of Qi and Q, are zero and one can again work with vectors representing the diagonals, i.e., q, and q/. Those are obtained by
qed = o Ø (a. -h + o g. of- @ (o u. + bobb).
where @ denotes an elementwise vector division. The update equation for the mean therefore sim- plifies
z = z,
wi 2 21
wt
where O denotes the elementwise vector product. The update equations for the individual parts of covariance are given by
o (1m
qi) O o o = (1m q:) @ o
of qf 3 0
where 1m denotes the m dimensional vector consisting of ones.
G ROBOTS AND DATA
The experiments are performed on data from two different real robots. The details of robots, data and data preprocessing is explained below:
G.1 MUSCULOOSKELETAL ROBOT ARM
Observation and Data Set: For this soft robot we have dimensional observation inputs(joint angles) and 8 dimensional action inputs(rressures). We collected the data of four DoF robot
18

---

Published as a conference paper at ICLR 2022
Figure 7:
The experiments are performed on data from robots with different actuator dynamics. From left to right these include: Pneumatically actuated artificial muscles (Büchler et al. 2016), electrically actuated Franka Emika Panda Robotic Arm.
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 model and is relevant to the soft robotics in general.
Training Procedure: For the fully observable case, trained HiP-RSSM for one-step ahead pre- diction 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.c., we train the models perform action conditional future predictions 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 context set of interaction histories during training/inferenee.
G.2 FRANKA EMIKA PANDA ROBOT ARM
Observation and Data Set: We collected the data from 7 DoF Franka Emika Panda manipulator during free motion and while manipulating loads with weights Okg (free motion), 0.5 kg, I kg. 1.5 kg, 2 kg and 2.5 kg. Data is sampled at high frequencies (IkHz). The training trajectories were motions with loads Okg(free motion), 1kg, 1.5kg, 2.5 kgs, while the testing trajectories contained motions with loads of 0.5kg and kgs. The observations for the forward model consist of the seven joint angles in radians, and the corresponding actions were werejoint 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, trained HiP-RSSM for one-step ahead pre- diction using an RMSE loss. Similar to the training procedure for partially observable case as in G.I 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 basclines (RKN, LSTM, GRU) are trained in similar fashion except that, we dont maintain a context set of interaction histories during training/inferenee.
G.3 WHEELEDD MOBILE ROBOT
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 high frequen- cies (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
19

---

Published as a conference paper at ICLR 2022
robot. The observation of the robot any time instance consists of the following features:.
= [z, cas( sin(a), cos
sin (3n cos(7), si)
where, 2, y/ z denote the global position of the Center of Mass he robot, a; y 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 context set and rest 150 is used for the recurrent time-series modelling.
For the fully observable case, we trained HiP-RSSM for one-step ahead pre- diction using an RMSE loss. Similar to the training procedure for partially observable G.T during training we randomly removed half the observations from the sequences and tasked the models with imputing those missing observations, only based on the knowledge of available actions/control commands.
Training Procedure:
in
Other recurrent basclines (RKN, LSTM, GRU) are trained in similar fashion except that, we dont maintain a context set of interaction histories during training/inferenee.
H HYPERPARAMETERS
H.1
PNEUMATIC MUSCULOSKELTAL ROBOT ARM
Recurrent Models
Hyperparameter
HiP-RSSM
RKN
LSTM
GRU
Learning Rate
8c-4
8c-4
le-3
lc-3
Latent Observation Dimension
15
15
15
15
Latent State Dimension
30
75
75
Latent Task Dimension
30
Context Encoder (HiP-RSSM): fully connected linear output (elu + 1)
Fully Connected 240, ReLU
Observation Encoder SHii-RSSSM.KNNLLSMM..CH fully connected linear output (elu + 1)
Fully Connected 120, ReLU
Observation Decoder (HiP-RSSM,,RKNLSSTM))) fully connected + linear output:
Fully Connected 1: 120, ReLU
Transition Model (HiP-RSSM.RKN) number of basis: 15
a<22): No hidden layers softmax output
Control Model (HiP-RSSM.RKKN) 3 fully connected + linear output
Fully Connected 1: 120, ReLU
Fully Connected 2: 120, ReLU
Fully Connected 120, ReLU
Neural Process (NP) Baseline
Latent Task Dimension: 30
Learning Rate: 9e-4
Optimizer Used: Adam Optimizer
Context Encoder fully connected + linear output (elu 1)
20

---

Published as a conference paper at ICLR 2022
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
FFNN Baseline
Learning Rate: 1e-3
Optimizer Used: Adam Optimizer
2 fully connected linear output
Fully Connected 1: 6000, ReLU
Fully Connected 2: 3000, ReLU
MAML Baseline
Meta Optimizer Learning Rate: 3e-3
Inner Optimizer Learning Rate: 0.4
Number Of Gradient Steps:
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
H.2
FRANKA ROBOT ARM WITH VARYING LOADS
Recurrent Models
Hyperparameter
HiP-RSSM
RKN
LSTM
GRU
Learning Rate
le-3
1e-3
3e-3
3e-3
Latent Observation Dimension
15
15
15
15
Latent State Dimension
30
30
75
75
Latent Task Dimension
Encoder (HI--BSMM.RNLLSTM..RR)) fully connected + linear output (elu
Fully Connected 30, ReL.U
Observation Decoder (HiP- CEI.RRSS..KLLSM.GRRRIII fully connected linear output:
Fully Connected 1: 30, ReLU
Transition Model (HiP-RSSM,RKN); number of basis: 32
a(zz)): No hidden layers - softmax output
Control Model (HiP-RSSM, RKN): 1 fully connected + linear output
Fully Connected 120, ReLU
21

---

Published as a conference paper at ICLR 2022
Neural Process (NP) Baseline
Latent Task Dimension: 30
Learning Rate: 1e-3
Optimizer Used: Adam Optimizer
Context Encoder:
fully connected + linear output (elu 1)
Fully Connected 240, ReLU
Linear output: 30
Decoder 3 fully connected + linear output
Fully Connected 512, ReLU
Fully Connected 2: 240, ReLU
Fully Connected 2: 120, ReLU
FENN Baseline
Learning Rate: 1e-3
Optimizer Used: Adam Optimizer
2 fully connected linear output
Fully Connected 6000, ReLU
Fully Connected 2: 3000, ReLU
MAML Baseline
Meta Optimizer Learning Rate: 3e-4
Inner Optimizer Learning Rate: 0.04
Number Of Gradient Steps:
Optimizer Used: Adam Optimizer (Meta Update) and SGD Optimizer (Inner Update)
Model
Feed Forward Neural Network with 3 fully connected linear output
Fully Connected 512, ReLU
Fully Connected 240, ReLU
Fully Connected 120, ReLU
H.3
WHEELED ROBOT TRAVERSING SLOPES OF DIFFERENT HEIGHIT
Recurrent Models
Hyperparameter
HiP-RSSM
RKN
LSTM
GRU
Learning Rate
9c-4
9e-4
le-2
lc-2
Latent Observation Dimension
30
30
15
15
Latent State Dimension
60
60
75
75
Latent Task Dimension
Encoder (HI-RRSSMRKKLLSTM.GRRU fully connected + linear output (elu 1)
Fully Connected 120, ReLU
Observation Decoder (HiF-RSSM.RKNNLSTM..RU)) fully connected linear output:
Fully Connected 1: 240, ReLU
Transition Model (HiP-RSSM,RKN) number of basis: 15
22

---

Published as a conference paper at ICLR 2022
(xz)): No hidden layers softmax output
Control Model (HiP-RSSM) RKN): 3 fully connected + linear output
Fully Connected 120, ReLU
Neural Process (NP) Baseline
Latent Task Dimension: 30
Learning Rate: le-3
Optimizer Used: Adam Optimizer
Context Encoder: fully connected + linear output (elu I)
Fully Connected I: 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 60, ReLU
FFNN Baseline
Learning Rate: 2e-3
Optimizer Used: Adam Optimizer
3 fully connected linear output
Fully Connected 512, ReLU
Fully Connected 2: 240, ReLU
Fully Connected 2: 120, ReLU
Fully Connected 60, ReLU
MAML Baseline
Meta Optimizer Learning Rate: 3e-4
Inner Optimizer Learning Rate: 0.4
Number Of Gradient Steps:
Optimizer Used: Adam Optimizer (Meta Update) and SGD Optimizer (Inner Update).
Model
Feed Forward Neural Network with 3 fully connected + linear output
Fully Connected 512, ReLU
Fully Connected 240, ReLU
Fully Connected 2: 120, ReLU
Fully Connected 3: 60, ReLU
23


Reviews:
[{"review_id": "lljZ43gMVcw", "reviewer": "Reviewer_8V7G", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "### Strengths\nExtending current state-space models to handle changing dynamics is important and the problem being addressed in this work is relevant to the community. The paper describes the approach reasonably well (up to the points below). A plus point is that the authors conducted experiments using real robots to demonstrate their approach can be applied to these settings.\n\n### Weaknesses\nAt present, technical novelty is unclear in that HiP-RSSM appears to be a variation of existing models. The primary equations are derived from the standard Kalman filter and there doesn't appear to be new insights or technical advances in the method's derivation. The experiments are not clearly described in the main text, which makes it difficult to validate the main claim that HiP-RSSM works better than existing approaches. Please see below for details. \n\n- The proposed model bears resemblance to the Kalman VAE (KVAE) [Fraccaro et al. 2017] which is not cited but highly related. The key difference is the introduction of the latent variable $l$. Initially, I expected that the transition matrices $A_t$ or control function to be dependent on $l$ but the variable enters only as an additive term in the updates. As an aside, I suggest that B.5 be brought forward into the main text since it is a core part of the method. \n- The proposed model is also rather confusing since the mean $\\mu_l$ is used in these updates rather than $r^l_n$? Can the authors better motivate the need to introduce the latent variable in this manner? How would it compare to a simple baseline that modifies the latent representation with an extra bias dimension and learns a time-varying transition matrix $A_t$? \n- The literature review also misses out on related work on switching dynamical systems, e.g., state-dependent or explicit duration models, which can capture the change in dynamics. This work appears to be a special case whereby the switching is not explicitly captured and technically contribution appears limited. \n- Is the training in the experiments done using RMSE or using eqn (9)? If I understand the paper, it seems like the method is training solely using RMSE which makes the proposed probabilistic setup appear unnecessary. \n- The prediction experiments results are for a one time-step prediction (only explained in the supplementary) and it's unclear how meaningful the results are. Typically, we are interested in predictions over long horizons, e.g., for MPC. This also motivates the need for capturing the change in dynamics (e.g., in the switching models) that is modelled by the HiP-RSSM.  \n\nReferences:\n\nFraccaro, Marco, et al. \"A disentangled recognition and nonlinear dynamics model for unsupervised learning.\" NeurIPS (2017).", "summary_of_the_paper": "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).", "summary_of_the_review": "While I do believe that the paper has some merits, the technical contribution appears limited and I am unable to recommend an accept. It would be helpful if the authors clarify differences to existing models, especially KVAE and switching state-space models.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "KBGtBVDTJPM", "reviewer": "Reviewer_VmRd", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The core idea leverages the model specification of a hidden-parameter-MDP (HiP-MDPs)  such that the hidden parameter encodes the time-varying aspect of the dynamics.  If I understand it right, the model considered overall consists of a latent variable (for e.g., motor friction, which is perhaps globally changing but locally fixed), a true underlying state (for e.g., exact location of the robot which changes locally as well), and an observation (for e.g., image of where the robot is).\n\n\n- Why are permutation invariant deep models used for encoding the history of past N steps? Why is it useful to discard the temporal information?\n\n- Section 3.2.2: It is mentioned that $z_{t-1},l, a$ are independent of each other as $z_t$ is unobserved. How are $z_{t-1}$ and $l$ independent. From Figure 2, does not $l$ directly influences $z_{t-1}$?\n\n- Section 3.2.2: what is the random variable being sampled from the third Gaussian distribution? Is it supposed to be $y$ instead of $v$?\n\n- Section 3.2.3: It might be better to provide some information in the main paper about how inverse on the term that depends on H is being done in a computationally efficient manner. \n\n\n- The key idea of the proposed work is to provide a wrapper around the Kalman filter. Exploiting the fact that the hidden-parameter is fixed throughout the task, the proposed method used deep nets to infer the hidden-parameter. This inferred hidden parameter is then used along with the standard state-space model, where the next state changes linearly in the past state, action, and the hidden parameter. \n\n- If I understand it right, in the proposed setup, deep nets were possible for inferring the hidden-parameter as the hidden parameters never change and thus do not need any posterior updates. Whereas, the state does evolve and thereby necessitating linear evolution structure. \n\n- Further, because the method builds around Kalman filtering, it is assumed that all the noises are isotropic Gaussian, such that covariances are zero and Kalman updates can be done easily. I am not sure why this assumption is reasonable in general.\n\n\n- As the paper title says \"changing dynamics scenario\", perhaps it is only reasonable to assume that the proposed method will handle that. While it is fine to assume that locally the hidden parameter is fixed (implying dynamics is fixed), ultimately it needs to be shown how at the global level the method can track the changing dynamics parameter as well. \n\n  - Currently, both in the model setup and the experiments, this setup is missing. The model assumes that the hidden parameter value is fixed throughout and does not consider at all _how the current hidden parameter influences the next value of the hidden parameter_. \n\n  - If I understand it right, even in the experiments constructed, each observation sequence is generated using a fixed value of the hidden parameter. Notice that this is not a fair comparison with the baseline methods. For example, even the basic GRU and LSTM models are more general as they can consider the latent variable to be a part of state-space that is evolving with time.\n\n- Further, the paper misses out on recent works by Xie et al. (2020) that build upon the work by Nagabandi et al. (2018a;b) and Finn et al. (2017). They look at a very similar HiP-MDP setup (and also goes beyond just modeling the time-varying dynamics and tries to do control). While they make the assumption of dynamics having a locally fixed hidden parameter like this work, they also consider how the hidden parameter is evolving more globally.\n\nDeep Reinforcement Learning amidst Lifelong Non-Stationarity.\nAnnie Xie, James Harrison, Chelsea Finn.\nInternational Conference on Machine Learning (ICML), 2021. \n\n- Under related works, it is mentioned that existing deep state-space models assume stationarity but their proposed method does not. I do not think this is true. If one considers the latent variable to be sampled from a starting distribution and held fixed throughout the episode (as what the proposed method has done both in experiments and the methodology), then it can be seen that everything is just stationary as well and warrants comparison with the existing deep state-space methods.\n\n-----\nMinor\n\n- The preliminary section builds up the notations as if arbitrary non-linear functions will be used later to model the dynamics, which is not completely true. \n\n- I think the paper can benefit a lot by clarifying the relations between $x$, $w$, $o$, and also between $\\eta$, $z$, $l$, \n\nTypos  Section 3.2.2: \n- Figure 1 => Figure 2?\n- Should $z_t^{-}$ depended on $z_{t-1}^-$ instead of $z_{t-1}$?\n- Is the matrix A time-varying or fixed? Notations are used inter-changeably.\n\n\n ", "summary_of_the_paper": "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.", "summary_of_the_review": "Overall, the proposed method claims to be designed for environments that are dynamically changing over time, but actually assumes the latent variable to be fixed throughout the entire episode (not just locally for a small part of the episode). As such it contradicts the main claim.\n\nLet me know if my understanding of the work is not accurate and I will be happy to reconsider the score.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "8IhjZTMaW15", "reviewer": "Reviewer_vXB5", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "# Strengths \n\nI am a huge fan of this work as I think a lot of nifty tricks are used. Firstly, I thought the introduction and the motivation of this work is top-notch! Secondly, the form of the generative model is very elegant as it allows for a global dynamics model to be learned that can be shared across tasks/environments while the context variable allows for specialization within a task (and is also very easy to compute at test-time). \nThe use of \"linear\" (I will have comments on this later) dynamics allows for everything to computed in closed-form as well. To avoid issues with likelihoods that are parameterized by NNs, the authors also use a pseudo-like likelihood where a separate encoder is used to project the observations to some feature space that is used to inform the latent states and while retaining the closed-form updates of Kalman filtering.\n\n# Weaknesses\n\nI think there are very important details that are swept under the rug. \nFirstly, the authors use a locally-linear transition model\n$$ A_{t-1} = \\sum_{k=1}^K \\alpha(z_{t-1}) A^{k} $$\nwhere $\\alpha(\\cdot)$ is a softmax parameterized by a NN. \nNote that while Kalman filtering can handle time-varying linear dynamics, they have to be *indepedent* of $z_{t-1}$.\nThis parameterization *prevents* the use of Kalman filtering as now $z_t$ nonlinearly depends on $z_{t-1}$! This is a well-known problem that has been attempted to be addressed in the switching linear dynamical systems literature i.e.,  (https://proceedings.mlr.press/v54/linderman17a.html, https://arxiv.org/abs/1811.12386). \n\nSecondly, the marginalization of the context variable, $l$, is only available in closed form if the $z_t$ depends linearly on $l$. But in some of the experiments, the authors pass the context variable through a NN, preventing closed-form marginalization; they instead take an approximation, passing in the mean and variance of the through the NN instead (I'm somewhat confused by the passing the variance through the NN. By the principle of propagation of uncertainty, it should be something like the derivative squared times the variance, https://en.wikipedia.org/wiki/Propagation_of_uncertainty).\nWhile I'm somewhat fine with these two approximations as it allows for ease of computations, the authors should explicitly state this and avoid the use of \"exact inference\". Moreover, I would love to see some sort of empirical evidence that taking these approximations is okay i.e. comparing their KF to an extended KF.\n\nLastly, one other detail that is missing is the discussion of the context sets. The description is somewhat confusing. On page 3, the authors state \"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\".\nBut on page 4, the parameterization and description of the context set makes it seem that one is inferring a sequence of $l$'s, \n\"for any time series $\\mathcal{T} = (o_t, a_t, o_{t+1}, ..., o_{t+N}, a_{t+N})$, the context set $C_l$ consists of tuples of\nthe current state/observation, current action and next state/observation for the previous N time steps, i.e. $C_l = (o^l_{t-n}, a_{t-n}^l, o_{t-n+1}^l)_{n=1}^N $\".\nWhat I am assuming is happening is that the data is split a priori into a context set and another set used for performing the Kalman filtering updates. \nI suggest using a separate subscript for the context sets.\nAlso, the constructions of the context sets don't seem to be discussed: what is $N$, how were the context sets collected, etc.\n\n# Comments/Questions\n1. On equation &, it should be $z_t^- = A_{t-1} z_{t-1} + Ba_t + C\\mu_l$\n2. Many models can perform decently well on 1-step ahead prediction (for instance, an LDS can perform well at one-step ahead prediction if the dynamics are slow, see Figures 2 and 3 in https://arxiv.org/abs/1811.12386). Have you investigated how your model fares on multi-step ahead prediction?\n\n\n------------------------\n# Update\nI think the authors have done a good job responding to my concerns and the concerns of the other reviewers. As such, I have raised my score to an 8.", "summary_of_the_paper": "This paper introduces a state-space model for non-stationary environments (where in this formulation the latent dynamics are changing while the observation model is fixed), which they call Hidden Parameter Recurrent State Space Models (HiP-RSSMs). In HiP-RSSMs, a context variable is inferred which is used to modify the latent dynamics to adapt towards the new setting. ", "summary_of_the_review": "I think the idea presented in the paper is great and well-motivated. The approach is elegant and the empirical results are compelling. Most of my concerns are in--important--missing details and details that are swept in the appendix, which I think can be easily fixed with some rewriting. For this reason, I will give a marginal acceptance. I am more than happy to increase my score if the authors address my concerns. \nAlso, if I misunderstood/missed something please let me know!", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "4aiMvm9jEHR", "reviewer": "Reviewer_nyLc", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "\n#### Pros\n- Moderate extension of existing recurrent state space models for changing dynamics for well-motivated problems.\n- Detailed experimental configurations are in the appendix.\n- Applications on different robot learning problems and settings are evaluated.\n\n#### Cons\n- For the experiments with partial observations and missing values, the authors may need to compare with some RNN baselines designed for handling partial observations/missing values.\n- The definition and scope of the context set $C_l$ and some related contents are quite ambiguous to the readers, which disrupts the understanding and evaluation of the proposed method. \n  - Based on the definition in Section 3.2.1, $C_l$s are different for different time step $t$; Meanwhile, $l$ is fixed for the duration of the task.\n  - In Section 3.2.1: $T$ has $N$ time steps. While in Section 3.3 $T$ is of length $M$. Please explain the differences between these variables.\n  - $t \\in T \\in R^{B \\times M \\times D}$. What is $t$?\n\n#### Details\n- In the second paragraph of the introduction, the authors motivate this work by listing both time-invariant (\"similar but not identical\") dynamics and time-variant dynamics (such as the friction changes and wear and tear). How are these two different types of dynamics captured in the proposed model, in the same way or with varying model formations? The experiments conducted in Section 5.1 are assumed to be for the time-variant dynamics, but more evidence about this assumption is needed.\n- However, the proposed model seems only to handle the former ones. Also, the experiments are not conducted with the time-changing dynamical systems.\n- It is unclear what the concept of latent task refers to (could the authors provide some examples?) and why multi-task models are suitable to compare as baselines.\n\n- Some definitions and concepts in Section 2.2 are not clear until the readers reach Section 3.2. Content rearrangements or clarifications are needed.\n- Similarly, it would be suggested to elaborate more model details about, e.g., locally linear transition.\n\n- Minur issues: \n  - Introduction: \"... like images was ...\" -> \"... like images and was ...\"\n  - Section 3.2: \"Each of these are three stages ...\" -> \"Each of these three stages ...\"\n  - Section 3.2.3: \"this a Gaussian conditioning layer\" -> \"this is a ...\"", "summary_of_the_paper": "The authors extend the recurrent Kalman networks with hidden parameters and context set features to model changing dynamics. The authors demonstrate the performance of the proposed method with a series of experiments on robot learning.", "summary_of_the_review": "This paper provides plausible and moderate extensions to the existing model to solve problems in a new setting. However, some necessary clarifications and important baselines are missing. Therefore, I'd like to put it borderline and am willing to adjust my scores if more information from the rebuttal and discussion phases is provided.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper addresses a few very important points on sequential latent-variable models, and introduce a different view on meta-RL.  Even  though the methods that this paper poses are incremental, it is such a hot-debated topic that I would prefer to see this published now.
Decision: Accept (Poster)