We finally visualize the latent task representations using TSNE in Figure <ref>. As seen in the plot, the HiP-SSM instances under similar terrain slopes cluster in the latent task space, indicating that the model correctly identifies the causal effects in the dynamics in an unsupervised fashion.

§.§ Ablation On Context Encoder For Inferring Task Abstraction

In table <ref>, we report the details of evaluating our Bayesian aggregation based context set encoder (discussed in <ref>) against a causal / recurrent encoder that takes into account the temporal structure of the context data owing to its sequential nature. We used a probabilistic recurrent encoder becker2019recurrent, whose mean and variance from the last time step is used to infer the posterior latent task distribution $p(\cvec{l}|\mathcal{C}_{\cvec{l}}) = \mathcal{N}(\cvec{\mu}_{\cvec{l}}, \textrm{diag}(\cvec{\sigma}_{\cvec{l}}^2))$. The dimensions of the latent task parameters obtained from both the set and recurrent encoders remain the same (60).

The reported experiments are conducted on data from a wheeled mobile robot discussed in Section <ref>. As reported in Table <ref> the permutation invariant set encoder outperforms the recurrent encoder by a good margin in terms of prediction accuracy for both fully and partially observable scenarios. Additionally the set encoder is far more efficient in terms of computational time required for training as seen from the time taken per epoch for each of these cases as it allows for efficient parallelization.

Comparison between the permutation invariant set encoder 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 5 runs.

No Imputation RMSE      $50 \%$ Imputation RMSE     Training Time Per Epoch

Set Encoder     $2.96 \pm 0.212$     $6.15 \pm 0.327$       $6.71$

Recurrent Encoder     $5.10 \pm0.041$     $10.12\pm 0.112$      14.13

§.§ Visualizing Changing Hidden Parameters At Test Time Over Trajectories With Varying Dynamics

An ablation on the performance of different task transformation models discussed in section <ref>.

We finally perform inference using the trained HiP-RSSM in a multi-task / changing dynamics scenario where the dynamics continuously changes over time. We use the inference procedure described in appendix <ref> based on a fluid definition for “task” as the local dynamics in a temporal segment. We plot the global variation in the latent task variable captured by each instance of the HiP-RSSM over these local temporal segments using the dimensionality reduction technique UMAP mcinnes2018umap. As seen in figures <ref> and <ref>, the latent task variable captures these causal factors of variations in an interpretable manner.

(a) and (b) shows how the one dimensional UMAP mcinnes2018umap embedding of the inferred latent task variable (top blue plots) changes at test time when an agent undergoes changes in its dynamics for the franka robot and 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.

§ CONCLUSION

We proposed HiP-RSSM, a probabilistically principled recurrent neural network architecture for modelling the world under non-stationary dynamics. We start by formalizing a new framework, HiP-SSM, to address the multi-task state-space modelling setting. HiP-SSM assumes a shared latent state and action space across tasks but additionally assumes latent structure in the dynamics.  We exploit the structure of the resulting Bayesian network to learn a universal dynamics model with latent parameter $l$ via 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 unsupervised fashion. Our experimental results on various robotic benchmarks show that HiP-RSSMs significantly outperform state of the art recurrent neural network architectures on dynamics modelling tasks. We believe that modelling the dynamics in the latent space which disentangles the state, action and task representations can benefit multiple future applications including planning/control in the latent space and causal factor identification.

CHAPTER: MULTI TIME SCALE SSMS: HIERARCHICAL WORLD MODELS AT MULTIPLE TEMPORAL ABSTRACTIONS

This chapter is based on "Multi Time Scale World Models" shaj2023multi.

This chapter of the thesis explores the concept of "temporal depth" within generative models of the world by structuring them through nested hierarchies that span across different spatio-temporal scales. This approach addresses a significant limitation found in the prevailing literature on world models: the reliance on a singular temporal depth or time scale. Present formulations in the fields of machine learning and artificial intelligence typically operate with a single, finely detailed temporal scale, such as milliseconds. Singular time scale often limits accurate prediction and planning over long hoizons lecun2022path. For efficient long-term prediction and planning, the model must predict at multiple levels of temporal abstractions sutton1995td,precup1997multi,lecun2022path.

Conceptual Depiction of Hierarchical Temporal Abstraction in World Models. At the lower level, represented by the array of yellow and blue squares, we observe a drone's movement patterns around a central tree—circular during the day (indicated by a yellow background) and elliptical at night (indicated by a blue background). The upper tier conveys a more abstracted temporal perspective, encapsulating broader informational trends without the granular detail. This abstraction captures the cyclical day-to-night transition and the drone's corresponding behavioral patterns: circular in daylight, elliptical in darkness. This model illustrates how complex, time-variant dynamics are synthesized into simplified trends to facilitate predictive insights within multi-scale temporal frameworks.

The world is rife with regularities. Day follows night, seasons follow one another, milk goes sour, faulty brakes are often followed by accidents, and so on. Regularities come at different time scales, ranging from tens of milliseconds to hundreds, to seconds, minutes, and upwards towards regularities or rules that are stable over weeks, months, and years.  For example, commuting to work involves planning at a higher level slow time scale (in minutes or hours) like which route to drive (or walk) and at a lower level fast time scale (in seconds or milliseconds) like how to brake/steer or choose precise skeletal muscle actions. Similarly, in robotic manipulation, the robot must be able to perform precise and coordinated movements to grasp and manipulate the object at a fast time scale while at a slower time scale the robot must also be able to recognize and utilize higher-level patterns and structures in the task, such as the shape, size, and location of objects, and the overall goal of the manipulation task. Generative models with "temporal depth" in the form of nested hierarchies offers several advantages in dealing with such scenarios lee2003hierarchical, hohwy2013predictive, lecun2022path as listed below:

* Reflection of Causal Depth: The world is inherently structured hierarchically hohwy2013predictive, with causes and effects nested within one another on various scales of time and space (for example, in autonomous driving applications, slow regularities like traffic congestion trends including peak and off-peak hours, holiday traffic patterns, and seasonal changes in road conditions affect the fast regularities like real-time decisions on route navigation and speed adjustments). A hierarchical model allows the representation of this causal depth, enabling intelligent agents to understand not just immediate sensory inputs but also the broader and interconnected causal relationships that govern the external world.

* Better Long Term Prediction: By organizing regularities in the world from faster, detailed levels to slower, more abstract levels, hierarchical models enable predictions at multiple scales of precision and time. This flexibility is fundamental to adaptive behavior, allowing for both immediate, detailed responses and longer-term planning based on abstracted patterns of regularity friston2008hierarchical,hohwy2013predictive. Fast changing regularities are good for detail; slower regularities are more general and abstract. This makes sense when we consider what regularities allow us to predict. If I want to predict something with great perceptual precision, then I cannot do it very far into the future, so I need to rely on a fast changing regularity. On the other hand, predictions further into the future come at a loss of precision and often detail.

* Adaptability and Transferability: As seen in Chapter <ref> hierarchical temporal abstractions can capture relevant task structures across dynamical systems under non-stationary conditions, which can be used to identify similarities and differences between tasks. Learning dynamics of these nested hierarchical abstractions allow for making "predictions about predictions" and subsequent transfer of knowledge learned from one task to another shanahan2022abstraction, lecun2022path. This adaptability is key to navigating a world that is constantly changing, allowing for the refinement of predictions and behaviors in response to new information.

This chapter of the dissertation attempt to come up with a principled probabilistic formalism for learning such multi-time scale world models as a hierarchical sequential latent variable model. We show that such models can better capture the complex, nonlinear dynamics of a system more efficiently and robustly than models that learn on a single timescale. This is exemplified in several challenging simulated and real-world prediction tasks such as the D4RL dataset, a simulated mobile robot, and real manipulators including data from heavy machinery excavators.

§ RELATED WORK

Multi Time Scale World Models One of the early works that enabled environment models at different temporal scales to be intermixed, producing temporally abstract world models was proposed by <cit.>. The work was limited to tabular settings but showed the importance of learning environment dynamics at multiple abstractions. However, there have been limited works that actually solve this problem at scale as discussed in <cit.>. A probabilistically principled formalism for these has been lacking in literature and this work is an early attempt to address this issue.
Deep State Space Models. Deep SSMs combine the benefits of deep neural nets and SSMs by offering tractable probabilistic inference and scalability to high-dimensional and large datasets. <cit.> use neural network architectures based on exact inference on SSMs and perform state estimation and dynamics prediction tasks. <cit.> extend these models to modelling non-stationary dynamics. <cit.> perform learning and inference on SSMs using variational approximations. However, most of these recurrent state-space models have been evaluated on very short-term prediction tasks in the range of a few milliseconds and model the dynamics at a single time scale.
Recent advancements in Transformers vaswani2017attention, radford2019language, brown2020language, which rely on attention mechanism, have demonstrated superior performance in capturing long-range dependency compared to RNN models in several domains including time series forecasting zhou2021informer,liu2022nonstationary and learning world models micheli2023transformers. <cit.> use transformer architectures based on a direct multistep loss zeng2022transformers and show promising results for long-term forecasting since they avoid the accumulation of errors from autoregression. On the other hand <cit.> uses a GPT-like autoregressive version of transformers to learn world models. These deterministic models, however, do not deal with temporal abstractions and uncertainty estimation in a principled manner. Nevertheless, we think Transformers that operate at multiple timescales based on our formalism can be a promising alternative research direction.

§ MULTI TIME SCALE STATE SPACE MODELS

Regularities in our world can be ordered hierarchically, from faster to slower hohwy2013predictive. Levels in the hierarchy can be connected such that certain slow regularities, at higher levels, pertain to relevant lower level, faster regularities. A complete such hierarchy would reveal the causal structure and depth of the world—the way causes interact and nest with each other across spatiotemporal scales.

Our goal is to learn a principled sequential latent variable model that can model the dynamics of artificial embodied agents (specifically robotic systems) under multiple levels of temporal abstractions based on observed sensory data. To do so, we introduce a new formalism, called Multi-Time Scale State Space (MTS3) Model, with the following desiderata:
i) It is capable of modeling dynamics at multiple time scales. ii) It allows for a single global model to be learned that can be shared across changing configurations of the environments. iii) It can provide accurate long-term predictions and uncertainty estimates. iv) It is probabilistically principled, yet scalable during learning and inference.

We start by introducing a formal definition of a 2 time scale Multi-time scale SSM in Section <ref> and then proceed to a general definition of an MTS3 with an arbitrary number of hierarchies/timescales in Section <ref>.

§.§ Formal Definition Of a 2-Level MTS3

An MTS3 model with 2 timescales is defined by two SSMs on a fast and a slow time scale respectively. Both SSMs are coupled via the latent state of the slow time scale SSM, which parameterizes / “reconfigures” the system dynamics of the fast time scale SSM. While the fast time scale SSM runs at the original time step $\Delta t$ of the dynamical system, the slow time scale SSM is only updated every $H$ step, i.e., the slow time scale time step is given by $H \Delta t$. We will derive closed-form Gaussian inference for obtaining the beliefs for both time scales, resulting in variations of the Kalman update rule which are also fully differentiable and used to back-propagate the error signal becker2019recurrent,haarnoja2016backprop. The definition with a 2-level MTS3 along with the inference and learning schemes that we propose is directly extendable to an arbitrary number of temporal abstractions by introducing additional feudal dayan1992feudal hierarchies with longer discretization steps and is further detailed in Section <ref>.

§.§.§ Fast time-scale SSM

The fast time-scale (fts) SSM is given by $\mathcal{S}_{\textrm{fast}} = ( \mathcal{Z}, \mathcal{A}, \mathcal{O}, f_{\cvec l}^\textrm{fts}, h^\textrm{fts}, \Delta t, \mathcal{L}).$ Here, $\cvec l \in \mathcal{L}$ is a task descriptor that parameterizes the dynamics model of the SSM and is held constant for H steps. We will denote the task descriptor for the $k$th time window of $H$ steps as $\cvec l_k$. The probabilistic dynamics and observation model of the fast time scale for the $t$th time step in the $k$th window can then be described as

\begin{align}
p(\cvec{z}_{k,t}|\cvec{z}_{k,t-1},\cvec{a}_{k,t-1}, \cvec l_k) & = \mathcal{N}(f^\textrm{fts}_{\cvec l}(\cvec{z}_{k,t-1},\cvec{a}_{k,t-1}, \cvec l_k), \cmat{Q}),
\textrm{ and } \nonumber\\
p(\cvec{o}_{k,t}|\cvec{z}_{k,t}) & = \mathcal{N}(h^\textrm{fts}(\cvec{z}_{k,t}), \cmat{R}).
\end{align}

Task-conditioned marginal transition model.
Moreover, we have to consider the uncertainty in the task descriptor (which will, in the end, be estimated by the slow time scale model), i.e., instead of considering a single task descriptor $\cvec l_k$, we have to consider a distribution over task descriptors $p(\cvec l_k)$ for inference in the fts-SSM. This distribution will be provided by the slow-time scale SSM for every time window $k$. We can further define the marginal task-conditioned transition model for the time window $k$ that is given by

\begin{align}p_{\cvec l_k}(\cvec{z}_{k,t}|\cvec{z}_{k,t-1},\cvec{a}_{k,t-1}) & = \int p(\cvec{z}_{k,t}|\cvec{z}_{k,t-1},\cvec{a}_{k,t-1}, \cvec l_k) p(\cvec l_k) d \cvec l_k
\end{align}

Latent observations.
Following <cit.>, we replace the observations by latent observations and their uncertainty, i.e., we use latent observation encoders to obtain $\cvec w_{k,t} = \textrm{enc}_w(\cvec o_{k,t})$ and an uncertainty encoder $\cvec \sigma_{k,t} = \textrm{enc}_\sigma(\cvec o_{k,t})$. The observation model is thus given by $p(\cvec w_{k,t}|\cvec z_{k,t}) = \mathcal{N}(h^\textrm{fts}(\cvec z_{k,t}), \textrm{diag}(\cvec \sigma_{k,t}))$.

§.§.§ Slow time-scale SSM

The slow time scale (sts) SSM only updates every time step H and uses the task parameter $\cvec l$ as a latent state representation. Formally, the SSM is defined as $\mathcal{S}_{\textrm{slow}} = ( \mathcal{L}, \mathcal{E}, \mathcal{T}, f^\textrm{sts}, h^\textrm{sts}, H \Delta t)$. It uses an abstract observation $\cvec \beta \in \mathcal{B}$ and abstract action $\cvec \alpha \in \mathcal{A}$ that summarize the observations and actions, respectively, throughout the current time window. The general dynamics model is hence given by
\begin{align}p(\cvec l_{k}|\cvec l_{k-1}, \cvec \alpha_k) = \mathcal{N}(f^{\textrm{sts}}(\cvec l_{k-1}, \cvec \alpha_k), \cvec S).\end{align}

Although there exist many ways to implement the abstraction of observations and actions of time windows, we choose to use a consistent formulation by fusing the information from all $H$ time steps of time window $k$ using Gaussian conditioning.

The graphical model corresponding to an MTS3 with 2 timescales. The latent task variable $\cvec l_k$ captures the slow changing dynamics using abstract observation inferred from $\{\cvec \beta_{k,t}\}_{t=1}^H$and abstract action $\cvec \alpha_k$ as described in Section <ref>. The inference in the fast time scale uses primitive observations $\cvec w_{k,t}$, primitive actions $\cvec a_{k,t}$ and the latent task descriptor $l_k$ which parameterizes the fast-changing dynamics of $\cvec z_{k,t}$ for a time window k as discussed in the section <ref>.

Observation abstraction.
In terms of the abstract observation model, we choose to model $H$ observations $\cvec \beta_{k,t}$, $t \in [1,H]$ for a single slow-scale time step $k$. All these observations can then be straightforwardly integrated into the belief state representation using incremental observation updates. The abstract observation and its uncertainty for time step $t$ is again obtained by an encoder architecture, i.e,

$$\cvec \beta_{k,t} = \textrm{enc}_{\beta}(\cvec o_{k,t}, t), \quad \cvec \nu_{k,t} = \textrm{enc}_{\nu}(\cvec o_{k,t}, t),$$

and $p(\cvec \beta_{k,t}|\cvec l_k) = \mathcal{N}(h^{\textrm{sts}}(\cvec l_k), \textrm{diag}(\cvec \nu_{k,t})).$
Hence, the abstract observation $\cvec \beta_{k,t}$ contains the actual observation $\cvec o_{k,t}$ at time step $t$ as well as a temporal encoding for the time-step. Although multiple Bayesian observation updates are permutation invariant, the temporal encoding preserves the relative time information between the observations, similar to current transformer architectures.

Action abstraction.
The abstract action $\cvec \alpha_k$ causes transitions to latent task $\cvec l_{k}$ from $\cvec l_{k-1}$. It should contain relevant information of all primitive actions $\cvec a_{k,t}$, $t \in [1,H]$ executed in the time window $k$. To do so, we again use Bayesian conditioning and latent action encoding. Each control action $\cvec a_{k,t}$ and the encoding of the time step $t$ are encoded in its latent representation and its uncertainty estimate, i.e.,

$$\cvec \alpha_{k,t} = \textrm{enc}_{\alpha}(\cvec a_{k,t}, t), \quad \cvec \rho_{k,t} = \textrm{enc}_{\rho}(\cvec a_{k,t}, t).$$

Generative model for the abstract action $\alpha_k$. The hollow arrows are deterministic transformations leading to implicit distribution $\alpha_{k,t}$ using an action set encoder.

Single latent actions $\cvec \alpha_{k,t}$ can be aggregated into a consistent representation $\cvec \alpha_{k}$ using Bayesian aggregation <cit.>. To do so, we use the likelihood $p(\cvec \alpha_{k,t}|\cvec \alpha_{k}) = \mathcal{N}(\cvec \alpha_{k}, \textrm{diag}( \cvec \rho_{k,t}))$ and obtain the posterior $p(\cvec \alpha_{k}|\cvec \alpha_{k,1:H}) = \mathcal{N}(\cvec \mu_{ \alpha_{k}}, \cvec \Sigma_{\alpha_k})$, which is obtained by following the standard Bayesian aggregation equations, see Appendix <ref>. Note that our abstract action representation also contains an uncertainty estimate which can be used to express different effects of the actions on the uncertainty of the prediction. Due to the Gaussian representations, we can compute the marginal transition model

\begin{align}p_{\cvec \alpha_k}(\cvec l_{k}|\cvec l_{k-1}, \cvec \alpha_{k,1:H}) = \int p_{\cvec \alpha_k}(\cvec l_{k}|\cvec l_{k-1}, \cvec \alpha_k) p(\cvec \alpha_k |\cvec \alpha_{k,1:H}) d\cvec \alpha_k.\end{align}

This transition model is used for inference and its parameters are learned.

§.§.§ Connecting both SSMs via inference

In the upcoming sections, we will devise Bayesian update rules to obtain the prior $p(\cvec l_k| \cvec \beta_{1:k-1},  \cvec \alpha_{1:k})$ and posterior $p(\cvec l_k| \cvec \beta_{1:k},  \cvec \alpha_{1:k})$ belief state for the sts-SSM as well as the belief states for the fts-SSM. The prior belief $p(\cvec l_k| \cvec \beta_{1:k-1},  \cvec \alpha_{1:k})$ contains all information up to time window $k-1$ and serves as a distribution over the task-descriptor of the fts-SSM, which connects both SSMs. This connection allows us to learn both SSMs jointly in an end-to-end manner.

The probabilistic graphical model of our MTS3 model is depicted in Figure <ref>. In the next section, we will present the detailed realization of each SSM to perform closed-form Gaussian inference and end-to-end learning on both time scales.

§.§ Inference in the Fast Time-Scale SSM

The fts-SSM performs inference for a given time window $k$ of horizon length $H$. To keep the notation uncluttered, we will also omit the time-window index $k$ whenever the context is clear.
We use a linear Gaussian task conditional transition model, i.e,
\begin{align}p(\cvec{z}_{t}|\cvec{z}_{t-1},\cvec{a}_{t-1},\cvec{l}_k) = \mathcal{N}\left(\cmat{A}\cvec{z}_{t-1} + \cmat{B}\cvec{a}_{t-1} + \cmat{C}\cvec{l}_k,  \cmat Q\right),\end{align}
where $\cvec A$, $\cvec B$, $\cvec C$ and $\cvec Q$ are state-independent but learnable parameters.
In our formulation, the task descriptor can only linearly modify the dynamics which was sufficient to obtain state-of-the-art performance in our experiments, but more complex parametrizations, such as locally linear models, would also be feasible. Following <cit.>, we split the latent state $\cvec z_t = [\cvec p_t, \cvec m_t]^T$ into its observable part $\cvec p_t$ and a part $\cvec m_t$ that needs to be observed over time.
We also use a linear observation model $p(\cvec w_{t}|\cvec z_{t}) = \mathcal{N}(\cvec H \cvec z_{t}, \textrm{diag}(\cvec \sigma_{t}))$ with $\cvec H = [\cvec I, \cvec 0]$.

We will assume that the distribution over the task descriptor is also given by a Gaussian distribution, i.e., $p(\cvec l_k) = \mathcal{N}(\cvec \mu_{\cvec l_k}, \cvec \Sigma_{\cvec l_k})$, which will be provided by the slow-time scale (sts) SSM, see Section <ref>. Given these modelling assumptions, the  task variable can now be integrated out in closed form, resulting in the following task-conditioned marginal transition model

\begin{align}p_{\cvec l_k}(\cvec{z}_{t}|\cvec{z}_{t-1},\cvec{a}_{t-1})
= \mathcal{N}\left(\cmat{A}\cvec{z}_{t-1} + \cmat{B}\cvec{a}_{t-1} + \cmat{C}\cvec{\mu}_{\cvec l_k},  \cmat{Q} + \cvec C \cvec \Sigma_{\cvec l_k} \cvec C^T \right), \end{align}

which will be used instead of the standard dynamics equations.

We follow the same factorization assumptions as in <cit.> and only estimate the diagonal elements of the block matrices of the covariance matrix of the belief, see Appendix B. The update equations for the Kalman prediction and observation updates are therefore equivalent to the RKN becker2019recurrent.

§.§ Inference in the Slow-Time Scale SSM

Prediction Update.
We follow the same Gaussian inference scheme as for the fts-SSM, i.e., we again employ a linear dynamics model
$p(\cvec l_{k}|\cvec l_{k-1}, \cvec \alpha_k) = \mathcal{N}(\cvec X \cvec l_{k-1} + \cvec Y \cvec \alpha_{k}, \cvec S),$ where $\cvec X$, $\cvec Y$ and $\cvec S$ are learnable parameters. The marginalized transition model for the abstract actions is then given by

\begin{align}p_{\cvec \alpha_k}(\cvec{l}_{k}|\cvec{l}_{k-1}) & = \int p(\cvec{l}_{k}|\cvec{l}_{k-1},\cvec{\alpha}_{k}) p(\cvec \alpha_k) d \cvec \alpha_k
= \mathcal{N}\left(\cmat{X}\cvec{l}_{k-1} + \cmat{Y}\cvec{\mu}_{\alpha_k}, \cmat{S} + \cvec Y \cvec \Sigma_{\cvec \alpha_k} \cvec Y^T \right). \end{align}

We can directly use this transition model to obtain the Kalman prediction update which computes the prior belief $p_{\cvec \alpha_{1:k}}(\cvec l_k | \cvec \beta_{1:k-1}) = \mathcal{N}(\cvec \mu_{l_k}^-, \cvec \Sigma_{l_k}^-)$ from the posterior belief  $p_{\cvec \alpha_{1:k-1}}(\cvec l_{k-1} | \cvec \beta_{1:k-1}) = \mathcal{N}(\cvec \mu_{l_{k-1}}^+, \cvec \Sigma_{l_{k-1}}^+)$ of the previous time window, see Appendix <ref>.

Observation/Task Update.
Similarly, we will use a linear observation model for the abstract observations
$p(\cvec \beta_{k,t}| \cvec l_k) = \mathcal{N}(\cmat H \cvec l_{k}, \textrm{diag}(\cvec \nu_{k,t}))$ with $\cmat H = [\cmat I, \cmat 0]$.
As can be seen from the definition of the observation matrix $\cmat H$, the latent space is also decomposed into its observable and unobservable part, i.e., $\cvec l_k = [\cvec u_{k}, \cvec v_{k}]$. In difference to the standard factorized Kalman observation update given in Section <ref>, we have to infer with a set of observations $\vec \beta_{k,t}$ with $t = 1 \dots H$ for a single time window $k$. While in principle, the Kalman observation update can be applied incrementally $H$ times to obtain the posterior $p_{\cvec \alpha_{1:k}}(\cvec l_{k} | \cvec \beta_{1:k}) = \mathcal{N}(\cvec \mu_{l_{k}}^+, \cvec \Sigma_{l_{k}}^+)$, such an update would be very slow and also cause numerical inaccuracies.[13]r0.55

Implementation of task update layer which performs posterior latent task inference in the sts-SSM.

Hence, we devise a new permutation invariant version of the update rule that allows parallel processing with set encoders zaheer2017deep. We found that this update rule is easier to formalize using precision matrices. Hence, we first transform the prior covariance vectors $\cvec \sigma_{l_{k}}^{u-}$, $\cvec \sigma_{l_{k}}^{l-}$ and $\cvec \sigma_{l_{k}}^{s-}$ to its corresponding precision representation $\cvec \lambda_{l_{k}}^{u-}$, $\cvec \lambda_{l_{k}}^{l-}$ and $\cvec \lambda_{l_{k}}^{s-}$ which can be performed using block-wise matrix inversions of $\cvec \Sigma_{l_{k}}^-$. Due to the factorization of the covariance matrix, this operation can be performed solely by scalar inversions.

As the update equations are rather lengthy, they are given in Appendix <ref>, <ref>. Subsequently, we compute the posterior precision, where only $\cvec \lambda_{l_{k}}^{u}$ is changed by

\begin{align}\cvec \lambda_{l_{k}}^{u+} = \cvec \lambda_{l_{k}}^{u-} + \sum_{t=1}^H \cvec 1 \oslash \cvec \nu_{k,t}\end{align}

while $\cvec \lambda_{l_{k}}^{l+} = \cvec \lambda_{l_{k}}^{l-}$ and $\cvec \lambda_{l_{k}}^{s+} = \cvec \lambda_{l_{k}}^{s-}$ remain constant. The operator $\oslash$ denotes the element-wise division. From the posterior precision, we can again obtain the posterior covariance vectors $\cvec \sigma_{l_{k}}^{u+}$, $\cvec \sigma_{l_{k}}^{l+}$ and $\cvec \sigma_{l_{k}}^{s+}$ using only scalar inversions, see Appendix <ref>, <ref>. The posterior mean $\cvec{\mu}_{l,k}^{+}$ can now be obtained from the prior mean $\cvec{\mu}_{l,k}^{-}$ as

\begin{equation}
\begin{aligned}[c]
\cvec{\mu}_{l,k}^{+}=\cvec{\mu}_{l,k}^{-} + \left[ \begin{array}{l}
\cvec{\sigma}_{l_k}^{u+} \\
\cvec{\sigma}_{l_k}^{s+} \\
\end{array}\right]  \odot \left[ \begin{array}{l}
\sum_{t=1}^{H}  \left(\cvec{\beta}_{k,t}-\cvec{\mu}^{\mathrm{u},-}_{l_k}\right) \oslash \cvec \nu_{k,t} \\
\sum_{t=1}^{H}  \left(\cvec{\beta}_{k,t}-\cvec{\mu}^{\mathrm{u},-}_{l_k}\right) \oslash \cvec \nu_{k,t} \\
\end{array}\right].   \hspace{1.5cm}\\
\end{aligned}
\end{equation}

For $H = 1$, i.e a single observation, the given equation is equivalent to the factorized Kalman observation update becker2019recurrent. Furthermore, the given rule constitutes a unification of the batch update rule for Bayesian aggregation volpp2020bayesian and the incremental Kalman update for our factorization of belief state representation becker2019recurrent detailed in Section <ref> and Appendix <ref>.

§.§ A General Definition For an  N-level MTS3

The human brain's internal representation of the real world is complex, comprising more than just two distinct causal hierarchies. Neuroscience research hohwy2013predictive,friston2008hierarchical,jiang2021predictive suggests that these systems exhibit multiple hierarchical structures. Consequently, we propose a formal definition for Multi-Hierarchy Theory of Mind (MTS3) that accommodates an arbitrary number (N) of hierarchies.

An N-level MTS3 can be defined as a family of N-state space models, $\{S_0, S_1, ..., S_{N-1}\}$. Each of the state space models $S_i$ is given by $S_i = (Z_i, A_i, O_i, f_i, h_i, H_i \Delta t, L_i)$, where $Z_i$ is the state space, $A_i$ the action space, and $O_i$ the observation space of the SSM. The parameter $H_i \Delta t$ denotes the discretization time step and $f_i$ and $h_i$ the dynamics and observation models, respectively. Here, $l_i \in L_i$ is a task descriptor that parameterizes the dynamics model of the SSM and is constant for a local window of steps $H_{i+1}$. $l_i$ is a function of the latent state of the SSM one level above it, i.e., $S_{i+1}$. The boundary cases can be defined as follows: for $i=0$, $H_0 = 1$. Similarly, for $i=N-1$, the latent task descriptor $L_i$ is an empty set. For all $i$, $H_i < H_{i+1}$.

§.§ Inference In N-Level MTS3

As seen from the definition, an N-level MTS3 is a set of N SSMs that are strictly feudal dayan1992feudal from top to bottom. Top-level SSM (managers) makes decisions/predictions independently of the bottom-level SSMs, while the bottom-level SSM (worker) is conditioned on its immediate manager. Each of these SSMs performs inference via two stages at every time-step, i.e,

* Gaussian Conditioning Updating the posterior belief based on incoming observations.

* Gaussian Marginalization Estimating the prior belief for the next step via learned linear dynamics. Note that this stage is conditioned on the current belief of an immediate manager if present.

The efficient Gaussian conditioning and marginalization processes for State-Space Models (SSMs) across all levels in the MTS3 model are specific instances of two broad Gaussian principles derived within the thesis. These principles are formally presented as Gaussian Identity <ref> and Gaussian Identity <ref>, as detailed below:

[Gaussian Conditioning]gigausscond

Consider the graphical model given in Figure <ref>, where a set of N conditionally i.i.d observations $\cvec{\Bar{r}} = \{\cvec{r}_i\}_{i=1}^N$ are generated by a latent variable $\cvec{l}$ and the observation model $p(\cvec{r}_i|\cvec{l}) = \mathcal{N}\left(\cvec{r}_i \mid \cmat{H} \cvec{l}, \textrm{diag}(\cmat{\sigma}_i^{obs})\right)$. Assuming an observation model $\cmat{H}=[\cmat{I},\cmat{0}]$, the mean ($\cvec{\mu}$) and precision matrix ($\cvec{\Lambda}$) of the posterior over the latent variable $\cvec{l}$, $p(\cvec{l}|\cvec{\Bar{r}}) = \mathcal{N}\left( \cvec{\mu}_l^{+}, \cmat{\Sigma}_l^{+} \right) = \mathcal{N}\left( \cvec{\mu}_l^{+}, (\cmat{\Lambda}_l^{+})^{-1} \right)$, given the prior $p_0(\cvec{l}) = \mathcal{N}\left( \cvec{\mu}_l^{-}, \cmat{\Sigma}_l^{-} \right) = \mathcal{N}\left( \cvec{\mu}_l^{-}, (\cmat{\Lambda}_l^{-})^{-1} \right)$ have the following permutation invariant closed form updates.

\begin{equation}
\begin{aligned}
\cmat{\Lambda}_l^{+} & = \cmat{\Lambda}_l^{-} + \left[ \begin{array}{ll}
\textrm{diag}(\sum_{i=1}^n\frac{1}{\cvec{\sigma}_i^{obs}}), &\cvec{0 }  \\
\cvec{0 },& \cvec{0} \\
\end{array}\right] \\
\cvec{\mu}_l^{+}&=\cvec{\mu}_l^{-} + \left[ \begin{array}{l}
\cvec{\sigma}_l^{u+} \\
\cvec{\sigma}_l^{s+} \\
\end{array}\right]  \odot \left[ \begin{array}{l}
\sum_{i=1}^{N}  \left(\cvec{r}_i-\cvec{\mu}^{\mathrm{u},-}_l\right) \odot \frac{1}{\cvec{\sigma}_i^{obs}} \\
\sum_{i=1}^{N}  \left(\cvec{r}_i-\cvec{\mu}^{\mathrm{u},-}_l\right) \odot \frac{1}{\cvec{\sigma}_i^{obs}
} \\
\end{array}\right]   \hspace{1.5cm}\\
\end{aligned}
\label{eq: ba2}
\end{equation}

The proof for the above identity is given in Appendix <ref>.

The closed form updates for the resulting posterior distribution $p(l|\Bar{r})$ is permutation invariant with respect to the observation set $\Bar{r}$.

Corollary <ref> is also a direct consequence of Theorem 1, which states that the sequential Bayesian update is permutation invariant. The permutation invariance/exchangeability property has both computational and theoretical implications in efficient posterior inference over latent task variables, which is further discussed in the section.

[Permutation Invariance Of Bayesian Inversion]thmperminv

For any conditionally i.i.d model where you have a global parameter $\theta$, and a set of observations $X= \{x_i\}_{i=1}^N$ drawn conditionally i.i.d from a distribution $p(X\mid\theta)$, then for any permutation $\pi$,  the posterior $p\left(\theta \mid x_1, \ldots, x_N\right) = p\left(\theta \mid \pi(x_1), \ldots, \pi(x_N)\right)$. Thus the posterior $p\left(\theta \mid X\right)$ is permutation invariant with respect to the set $X$.

The proof of the above theorem can be found in Appendix <ref>.

The derived update equations can be coded as a layer in the neural network architecture as shown in Figure <ref>.

Gaussian Conditioning Process and Implementation. Left: A graphical model illustrating the generative process underlying Gaussian conditioning, utilized for Bayesian inversion to infer distributions over model parameters (\(l\)) from observed data (\(\cvec{\Bar{o}} = \{\cvec{o}_i\}_{i=1}^N\)). This model supports various operations including observation updates, task updates, and abstract action updates, each tailored to specific conditions such as different numbers of observations (\(N\)) or observation models (\(H\)). Right: The Gaussian conditioning within neural network architectures across the thesis is implemented as a network layer to perform dynamic parameter updating and end-to-end learning using loss functions of choice.

[Linear Combination Gaussian Marginalization]gigaussmarg

Consider the graphical model in Figure <ref>, where a set of N normally distributed independent random variables $\Bar{u} = \{ \cvec{u_i} \sim \mathcal{N}(\cvec{\mu}_{\cvec{u_i}},\cmat{\Sigma}_{\cvec{u_i}})\}_{i=0}^N$ forms a “common effect/V Structure” with a latent variable $\cvec{y}$. If the conditional distribution $p\left(\cvec{y}|\cvec{u_1}, \cvec{u_2}, . ., \cvec{u_N}\right) = \mathcal{N}( \sum_{i=0}^N \cmat{A_i}\cvec{u_i}, \cmat{\Sigma})$, then marginal $p(\cvec{y}) = \int p(\cvec{y}|\cvec{u_1}, \cvec{u_2}, . ., \cvec{u_N})\prod_{i=1}^Np(\cvec{u_i})d\cvec{u_i} = \mathcal{N}(\sum_{i=1}^N \cmat{A_i}\cvec{\mu}_{\cvec{u_i}} , \cmat{\Sigma} + \sum_{i=1}^N \cmat{A_i}\cmat{\Sigma}_{u_i}\cmat{A_i}^T) $.

The proof for the above identity is given in Appendix <ref>.

Gaussian Marginalization Process and Implementation: Left: A graphical model illustrating the common effect of \(N\) normally distributed independent causes \(\{\cvec{c_i}\}\) on a latent variable \(\cvec{e}\). These causes are marginalized out to obtain \(p(\cvec{e})\). The marginalization can be performed in closed form as per Identity <ref>. The “predict” step across time scales in every SSM formalism proposed in the thesis is an instance of this operation. Right: The Gaussian marginalization within neural network architectures across the thesis is implemented as a neural network layer (with learnable parameters) to perform Gaussian marginalization/model averaging and end-to-end learning.

Even though our experiments focus on MTS3 models with 2 hierarchies, extensive experimentation with more hierarchies can be taken as future work.

§ MTS3 AS A HIERARCHICAL WORLD MODEL

Using hierarchical internal world models to envision the future: Our predictions
about future world states rely on top-down predictions at multiple time scales and abstractions. Initially, the model generates abstract predictions at the highest level, denoted $l_k$, using abstract actions ($\alpha_k$): this process is symbolized by blue arrows. Subsequently, each higher-level abstract state adjusts the more detailed, lower-level granular states ($z_{k,t}$), represented by red arrows, over a period determined by the time scale parameter $H$. This mechanism allows the higher level to effectively make "predictions about predictions" concerning the lower level.

MTS3 allows for a natural way to build world models that can deal with partial observability, nonstationarity, and uncertainty in long-term predictions, properties which are critical for model-based control and planning. Furthermore, introducing several levels of latent variables, each working at a  different time scale allows us to learn world models that can make action conditional predictions/“dreams” at multiple time scales and multiple levels of state and action abstractions.

§.§ Conditional Multi Time Predictions With World Model

Conditional multistep ahead predictions involve estimating plausible future states of the world resulting from a sequence of actions. Our principled formalism allows for action-conditional future predictions at multiple levels of temporal abstraction. The prediction update for the sts-SSM makes prior estimates of future latent task variables conditioned on the abstract action representations. However, the task conditional prediction update in the fts-SSM estimates the future prior latent states, conditioned on primitive actions and the inferred latent task priors, which are decoded to reconstruct future observations.
To initialize the prior belief $p(\cvec z_{k,1})$ for the first time step of the time window $k$, we use the prior belief $p(\cvec z_{k-1, H+1})$ for the last time step of the time window $k-1$.

§.§ Optimizing the Predictive Log-Likelihood

The training objective for the MTS3 involves maximizing the posterior predictive log-likelihood which is given below for a single trajectory, i.e.,

\begin{align}
\label{eq:objective}
L & = \sum_{k=1}^N \sum_{t=1}^H \log p(\cvec{o}_{k,t+1}|\cvec{\beta}_{1:k-1},\cvec{\alpha}_{1:k},\cvec{w}_{k,1:t}, \cvec{a}_{k,1:t}) \nonumber \\ &
= \sum_{k=1}^N \sum_{t=1}^H \log \iint p(\cvec{o}_{k,t+1}|\cvec{z}_{k,t+1})  p(\cvec{z}_{k,t+1}|\cvec{w}_{k,1:t}, \cvec a_{k,1:t}, \cvec l_k) \\ & \hspace{4cm}p(\cvec l_k|\cvec{\beta}_{1:k-1},\cvec{\alpha}_{1:k})d\cvec{z}_{k,t+1} d \cvec l_k \nonumber \\
& =  \sum_{k=1}^N \sum_{t=1}^H \log \int p(\cvec{o}_{k,t+1}|\cvec{z}_{k,t+1})  p_{\cvec l_k}(\cvec{z}_{k,t+1}|\cvec{w}_{k,1:t}, \cvec a_{k,1:t}) d\cvec{z}_{k,t+1}.
\end{align}

The extension to multiple trajectories is straightforward and was omitted to keep the notation uncluttered. Here, $\cvec{o}_{k,t+1}$ is the ground-truth observations at time step $t+1$ and time window $k$ that needs to be predicted from all (latent and abstract) observations up to time step $t$. The prior belief corresponding to the latent state $p_{\cvec l_k}(\cvec{z}_{k,t+1}|\cvec{w}_{k,1:t}, \cvec a_{k,1:t})$ has a closed form solution as discussed in Section <ref>.

We employ a Gaussian approximation of the posterior predictive log-likelihood of the form $ p(\cvec{o}_{k,t+1}|\cvec{\beta}_{1:k-1},\cvec{\alpha}_{1:k},\cvec{w}_{k,1:t}, \cvec{a}_{k,1:t}) \approx \mathcal{N}(\cvec{\mu}_{\cvec{o}_{k,t+1}},\textrm{diag}(\cvec{\sigma}_{\cvec{o}_{k,t+1}}))$ where we use the mean of the prior belief $\cvec{\mu}_{z_{k,t+1}}^-$ to decode the predictive mean, i.e, $\cvec{\mu}_{\cvec{o}_{k,t+1}} =  \textrm{dec}_{\cvec{\mu}}(\cvec{\mu}_{z_{k,t+1}}^{-})$ and the variance estimate of the prior belief to decode the observation variance, i.e., $\cvec{\sigma}_{o_{k,t+1}} = \textrm{dec}_{\sigma}(\cmat{\Sigma}_{z_{k,t+1}}^{-})$. This approximation can be motivated by a moment matching perspective and allows for end-to-end optimization of the logarithmic likelihood without using auxiliary objectives such as the ELBO <cit.>.

Gradients are computed using (truncated) backpropagation over time (BPTT) werbos1990backpropagation and clipped.
We optimize the objective using the Adam kingma2014adam stochastic gradient descent optimizer with default parameters. We refer to Appendix A for more details. For training, we also initialize the prior belief $p(\cvec z_{k,1})$ with the prior belief $p_{\cvec l_{k-1}}(\cvec z_{k-1,H+1}|\cvec w_{k-1,1:H},\cvec a_{k-1,1:H})$ from the previous time window $k-1$. However, we cut the gradients for the fast time scale between time windows as this avoids vanishing gradients, and we observed a more stable learning behavior. Yet, the gradients can still flow between time windows for the fts-SSM via the sts-SSM.

Schematic of a 2-Level MTS3 Architecture. Inference in MTS3 takes place via closed-form equations derived using exact inference, spread across two-time scales. For the fast time scale (fts) SSM, these include the task conditional state predict and observation update stages as discussed in Section <ref> of the main paper. Whereas, for the slow time scale (sts) SSM, these include the task prediction and task update stages which are described in Section <ref>. More implementation details can also be found in Appendix <ref>.

§.§ Imputation Based Self Supervised Training For Long Term Prediction

Using the given training loss results in models that are good in one-time step prediction, but typically perform poorly in long-term predictions as the loss assumes that observations are always available up to time step $t$. To increase the performance of the long-term prediction, we can treat the long-term prediction problem as a case of the problem of “missing value”, where the missing observations occur in future time steps. Thus, to train our model for long-term prediction, we randomly mask a fraction of observations and explicitly task the network to impute the missing observations, resulting in a strong self-supervised learning signal for long-term prediction with varying prediction horizon length. This imputation scheme is applied at both time scales, masking out single time steps or whole time windows of length H. The imputation mask is also randomly resampled for every minibatch.

Self-supervised training for multi-step ahead predictions by tasking the model to "fill in" the masked observations. Here masking is done at both time scales based on the available observation/set of observations in each time window.

§ EXPERIMENTS

In this section, we evaluate our approach to a diverse set of simulated and real-world dynamical systems for long-horizon prediction tasks. Our experiments are designed to answer the following questions. (a) Can MTS3 make accurate long-term deterministic predictions (mean estimates)? (b) Can MTS3 make accurate long-term probabilistic predictions (variance estimates)? (c) How important are the modelling assumptions and training scheme?

§.§ Baseline Dynamics Models

While a full description of our baselines can be found in Appendix <ref>, a brief description of them is given here: (a) RNNs - We compare our method to two widely used recurrent neural network architectures, LSTMs lstm and GRUs gru.
(b) RSSMs - Among several RSSMs from the literature, we chose RKN becker2019recurrent and HiP-RSSM shaj2022hidden as these have shown excellent performance for dynamics learning for short-term predictions and rely on exact inference as in our case. (c) Transformers - We also compare with two state-of-the-art Transformer vaswani2017attention variants. The first variant (AR-Transformer) relies on a GPT-like autoregressive prediction radford2019language,brown2020language. Whereas the second variant (Multi-Transformer) uses direct multi-step loss zeng2022transformers from recent literature on long horizon time-series forecasting zhou2021informer,liu2022nonstationary,nie2023a. Here, multistep ahead predictions are performed using a single shot given the action sequences.

§.§ Environments and Datasets

D4RL Environments: (left) HalfCheetah, (middle) Franka Kitchen, (right) Maze2D-Medium(left) JCB Hydradig 110W Excavator (right) Franka Emika Panda Robot

We experiment with three broad datasets. While full descriptions of these datasets, dataset creation procedure, and overall statistics are given in Appendix E, a brief description of them is as follows.
(a) D4RL Datasets - We use a set of 3 different environments/agents from D4RL dataset fu2020d4rl, which includes the HalfCheetah, Medium Maze and Franka Kitchen environment. Each of these was chosen because of their distinct properties like sub-optimal trajectories (HalfCheetah), realistic domains / human demonstrations (Kitchen), multi-task trajectories, non-markovian collection policies (Kitchen and Maze) and availability of long horizon episodes (all three). (b) Manipulation Datasets - We use 2 datasets collected from a real excavator arm and a Panda robot. The highly non-linear non-markovian dynamics due to hydraulic actuators in the former and non-stationary dynamics owing to different payloads in the latter make them challenging benchmarks. Furthermore, accurate modelling of the dynamics of these complex systems is important since learning control policies for automation directly on large excavators is economically infeasible and potentially hazardous. (c) Mobile Robotics Dataset - We set up a simulated four-wheeled mobile robot traversing a highly uneven terrain of varying steepness generated by a mix of sinusoidal functions. This problem is challenging due to the highly non-linear dynamics involving wheel-terrain interactions and non-stationary dynamics introduced by varying steepness levels. In all datasets, we only use information about agent/object positions and we mask out velocities to create a partially observable setting.

Comparison with baselines in terms of RMSE for long horizon predictions (in seconds) as discussed in Section <ref>.

§.§ Can MTS3 make accurate long-term deterministic predictions (mean estimates)?

Here we evaluate the quality of the mean estimates for long-term prediction using our approach. The results are reported in terms of "sliding window RMSE" in Figure <ref>. We see that MTS3 gives consistently good long-term action conditional future predictions on all 6 datasets. Deep Kalman models becker2019recurrent,shaj2022hidden which operate on a single time scale fail to give meaningful mean estimates beyond a few milliseconds. Similarly, widely used RNN baselines lstm,gru which form the backbone of several world models ha2018world,hafner2019learning give poor action conditional predictions over long horizons. AR-Transformers also fail possibly due to error accumulation caused by the autoregression. However, Multi-Transformers are a strong baseline that outperforms MTS3 in the Medium Maze and Panda dataset by a small margin. However, on more complex tasks like the Kitchen task, which requires modelling multi-object, multi-task interactions gupta2019relay, MTS3 is the only model that gives meaningful long horizon predictions. A detailed description of the metric "sliding window RMSE" is given in Appendix <ref>. A visualization of the predicted trajectories vs. ground truth is given in Appendix <ref>.

Visualizations of the predicted trajectories vs ground truth for MTS3 and GPT like autoregressive transformer (AR-Transformer). More visualizations across algorithms and datasets can be found in Appendix <ref>.

§.§ Can MTS3 make accurate long-term probabilistic predictions (variance estimates)?

Next, we examine the question of whether the principled probabilistic inference translates to accurate uncertainty quantification during long-horizon predictions. We trained all the baselines with a negative
log-likelihood loss and used the same as a metric to quantify the quality of uncertainty estimates. During the evaluation we rely on a sliding window approach (see Appendix <ref>) and report the results for the last timestep in Table <ref>. As seen in Table <ref>, MTS3 gives the most accurate uncertainty estimates in all datasets except Medium Maze, where it is outperformed by Multi-Transformer. Also, notably, AR-Transformers and deep Kalman models fail to learn any meaningful uncertainty representation when it comes to long-term predictions.

Prediction     7c|Algorithm
Horizon     MTS3     Multi-Trans     AR-Trans     LSTM     GRU     RKN     HiP-RSSM
Half Cheetah     6 s     $\mathbf{-2.80 \pm 0.30}$     $0.25 \pm 0.05$         $7.34 \pm 0.06$     $7.49 \pm 0.04$
Kitchen     2.5 s     $\mathbf{-25.74 \pm 0.12}$     $-7.3 \pm 0.2$         $32.45 \pm 1.64$     $32.72 \pm 0.65$
Medium Maze     4 s     $-0.21 \pm 0.022$     $\mathbf{-0.88 \pm 0.02}$         $4.03 \pm 0.32$     $7.76 \pm 0.07$
Panda     1.8 s     $\mathbf{2.79 \pm 0.32}$     $3.77 \pm 0.33$         $7.94 \pm 0.39$     $7.91 \pm 0.23$
Hydraulic     12 s     $\mathbf{-2.64 \pm 0.12}$     $-2.46 \pm 0.03$         $7.35 \pm 0.061$     $7.35 \pm 0.06$
Mobile Robot     3 s     $\mathbf{-6.47 \pm 0.71}$     $-5.17 \pm 0.23$         $11.27 \pm 2.3$     $14.55 \pm 5.6$

Comparison in terms of Negative Log Likelihood (NLL) for long horizon predictions (in seconds). Here bold numbers indicate the top methods and  denotes very high/nan values resulting from the highly divergent mean/variance long-term predictions.

§.§ How important are the modelling assumptions and training scheme?

Now, we look at three important modelling and training design choices: (i) splitting the latent states to include an unobservable “memory” part using observation model $h^{sts}=h^{fts}=\cmat H = [\cmat I, \cmat 0]$ as discussed in Sections <ref> and <ref>, (ii) action abstractions in Section <ref>, (iii) training by imputation.

Ablation on discretization step $H. \Delta t$ (a) The long-term prediction results in terms of RMSE, with different $H$ values as discussed in Section <ref> on the hydraulics dataset.

To analyze the importance of the memory component, we derived and implemented an MTS3 variant with an observation model of $h^{sts}=h^{fts}=\cmat I$ and a pure diagonal matrix representation for the covariance matrices. As seen in Figure <ref>, this results in worse long-term predictions, suggesting that splitting the latent states in its observable and unobservable part in MTS3 is critical for learning models of non-markovian dynamical systems. Regarding (ii), we further devised another variant where MTS3 only had access to observations, primitive actions and observation abstractions, but no action abstractions. As seen in our ablation studies, using the action abstraction is crucial for long-horizon predictions.

Our final ablation (iii) shows the importance of an imputation-based training scheme discussed in Section <ref>. As seen in Figure <ref> when trained for 1 step ahead predictions without imputation, MTS3 performs significantly worse for long-term prediction suggesting the importance of this training regime.

Ablation on discretization step $H. \Delta t$. The predictions by MTS3 variants with different values of timescale parameter $H.\Delta t$ on a trajectory picked from the hydraulics excavator dataset. The top images/mts3/ are for $H=3$ and
$H=10$. Bottom images/mts3/ are for $H=30$ and $H=75$. Note that the results reported in the paper are with $H=30$.

§.§ What is the role of the discretization step $H.\Delta t$?

Finally, we perform ablation for different values of $H.\Delta t$, which controls the time scale of the task dynamics. The results reported are for the hydraulics dataset. The higher the value of H, the slower the timescale of the task dynamics relative to the state dynamics. As seen in Figure Figure <ref>, smaller values of $H$ (2,3,5 and 10) give significantly worse performance. Very large values of $H$ (like 75) also result in degradation of performance. To further get an intuitive understanding of the MTS3's behaviour under different timescales, we plot the predictions given by MTS3 for different values of $H$ on a trajectory handpicked from the hydraulics excavator dataset. As seen in Figure <ref> for large values of $H$ like 30 and 75, we notice that the slow-changing task dynamics "reconfigures" the fast dynamics every 30 and 75-step window respectively, by conditioning the lower level dynamics with the newly updated task prior. This effect is noticeable as periodic jumps or discontinuities in the predictions, occurring at 30 and 75-step intervals. Also, for a very large $H$ like 75, the fast time scale ssm has to make many more steps in a longer window resulting in error accumulation and poor predictions. A similar trend was observed for the mobile robot dataset. As seen in Figure <ref>, smaller values of H (like 2,3,5 and 10) give significantly worse performance. Very large values of H (like 150) also result in degradation of performance. In the paper, we used a value of H=75.

Ablation on discretization step $H. \Delta t$. The long-term prediction results in terms of RMSE, with different $H$ values as discussed in Section <ref> on (left) the hydraulics dataset and (right) the mobile robot dataset.

§ CONCLUSION AND FUTURE WORK

In this work, we introduce MTS3, a probabilistic formalism for learning the dynamics of complex environments at multiple time scales. By modelling the dynamics of the world at multiple levels of temporal abstraction we capture both the slow-changing long-term trends and fast-changing short-term trends in data, leading to highly accurate predictions spanning several seconds into the future. Our experiments demonstrate that simple linear models with principled modelling assumptions can compete with large transformer model variants that require several times more parameters. Furthermore, our inference scheme also allows for principled uncertainty propagation over long horizons across multiple time scales which capture the stochastic nature of environments. We believe our formalism can benefit multiple future applications including hierarchical planning/control. We discuss the limitations and broader impacts of our work in Chapter <ref>.

CHAPTER: OUTLOOK

This thesis questioned the current formalisms for learning-world models due to their inability to capture some critical criteria for a foundational world model. Specifically, we looked at three aspects (i) scalable probabilistic modelling that captures causal relations in our world (Chapter <ref>, <ref> and <ref>), (ii) adaptability to changing tasks (Chapter <ref>) and (iii) hierarchical modelling at multiple time scales and temporal abstractions (Chapter <ref>).
In this context the thesis proposed two new formalisms Hidden Parameter State Space Models (HiP-SSM in Chapter <ref>) and Multi Time Scale State Space Model (MTS3 in Chapter <ref>) besides making an existing formalism of SSMs more robust for performing interventions/counterfactuals with action/control signals (Ac-SSM in Chapter <ref>).

The proposed formalisms are in line with related theory in computational neuroscience called predictive processing and Bayesian brain hypothesis (Chapter <ref>). In the proposed generative models, Bayesian inversion is employed to deduce the underlying causes from the incoming sensory data at various levels of abstraction. The Bayesian inversion results in update rules that use "precision weighting" friston2009free,hohwy2013predictive,seth2014cybernetic as a mechanism of directing "attention" to the relevant sensory information. The idea of "precision estimation" by learned sensory encoders is also in line with discussions in the broader neuroscience community. The world model formalisms proposed are adaptable to changing situations (embodying cognitive flexibility) and can make top-down long-horizon predictions based on nested causal hierarchies.

Machines that can replicate human intelligence and type 2 reasoning capabilities should be able to reason at multiple
levels of temporal abstraction using internal world models lecun2022path,gupta2024essential. The work represents a hopeful step towards developing embodied AI systems with a deeper, more intuitive grasp of the world, aligning machine learning processes more closely with natural intelligence. We believe that the findings of this dissertation will contribute to ongoing research in the creation of robust, principled, and scalable world models, while acknowledging the vast scope for further research and development in this area.

Limitations and Future Work
The thesis specifically concentrated on the aspect of prediction rather than control or planning. Addressing hierarchical planning, much like hierarchical prediction, remains a significantly underexplored challenge. A logical progression for future research would be to extend the concept of planning within the exact inference framework botvinick2012planning,watson2020stochastic. This approach necessitates the derivation of backward messages and the use of the Kalman duality principle within the learned latent spaces of a deep encoder. Furthermore, investigating variational approaches to the hierarchical models presented, and considering hierarchical planning/control as a problem of approximate inference toussaint2009probabilistic,friston2009free,millidge2020relationship, presents another intriguing avenue for research.

A further limitation noted in the thesis is its reliance on proprioceptive sensors for experimentation, without empirical validation using high-dimensional sensory information, such as vision. An important avenue for future research would involve evaluating the usefulness and characteristics of abstractions derived through Bayesian inversion/aggregation when applied to high-dimensional and  multimodal sensory inputs. Conducting experiments with image-based data and exploring "non-reconstruction" based loss functions, as suggested by <cit.>, represent promising directions for upcoming work.

Though our hierarchical formalism outperform transformer variants on long horizon predictions on several tasks, a limitation for these to be used as foundational world models is the computational bottleneck due to the sequential nature of dynamics. However, there have been exciting progress recently in deterministic linear SSMs gu2021efficiently,smith2022simplified,mondal2023efficient that allow efficient parallelization during training. Since our dynamics is linear similar to these approaches, parallelization using similar techniques as employed by deterministic SSMs can be another important direction of future research.

Our experimentation was limited to the hierarchical models with just two levels of abstraction, as this configuration proved to be adequate for numerous tasks we carried out. Nevertheless, there is room for deeper exploration.The intersection of machine learning with broader cognitive theories, particularly the concept of machine consciousness, presents a novel yet underexplored frontier in the field. Traditionally, consciousness has been a topic more familiar to the realms of philosophy and psychology than to machine learning and computer science. This often makes it a delicate subject within our discipline.

Perception and consciousness as hierarchical inference from <cit.>.

As I write this thesis in the Spring of 2024, the machine learning field finds itself at a crossroads, significantly influenced by the impressive capabilities of attention/transformer vaswani2017attention based large language models (LLMs). These advancements suggests a pressing need for mathematical models or analytical frameworks to better understand machine consciousness bengio2017consciousness,chalmers2023could. Numerous theories on consciousness highlight the critical role of recurrent processing chalmers2023could,lamme2010neuroscience, a feature notably absent in current LLM architectures. In this context, I believe that hierarchical generative models, particularly those extending beyond the 2-level temporal depth as explored in this thesis, represent a promising research direction for embedding consciousness priors friston2013consciousness,bengio2017consciousness into machine learning systems. This approach could offer profound insights into the mechanisms underpinning consciousness and its potential replication in artificial systems.

Missing 'biblatex' package
The bibliography requires the 'biblatex' package.

booktitleInternational Conference on Artificial Intelligence and Statistics
titleExplore the Context: Optimal Data Collection for Context-Conditional Dynamics Models

titleGPyOpt: A Bayesian Optimization framework in python

Nature Publishing Group UK London

titleParallel visual computation

journaltitleRobotics and Autonomous Systems
titleActive learning of inverse models with intrinsically motivated goal exploration in robots

journaltitlearXiv preprint arXiv:2210.09256
titleOn Uncertainty in Deep State Space Models for Model-Based Reinforcement Learning

booktitleInternational Conference on Machine Learning
titleRecurrent kalman networks: Factorized inference in high-dimensional deep feature spaces

booktitleInternational Conference on Machine Learning
titleSwitching linear dynamics for variational bayes filtering

journaltitlearXiv preprint arXiv:1709.08568
titleThe consciousness prior

journaltitleMachine learning
titlePattern recognition

journaltitleTrends in cognitive sciences
titlePlanning as inference

journaltitleAdvances in neural information processing systems
titleLanguage models are few-shot learners

journaltitlearXiv preprint arXiv:2006.05935
titleLearning to Play Table Tennis From Scratch using Muscular Robots

booktitle2016 IEEE International Conference on Robotics and Automation (ICRA)
titleA lightweight robotic arm with pneumatic muscles for robot learning

journaltitleIEEE Robotics & Automation Magazine
titleLearning and reproduction of gestures by imitation

journaltitlearXiv preprint arXiv:2303.07103
titleCould a large language model be conscious?

Association for Computational Linguistics

booktitleProceedings of SSST-8, Eighth Workshop on Syntax, Semantics and Structure in Statistical Translation
titleOn the Properties of Neural Machine Translation: Encoder–Decoder Approaches

journaltitleAdvances in neural information processing systems
titleDeep reinforcement learning in a handful of trials using probabilistic dynamics models

Cambridge University Press

journaltitleBehavioral and brain sciences
titleWhatever next? Predictive brains, situated agents, and the future of cognitive science

titlePybullet, a python module for physics simulation for games, robotics and machine learning

titleThinking, fast and slow

journaltitleAdvances in neural information processing systems
titleFeudal reinforcement learning

MIT Press One Rogers Street, Cambridge, MA 02142-1209, USA journals-info …

journaltitleNeural computation
titleThe helmholtz machine

booktitleProceedings of the 28th International Conference on machine learning (ICML-11)
titlePILCO: A model-based and data-efficient approach to policy search

journaltitleIEEE transactions on pattern analysis and machine intelligence
titleGaussian processes for data-efficient learning in robotics and control

journaltitlearXiv preprint arXiv:1605.07127
titleLearning and policy search in stochastic dynamical systems with bayesian neural networks

journaltitlearXiv preprint arXiv:1810.04805
titleBert: Pre-training of deep bidirectional transformers for language understanding

booktitleInternational Conference on Machine Learning
titleCollapsed amortized variational inference for switching nonlinear dynamical systems

NIH Public Access

booktitleIJCAI: proceedings of the conference
titleHidden parameter markov decision processes: A semiparametric regression approach for discovering latent task parametrizations

booktitle2007 IEEE International Symposium on Intelligent Signal Processing
titleSimulation control of a biped robot with support vector regression

booktitleInternational Conference on Machine Learning
titleModel-agnostic meta-learning for fast adaptation of deep networks

journaltitleAdvances in neural information processing systems
titleUnsupervised learning for physical interaction through video prediction

journaltitlearXiv preprint arXiv:1710.05741
titleA disentangled recognition and nonlinear dynamics model for unsupervised learning

The Royal Society London

journaltitlePhilosophical transactions of the Royal Society B: Biological sciences
titleA theory of cortical responses

Public Library of Science San Francisco, USA

journaltitlePLoS computational biology
titleHierarchical models in the brain

journaltitleTrends in cognitive sciences
titleThe free-energy principle: a rough guide to the brain?

Taylor & Francis

titleConsciousness and hierarchical inference

journaltitlearXiv preprint arXiv:2004.07219
titleD4rl: Datasets for deep data-driven reinforcement learning

booktitle2016 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)
titleOne-shot learning of manipulation skills with online dynamics adaptation and neural network priors

booktitleData-efficient machine learning workshop, ICML
titleImproving PILCO with Bayesian neural network dynamics models

journaltitlearXiv preprint arXiv:1807.01622
titleNeural processes

journaltitlearXiv preprint arXiv:1805.09921
titleMeta-learning probabilistic inference for prediction

titleExplicit stochastic predictive control of combustion plants based on Gaussian process models

The Royal Society London

journaltitleProceedings of the Royal Society of London. Series B. Biological Sciences
titlePerceptual illusions and brain models

The Royal Society London

journaltitlePhilosophical Transactions of the Royal Society of London. B, Biological Sciences
titlePerceptions as hypotheses

booktitleInternational Conference on Learning Representations
titleEfficiently Modeling Long Sequences with Structured State Spaces

journaltitlearXiv preprint arXiv:1910.11956
titleRelay policy learning: Solving long-horizon tasks via imitation and reinforcement learning

journaltitlearXiv preprint arXiv:2402.06665
titleThe Essential Role of Causality in Foundation World Models for Embodied AI

journaltitlearXiv preprint arXiv:1803.10122
titleWorld models

booktitleAdvances in neural information processing systems
titleBackprop kf: Learning discriminative deterministic state estimators

booktitleInternational Conference on Machine Learning
titleLearning latent dynamics for planning from pixels

MIT Press

journaltitleNeural computation
titleMosaic model for sensorimotor learning and control

titleConcerning the perceptions in general, 1867.

MIT Press

journaltitleNeural Comput.
titleLong Short-Term Memory

OUP Oxford

titleThe predictive mind

Neural Comput

journaltitleNeural computation
titleBayesian Filtering with Multiple Internal Models: Toward a Theory of Social Intelligence

journaltitlearXiv preprint arXiv:2112.10048
titlePredictive coding theories of cortical function

Institute of Mathematical Statistics

journaltitleStatistical science
titleGraphical models

titleNew results in linear filtering and prediction theory

titleA new approach to linear filtering and prediction problems

booktitleInternational conference on artificial intelligence and statistics
titleData-efficient reinforcement learning with probabilistic model predictive control

journaltitlearXiv preprint arXiv:1605.06432
titleDeep variational bayes filters: Unsupervised learning of state space models from raw data

IOP Publishing

journaltitleNetwork: computation in neural systems
titleA forward-inverse optics model of reciprocal connections between visual cortical areas

Annual Reviews

journaltitleAnnu. Rev. Psychol.
titleObject perception as Bayesian inference

Cambridge University Press

titleIdentifiable Parameters and Optimum Configurations for Robots Calibration

journaltitleIEEE Transactions on Robotics
titleLearning stable nonlinear dynamical systems with gaussian mixture models

journaltitlearXiv preprint arXiv:1612.00475
titleTransfer learning across patient variations with hidden parameter markov decision processes

journaltitlearXiv preprint arXiv:1412.6980
titleAdam: A method for stochastic optimization

journaltitleTRENDS in Neurosciences
titleThe Bayesian brain: the role of uncertainty in neural coding and computation

booktitleProceedings 2007 ieee international conference on robotics and automation
titleGaussian processes and reinforcement learning for identification and control of an autonomous blimp

booktitleProceedings of the 2004 American control conference
titleGaussian process model based predictive control

journaltitleInformation Fusion
titleChatGPT: Jack of all trades, master of none

MIT press

titleProbabilistic graphical models: principles and techniques

MIT press Cambridge, MA

journaltitleIntroduction to statistical relational learning
titleGraphical models in a nutshell

National Acad Sciences

journaltitleProceedings of the National Academy of Sciences
titleHamiltonian systems and transformation in Hilbert space

booktitleProceedings of the AAAI Conference on Artificial Intelligence
titleStructured inference networks for nonlinear state space models

Taylor & Francis

journaltitleCognitive neuroscience
titleHow neuroscience will change our view on consciousness

journaltitleOpen Review
titleA path towards autonomous machine intelligence version 0.9. 2, 2022-06-27

Optica Publishing Group

journaltitleJOSA a
titleHierarchical Bayesian inference in the visual cortex

Rome, Italy

booktitleRobotics: Science and Systems
titleDeepMPC: Learning deep latent features for model predictive control.

booktitleArtificial Intelligence and Statistics
titleBayesian learning and inference in recurrent switching linear dynamical systems

booktitleAdvances in Neural Information Processing Systems
titleNon-stationary Transformers: Exploring the Stationarity in Time Series Forecasting

journaltitleAutomata studies
titleThe epistemological problem for automata

journaltitlearXiv preprint arXiv:1802.03426
titleUmap: Uniform manifold approximation and projection for dimension reduction

booktitleThe Eleventh International Conference on Learning Representations
titleTransformers are Sample-Efficient World Models

booktitleActive Inference: First International Workshop, IWAI 2020, Co-located with ECML/PKDD 2020, Ghent, Belgium, September 14, 2020, Proceedings 1
titleOn the relationship between active inference and control as inference

journaltitlearXiv preprint arXiv:2306.11941
titleEfficient Dynamics Modeling in Interactive Environments with Koopman Theory

journaltitlearXiv preprint arXiv:1803.11347
titleLearning to adapt in dynamic, real-world environments through meta-reinforcement learning

journaltitlearXiv preprint arXiv:1812.07671
titleDeep online learning via meta-learning: Continual adaptation for model-based rl

booktitle2018 IEEE International Conference on Robotics and Automation (ICRA)
titleNeural network dynamics for model-based deep reinforcement learning with model-free fine-tuning

Psychology press

titleCognitive psychology: Classic edition

booktitle2010 IEEE International Conference on Robotics and Automation
titleUsing model knowledge for learning inverse dynamics

titleIncremental online sparsification for model learning in real-time robot control

journaltitleAdvances in neural information processing systems
titleLocal Gaussian process regression for real time online model learning

Taylor & Francis

journaltitleAdvanced Robotics
titleModel learning with local gaussian process regression

booktitleThe Eleventh International Conference on Learning Representations
titleA Time Series is Worth 64 Words: Long-term Forecasting with Transformers

booktitleAdvances in neural information processing systems
titleAction-conditional video prediction using deep networks in atari games

journaltitleTrends in Cognitive Sciences
titleThe repurposed social brain

journaltitleTrends in cognitive sciences
titleHierarchical active inference: a theory of motivated control

booktitle2015 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)
titleReal-time deep learning of robotic manipulator inverse dynamics

journaltitleAdvances in neural information processing systems
titleMulti-time models for temporally abstract planning

booktitle2015 IEEE International Conference on Robotics and Automation (ICRA)
titleDeep learning helicopter dynamics models

journaltitleOpenAI blog
titleLanguage models are unsupervised multitask learners

journaltitleVision research
titleAn optimal estimation approach to visual perception and learning

booktitleSummer School on Machine Learning
titleGaussian processes in machine learning

journaltitlearXiv preprint arXiv:2205.13804
titleEnd-to-end learning of hybrid inverse dynamics models for precise and compliant impedance control

booktitle2017 IEEE-RAS 17th International Conference on Humanoid Robotics (Humanoids)
titleLearning Inverse Dynamics Models in O (n) time with LSTM networks

journaltitlearXiv preprint arXiv:1803.07551
titleMeta reinforcement learning with latent variable gaussian processes

journaltitleApplied Intelligence
titleScalable techniques from nonparametric statistics for real time robot learning

booktitleInternational Conference on Machine Learning
titlePlanning to explore via self-supervised world models

titleBeing you: A new science of consciousness

Open MIND. Frankfurt am Main: MIND Group

booktitleOpen mind
titleThe cybernetic Bayesian brain

journaltitleConference On Robot Learning
titleAction-Conditional Recurrent Kalman Networks For Forward and Inverse Dynamics Learning

journaltitleInternational Conference On Learning Representations
titleHidden Parameter Recurrent State Space Models For Changing Dynamics Scenarios

journaltitleAdvances in Neural Information Processing Systems
titleMulti Time Scale World Models

journaltitlearXiv preprint arXiv:2202.05839
titleAbstraction for deep reinforcement learning

booktitleThe Eleventh International Conference on Learning Representations
titleSimplified State Space Layers for Sequence Modeling

journaltitleIEEE Robotics and Automation Letters
titleAdding Terrain Height to Improve Model Learning for Path Tracking on Uneven Terrain by a Four Wheel Robot

journaltitleIEEE/ASME Transactions on Mechatronics
titleInertia Tensor Properties in Robot Dynamics Identification: A Linear Matrix Inequality Approach

booktitleMachine Learning Proceedings 1995
titleTD models: Modeling the world at a mixture of time scales

journaltitleTrends in ecology & evolution
titleSocial competence: an evolutionary approach

journaltitleControl Engineering Practice
titleState-dependent control of a hydraulically actuated nuclear decommissioning robot

booktitleProceedings of the 26th annual international conference on machine learning
titleFactored conditional restricted Boltzmann machines for modeling motion style

journaltitleKünstliche Intell.
titleProbabilistic inference as a model of planned behavior.

journaltitleJournal of machine learning research
titleVisualizing data using t-SNE.

journaltitleAdvances in neural information processing systems
titleAttention is all you need

booktitleInternational Conference on Learning Representations
titleBayesian Context Aggregation for Neural Processes

Courier Corporation

titleTreatise on Physiological Optics, volume III

booktitleConference on Robot Learning
titleStochastic optimal control as approximate input inference

journaltitleProceedings of the IEEE
titleBackpropagation through time: what it does and how to do it

booktitle2017 IEEE international conference on robotics and automation (ICRA)
titleInformation theoretic mpc for model-based reinforcement learning

journaltitlearXiv preprint arXiv:1703.06114
titleDeep sets

journaltitlearXiv preprint arXiv:2308.13067
titleCausal parrots: Large language models may talk causality but are not causal

journaltitlearXiv preprint arXiv:2205.13504
titleAre transformers effective for time series forecasting?

booktitleProceedings of the AAAI conference on artificial intelligence
titleInformer: Beyond efficient transformer for long sequence time-series forecasting

CHAPTER: APPENDIX: ACTION CONDITIONAL SSM

§ INVERSE DYNAMICS LEARNING WITH ACTION CONDITIONAL SSM

For the inverse dynamics case  we want to learn a model  where $\vec o_{t+1}$ is the desired next observation for time step $t + 1$ and $\hat{\vec{a}}_{t}$ is the predicted action, i.e., the one to be applied. We introduce an action decoder which decodes the latent posterior $\left(\vec{z}_{t}^+, \vec{\Sigma}_{t}^+ \right)$ and estimates the action required to move to the desired next observation. The action decoder also gets information regarding the next observation as an input.
During training this corresponds to the next observation in the data.
For control a desired next observation is used to obtain the action required to reach that observation. The joint angles and velocities are treated as observations in our inverse dynamics experiments.

Schematic diagram of the inverse dynamics learning architecture. Here the posterior $\left(\vec{z}_{t}^+, \vec{\Sigma}_{t}^+ \right)$, at the current time step is fed to an action decoder along with the desired target. In the next time-step, the executed action $\vec a_t$ is fed to the predict stage of RKN cell.
Further, the next predicted prior $\left(\vec{z}_{t+1}^-, \vec{\Sigma}_{t+1}^- \right)$ is decoded to get the next state $\hat{\vec{o}}_{t + 1}$.

As seen in Figure <ref>, instead of merely learning the inverse dynamics, our approach learns inverse and forward dynamics simultaneously by feeding back the executed action to the action-conditional predict stage, which predicts forward in latent space.

Note that the action decoder also gets the same action as the target during training.
However, the model sees the true action only after the prediction since the action decoder works with the posterior estimate $\left(\vec{z}_t^+, \vec{\Sigma}_t^+ \right)$.

Hence, the prediction of the inverse dynamics model for the current time step is made independent of this feedback.
We found that enforcing this causal feedback, a necessary structural component of the Bayesian network of the underlying latent dynamical system, improves the performance of the inverse model, as seen in Figure <ref>.

Loss And Training.
The dual output architecture for our inverse model leads to two different loss functions for the action and observation decoders which are optimized jointly.
Our experiments showed that the loss of the forward model is an excellent auxiliary loss function for the inverse model and learning this implicit forward model jointly with the inverse model results in much better performance for the inverse model.
We assume the reasons for this effect is that the forward model loss is providing crucial gradient information to form an informative latent state representation that is also useful for the inverse model.

In order to deal with high frequency data we again chose to predict the normalized differences to the previously executed action, i.e. $\hat{\vec{a}}_t = \vec{a}_{t-1} + \textrm{dec}_\textrm{action}(\vec{z}_t^+)$.
Here $\vec{a}_{t-1}$ is the executed action at $t-1$ and $\textrm{dec}_\textrm{action}(\vec{z}_t^+)$ the output of the actual action decoder network.
The combined loss function for the inverse dynamic case is given by

\begin{align*}
\mathcal{L}_\textrm{inv} = \sqrt{\dfrac{1}{T} \sum_{i=1}^T \parallel \left(\vec{a}_{t+1} - \vec{a}_{t}\right) - \textrm{dec}_\textrm{action}\left(\vec{z}_t^+\right) \parallel^2}    +
\lambda \mathcal{L}_\textrm{fwd} %\sqrt{ \dfrac{1}{T} \sum_{t=1}^T  \parallel \vec{o}_{t+1} - \left( \vec{o}_{t} + \textrm{dec}\left(\vec{z}_t^- \right)  \right)
\end{align*}

$\lambda$ chooses the trade-off between the inverse model loss and the forward model loss. The value of $\lambda$ is chosen via hyperparameter optimization using GPyOpt gpyopt2016.
Note that for the inverse model,  the actions are decoded based on the posterior mean of the current time step while for the forward model the observations are decoded from the prior mean of the next time step.

We evaluated the performance of the proposed method for inverse dynamics learning on two real robots, Franka Emika Panda and Barrett WAM. Barret WAM is a robot with direct cable drives. The direct cable drives produce high torques, generating fast and dexterous movements but yield complex dynamics. Rigid-body dynamics cannot model this complex dynamics accurately due to the variable stiffness and lengths of the cables nguyen2011incremental.

Franka Panda

Barret WAM

Action Feedback

(a) and (b) Joint torque prediction RMSE values in NM of Action-Conditional RKN, LSTM and FFNN for Panda and Barret WAM. A comparison is also provided with the analytical (RBD) model of Panda. (c) Comparison of ac-RKN for inverse dynamics learning with and without the action feedback as discussed in Section <ref>.

Predicted joint torques(normalized) for first 3 joints of the Panda robot arm. The learned inverse dynamics model by ac-RKN match closely with the ground truth data while the rigid body dynamics model can not capture the high-frequency variations in the data.

Joint Torque Prediction Task. We benchmark the representational capability of the latent state posterior, $\vec z_t^+$, of ac-RKN in accurately modelling the inverse dynamics of these robots in comparison to deterministic models like LSTMs and FFNN. It is clear from <ref> and <ref> that ac-RKN learns highly accurate models of these in contrast with other data-driven methods. This highly precise modelling is often a requirement for high fidelity and compliant robotic control.

Impact of Action Feedback. We also perform an ablation study with and without the action feedback for the prediction step in the latent dynamics. As seen in the Figure <ref>, the action feedback always results in better representational capability as this helps the implicit forward model in making better predictions by taking into account its causal effect on the state transitions.

Comparison to Analytical Models. Finally, we make a comparison with the analytical model of the Panda robot for inverse dynamics. Please refer to Appendix C for more details on the analytical model. As evident from Figure <ref> analytical models gave predictions with much lesser accuracy in comparison to ac-RKN, as it does not consider unmodelled effects such as joint and link flexibilities, backlash, stiction and actuator dynamics.  In such cases, the robot would continuously have to track its current position and compensate the errors with high-gain feedback control thus making it dangerous to interact with the real world and impossible to work in human-centred environments.

§ IMPLEMENTATION DETAILS

§.§ Locally Linear Transition Model
The state transitions in the predict stage of the Kalman filter is governed by a locally linear transition model. To obtain a locally linear transition model, the RKN learns $K$ constant transition matrices $\vec{A}^{(k)}$ and combines them using state dependent coefficients $\alpha^{(k)}(\vec{z_t})$, i.e.,
$ \vec{A}_t = \sum_{k=0}^K \alpha^{(k)}(\vec{z_t}) \vec{A}^{(k)}. $
A small neural network with softmax output is used to learn $\alpha^{(k)}$. Each $\vec{A}^{(k)}$ is designed to consist of four band matrices as in <cit.> in order to reduce the number parameters without affecting the performance.

§ DETAILS OF RIGID BODY DYNAMICS MODEL

The analytical model for Franka Emika Panda is a rigid-body dynamics model that was identified in its so-called base parameters Khalil1991. Due to the friction compensation in the joints, we observed that the viscous friction is negligible, whereas the observed Coulomb friction is very small yet included in our parameterization. Which results is a model with 50 parameters.
The base parameterization is computed based on provided kinematic properties of the robotic arm and provides a linear relation between the base parameters and joint torques for a given set of joint positions, velocities and accelerations.

Due to this linearity, the regression problem can be solved using a linear least-squares method, although additional linear matrix inequality constraints must be fullfilled to ensure that the resulting parameters are physically realizable Sousa2019,reuss2022end.
In order to perform forward simulation of the robot dynamics, we numerically solve an initial value problem for the implicit set of differential equations defined by base parameterization of the rigid-body model.
Note that, the model does not parameterize actuator dynamics, nor does it model joint flexibilities, link flexibilities, or stiction. The focus here is to provide a reference baseline to show which effects the acRKN captures in comparison to a text-book robot model.

CHAPTER: APPENDIX: HIDDEN PARAMETER SSM

§ PROOF FOR GAUSSIAN IDENTITY <REF>

This section provides a proof for Gaussian Identity <ref>. First we derive an expression for the joint distribution $p(\cvec{u},\cvec{v},\cvec{y})$.

[Joint Gaussian Distribution]gijointgauss2
If $\cvec{u} \sim \mathcal{N}(\cvec{\mu}_u + b,\cvec{\Sigma}_u)$ and $v \sim \mathcal{N}(\cvec{\mu}_v,\cvec{\Sigma}_v)$ are normally distributed independent random variables and if conditional distribution $p(\cvec{y}|\cvec{u},\cvec{v}) = \mathcal{N}(\cvec{A}\cvec{u} + b + \cvec{B}v, \cvec{\Sigma})$, the joint distribution has an expression as follows:

$\left(\begin{array}{c}\mathbf{u} \\ \mathbf{v} \\ \mathbf{y}\end{array}\right)  \sim \mathcal{N}\left(\left(\begin{array}{c}\cvec{\mu}_{u} \\ \cvec{\mu}_{v} \\ \mathbf{A} \mu_{u}+ b + \mathbf{B} \mu_{v} \end{array}\right),\left(\begin{array}{cccc}\boldsymbol{\Sigma}_{u} & 0& \boldsymbol{\Sigma}_{u} \mathbf{A}^{\top} \\ 0 &\boldsymbol{\Sigma}_{v}  &\boldsymbol{\Sigma}_{v} \mathbf{B}^{\top} \\ \mathbf{A} \boldsymbol{\Sigma}_{u}^{\top} & \mathbf{B} \boldsymbol{\Sigma}_{v}^{\top} & \mathbf{A} \boldsymbol{\Sigma}_{u} \mathbf{A}^{\top}+\mathbf{B} \boldsymbol{\Sigma}_{v} \mathbf{B}^{\top} + \boldsymbol{\Sigma} \end{array}\right)\right)$

Proof for Identity <ref>

Let displacement of a variable  $\mathbf{u}$ be denoted by $\Delta \mathbf{u}=\mathbf{u}-\langle\mathbf{u}\rangle $.

Since $\mathbf{u}$ and $\mathbf{v}$ are independent, covariances $\left\langle\Delta \mathbf{u} \Delta \mathbf{v}^{\top}\right\rangle=0$.

We can write $\mathbf{y}=\mathbf{A}\cvec{u} + b + \mathbf{B}\cvec{v} + \cvec{\epsilon}$, where $\cvec{\epsilon} \sim \mathcal{N}(\cvec{0},\cvec{\Sigma})$ and $b$ is a constant. Then we have covariance $\left\langle\Delta \mathbf{u} \Delta \mathbf{y}^{\top}\right\rangle=\left\langle\Delta \mathbf{u}(\mathbf{A} \Delta \mathbf{u}+\mathbf{B} \Delta \mathbf{v} + \Delta \epsilon)^{\top}\right\rangle=$ $\left\langle\Delta \mathbf{u} \Delta \mathbf{u}^{\top}\right\rangle \mathbf{A}^{\top}+\left\langle\Delta \mathbf{u} \Delta \mathbf{v}^{\top}\right\rangle \mathbf{B}^{\top} +  \left\langle\Delta \mathbf{u} \Delta \mathbf{\epsilon}^{\top}\right\rangle$. Since $\left\langle\Delta \mathbf{u} \Delta \mathbf{v}^{\top}\right\rangle=\left\langle\Delta \mathbf{u} \Delta \mathbf{\epsilon}^{\top}\right\rangle=0$ we therefore have $\left\langle\Delta \mathbf{u} \Delta \mathbf{y}^{\top}\right\rangle=\boldsymbol{\Sigma}_{u} \mathbf{A}^{\top}.$

The derivations for covariances $\left\langle\Delta \mathbf{v} \Delta \mathbf{y}^{\top}\right\rangle$ follows similarly and the corresponding covariance has the expression, $\left\langle\Delta \mathbf{v} \Delta \mathbf{y}^{\top}\right\rangle=\boldsymbol{\Sigma}_{v} \mathbf{B}^{\top}$.

Similarly, $\left\langle\Delta \mathbf{y} \Delta \mathbf{y}^{\top}\right\rangle=$ $\left\langle(\mathbf{A} \Delta \mathbf{u}+\mathbf{B} \Delta \mathbf{v} + \Delta \mathbf{\epsilon})(\mathbf{A} \Delta \mathbf{u}+\mathbf{B} \Delta \mathbf{v} + \Delta \mathbf{\epsilon})^{\top}\right\rangle=\mathbf{A}\left\langle\Delta \mathbf{u} \Delta \mathbf{u}^{\top}\right\rangle \mathbf{A}^{\top}+\mathbf{B}\left\langle\Delta \mathbf{v} \Delta \mathbf{v}^{\top}\right\rangle \mathbf{B}^{\top} + \left\langle\Delta \mathbf{\epsilon} \Delta \mathbf{\epsilon}^{\top}\right\rangle = \mathbf{A} \boldsymbol{\Sigma}_{u} \mathbf{A}^{\top}+\mathbf{B} \boldsymbol{\Sigma}_{v} \mathbf{B}^{\top} +\Sigma $. The result follows.

[Gaussian Marginalization]gigeneralgaussmarg

\left(\begin{array}{l}
\mathbf{\cvec{u}} \\
\mathbf{\cvec{v}} \\
\mathbf{\cvec{y}}
\end{array}\right) \sim \mathcal{N}\left(\left(\begin{array}{l}
\cvec{\mu}_{u} \\
\cvec{\mu}_{v} \\
\cvec{\mu}_{z}
\end{array}\right),\left(\begin{array}{cccc}
\boldsymbol{\Sigma}_{uu} & \boldsymbol{\Sigma}_{uv} & \boldsymbol{\Sigma}_{uy}\\
\mathbf{\Sigma}_{uv}^{\top} & \boldsymbol{\Sigma}_{vv} & \mathbf{\Sigma}_{v y}\\
\boldsymbol{\Sigma}_{uy}^{\top}  & \boldsymbol{\Sigma}_{vy}^{\top}   & \boldsymbol{\Sigma}_{yy}
\end{array}\right)\right)

then marginal over y is given as $p(\cvec{y}) = \int_{\cvec{u},\cvec{v}} p(\cvec{y}|\cvec{u},\cvec{v})p(\cvec{u})p(\cvec{v})d\cvec{u}d\cvec{v} = \mathcal{N}\left(\cvec{\mu}_{y}, \mathbf{\Sigma}_{yy}\right)$

Proof For Identity <ref>

We refer to <cit.> for the derivation, which requires calculation of the Schur complement as well as completing the square of the Gaussian p.d.f. to integrate out the variable. The given derivation bishop2006pattern for two variable multivariate Gaussians can be extended to 3 variable case WLOG.

Proof For Identity <ref> is immediate from Identity <ref> and Identity <ref>.

§ ADDITIONAL EXPERIMENTS

§.§ Comparison To Soft Switching Baseline

We also implement a soft-switching baseline similar to <cit.>, where a soft mixture of dynamics is implemented in the latent transition dynamics (Kalman time update).  Similar to <cit.>, we now globally learn $K$ constant transition matrices $\boldsymbol{A}^{(k)}$ and control matrices $\boldsymbol{B}^{(k)}$. An interpolation is done between these using a “dynamics parameter network" fraccaro2017disentangled $\alpha_{t}=$ $\boldsymbol{\alpha}_{t}\left(\mathbf{w}_{0: t-1}\right)$. The dynamics parameter network is implemented with a recurrent neural network with LSTM cells that takes at each time step the mean of the encoded observation $w_t$ as input and recurses $\mathbf{d}_{t}=$ $\operatorname{LSTM}\left(\mathbf{w}_{t-1}, \mathbf{d}_{t-1}\right)$ and $\boldsymbol{\alpha}_{t}=\operatorname{softmax}\left(\mathbf{d}_{t}\right)$. The output of the dynamics parameter network is weights that sum to one, $\sum_{k=1}^{K} \alpha_{t}^{(k)}\left(\mathbf{w}_{0: t-1}\right)=1$.
These weights choose and interpolate between $K$ different operating modes:

$$\mathbf{A}_{t}=\sum_{k=1}^{K} \alpha_{t}^{(k)}\left(\mathbf{w}_{0: t-1}\right) \mathbf{A}^{(k)}, \quad \mathbf{B}_{t}=\sum_{k=1}^{K} \alpha_{t}^{(k)}\left(\mathbf{w}_{0: t-1}\right) \mathbf{B}^{(k)}$$

Authors interpret the weighted sum as a soft mixture of $K$ different Linear Gaussian SSMs whose time-invariant matrices are combined using the time varying weights $\boldsymbol{\alpha}_{t}$. In practice, each of the $K$ sets $\left\{\mathbf{A}^{(k)}, \mathbf{B}^{(k)}\right\}$ models different/changing dynamics, that will dominate when the corresponding $\alpha_{t}^{(k)}$ is high.

Figure <ref> compares HiP-RSSM with different recurrent architectures including the soft-switching baseline. As seen in figure <ref>, HiP-RSSM clearly outperforms the soft-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 <ref>.

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.

§.§ Multi-Step Ahead Predictions

In figure <ref>, we compare the results of multi-step ahead predictions for upto 50 steps of HiP-RSSM with recurrent baselines like RKNbecker2019recurrent and LSTM. Inorder to train the recurrent models for multi step ahead predictions, we removed three-quarters of the observations from the temporal sequence and tasked the models with imputing those missing observations, only based on the knowledge of available actions/control commands, i.e., we train the models to perform action conditional future predictions to impute missing observations. The imputation employs the model for multi-step ahead predictions in a convenient way shaj2020action. One could instead also go for a dedicated multi-step loss function as in approaches like <cit.>.

As seen in figure <ref>, HiP-RSSM clearly outperforms contemporary recurrent models for multi-step ahead prediction tasks since it takes into account additional causal factors of variation (slopes of the terrain for this robot) in the latent dynamics in an unsupervised manner.

§ IMPLEMENTATION DETAILS

§.§ Context Set Encoder and Latent Task Representation
The HiP-RSSM maps a set of previous interaction histories, $ \{\cvec{o}_{t,n}^l,\cvec{a}_{t,n}^l,\cvec{o}_{t+1,n}^l\}_{n=1}^N$, to a set of latent features and an estimate of uncertainty in those features, $
\{\cvec{r}_{n}^l,(\cvec{\sigma}_n^l)^2\}_{n=1}^N$, using a context encoder. We use a feed-forward neural network as the encoder in all of our experiements since we deal with high dimensional vectors as our observations. However, depending upon the nature of observations, we could use different encoder architectures.

The set of latent features and uncertainties, $
\{\cvec{r}_{n}^l,\left(\cvec{\sigma}_n^{\cvec{l}}\right)^2\}_{n=1}^N$, are further aggregated in a probabilistically principled manner using bayesian aggregation operator discussed in <ref> to get a gaussian latent task variable, with a mean ($\cvec{\mu}_l$) and diagonal covariance ($\cvec{\sigma}_l$). Intuitively the context encoder learns to weight the contribution from each observation in the context set based on bayesian priniciples and emits a probabilistic representation of the latent task.

§.§ Latent Task Transformation Model

To achieve latent task conditioning within the recurrent cell, we include  a task transformation model $(\boldsymbol{c})$, in addition to the locally linear transition model $\boldsymbol{A}_t$ and control model $\boldsymbol{b}$ in the time update stage (section <ref>).
Though in section <ref>, we used the notation for a linear task transformation matrix, $\boldsymbol{C}$, to motivate the additive interaction of latent task variables, $\cvec{\mu}_l$ and $\cvec{\sigma}_l$ in the latent space, the task transformation function can be designed in several ways, i.e.:

(i) Linear: $\boldsymbol{c} =\boldsymbol{C}$, where $\boldsymbol{C} $ is a linear transformation matrix. The corresponding time update equations are given below:

\begin{align*}
\cvec{z}_t^- &=\cmat{A}_{t-1}\cvec{z}_{t-1}^+ + \cmat{b}(\cvec{a}_t) + \cmat{C}\cvec{\mu_l}, \\
\cmat{\Sigma}_t^- &= \cmat{A}_{t-1}\cmat{\Sigma}_{t-1}^+\cmat{A}_{t-1}^T + \cmat{C}(\cmat{I}\cdot\cmat{\sigma}_{\cvec{l}})\cmat{C}^T + \cmat{\Sigma}_{\textrm{trans}}.
\end{align*}

(ii) Locally-Linear: $\boldsymbol{c} = \boldsymbol{C}_t $, where $\boldsymbol{C}_t = \sum_{k=0}^K \beta^{(k)}(\boldsymbol{z_t}) \boldsymbol{C}^{(k)}$ is a linear combination of k linear control models $\boldsymbol{C}^{(k)}$. A  small  neural  network  with softmax  output  is  used  to  learn $\beta^{(k)}$. The corresponding time update equations are given below:

\begin{align*}
\cvec{z}_t^- &=\cmat{A}_{t-1}\cvec{z}_{t-1}^+ + \cmat{b}(\cvec{a}_t) + \cmat{C_t}\cvec{\mu_l}, \\
\cmat{\Sigma}_t^- &= \cmat{A}_{t-1}\cmat{\Sigma}_{t-1}^+\cmat{A}_{t-1}^T + \cmat{C_t}(\cmat{I}\cdot\cmat{\sigma}_{\cvec{l}})\cmat{C_t}^T + \cmat{\Sigma}_{\textrm{trans}}.
\end{align*}

(iii) Non-Linear: $\boldsymbol{c} = \boldsymbol{f}$, where $\boldsymbol{f}(.)$ can be any non-linear function approximator. We use a multi-layer neural network regressor with ReLU activations, which transforms the latent task moments $\cvec{\mu_l}$ and $\cvec{\sigma_l}$ directly into the latent space of the state space model via additive interactions. The corresponding time update equations are given below:

\begin{align*}
\cvec{z}_t^- &=\cmat{A}_{t-1}\cvec{z}_{t-1}^+ + \cmat{b}(\cvec{a}_t) + \cmat{f}(\cvec{\mu}_l), \\
\cmat{\Sigma}_t^- &= \cmat{A}_{t-1}\cmat{\Sigma}_{t-1}^+\cmat{A}_{t-1}^T + \cmat{f}(\cvec{\sigma}_{\cvec{l}})+ \cmat{\Sigma}_{\textrm{trans}}.
\end{align*}

In our ablation study (Figure <ref>), for the linear and locally linear task transformation models, we assume that the dimension of the latent context variable $\cvec{l}$ and the latent state space $\cvec{z}_t$ are equal. This allows us to work with square matrices which are more convenient. For the non-linear transformation we are free to choose the size of the latent context variable.  However for a fairer comparison we keep the dimension of latent task variable to be similar in all three cases. We choose the non-linear task transformation model in HiP-RSSM architecture as this gave the best performance in practice.

CHAPTER: APPENDIX: MULTI TIME SCALE SSM

§ PROOFS AND DERIVATIONS

Graphical Model For Bayesian conditioning with $N$ observations.

In the following sections vectors are denoted by a lowercase letter in bold, such as "$\cvec{\mathrm{v}}$", while Matrices as an uppercase letter in bold, such as "$\cmat{M}$". $\cmat{I}$ denotes identity matrix and $\cmat{0}$ represents a matrix filled with zeros. For any matrix $\cmat{M}$, $\cvec{m}$ denotes the corresponding vector of diagonal entries. Also, $\odot$ denotes the elementwise vector product and  $\oslash$ denotes an elementwise vector division.

§.§ Proof For Bayesian Conditioning As Permutation Invariant Set Operations (Identity <ref>)

Note that $\cvec \Sigma_l$ is the covariance matrix which is the inverse of the precision matrix $\cvec \Lambda_l$. Due to the observation model assumption $\cmat{H}=[\cmat{I},\cmat{0}]$, they take block diagonal form,
$$\cvec \Sigma_l = \left[\begin{array}{cc}\cvec \Sigma_l^u & \cvec \Sigma^{s}_l \\ \cvec \Sigma^{s}_l & \cvec \Sigma^{l}_l  \end{array} \right], \textrm{ with } \cvec \Sigma_u = \textrm{diag}(\cvec \sigma_l^u), \; \cvec \Sigma_l = \textrm{diag}(\cvec \sigma_l^l) \textrm{ and } \cvec \Sigma_{s} = \textrm{diag}(\cvec \sigma_{l}^s).$$

Case 1 (Single Observation): Before deriving the update rule for $N$ conditionally iid observations, let us start with a simpler case consisting of a single observation $\cvec r$. If the marginal Gaussian distribution for the latent variable $\cvec l$ takes the form $ p(\mathbf{l}) =\mathcal{N}\left(\mathbf{l} \mid \boldsymbol{\mu}, \boldsymbol{\Lambda}^{-1}\right)$ and the conditional Gaussian distribution for the single observation $\cvec{r}$ given $\cvec{l}$ has the form
, $p(\mathbf{r} \mid \mathbf{l}) =\mathcal{N}\left(\mathbf{r} \mid \mathbf{H} \mathbf{l}+\mathbf{b}, \mathbf{L}^{-1}\right)$. Then the posterior distribution over $\mathbf{l}$ can be obtained in closed form as,

\begin{equation}
\label{eq:bishop1}
\begin{aligned}
p(\mathbf{l} \mid \mathbf{r}) & =\mathcal{N}\left(\mathbf{l} \mid \mathbf{\Sigma}\left\{\mathbf{H}^{\mathrm{T}} \mathbf{L}(\mathbf{r}-\mathbf{b})+\boldsymbol{\Lambda} \boldsymbol{\mu}\right\}, \boldsymbol{\Lambda}^{-1}\right) ,\textrm{where } \boldsymbol{\Lambda}=\left(\boldsymbol{\Lambda}+\mathbf{H}^{\mathrm{T}} \mathbf{L} \mathbf{H}\right).
\end{aligned}
\end{equation}

We refer to Section 2.3.3 of <cit.>, to the proof for this standard result.

Case 2 (Set Of Observations): Now instead of a single observation, we wish to derive a closed form solution for the posterior over latent variable $\boldsymbol{l}\in \mathbb{R}^{2d}$, given a set of N conditionally i.i.d observations $\Bar{r} = \{r_i\}_{i=1}^N$. Here each element $\boldsymbol{r_i} \in \mathbb{R}^d$ of the set $\Bar{r}$ is assumed to to have an observation model $\cmat{H}=[\cmat{I},\cmat{0}]$. In the derivation, we represent the set of N observations as a random vector
$$\Bar{r} = \left[ \begin{array}{ll}
\cvec{r_1}  \\
\cvec{r_2} \\
\cvec{.} \\
\cvec{.} \\
\cvec{r_N}
\end{array}\right]_{Nd \times 1}.$$

Since each observation in the set $\Bar{r}$ are conditionally independent, we denote the conditional distribution over the context set as $\Bar{r}  \mid \mathbf{l} \sim \mathcal{N}\left(\bar{H} \mathrm{l},  \cmat{\Sigma}_{r}\right)$, where the diagonal covariance matrix has the following form:
$$\cmat{\Sigma}_{r}= \left[ \begin{array}{llllll}
\textrm{diag}(\cvec \sigma_{r_1}),&0,&0,&..,&0  \\
0,&\textrm{diag}(\cvec \sigma_{r_2}),&0,&..,&0 \\
\boldsymbol{.},&.,&.,&..,&. \\
\boldsymbol{. },&.,&.,&..,&.\\
0,&0,&0,&..,&\textrm{diag}(\cvec \sigma_{r_N})
\end{array}\right]_{Nd \times Nd}.

The corresponding observation model  $\bar{\cmat H}$ is

\bar{\cmat H} =  \left[ \begin{array}{ll}
\cmat{H}  \\
\cmat{H} \\
\cmat{.} \\
\cmat{.} \\
\cmat{H}
\end{array}\right]_{Nd \times 2d} = \left[ \begin{array}{ll}
\boldsymbol{I }, \boldsymbol{0 }  \\
\boldsymbol{I }, \boldsymbol{0} \\
\boldsymbol{.}, \boldsymbol{.} \\
\boldsymbol{. }, \boldsymbol{.} \\
\boldsymbol{I }, \boldsymbol{0}
\end{array}\right]_{Nd \times 2d}.$$

Now given the prior over the latent task variable $\mathrm{l} \sim \mathcal{N}\left(\mu_l^-, \boldsymbol{\Sigma}_l^-\right)$, the parameters of the posterior distribution over the task variable, $p(l|\Bar{r}) \sim \mathcal{N}\left(\mu_l^+, \boldsymbol{\Lambda}_l^+\right)$, can be obtained in closed-form substituting in Equation (<ref>) as follows.

\begin{equation}
\begin{aligned}
\Lambda_l^+ & =  (\Sigma_l^+)^{-1}  \\ &= \boldsymbol{\Sigma}_l^{-1}+ \bar{\cmat {H}} ^T \cmat{\Sigma}_{r} \bar{\cmat {H}}  \\
& = \boldsymbol{\Sigma}_l^{-1}+ \left[ \begin{array}{ll}
\textrm{diag}(\cvec \sigma_{r_1}),  \textrm{diag}(\cvec \sigma_{r_2}),  \textrm{diag}(\cvec \sigma_{r_3}), . , . ,   \textrm{diag}(\cvec \sigma_{r_N})  \\
\boldsymbol{0 }, \quad \quad \boldsymbol{0}, \quad \quad \boldsymbol{0}, \quad. , . ,\quad  \boldsymbol{0 }  \\
\end{array}\right]_{2d\times nd} \bar{ \cmat{H}} \\
& =\boldsymbol{\lambda}_l^{-} + \left[ \begin{array}{ll}
\mathrm{diag}(\sum_{i=1}^n\frac{1}{\cvec \sigma_{r_i}}), &\boldsymbol{0 }  \\
\boldsymbol{0 },& \boldsymbol{0} \\
\end{array}\right]_{2d\times 2d}  \\[15pt]
\mu_l^+ & = \boldsymbol{\mu}_l^{-} + (\Lambda^+)^{-1}\bar{\cmat {H}} ^T \left( \sigma_{\boldsymbol{r}}^{-2} \boldsymbol{I}\right)\left(\boldsymbol{y}-\bar{\cmat {H}}  \boldsymbol{\mu}_{\boldsymbol{x}}\right) \\
& = \boldsymbol{\mu}_l^{-}+ \Sigma^+ \bar{\cmat {H}} \left( \sigma_{\boldsymbol{r}}^{-2} \boldsymbol{I}\right) \left(\boldsymbol{y}- \bar{\cmat {H}}  \boldsymbol{\mu}_{\boldsymbol{x}}\right) \\
& = \boldsymbol{\mu}_l^{-}+ \Sigma^+\left[ \begin{array}{ll}
\sigma_{\boldsymbol{r_1}}^{-2}\boldsymbol{I }, \sigma_{\boldsymbol{r_2}}^{-2}\boldsymbol{I}, \sigma_{\boldsymbol{r_3}}^{-2}\boldsymbol{I }, . , . ,  \sigma_{\boldsymbol{r_n}}^{-2}\boldsymbol{I }  \\
\boldsymbol{0 }, \quad \quad \boldsymbol{0}, \quad \quad \boldsymbol{0}, \quad. , . ,\quad  \boldsymbol{0 }  \\
\end{array}\right] \left(\boldsymbol{y}-\bar{\cmat {H}}  \boldsymbol{\mu}_{\boldsymbol{x}}\right) \\
& =\boldsymbol{\mu}_l^{-} + \left[ \begin{array}{ll}
\boldsymbol{\sigma_l^{u+}}, &\boldsymbol{\sigma_l^{s+}}  \\
\boldsymbol{\sigma_l^{s+}} ,& \boldsymbol{\sigma_l^{l+}}  \\
\end{array}\right]  \left[ \begin{array}{l}
\sum_{n=1}^{N}  \left(\mathbf{r_n}-\mathbf{\mu}^{\mathrm{u},-}_l\right) \odot \frac{1}{\sigma_i} \\
\boldsymbol{0 } \\
\end{array}\right]  \\
& =\boldsymbol{\mu}_l^{-} + \left[ \begin{array}{l}
\boldsymbol{\sigma_l^{u+}} \\
\boldsymbol{\sigma_l^{s+}}  \\
\end{array}\right]  \odot \left[ \begin{array}{l}
\sum_{i=1}^{N}  \left(\mathbf{r_i}-\mathbf{\mu}^{\mathrm{u},-}_l\right) \odot \frac{1}{\cvec \sigma_{r_i}} \\
\sum_{i=1}^{N}  \left(\mathbf{r_n}-\mathbf{\mu}^{\mathrm{u},-}_l\right) \odot \frac{1}{\cvec \sigma_{r_i}} \\
\end{array}\right]  \\
\end{aligned}
\end{equation}

Here $\mu_l^+$ is the posterior mean and $\boldsymbol{\Lambda}_l^+$ is the posterior precision matrix.

§.§ Derivation For Matrix Inversions as Scalar Operations

Consider a block  matrix of the following form $\cvec A = \left[\begin{array}{cc} \mathrm{diag}(\cvec a^u) & \mathrm{diag}(\cvec a^s) \\  \mathrm{diag}(\cvec a^s) &  \mathrm{diag}(\cvec a^l)  \end{array}\right]$. Then inverse $A^{-1} = \cvec B $ can be calculated using scalar operations and is given as, $\cvec B = \left[\begin{array}{cc} \mathrm{diag}(\cvec b^u) & \mathrm{diag}(\cvec b^s) \\  \mathrm{diag}(\cvec b^s) &  \mathrm{diag}(\cvec b^l)  \end{array}\right]$ where,

\begin{align}
\label{eq: matscalar}
\begin{split}
\cvec b^u & =\cvec a_l \oslash{ ( \cvec a_u \odot \cvec a_l- \cvec a_s \odot \cvec a_s )}  \\
\cvec b^s & = - \cvec a_s \oslash{ ( \cvec a_u \odot \cvec a_l- \cvec a_s \odot \cvec a_s )}  \\
\cvec b^l & =\cvec a_u \oslash{ ( \cvec a_u \odot \cvec a_l- \cvec a_s \odot \cvec a_s )}
\end{split}
\end{align}

Proof: To prove this we will use the following matrix identity of a partitioned matrix from <cit.>, which states

\begin{align}
\label{eq: matide}
\left(\begin{array}{ll}
\mathrm{A} & \mathrm{B} \\
\mathrm{C} & \mathrm{D}
\end{array}\right)^{-1}=\left(\begin{array}{cc}
\mathrm{M} & -\mathrm{MBD}^{-1} \\
-\mathrm{D}^{-1} \mathrm{CM} & \mathrm{D}^{-1}+\mathrm{D}^{-1} \mathrm{CMBD}^{-1}
\end{array}\right)
\end{align}

where M is defined as

\mathrm{M}=\left(\mathrm{A}-\mathrm{BD}^{-1} \mathrm{C}\right)^{-1}.
Here M is called the Schur complement of the Matrix on the left side of Equation <ref>. The algebraic manipulations to arrive at scalar operations in Equation <ref> are straightforward.

§.§ Proof for Permutation Invariance (Theorem <ref>)

The posterior distribution $p(\theta \mid X)$, defined by Bayes' theorem as $p(\theta \mid X) = \frac{p(X \mid \theta)p(\theta)}{p(X)}$, relies on the likelihood $p(X \mid \theta) = \prod_{i=1}^N p(x_i \mid \theta)$. Under any permutation $\pi$, the permuted set $X^\pi = \{\pi(x_1), \ldots, \pi(x_N)\}$ preserves the likelihood since $p(X^\pi \mid \theta) = \prod_{i=1}^N p(\pi(x_i) \mid \theta) = \prod_{i=1}^N p(x_i \mid \theta) = p(X \mid \theta)$. The prior $p(\theta)$ and the marginal likelihood $p(X)$ are invariant under permutation. Therefore, $p(\theta \mid X^\pi) = p(\theta \mid X)$, establishing the permutation invariance of the posterior.

§.§ Proof for Gaussian Marginalization (Identity <ref>)

The proof is a staright forward extension to the Identity <ref> derived in Chapter <ref>.  First, we derive an expression for the joint distribution $p(\cvec{u_1},\cvec{u_2},..,\cvec{u_N},\cvec{y})$.

[Joint Distribution]gigaussjoint
If $\cvec{u_1} \sim \mathcal{N}(\cvec{\mu}_1,\cvec{\Sigma}_1)$ and $\cvec{u_2} \sim \mathcal{N}(\cvec{\mu}_2,\cvec{\Sigma}_2)$ are normally distributed independent random variables and if conditional distribution $p(\cvec{y}|\cvec{u_1},\cvec{u_2}) = \mathcal{N}(\cvec{A_1}\cvec{u_1} + \cvec{A_2}\cvec{u_2}, \cvec{\Sigma})$, the joint distribution has an expression as follows:

\begin{equation}
\left(\begin{array}{c}
\mathbf{u_1} \\
\mathbf{u_2} \\
\vdots \\
\mathbf{u_N} \\
\mathbf{y}
\end{array}\right) \sim \mathcal{N}\left(
\left(\begin{array}{c}
\cvec{\mu}_1 \\
\cvec{\mu}_2 \\
\vdots \\
\cvec{\mu}_N \\
\sum_{n=1}^N \mathbf{A}_n \cvec{\mu}_n
\end{array}\right),
\left(\begin{array}{cccc}
\mathbf{\Sigma}_1 & 0 & \cdots & \mathbf{A}_1^\top \mathbf{\Sigma}_1 \\
0 & \mathbf{\Sigma}_2 & \cdots & \mathbf{A}_2^\top \mathbf{\Sigma}_2 \\
\vdots & \vdots & \ddots & \vdots \\
\mathbf{\Sigma}_1 \mathbf{A}_1 & \mathbf{\Sigma}_2 \mathbf{A}_2 & \cdots & \sum_{n=1}^N \mathbf{A}_n \mathbf{\Sigma}_n \mathbf{A}_n^\top + \mathbf{\Sigma}
\end{array}\right)
\right)
\end{equation}

We can write $\mathbf{y}= \sum_{n=1}^N \cmat{A_n}\cvec{u_n} + \cvec{\epsilon}$, where $\cvec{\epsilon} \sim \mathcal{N}(\cvec{0},\cvec{\Sigma})$ and $b$ is a constant.

Let displacement of a variable  $\mathbf{u}$ be denoted by $\Delta \mathbf{u}=\mathbf{u}-\langle\mathbf{u}\rangle $.

Since $\mathbf{u_i}$ and $\mathbf{u_j}$ are independent $\forall i \neq j$, the covariances

\begin{equation}
\label{eq:gu2indep1}
\begin{aligned}
\left\langle\Delta \mathbf{u_i} \Delta \mathbf{u_j}^{\top}\right\rangle=0, \forall i \neq j .
\end{aligned}
\end{equation}

\begin{equation}
\label{eq:gu2indep2}
\begin{aligned}
\left\langle\Delta \mathbf{u_i} \Delta \mathbf{\epsilon}^{\top}\right\rangle=0, \forall i.
\end{aligned}
\end{equation}

For any \(i\), we have the covariance

\begin{equation}
\begin{aligned}
\left\langle \Delta \mathbf{u}_i \Delta \mathbf{y}^\top \right\rangle & = \left\langle \Delta \mathbf{u}_i \left(\sum_{n=1}^N \mathbf{A}_n \Delta \mathbf{u}_n + \Delta \epsilon \right)^\top \right\rangle\\ & = \left\langle \Delta \mathbf{u}_i \Delta \mathbf{u}_i^\top \right\rangle \mathbf{A}_i^\top + \sum_{j=1, j \neq i}^N \left\langle \Delta \mathbf{u}_i \Delta \mathbf{u}_j^\top \right\rangle \mathbf{A}_j^\top + \left\langle \Delta \mathbf{u}_i \Delta \epsilon^\top \right\rangle.
\end{aligned}
\end{equation}

Using equations \(\ref{eq:gu2indep1}\) and \(\ref{eq:gu2indep2}\), we therefore derive an expression for the corresponding covariance as:

\[
\left\langle \Delta \mathbf{u}_i \Delta \mathbf{y}^\top \right\rangle = \mathbf{\Sigma}_i \mathbf{A}_i^\top.
\]

\begin{align*}
\left\langle \Delta \mathbf{y} \Delta \mathbf{y}^\top \right\rangle
&= \left\langle \left(\sum_{n=1}^N \mathbf{A}_n \Delta \mathbf{u}_n + \Delta \epsilon \right) \left(\sum_{n=1}^N \mathbf{A}_n \Delta \mathbf{u}_n + \Delta \epsilon\right)^\top \right\rangle \\
&= \sum_{n=1}^N \mathbf{A}_n \left\langle \Delta \mathbf{u}_n \Delta \mathbf{u}_n^\top \right\rangle \mathbf{A}_n^\top + \left\langle \Delta \epsilon \Delta \epsilon^\top \right\rangle \\
&= \sum_{n=1}^N \mathbf{A}_n \mathbf{\Sigma}_n \mathbf{A}_n^\top + \mathbf{\Sigma}_\epsilon.
\end{align*}

The result follows.

\[
\left(\begin{array}{c}
\cvec{u}_1 \\
\cvec{u}_2 \\
\vdots \\
\cvec{u}_N \\
\cvec{y}
\end{array}\right) \sim \mathcal{N}\left(\left(\begin{array}{c}
\cvec{\mu}_1 \\
\cvec{\mu}_2 \\
\vdots \\
\cvec{\mu}_N \\
\cvec{\mu}_y
\end{array}\right),\left(\begin{array}{ccccc}
\cmat{\Sigma}_{11} & \cmat{\Sigma}_{12} & \cdots & \cmat{\Sigma}_{1N} & \cmat{\Sigma}_{1y} \\
\cmat{\Sigma}_{21} & \cmat{\Sigma}_{22} & \cdots & \cmat{\Sigma}_{2N} & \cmat{\Sigma}_{2y} \\
\vdots & \vdots & \ddots & \vdots & \vdots \\
\cmat{\Sigma}_{N1} & \cmat{\Sigma}_{N2} & \cdots & \cmat{\Sigma}_{NN} & \cmat{\Sigma}_{Ny} \\
\cmat{\Sigma}_{y1} & \cmat{\Sigma}_{y2} & \cdots & \cmat{\Sigma}_{yN} & \cmat{\Sigma}_{yy}
\end{array}\right)\right)
\]

then marginal over y is given as $p(\cvec{y}) = \int p(\cvec{y}, \cvec{u_1}, \cvec{u_2}, . ., \cvec{u_N})\prod_{i=1}^N \cvec{u_i} = \mathcal{N}\left(\cvec{\mu}_{y}, \mathbf{\Sigma}_{yy}\right)$

Proof For Identity <ref>
We refer to <cit.> for the derivation, which requires calculation of the Schur complement as well as completing the square of the Gaussian p.d.f. to integrate out the variable. The given derivation bishop2006pattern for two variable multivariate Gaussians can be extended to the case of N variables WLOG.

Proof For Identity <ref> is immediate from Identity <ref> and Identity <ref>.

§ IMPLEMENTATION DETAILS

§.§ Inference In Slow Time Scale SSM

§.§.§ Inferring Action Abstraction (sts-SSM)

Given a set of encoded primitive actions and their corresponding variances $ \{ \cvec{\alpha_{k,t}}, \cvec{ \rho_{k,t}} \}_{t=1}^H $, using the prior and observation model assumptions in Section 3.1.2 of main paper, we infer the latent abstract action $p(\cvec \alpha_{k}|\cvec \alpha_{k,1:H}) = \mathcal{N}(\cvec \mu_{ \alpha_{k}}, \cvec \Sigma_{\alpha_k})  = \mathcal{N}(\cvec \mu_{ \alpha_{k}}, \textrm{diag}\cvec(\sigma_{\alpha_k}))$ as a Bayesian aggregation <cit.> of these using the following closed-form equations:

Generative model for the abstract action $\alpha_k$. The hollow arrows are deterministic transformations leading to implicit distribution $\alpha_{k,t}$ using an action set encoder.

\begin{align*}
\cvec{\sigma}_{\alpha_{k}} & = \left( \left( \cvec{\sigma}_{0} \right)^\ominus + \sum_{n=1}^N \left(\left(\cvec{\rho_{k,t}}\right)^\ominus\right) \right)^\ominus, \\
\cvec{\mu}_{\alpha_{k}} & = \cvec{\mu}_{0} + \cvec{\sigma}_{\alpha_{k}} \odot \sum_{n=1}^N \left(\cvec{\alpha_{k,t}} - \cvec{\mu}_{0} \right) \oslash \cvec{\rho_{k,t}}
\end{align*}

Here, $\ominus$, $\odot$ and $\oslash$ denote element-wise inversion, product, and division, respectively. The update equation is coded as the “abstract action inference” neural network layer as shown in Figure <ref>.

§.§.§ Task Prediction (sts-SSM)

The goal of this step is to update the prior marginal over the latent task variable $\cvec{l}_k$, $p(\cvec{l}_k|\cvec{\beta}_{1:k-1},\cvec{\alpha}_{1:k})$, given the posterior beliefs from the time window $k-1$ and abstract action $\cvec \alpha_{k}$.

Using the linear dynamics model assumptions from Section 3.3, we can use the following closed-form update equations to compute,
$p(\cvec{l}_k|\cvec{\beta}_{1:k-1},\cvec{\alpha}_{1:k}) = \mathcal{N}(\cvec{\mu}_{l_{k}}^-,\cvec{\Sigma}_{l_{k}}^-)$, where

\begin{equation}
\label{eq:task-predict}
\begin{aligned}[c]
\cvec{\mu}_{l_{k}}^- &=\cmat{X}\cvec{\mu}_{l_{k-1}}^+ + \cmat{Y}\cvec{\alpha}_{k}\\
\cvec{\Sigma}_{l_{k}}^- &= \cmat{X}\cvec{\Sigma}_{l_{k-1}}^+\cmat{X}^T + \cmat{Y}\cvec{\Sigma}_{\alpha_{k}}\cmat{Y}^T + \cmat S.
\end{aligned}
\end{equation}

These closed-form equations are coded as the “task predict” neural net layer as shown in Figure <ref>.

§.§.§ Task Update (sts-SSM)

In this stage, we update the prior over $l_k$ using an abstract observation set $\{\cvec \beta_{k,t}\}_{t=1}^H$, to obtain the latent task the posterior $\mathcal{N}(\cvec{\mu}_{z_{k,t}}^+,\cmat{\Sigma}_{z_{k,t}}^+) = \mathcal{N}(\left[\begin{array}{c}\cvec{\mu}_{t}^{u+} \\
\cvec{\mu}_{t}^{l+} \end{array}\right],\left[\begin{array}{cc}\cvec \Sigma_t^u & \cvec \Sigma^{s}_t \\ \cvec \Sigma^{s}_t & \cvec \Sigma^{l}_t  \end{array}\right]^+)$, with $\cvec \Sigma_{l_k}^u = \textrm{diag}(\cvec \sigma_{l_k}^u), \; \cvec \Sigma_{l_k}^l = \textrm{diag}(\cvec \sigma_{l_k}^l)$  and  $\cvec \Sigma_{l_k}^{s} = \textrm{diag}(\cvec \sigma_{l_k}^s)$.

To do so we first invert the prior covariance matrix $\left[\begin{array}{cc}\cvec \Sigma_{l_k}^u & \cvec \Sigma_{l_k}^{s} \\ \cvec \Sigma_{l_k}^{s} & \cvec \Sigma_{l_k}^{l}  \end{array}\right]^+$ to the precision matrix $\left[\begin{array}{cc}\cvec \lambda_{l_k}^u & \cvec \lambda_{l_k}^{s} \\ \cvec \lambda_{l_k}^{s} & \cvec \lambda_{l_k}^{l}  \end{array}\right]^+$ for permutation invariant parallel processing. The posterior precision is then computed using scalar operations are follows, where only $\cvec \lambda_{l_{k}}^{u}$ is changed by

\begin{align}\cvec \lambda_{l_{k}}^{u+} = \cvec \lambda_{l_{k}}^{u-} + \sum_{t=1}^H \cvec 1 \oslash \cvec \nu_{k,t}\end{align}

while $\cvec \lambda_{l_{k}}^{l+} = \cvec \lambda_{l_{k}}^{l-}$ and $\cvec \lambda_{l_{k}}^{s+} = \cvec \lambda_{l_{k}}^{s-}$ remain constant. The operator $\oslash$ denotes the element-wise division. The posterior precision is inverted back to the posterior covariance vectors $\cvec \sigma_{l_{k}}^{u+}$, $\cvec \sigma_{l_{k}}^{l+}$ and $\cvec \sigma_{l_{k}}^{s+}$. Now, the posterior mean $\cvec{\mu}_{l,k}^{+}$ can be obtained from the prior mean $\cvec{\mu}_{l,k}^{-}$ as

\begin{equation}
\begin{aligned}[c]
\cvec{\mu}_{l,k}^{+}=\cvec{\mu}_{l,k}^{-} + \left[ \begin{array}{l}
\cvec{\sigma}_{l_k}^{u+} \\
\cvec{\sigma}_{l_k}^{s+} \\
\end{array}\right]  \odot \left[ \begin{array}{l}
\sum_{t=1}^{H}  \left(\cvec{\beta}_{k,t}-\cvec{\mu}^{\mathrm{u},-}_{l_k}\right) \oslash \cvec \nu_{k,t} \\
\sum_{t=1}^{H}  \left(\cvec{\beta}_{k,t}-\cvec{\mu}^{\mathrm{u},-}_{l_k}\right) \oslash \cvec \nu_{k,t} \\
\end{array}\right].   \hspace{1.5cm}\\
\end{aligned}
\end{equation}

The inversion between the covariance matrix and precision matrix can be done via scalar operations leveraging block diagonal structure as derived in Appendix <ref>. Figure <ref> shows the schematic of the task update layer.

§.§ Inference In Fast Time Scale SSM

The inference in fts-SSM for a time-window $k$ involves two stages as illustrated in Figure <ref>, calculating the prior and posterior over the latent state variable $z_t$. To keep the notation uncluttered, we will also omit the time-window index $k$ whenever the context is clear as in Section 3.2.

§.§.§ Task Conditional State Prediction (fts-SSM)

Following the assumptions of a task conditional linear dynamics as in Section 3.2 of the main paper, we obtain the prior marginal for $p(\cvec{z}_{k,t}|\cvec{w}^k_{1:t-1},\cvec{a}^k_{1:t-1},\cvec{\beta}_{1:k-1},\cvec{\alpha}_{1:k-1}) = \mathcal{N}(\cvec{\mu}_{z_{k,t}}^-,\cmat{\Sigma}_{z_{k,t}}^-)$ in closed form, where

\begin{equation}
\label{eq:cond-predict}
\centering
\begin{aligned}
\cvec{\mu}_{z_{k,t}}^- &=\cmat{A}\cvec{\mu}_{z_{k,t-1}}^- + \cmat{B}\cvec{a}_{k,t-1} + \cmat{C}\cvec{\mu}_{l_{k}}^-, \\
\cmat{\Sigma}_{k,t}^- &= \cmat{A}\cmat{\Sigma}_{k,t-1}^+\cmat{A}^T + \cmat{C}\cvec{\Sigma}_{l_{k}}^-\cmat{C}^T + \cmat{Q}.
\end{aligned}
\end{equation}

§.§.§ Observation Update (fts-SSM)

In this stage, we compute the posterior belief $p(\cvec{z}_{k,t}|\cvec{w}^k_{1:t},\cvec{a}^k_{1:t},\cvec{\beta}_{1:k},\cvec{\alpha}_{1:k-1}) = \mathcal{N}(\cvec{\mu}_{z_{k,t}}^-,\cmat{\Sigma}_{z_{k,t}}^-)$.  using the same closed-form update as in  <cit.>. The choice of the special observation model splits the state into two parts, an upper $\boldsymbol{z}_t^\textrm{u}$ and a lower part $\boldsymbol{z}_t^\textrm{l}$, resulting in the posterior belief $\mathcal{N}(\cvec{\mu}_{z_{k,t}}^-,\cmat{\Sigma}_{z_{k,t}}^-) = \mathcal{N}(\left[\begin{array}{c}\cvec{\mu}_{t}^{u+} \\
\cvec{\mu}_{t}^{l+} \end{array}\right], \left[\begin{array}{cc}\cvec \Sigma_t^u & \cvec \Sigma^{s}_t \\ \cvec \Sigma^{s}_t & \cvec \Sigma^{l}_t  \end{array}\right]^+)$, with $\cvec \Sigma_t^u = \textrm{diag}(\cvec \sigma_t^s), \; \cvec \Sigma_t^l = \textrm{diag}(\cvec \sigma_t^l)$ and $\cvec \Sigma_t^{s} = \textrm{diag}(\cvec \sigma_{t}^s)$. Thus, the factorization allows for only the diagonal and one off-diagonal vector of the covariance to be computed and simplifies the calculation of the mean and posterior to simple scalar operations.

The closed-form equations for the mean can be expressed as the following scalar equations,

\begin{align*}
\boldsymbol{z}_t^+ = \boldsymbol{z}_t^- +
\left[\begin{array}{c} \boldsymbol{\sigma}^\mathrm{u,-}_t \\
\boldsymbol{\sigma}^\mathrm{l,-}_t \end{array}\right]
\odot
\left[\begin{array}{c}\boldsymbol{w}_t - \boldsymbol{z}^{\mathrm{u},-}_t \\
\boldsymbol{w}_t - \boldsymbol{z}^{\mathrm{u},-}_t  \end{array}\right]
\oslash
\left[\begin{array}{c}  \boldsymbol{\sigma}_t^{\mathrm{u},-} + \boldsymbol{\sigma}_t^\mathrm{obs}
\\  \boldsymbol{\sigma}_t^{\mathrm{u},-} + \boldsymbol{\sigma}_t^\mathrm{obs} \end{array}\right],
\end{align*}

The corresponding equations for the variance update can be expressed as the following scalar operations,

\begin{align*}
\boldsymbol{\sigma}^{\mathrm{u},+}_t &= \boldsymbol{\sigma}^{\mathrm{u},-}_t \odot \boldsymbol{\sigma}^{\mathrm{u},-}_t \oslash \left( \boldsymbol{\sigma}_t^{\mathrm{u},-} + \boldsymbol{\sigma}_t^\mathrm{obs} \right),  \\
\boldsymbol{\sigma}^{\mathrm{s},+}_t &= \boldsymbol{\sigma}^{\mathrm{u},-}_t \odot \boldsymbol{\sigma}^{\mathrm{s},-}_t \oslash \left( \boldsymbol{\sigma}_t^{\mathrm{u},-} + \boldsymbol{\sigma}_t^\mathrm{obs} \right), \\
\boldsymbol{\sigma}^{\mathrm{l},+}_t &= \boldsymbol{\sigma}^{\mathrm{l}, -}_t - \boldsymbol{\sigma}^{\mathrm{s},-}_t \odot \boldsymbol{\sigma}^{\mathrm{s},-}_t \oslash \left( \boldsymbol{\sigma}_t^{\mathrm{u},-} + \boldsymbol{\sigma}_t^\mathrm{obs} \right),
\end{align*}
where $\odot$ denotes the elementwise vector product and  $\oslash$ denotes an elementwise vector division.

§.§ Modelling Assumptions

§.§.§ Control Model
To achieve action conditioning within the recurrent cell of fts-SMM, we include a control model $b(a_{k,t})$ in addition to the linear transition model $A_t$. $b(a_{k,t}) =f(a_{k,t})$, where $f(.)$ can be any non-linear function approximator.  We use a multi-layer neural network regressor with ReLU activations <cit.>.

However, unlike the fts-SSM where actions are assumed to be known and subjected to no noise, in the sts-SSM, the abstract action is an inferred latent variable with an associated uncertainty estimate. Hence we use a linear control model $Y$, for principled uncertainty propagation.

§.§.§ Transition Noise

We assume the covariance of the transition noise $Q$ and $S$ in both timescales to be diagonal. The noise is learned and is independent of the latent state.

§ METRICS USED FOR MEASURING LONG HORIZON PREDICTIONS

§.§ Sliding Window RMSE

The sliding window RMSE (Root Mean Squared Error) metric is computed for a predicted trajectory in comparison to its ground truth. At each time step, the RMSE for each trajectory is determined by taking the root mean square of the differences between the ground truth and predicted values within a sliding window that terminates at the current time step. This sliding window, with a specified size, provides a smoothed localized assessment of prediction accuracy over the entire prediction length.
Mathematically, the sliding window RMSE at time step $t$ is given by:

\operatorname{RMSE}(t)=\sqrt{\frac{1}{W} \sum_{i=t-W+1}^t\left(\mathrm{gt}_i-\operatorname{pred}_i\right)^2}

where $t$ is the current time step, $W$ is the window size, and $\mathrm{gt}_i$ and $\operatorname{pred}_i$ are the ground truth and predicted values at time step $i$, respectively. The extension to multiple trajectories is straightforward and omitted to keep the notation uncluttered.

§.§ Sliding Window NLL

The sliding window NLL (Negative Log-Likelihood) metric is computed for a predicted probability distribution against the true distribution. At each time step, the NLL is determined by summing the negative log-likelihood values within a sliding window that terminates at the current time step. This sliding window, with a specified size, provides a smoothed localized evaluation of prediction accuracy across the entire sequence.

Mathematically, the sliding window NLL at time step $t$ is given by:

\mathrm{NLL}(t)=-\frac{1}{W} \sum_{i=t-W+1}^t \log \mathcal{N}\left(\mathrm{gt}_i \mid \mathrm{predMean}_i, \mathrm{predVar}_i\right)

where $t$ is the current time step, $W$ is the window size. $\mathrm{predMean}_i$, $\mathrm{predVar}_i$, and $\mathrm{gt}_i$ represent the predicted mean, predicted variance, and the ground truth at time step $i$.

§ VISUALIZATION OF PREDICTIONS GIVEN BY DIFFERENT MODELS.

In this section, we plot the multistep ahead predictions (mean and variance) by different models on 3 datasets on normalized test trajectories. Not that we omit NaN values in predictions while plotting.

§.§.§ Franka Kitchen

Multi-step ahead mean and variance predictions for a particular joint (joint 1) of Franka Kitchen Environment. The multi-step ahead prediction starts from the first red dot, which indicates masked observations. MTS3 gives the most reliable mean and variance estimates.

§.§.§ Hydraulic Excavator

Multi-step ahead mean and variance predictions for a particular joint (joint 1) of Excavator Dataset. The multi-step ahead prediction starts from the first red dot, which indicates masked observations. MTS3 gives the most reliable mean and variance estimates even up to 12 seconds into the future. Another interesting observation can also be seen in the predictions for MTS3, where after every window k of sts-SSM, which is 0.3 seconds (30 timesteps) long, the updation of the higher-level abstractions helps in grounding the lower-level predictions thus helping in the long horizon yet fine-grained predictions.

§.§.§ Mobile Robot

Multi-step ahead mean and variance predictions for a particular joint (joint 7) of Mobile Robot Dataset. The multi-step ahead prediction starts from the first red dot, which indicates masked observations. MTS3 gives the most accurate mean and variance estimates among all algorithms.

CHAPTER: APPENDIX: ROBOTS AND DATASET DETAILS

§ ROBOTS AND DATA USED FOR IN CHAPTER <REF> ON AC-SSM

The experiments are performed on data from four different robots. The details of robots, data, and data preprocessing are explained below:

§.§ Hydraulic Brokk 40 Robot Arm

Observation and Data Set: The data was obtained from a HYDROLEK–7W 6 degree–of–freedom manipulator with a continuous (360 degree) jaw rotation mechanism. We actuate the joints via hydraulic pistons, which are powered via an auxiliary output from the hydraulic pump. Thus learning the forward model is difficult due to inherent hysteresis associated with hydraulic control. For this robot, only one joint is moved at a time, so we have independent time series per joint. The joint data consists of measured joint positions and the input current to the controller of the joint sampled at 100Hz.

Training Procedure: During training, we work with sequences of length 500. For the first 300 time steps those sequences consist of the full observation, i.e., the joint position and current. We give only the current signals in the remaining 200 time steps. The models have to impute the missing joint positions in an uninformed fashion, i.e., we only indicate the absence of a position by unrealistically high values.

§.§ Musculoskeletal Robot Arm

Observation and Data Set:  For this soft robot we have 4 dimensional observation inputs(joint angles) and 8 dimensional action inputs(pressures). We collected the data of a four DoF robot actuated by Pneumatic Artificial Muscles
(PAMs). The robot arm has eight PAMs in total with each DoF actuated by an antagonistic pair. The robot arm reaches high joint angle accelerations of up to 28, 000deg/s2 while avoiding dangerous joint limits thanks to the antagonistic actuation and limits on the air pressure ranges. The data
consists of trajectories collected while training with a model-free reinforcement learning algorithm to hit balls while playing table tennis. We sampled the data at 100Hz. The hysteresis associated with the pneumatic actuators used in this robot is challenging to model and is relevant to the soft robotics in general.

Training Procedure: During training, we randomly removed three-quarters of the states from the sequences and tasked the models with imputing those missing states, only based on the knowledge of available actions/control commands, i.e., we train the models to perform action conditional future predictions to impute missing states. The imputation employs the model for multi-step ahead predictions in a convenient way. One could instead go for a dedicated loss function as in approaches like <cit.>, <cit.> for long term predictions.