# Multi Time Scale World Models

Vaisakh Shaj1 Saleh Gholam Zadeh1,2 Ozan Demir3 Luiz Ricardo Douat3
Gerhard Neumann1
1Karlsruhe Institute Of Technology (KIT), Germany
2SAP SE, Germany
3Robert Bosch GmbH, Germany
Corresponding author. Email to<EMAIL_ADDRESS>
###### Abstract

Intelligent agents use internal world models to reason and make predictions
about different courses of their actions at many scales [22]. Devising
learning paradigms and architectures that allow machines to learn world models
that operate at multiple levels of temporal abstractions while dealing with
complex uncertainty predictions is a major technical hurdle [18]. In this
work, we propose a probabilistic formalism to learn multi-time scale world
models which we call the Multi Time Scale State Space (MTS3) model. Our model
uses a computationally efficient inference scheme on multiple time scales for
highly accurate long-horizon predictions and uncertainty estimates over
several seconds into the future. Our experiments, which focus on action
conditional long horizon future predictions, show that MTS3 outperforms recent
methods on several system identification benchmarks including complex
simulated and real-world dynamical systems. Code is available at this
repository: https://github.com/ALRhub/MTS3.

##  1 Introduction

World models attempt to learn a compact and expressive representation of the
environment dynamics from observed data. These models can predict possible
future world states as a function of an imagined action sequence and are a key
ingredient of model-predictive control [4] and model-based reinforcement
learning (RL). One important dimension of world models is the level of
temporal granularity or the time scale at which the model operates. Existing
literature on world models operates at a single level of temporal abstraction,
typically at a fine-grained level such as milliseconds. One drawback of
single-time scale world models is that they may not capture longer-term trends
and patterns in the data [18].

For efficient long-horizon prediction and planning, the model needs to predict
at multiple levels of temporal abstractions [29, 23]. Intuitively, low-level
temporal abstractions should contain precise details about the input so as to
predict accurately in the short term, while high-level, abstract
representations should simplify accurate long-term predictions. Both
abstractions must also interrelate with each other at least in the sense that
the higher-level predictions/plans can be turned into low-level moment-by-
moment predictions. For example, 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.

Furthermore, temporal abstractions can capture relevant task structures across
dynamical systems under non-stationary which can be used to identify the
similarities and differences between tasks, allowing the robot to transfer
knowledge learned from one task to another [27, 18].

In this work, we 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, non-linear 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.

##  2 Preliminaries

State space models (SSMs) are Bayesian probabilistic graphical models [16, 13]
that are popular for learning patterns and predicting behaviour in sequential
data and dynamical systems. Formally, we define a state space model as a tuple
$(\mathcal{Z},\mathcal{A},\mathcal{O},f,h,\Delta t)$, where $\mathcal{Z}$ is
the state space, $\mathcal{A}$ the action space and $\mathcal{O}$ the
observation space of the SSM. The parameter $\Delta t$ denotes the
discretization time-step and $f$ and $h$ the dynamics and observation models
respectively. We will consider the Gaussian state space model that is
represented using the following equations

$\begin{aligned}
\boldsymbol{z}_{t}&=f(\boldsymbol{z}_{t-1},\boldsymbol{a}_{t-1})+\boldsymbol{\epsilon}_{t},\quad\boldsymbol{\epsilon}_{t}\sim\mathcal{N}(\boldsymbol{0},\mathbf{\Sigma_{z}}),\end{aligned}\textrm{and}\quad\begin{aligned}
\boldsymbol{o}_{t}&=h(\boldsymbol{z}_{t})+\boldsymbol{v}_{t},\quad\boldsymbol{v}_{t}\sim\mathcal{N}(\boldsymbol{0},\mathbf{\Sigma_{o}}).\end{aligned}$

Here $\boldsymbol{z}_{t}\in\mathcal{Z}$, $\boldsymbol{a}_{t}\in\mathcal{A}$
and $\boldsymbol{o}_{t}\in\mathcal{O}$ are the latent states, actions and
observations at time t. The vectors $\boldsymbol{\epsilon}_{t}$ and
$\boldsymbol{v}_{t}$ denote zero-mean Gaussian noise. When $f$ and $h$ are
linear/locally linear, inference can be performed efficiently via exact
inference. There have been several works recently [10, 1, 26] where these
closed-form solutions are coded as layers of a neural network in deep-state
space model literature, i.e, the architecture of the network is informed by
the structure of the probabilistic state estimator. Following this line of
work, we propose a multi-time scale linear Gaussian state space model, whose
inference can be performed via closed-form solutions.

##### (Locally-)Linear Gaussian SSMs.

To perform inference in SSMs, we follow [1]. They use a (locally-)linear
dynamics model. Moreover, they replace the observations $\boldsymbol{o}$ with
a latent observation $\boldsymbol{w}$. This latent observation is obtained by
an encoder $\boldsymbol{w}_{o_{t}}=\textrm{enc}_{w}(\boldsymbol{o}_{t})$ along
with the uncertainty of this observation, i.e.,
$\boldsymbol{\sigma}_{o_{t}}=\textrm{enc}_{\sigma}(\boldsymbol{o}_{t}).$ Due
to the non-linear observation encoder, a simplified linear observation model
can now be used. Hence, the dynamics and observation models can be described
as

$p(\boldsymbol{z}_{t+1}|\boldsymbol{z}_{t},\boldsymbol{a}_{t})=\mathcal{N}(\boldsymbol{A}_{t}\boldsymbol{z}_{t}+\boldsymbol{c}_{t}+\boldsymbol{B}_{t}\boldsymbol{a}_{t},\textrm{diag}(\boldsymbol{\sigma}_{z})),\textrm{
and
}p(\boldsymbol{w}_{o_{t}}|\boldsymbol{z}_{t})=\mathcal{N}(\boldsymbol{H}\boldsymbol{z}_{t},\textrm{diag}(\boldsymbol{\sigma}_{o_{t}})),$

where a simple observation matrix of
$\boldsymbol{H}=[\boldsymbol{I},\boldsymbol{0}]$ is used. The underlying
assumption behind this observation model is that the latent state
$\boldsymbol{z}_{t}=[\boldsymbol{p}_{t}^{T},\boldsymbol{d}_{t}^{T}]^{T}$ has
twice the dimensionality of the latent observation $\boldsymbol{w}_{t}$ and
only the first half of the latent state, i.e., $\boldsymbol{p}_{t}$, can be
observed. The second half of the latent state, i.e., $\boldsymbol{d}_{t}$,
serves as derivative or velocity units that can be used by the model to
estimate the change of the observable part of the latent state.

##### Factorized Inference in Linear Gaussian SSMs.

Inference in the introduced linear Gaussian SSM is straightforward and can be
performed using Kalman prediction and observation updates. However, these
updates involve high dimensional matrix inversions that are expensive to
evaluate and hard to backpropagate for end-to-end learning. Hence, [1]
introduce a factorization of the belief
$p(\boldsymbol{z}_{t}|\boldsymbol{o}_{1:t},\boldsymbol{a}_{1:t-1})=\mathcal{N}(\boldsymbol{\mu}_{t},\boldsymbol{\Sigma}_{t})$
such that only the diagonal and one off-diagonal vector of the covariance need
to be computed, i.e.

$\boldsymbol{\Sigma}_{t}=\left[\begin{array}[]{cc}\boldsymbol{\Sigma}_{t}^{u}&\boldsymbol{\Sigma}^{s}_{t}\\\
\boldsymbol{\Sigma}^{s}_{t}&\boldsymbol{\Sigma}^{l}_{t}\end{array}\right],\textrm{
with
}\boldsymbol{\Sigma}_{u}=\textrm{diag}(\boldsymbol{\sigma}_{t}^{s}),\;\boldsymbol{\Sigma}_{l}=\textrm{diag}(\boldsymbol{\sigma}_{t}^{l})\textrm{
and }\boldsymbol{\Sigma}_{s}=\textrm{diag}(\boldsymbol{\sigma}_{t}^{s}).$

Using this factorization assumption, closed-form Gaussian inference can be
performed using only scalar divisions which are fast and easy to back-
propagate. These factorization assumptions form the basis for the inference
update in our MTS3 model.

##### Bayesian Aggregation.

To aggregate information from several observations into a consistent
representation, [31] introduce Bayesian aggregation in the context of Meta-
Learning. They again use an encoder to obtain a latent observation vector
$\boldsymbol{r}_{o_{t}}$ and its uncertainty vector
$\boldsymbol{\sigma}_{o_{t}}$. Given the observation model
$p(\boldsymbol{r}_{o_{t}}|\boldsymbol{z})=\mathcal{N}(\boldsymbol{H}\boldsymbol{z},\textrm{diag}(\boldsymbol{\sigma}_{o_{t}}))$
with $\boldsymbol{H}=\boldsymbol{I}$ and a prior
$p(\boldsymbol{z})=\mathcal{N}(\boldsymbol{\mu}_{0},\textrm{diag}(\boldsymbol{\sigma}_{0}))$,
the posterior $p(\boldsymbol{z}|\boldsymbol{r}_{o_{1}:o_{t}})$ can again be
effectively computed by Gaussian inference that involve only scalar
inversions. Note that computing this posterior is a simplified case of the
Kalman update rule used in Gaussian SSMs [1], with no memory units,
$\boldsymbol{H}=\boldsymbol{I}$ and no dynamics. To increase efficiency, the
update rule can be formulated in a batch manner for parallel processing
instead of an incremental update [31].

##  3 Multi Time Scale State Space Models

Figure 1: MTS3 captures slow-moving long-term trends as the latent task
dynamics and the fast-moving short-term trends as the latent state dynamics.

Our goal is to learn a principled sequential latent variable model that can
model the dynamics of partially observable robotic systems under multiple
levels of temporal abstractions. 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 modelling 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 give accurate
long-term predictions and uncertainty estimates. iv) It is probabilistically
principled yet scalable during learning and inference.

###  3.1 General Definition

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 parametrizes/“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 [1, 10]. 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 [6]
hierarchies with longer discretization steps and is further detailed in
Section 3.4.

####  3.1.1 Fast time-scale SSM

The fast time-scale (fts) SSM is given by
$\mathcal{S}_{\textrm{fast}}=(\mathcal{Z},\mathcal{A},\mathcal{O},f_{\boldsymbol{l}}^{\textrm{fts}},h^{\textrm{fts}},\Delta
t,\mathcal{L}).$ Here, $\boldsymbol{l}\in\mathcal{L}$ is a task descriptor
that parametrizes 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 $\boldsymbol{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

$\displaystyle
p(\boldsymbol{z}_{k,t}|\boldsymbol{z}_{k,t-1},\boldsymbol{a}_{k,t-1},\boldsymbol{l}_{k})$
$\displaystyle=\mathcal{N}(f^{\textrm{fts}}_{\boldsymbol{l}}(\boldsymbol{z}_{k,t-1},\boldsymbol{a}_{k,t-1},\boldsymbol{l}_{k}),\mathbf{Q}),\textrm{
and }$ $\displaystyle p(\boldsymbol{o}_{k,t}|\boldsymbol{z}_{k,t})$
$\displaystyle=\mathcal{N}(h^{\textrm{fts}}(\boldsymbol{z}_{k,t}),\mathbf{R}).$
(1)

##### 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 $\boldsymbol{l}_{k}$, we have to consider
a distribution over task-descriptors $p(\boldsymbol{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 time window $k$ which is given by

$\displaystyle
p_{\boldsymbol{l}_{k}}(\boldsymbol{z}_{k,t}|\boldsymbol{z}_{k,t-1},\boldsymbol{a}_{k,t-1})$
$\displaystyle=\int
p(\boldsymbol{z}_{k,t}|\boldsymbol{z}_{k,t-1},\boldsymbol{a}_{k,t-1},\boldsymbol{l}_{k})p(\boldsymbol{l}_{k})d\boldsymbol{l}_{k}$
(2)

##### Latent observations.

Following [1], we replace the observations by latent observations and their
uncertainty, i.e., we use latent observation encoders to obtain
$\boldsymbol{w}_{k,t}=\textrm{enc}_{w}(\boldsymbol{o}_{k,t})$ and an
uncertainty encoder
$\boldsymbol{\sigma}_{k,t}=\textrm{enc}_{\sigma}(\boldsymbol{o}_{k,t})$. The
observation model is hence given by
$p(\boldsymbol{w}_{k,t}|\boldsymbol{z}_{k,t})=\mathcal{N}(h^{\textrm{fts}}(\boldsymbol{z}_{k,t}),\textrm{diag}(\boldsymbol{\sigma}_{k,t}))$.

####  3.1.2 Slow time-scale SSM

The slow time-scale (sts) SSM only updates every H time step and uses the task
parameter $\boldsymbol{l}$ as 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 $\boldsymbol{\beta}\in\mathcal{B}$ and
abstract action $\boldsymbol{\alpha}\in\mathcal{A}$ that summarize the
observations and actions respectively throughout the current time window. The
general dynamics model is hence given by

$\displaystyle
p(\boldsymbol{l}_{k}|\boldsymbol{l}_{k-1},\boldsymbol{\alpha}_{k})=\mathcal{N}(f^{\textrm{sts}}(\boldsymbol{l}_{k-1},\boldsymbol{\alpha}_{k}),\boldsymbol{S}).$
(3)

While there exist many ways to implement the abstraction of observations and
actions of the 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.

${z_{1,1}}$${z_{1,2}}$${z_{1,3}}$${w_{1,1}}$${w_{1,2}}$${w_{1,3}}$${a_{1,1}}$${a_{1,2}}$${l_{1}}$$\beta_{{1,t}}$$t=1..H$$\alpha_{{1}}$${z_{2,1}}$${z_{2,2}}$${z_{2,3}}$${w_{2,1}}$${w_{2,2}}$${w_{2,3}}$${a_{2,1}}$${a_{2,2}}$${l_{2}}$$\beta_{{2,t}}$$t=1..H$$\alpha_{{2}}$${z_{3,1}}$${z_{3,2}}$${z_{3,3}}$${w_{3,1}}$${w_{3,2}}$${w_{3,3}}$${a_{3,1}}$${a_{3,2}}$${l_{3}}$$\beta_{{3,t}}$$\alpha_{{3}}$$t=1..H$

Figure 2: The graphical model corresponding to an MTS3 with 2 timescales. The
latent task variable $\boldsymbol{l}_{k}$ captures the slow-changing dynamics
using abstract observation inferred from
$\\{\boldsymbol{\beta}_{k,t}\\}_{t=1}^{H}$and abstract action
$\boldsymbol{\alpha}_{k}$ as described in section 3.1.2. The inference in the
fast time scale uses primitive observations $\boldsymbol{w}_{k,t}$, primitive
actions $\boldsymbol{a}_{k,t}$ and the latent task descriptor $l_{k}$ which
parameterizes the fast-changing dynamics of $\boldsymbol{z}_{k,t}$ for a time
window k as discussed in the section 3.2.

##### Observation abstraction.

In terms of the abstract observation model, we choose to model $H$
observations $\boldsymbol{\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,

$\boldsymbol{\beta}_{k,t}=\textrm{enc}_{\beta}(\boldsymbol{o}_{k,t},t),\quad\boldsymbol{\nu}_{k,t}=\textrm{enc}_{\nu}(\boldsymbol{o}_{k,t},t),$

and
$p(\boldsymbol{\beta}_{k,t}|\boldsymbol{l}_{k})=\mathcal{N}(h^{\textrm{sts}}(\boldsymbol{l}_{k}),\textrm{diag}(\boldsymbol{\nu}_{k,t})).$
Hence, the abstract observation $\boldsymbol{\beta}_{k,t}$ contains the actual
observation $\boldsymbol{o}_{k,t}$ at time step $t$ as well as a temporal
encoding for the time-step. While 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 $\boldsymbol{\alpha}_{k}$ causes the transitions to the
latent task $\boldsymbol{l}_{k}$ from $\boldsymbol{l}_{k-1}$. It should
contain the relevant information of all primitive actions
$\boldsymbol{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 $\boldsymbol{a}_{k,t}$ and the encoding of time-step $t$ is encoded
into its latent representation and its uncertainty estimate, i.e.,

$\boldsymbol{\alpha}_{k,t}=\textrm{enc}_{\alpha}(\boldsymbol{a}_{k,t},t),\quad\boldsymbol{\rho}_{k,t}=\textrm{enc}_{\rho}(\boldsymbol{a}_{k,t},t).$

The single latent actions $\boldsymbol{\alpha}_{k,t}$ can be aggregated into a
consistent representation $\boldsymbol{\alpha}_{k}$ using Bayesian aggregation
[31]. To do so, we use the likelihood
$p(\boldsymbol{\alpha}_{k,t}|\boldsymbol{\alpha}_{k})=\mathcal{N}(\boldsymbol{\alpha}_{k},\textrm{diag}(\boldsymbol{\rho}_{k,t}))$
and obtain the posterior
$p(\boldsymbol{\alpha}_{k}|\boldsymbol{\alpha}_{k,1:H})=\mathcal{N}(\boldsymbol{\mu}_{\alpha_{k}},\boldsymbol{\Sigma}_{\alpha_{k}})$,
which is obtained by following the standard Bayesian aggregation equations,
see Appendix A. 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

$\displaystyle
p_{\boldsymbol{\alpha}_{k}}(\boldsymbol{l}_{k}|\boldsymbol{l}_{k-1},\boldsymbol{\alpha}_{k,1:H})=\int
p_{\boldsymbol{\alpha}_{k}}(\boldsymbol{l}_{k}|\boldsymbol{l}_{k-1},\boldsymbol{\alpha}_{k})p(\boldsymbol{\alpha}_{k}|\boldsymbol{\alpha}_{k,1:H})d\boldsymbol{\alpha}_{k}.$
(4)

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

####  3.1.3 Connecting both SSMs via inference

In the upcoming sections, we will devise Bayesian update rules to obtain the
prior
$p(\boldsymbol{l}_{k}|\boldsymbol{\beta}_{1:k-1},\boldsymbol{\alpha}_{1:k})$
and posterior
$p(\boldsymbol{l}_{k}|\boldsymbol{\beta}_{1:k},\boldsymbol{\alpha}_{1:k})$
belief state for the sts-SSM as well as the belief states for the fts-SSM. The
prior belief
$p(\boldsymbol{l}_{k}|\boldsymbol{\beta}_{1:k-1},\boldsymbol{\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 2.
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.

###  3.2 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,

$\displaystyle
p(\boldsymbol{z}_{t}|\boldsymbol{z}_{t-1},\boldsymbol{a}_{t-1},\boldsymbol{l}_{k})=\mathcal{N}\left(\mathbf{A}\boldsymbol{z}_{t-1}+\mathbf{B}\boldsymbol{a}_{t-1}+\mathbf{C}\boldsymbol{l}_{k},\mathbf{Q}\right),$
(5)

where $\boldsymbol{A}$, $\boldsymbol{B}$, $\boldsymbol{C}$ and
$\boldsymbol{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 [1], we split the latent state
$\boldsymbol{z}_{t}=[\boldsymbol{p}_{t},\boldsymbol{m}_{t}]^{T}$ into its
observable part $\boldsymbol{p}_{t}$ and a part $\boldsymbol{m}_{t}$ that
needs to be observed over time. We also use a linear observation model
$p(\boldsymbol{w}_{t}|\boldsymbol{z}_{t})=\mathcal{N}(\boldsymbol{H}\boldsymbol{z}_{t},\textrm{diag}(\boldsymbol{\sigma}_{t}))$
with $\boldsymbol{H}=[\boldsymbol{I},\boldsymbol{0}]$.

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

$\displaystyle
p_{\boldsymbol{l}_{k}}(\boldsymbol{z}_{t}|\boldsymbol{z}_{t-1},\boldsymbol{a}_{t-1})=\mathcal{N}\left(\mathbf{A}\boldsymbol{z}_{t-1}+\mathbf{B}\boldsymbol{a}_{t-1}+\mathbf{C}\boldsymbol{\mu}_{\boldsymbol{l}_{k}},\mathbf{Q}+\boldsymbol{C}\boldsymbol{\Sigma}_{\boldsymbol{l}_{k}}\boldsymbol{C}^{T}\right),$
(6)

which will be used instead of the standard dynamics equations. We follow the
same factorization assumptions as in [1] 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 [1].

###  3.3 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(\boldsymbol{l}_{k}|\boldsymbol{l}_{k-1},\boldsymbol{\alpha}_{k})=\mathcal{N}(\boldsymbol{X}\boldsymbol{l}_{k-1}+\boldsymbol{Y}\boldsymbol{\alpha}_{k},\boldsymbol{S}),$
where $\boldsymbol{X}$, $\boldsymbol{Y}$ and $\boldsymbol{S}$ are learnable
parameters. The marginalized transition model for the abstract actions is then
given by

$\displaystyle
p_{\boldsymbol{\alpha}_{k}}(\boldsymbol{l}_{k}|\boldsymbol{l}_{k-1})$
$\displaystyle=\int
p(\boldsymbol{l}_{k}|\boldsymbol{l}_{k-1},\boldsymbol{\alpha}_{k})p(\boldsymbol{\alpha}_{k})d\boldsymbol{\alpha}_{k}=\mathcal{N}\left(\mathbf{X}\boldsymbol{l}_{k-1}+\mathbf{Y}\boldsymbol{\mu}_{\alpha_{k}},\mathbf{S}+\boldsymbol{Y}\boldsymbol{\Sigma}_{\boldsymbol{\alpha}_{k}}\boldsymbol{Y}^{T}\right).$
(7)

We can directly use this transition model to obtain the Kalman prediction
update which computes the prior belief
$p_{\boldsymbol{\alpha}_{1:k}}(\boldsymbol{l}_{k}|\boldsymbol{\beta}_{1:k-1})=\mathcal{N}(\boldsymbol{\mu}_{l_{k}}^{-},\boldsymbol{\Sigma}_{l_{k}}^{-})$
from the posterior belief
$p_{\boldsymbol{\alpha}_{1:k-1}}(\boldsymbol{l}_{k-1}|\boldsymbol{\beta}_{1:k-1})=\mathcal{N}(\boldsymbol{\mu}_{l_{k-1}}^{+},\boldsymbol{\Sigma}_{l_{k-1}}^{+})$
of the previous time window, see Appendix A.

##### Observation Update.

Similarly, we will use a linear observation model for the abstract
observations
$p(\boldsymbol{\beta}_{k,t}|\boldsymbol{l}_{k})=\mathcal{N}(\mathbf{H}\boldsymbol{l}_{k},\textrm{diag}(\boldsymbol{\nu}_{k,t}))$
with $\mathbf{H}=[\mathbf{I},\mathbf{0}]$. As can be seen from the definition
of the observation matrix $\mathbf{H}$, the latent space is also decomposed
into its observable and unobservable part, i.e.,
$\boldsymbol{l}_{k}=[\boldsymbol{u}_{k},\boldsymbol{v}_{k}]$. In difference to
the standard factorized Kalman observation update given in Appendix A, 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_{\boldsymbol{\alpha}_{1:k}}(\boldsymbol{l}_{k}|\boldsymbol{\beta}_{1:k})=\mathcal{N}(\boldsymbol{\mu}_{l_{k}}^{+},\boldsymbol{\Sigma}_{l_{k}}^{+})$,
such an update would be very slow and also cause numerical inaccuracies.
Hence, we devise a new permutation invariant version of the update rule that
allows parallel processing with set encoders [33]. We found that this update
rule is easier to formalize using precision matrices. Hence, we first
transform the prior covariance vectors $\boldsymbol{\sigma}_{l_{k}}^{u-}$,
$\boldsymbol{\sigma}_{l_{k}}^{l-}$ and $\boldsymbol{\sigma}_{l_{k}}^{s-}$ to
its corresponding precision representation
$\boldsymbol{\lambda}_{l_{k}}^{u-}$, $\boldsymbol{\lambda}_{l_{k}}^{l-}$ and
$\boldsymbol{\lambda}_{l_{k}}^{s-}$ which can be performed using block-wise
matrix inversions of $\boldsymbol{\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 A, B. Subsequently, we compute the posterior precision,
where only $\boldsymbol{\lambda}_{l_{k}}^{u}$ is changed by

$\displaystyle\boldsymbol{\lambda}_{l_{k}}^{u+}=\boldsymbol{\lambda}_{l_{k}}^{u-}+\sum_{t=1}^{H}\boldsymbol{1}\oslash\boldsymbol{\nu}_{k,t}$
(8)

while $\boldsymbol{\lambda}_{l_{k}}^{l+}=\boldsymbol{\lambda}_{l_{k}}^{l-}$
and $\boldsymbol{\lambda}_{l_{k}}^{s+}=\boldsymbol{\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 $\boldsymbol{\sigma}_{l_{k}}^{u+}$, $\boldsymbol{\sigma}_{l_{k}}^{l+}$
and $\boldsymbol{\sigma}_{l_{k}}^{s+}$ using only scalar inversions, see
Appendix A, B. The posterior mean $\boldsymbol{\mu}_{l,k}^{+}$ can now be
obtained from the prior mean $\boldsymbol{\mu}_{l,k}^{-}$ as

$\displaystyle\boldsymbol{\mu}_{l,k}^{+}=\boldsymbol{\mu}_{l,k}^{-}+\left[\begin{array}[]{l}\boldsymbol{\sigma}_{l_{k}}^{u+}\\\
\boldsymbol{\sigma}_{l_{k}}^{s+}\\\
\end{array}\right]\odot\left[\begin{array}[]{l}\sum_{t=1}^{H}\left(\boldsymbol{\beta}_{k,t}-\boldsymbol{\mu}^{\mathrm{u},-}_{l_{k}}\right)\oslash\boldsymbol{\nu}_{k,t}\\\
\sum_{t=1}^{H}\left(\boldsymbol{\beta}_{k,t}-\boldsymbol{\mu}^{\mathrm{u},-}_{l_{k}}\right)\oslash\boldsymbol{\nu}_{k,t}\\\
\end{array}\right].$ (9)

Note that for $H=1$, i.e a single observation, the given equation is
equivalent to the RKN updates. Moreover, the given rule constitutes a
unification of the batch update rule for Bayesian aggregation [31] and the
incremental Kalman update for our factorization of the belief state
representation [1] detailed in Appendix A.

###  3.4 A General Definition For an N-level MTS3

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 model $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
parametrizes the dynamics model of the SSM and is held constant for a local
window of $H_{i+1}$ steps. $l_{i}$ is a function of the latent state of 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}$.

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

##  4 MTS3 as a Hierarchical World Model

MTS3 allows for a natural way to build world models that can deal with partial
observability, non-stationarity 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.

###  4.1 Conditional Multi Time Predictions With World Model

Conditional multi-step 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 abstractions. The prediction update for the sts-SSM makes prior
estimates about future latent task variables conditioned on the abstract
action representations. Whereas, 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. For initializing the prior belief
$p(\boldsymbol{z}_{k,1})$ for the first time step of the time window $k$, we
use the prior belief $p(\boldsymbol{z}_{k-1,H+1})$ of the last time step of
the time window $k-1$.

###  4.2 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.,

$\displaystyle L$ $\displaystyle=\sum_{k=1}^{N}\sum_{t=1}^{H}\log
p(\boldsymbol{o}_{k,t+1}|\boldsymbol{\beta}_{1:k-1},\boldsymbol{\alpha}_{1:k},\boldsymbol{w}_{k,1:t},\boldsymbol{a}_{k,1:t})$
$\displaystyle=\sum_{k=1}^{N}\sum_{t=1}^{H}\log\iint
p(\boldsymbol{o}_{k,t+1}|\boldsymbol{z}_{k,t+1})p(\boldsymbol{z}_{k,t+1}|\boldsymbol{w}_{k,1:t},\boldsymbol{a}_{k,1:t},\boldsymbol{l}_{k})p(\boldsymbol{l}_{k}|\boldsymbol{\beta}_{1:k-1},\boldsymbol{\alpha}_{1:k})d\boldsymbol{z}_{k,t+1}d\boldsymbol{l}_{k}$
$\displaystyle=\sum_{k=1}^{N}\sum_{t=1}^{H}\log\int
p(\boldsymbol{o}_{k,t+1}|\boldsymbol{z}_{k,t+1})p_{\boldsymbol{l}_{k}}(\boldsymbol{z}_{k,t+1}|\boldsymbol{w}_{k,1:t},\boldsymbol{a}_{k,1:t})d\boldsymbol{z}_{k,t+1}.$
(10)

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

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

Gradients are computed using (truncated) backpropagation through time (BPTT)
[32] and clipped. We optimize the objective using the Adam [15] stochastic
gradient descent optimizer with default parameters. We refer to Appendix A for
more details. For training, we also initialize the prior belief
$p(\boldsymbol{z}_{k,1})$ with the prior belief
$p_{\boldsymbol{l}_{k-1}}(\boldsymbol{z}_{k-1,H+1}|\boldsymbol{w}_{k-1,1:H},\boldsymbol{a}_{k-1,1:H})$
from of 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 behaviour. Yet, the gradients can still flow
between time windows for the fts-SSM via the sts-SSM.

###  4.3 Imputation Based 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 long-term prediction performance, we can treat the long-term
prediction problem as a case of the “missing value” problem, where the missing
observations are at the 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 mini-batch.

##  5 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 [29]. 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 [18]. 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. [10, 1, 25] use neural network
architectures based on exact inference on SSMs and perform state estimation
and dynamics prediction tasks. [26] extend these models to modelling non-
stationary dynamics. [17, 14, 11] 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.

Transformers Recent advancements in Transformers [30, 24, 3], 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 [35, 19] and learning world models [20]. [35, 19, 21] use
transformer architectures based on a direct multistep loss [34] and show
promising results for long-term forecasting since they avoid error
accumulation from autoregression. On the other hand [20] 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.

##  6 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?

###  6.1 Baseline Dynamics Models

While a full description of our baselines can be found in Appendix E, a brief
description of them is given here: (a) RNNs \- We compare our method to two
widely used recurrent neural network architectures, LSTMs [12] and GRUs [5].
(b) RSSMs \- Among several RSSMs from the literature, we chose RKN [1] and
HiP-RSSM [26] 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 [30]
variants. The first variant (AR-Transformer) relies on a GPT-like
autoregressive prediction [24, 3]. Whereas the second variant (Multi-
Transformer) uses direct multi-step loss [34] from recent literature on long
horizon time-series forecasting [35, 19, 21]. Here, multistep ahead
predictions are performed using a single shot given the action sequences.

###  6.2 Environments and Datasets

We experiment with three broad datasets. While full descriptions of these
datasets, dataset creation procedure, and overall statistics are given in
Appendix D, a brief description of them is as follows. (a) D4RL Datasets \- We
use a set of 3 different environments/agents from D4RL dataset [7], 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.

(a)

(b)

(c)

Figure 3: (a) and (b) Comparison with baselines in terms of RMSE for long
horizon predictions (in seconds) as discussed in Section 6.3. (c) Ablation
with different MTS3 variants is discussed in 4.

###  6.3 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 RMSE in Figure 3. We
see that MTS3 gives consistently good long-term action conditional future
predictions on all 6 datasets. Deep Kalman models [1, 26] which operate on a
single time scale fail to give meaningful mean estimates beyond a few
milliseconds. Similarly, widely used RNN baselines [12, 5] which form the
backbone of several world models [9, 11] 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 [8], MTS3
is the only model that gives meaningful long horizon predictions. A
visualization of the predicted trajectories vs. ground truth is given in
Appendix C.

###  6.4 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. As seen in table 1, 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 | 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$ | ✗ | ✗

Table 1: 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.

###  6.5 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}=\mathbf{H}=[\mathbf{I},\mathbf{0}]$ as
discussed in Sections 3.3 and 3.2, (ii) action abstractions discussed in
Section 3.1.2, (iii) training by imputation.

To analyze the importance of the memory component, we derived and implemented
an MTS3 variant with an observation model of $h^{sts}=h^{fts}=\mathbf{I}$ and
a pure diagonal matrix representation for the covariance matrices. As seen in
Figure 3(c), 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 4.3. As seen in Figure 3(c) when trained for 1
step ahead predictions without imputation, MTS3 performs significantly worse
for long-term prediction suggesting the importance of this training regime.

###  6.6 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 4, 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 4, 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.

(d)

(e)

Figure 4: 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 6.6 on the hydraulics dataset. (b) 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 figures are for $H=3$
and $H=10$. Bottom figures are for $H=30$ and $H=75$. Note that the results
reported in the paper are with $H=30$.

##  7 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 Appendix F and G.

##  8 Acknowledgement

We thank the anonymous reviewers for the valuable remarks and discussions
which greatly improved the quality of this paper. This work was supported by
funding from the pilot program Core Informatics of the Helmholtz Association
(HGF). The authors acknowledge support by the state of Baden-Württemberg
through bwHPC, as well as the HoreKa supercomputer funded by the Ministry of
Science, Research and the Arts Baden-Württemberg and by the German Federal
Ministry of Education and Research.

## References

  * Becker et al. [2019] Philipp Becker, Harit Pandya, Gregor Gebhardt, Cheng Zhao, C James Taylor, and Gerhard Neumann.  Recurrent kalman networks: Factorized inference in high-dimensional deep feature spaces.  In _International Conference on Machine Learning_ , pages 544–552. PMLR, 2019.
  * Bishop [2006] Christopher M Bishop.  Pattern recognition.  _Machine learning_ , 128(9), 2006.
  * Brown et al. [2020] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al.  Language models are few-shot learners.  _Advances in neural information processing systems_ , 33:1877–1901, 2020.
  * Camacho and Alba [2013] Eduardo F Camacho and Carlos Bordons Alba.  _Model predictive control_.  Springer science & business media, 2013.
  * Cho et al. [2014] Kyunghyun Cho, Bart van Merriënboer, Dzmitry Bahdanau, and Yoshua Bengio.  On the properties of neural machine translation: Encoder–decoder approaches.  In _Proceedings of SSST-8, Eighth Workshop on Syntax, Semantics and Structure in Statistical Translation_ , pages 103–111. Association for Computational Linguistics, October 2014.
  * Dayan and Hinton [1992] Peter Dayan and Geoffrey E Hinton.  Feudal reinforcement learning.  _Advances in neural information processing systems_ , 5, 1992.
  * Fu et al. [2020] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine.  D4rl: Datasets for deep data-driven reinforcement learning.  _arXiv preprint arXiv:2004.07219_ , 2020.
  * Gupta et al. [2019] Abhishek Gupta, Vikash Kumar, Corey Lynch, Sergey Levine, and Karol Hausman.  Relay policy learning: Solving long-horizon tasks via imitation and reinforcement learning.  _arXiv preprint arXiv:1910.11956_ , 2019.
  * Ha and Schmidhuber [2018] David Ha and Jürgen Schmidhuber.  World models.  _arXiv preprint arXiv:1803.10122_ , 2018.
  * Haarnoja et al. [2016] Tuomas Haarnoja, Anurag Ajay, Sergey Levine, and Pieter Abbeel.  Backprop kf: Learning discriminative deterministic state estimators.  In _Advances in neural information processing systems_ , pages 4376–4384, 2016.
  * Hafner et al. [2019] Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson.  Learning latent dynamics for planning from pixels.  In _International Conference on Machine Learning_ , pages 2555–2565. PMLR, 2019.
  * Hochreiter and Schmidhuber [1997] Sepp Hochreiter and Jürgen Schmidhuber.  Long short-term memory.  _Neural Comput._ , 9(8):1735–1780, November 1997.  ISSN 0899-7667.
  * Jordan [2004] Michael I Jordan.  Graphical models.  _Statistical science_ , 19(1):140–155, 2004.
  * Karl et al. [2016] Maximilian Karl, Maximilian Soelch, Justin Bayer, and Patrick Van der Smagt.  Deep variational bayes filters: Unsupervised learning of state space models from raw data.  _arXiv preprint arXiv:1605.06432_ , 2016.
  * Kingma and Ba [2014] Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Koller and Friedman [2009] Daphne Koller and Nir Friedman.  _Probabilistic graphical models: principles and techniques_.  MIT press, 2009.
  * Krishnan et al. [2017] Rahul Krishnan, Uri Shalit, and David Sontag.  Structured inference networks for nonlinear state space models.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 31, 2017.
  * LeCun [2022] Yann LeCun.  A path towards autonomous machine intelligence version 0.9. 2, 2022-06-27.  _Open Review_ , 62, 2022.
  * Liu et al. [2022] Yong Liu, Haixu Wu, Jianmin Wang, and Mingsheng Long.  Non-stationary transformers: Exploring the stationarity in time series forecasting.  In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, _Advances in Neural Information Processing Systems_ , 2022.  URL https://openreview.net/forum?id=ucNDIDRNjjv.
  * Micheli et al. [2023] Vincent Micheli, Eloi Alonso, and François Fleuret.  Transformers are sample-efficient world models.  In _The Eleventh International Conference on Learning Representations_ , 2023.  URL https://openreview.net/forum?id=vhFu1Acb0xb.
  * Nie et al. [2023] Yuqi Nie, Nam H Nguyen, Phanwadee Sinthong, and Jayant Kalagnanam.  A time series is worth 64 words: Long-term forecasting with transformers.  In _The Eleventh International Conference on Learning Representations_ , 2023.  URL https://openreview.net/forum?id=Jbdc0vTOcol.
  * Pezzulo and Cisek [2016] Giovanni Pezzulo and Paul Cisek.  Navigating the affordance landscape: feedback control as a process model of behavior and cognition.  _Trends in cognitive sciences_ , 20(6):414–424, 2016.
  * Precup and Sutton [1997] Doina Precup and Richard S Sutton.  Multi-time models for temporally abstract planning.  _Advances in neural information processing systems_ , 10, 1997.
  * Radford et al. [2019] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al.  Language models are unsupervised multitask learners.  _OpenAI blog_ , 1(8):9, 2019.
  * Shaj et al. [2020] Vaisakh Shaj, Philipp Becker, Dieter Buchler, Harit Pandya, Niels van Duijkeren, C James Taylor, Marc Hanheide, and Gerhard Neumann.  Action-conditional recurrent kalman networks for forward and inverse dynamics learning.  _Conference On Robot Learning_ , 2020.
  * Shaj et al. [2022] Vaisakh Shaj, Dieter Buchler, Rohit Sonker, Philipp Becker, and Gerhard Neumann.  Hidden parameter recurrent state space models for changing dynamics scenarios.  _International Conference On Learning Representations_ , 2022.
  * Shanahan and Mitchell [2022] Murray Shanahan and Melanie Mitchell.  Abstraction for deep reinforcement learning.  _arXiv preprint arXiv:2202.05839_ , 2022.
  * Sonker and Dutta [2020] Rohit Sonker and Ashish Dutta.  Adding terrain height to improve model learning for path tracking on uneven terrain by a four wheel robot.  _IEEE Robotics and Automation Letters_ , 6(1):239–246, 2020.
  * Sutton [1995] Richard S Sutton.  Td models: Modeling the world at a mixture of time scales.  In _Machine Learning Proceedings 1995_ , pages 531–539. Elsevier, 1995.
  * Vaswani et al. [2017] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  _Advances in neural information processing systems_ , 30, 2017.
  * Volpp et al. [2020] Michael Volpp, Fabian Flürenbrock, Lukas Grossberger, Christian Daniel, and Gerhard Neumann.  Bayesian context aggregation for neural processes.  In _International Conference on Learning Representations_ , 2020.
  * Werbos [1990] Paul J Werbos.  Backpropagation through time: what it does and how to do it.  _Proceedings of the IEEE_ , 78(10):1550–1560, 1990.
  * Zaheer et al. [2017] Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Ruslan Salakhutdinov, and Alexander Smola.  Deep sets.  _arXiv preprint arXiv:1703.06114_ , 2017.
  * Zeng et al. [2022] Ailing Zeng, Muxi Chen, Lei Zhang, and Qiang Xu.  Are transformers effective for time series forecasting?  _arXiv preprint arXiv:2205.13504_ , 2022.
  * Zhou et al. [2021] Haoyi Zhou, Shanghang Zhang, Jieqi Peng, Shuai Zhang, Jianxin Li, Hui Xiong, and Wancai Zhang.  Informer: Beyond efficient transformer for long sequence time-series forecasting.  In _Proceedings of the AAAI conference on artificial intelligence_ , volume 35, pages 11106–11115, 2021.

##  Appendix A Implementation Details

Figure 5: 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 3.2 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 3.3.

###  A.1 Inference In Slow Time Scale SSM

####  A.1.1 Inferring Action Abstraction (sts-SSM)

$\alpha_{k,t}$$t$$a_{t,k}$$\alpha_{k}$$N$

Figure 6: 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.

Given a set of encoded primitive actions and their corresponding variances
$\\{\boldsymbol{\alpha_{k,t}},\boldsymbol{\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(\boldsymbol{\alpha}_{k}|\boldsymbol{\alpha}_{k,1:H})=\mathcal{N}(\boldsymbol{\mu}_{\alpha_{k}},\boldsymbol{\Sigma}_{\alpha_{k}})=\mathcal{N}(\boldsymbol{\mu}_{\alpha_{k}},\textrm{diag}\boldsymbol{(}\sigma_{\alpha_{k}}))$
as a Bayesian aggregation [31] of these using the following closed-form
equations:

$\displaystyle\boldsymbol{\sigma}_{\alpha_{k}}$
$\displaystyle=\left(\left(\boldsymbol{\sigma}_{0}\right)^{\ominus}+\sum_{n=1}^{N}\left(\left(\boldsymbol{\rho_{k,t}}\right)^{\ominus}\right)\right)^{\ominus},$
$\displaystyle\boldsymbol{\mu}_{\alpha_{k}}$
$\displaystyle=\boldsymbol{\mu}_{0}+\boldsymbol{\sigma}_{\alpha_{k}}\odot\sum_{n=1}^{N}\left(\boldsymbol{\alpha_{k,t}}-\boldsymbol{\mu}_{0}\right)\oslash\boldsymbol{\rho_{k,t}}$

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 5.

####  A.1.2 Task Prediction (sts-SSM)

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

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

$\displaystyle\boldsymbol{\mu}_{l_{k}}^{-}$
$\displaystyle=\mathbf{X}\boldsymbol{\mu}_{l_{k-1}}^{+}+\mathbf{Y}\boldsymbol{\alpha}_{k}$
(11) $\displaystyle\boldsymbol{\Sigma}_{l_{k}}^{-}$
$\displaystyle=\mathbf{X}\boldsymbol{\Sigma}_{l_{k-1}}^{+}\mathbf{X}^{T}+\mathbf{Y}\boldsymbol{\Sigma}_{\alpha_{k}}\mathbf{Y}^{T}+\mathbf{S}.$

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

####  A.1.3 Task Update (sts-SSM)

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

To do so we first invert the prior covariance matrix
$\left[\begin{array}[]{cc}\boldsymbol{\Sigma}_{l_{k}}^{u}&\boldsymbol{\Sigma}_{l_{k}}^{s}\\\
\boldsymbol{\Sigma}_{l_{k}}^{s}&\boldsymbol{\Sigma}_{l_{k}}^{l}\end{array}\right]^{+}$
to the precision matrix
$\left[\begin{array}[]{cc}\boldsymbol{\lambda}_{l_{k}}^{u}&\boldsymbol{\lambda}_{l_{k}}^{s}\\\
\boldsymbol{\lambda}_{l_{k}}^{s}&\boldsymbol{\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
$\boldsymbol{\lambda}_{l_{k}}^{u}$ is changed by

$\displaystyle\boldsymbol{\lambda}_{l_{k}}^{u+}=\boldsymbol{\lambda}_{l_{k}}^{u-}+\sum_{t=1}^{H}\boldsymbol{1}\oslash\boldsymbol{\nu}_{k,t}$
(12)

while $\boldsymbol{\lambda}_{l_{k}}^{l+}=\boldsymbol{\lambda}_{l_{k}}^{l-}$
and $\boldsymbol{\lambda}_{l_{k}}^{s+}=\boldsymbol{\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
$\boldsymbol{\sigma}_{l_{k}}^{u+}$, $\boldsymbol{\sigma}_{l_{k}}^{l+}$ and
$\boldsymbol{\sigma}_{l_{k}}^{s+}$. Now, the posterior mean
$\boldsymbol{\mu}_{l,k}^{+}$ can be obtained from the prior mean
$\boldsymbol{\mu}_{l,k}^{-}$ as

$\displaystyle\boldsymbol{\mu}_{l,k}^{+}=\boldsymbol{\mu}_{l,k}^{-}+\left[\begin{array}[]{l}\boldsymbol{\sigma}_{l_{k}}^{u+}\\\
\boldsymbol{\sigma}_{l_{k}}^{s+}\\\
\end{array}\right]\odot\left[\begin{array}[]{l}\sum_{t=1}^{H}\left(\boldsymbol{\beta}_{k,t}-\boldsymbol{\mu}^{\mathrm{u},-}_{l_{k}}\right)\oslash\boldsymbol{\nu}_{k,t}\\\
\sum_{t=1}^{H}\left(\boldsymbol{\beta}_{k,t}-\boldsymbol{\mu}^{\mathrm{u},-}_{l_{k}}\right)\oslash\boldsymbol{\nu}_{k,t}\\\
\end{array}\right].$ (13)

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

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

###  A.2 Inference In Fast Time Scale SSM

The inference in fts-SSM for a time-window $k$ involves two stages as
illustrated in Figure LABEL:fig:_schematic, 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.

####  A.2.1 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(\boldsymbol{z}_{k,t}|\boldsymbol{w}^{k}_{1:t-1},\boldsymbol{a}^{k}_{1:t-1},\boldsymbol{\beta}_{1:k-1},\boldsymbol{\alpha}_{1:k-1})=\mathcal{N}(\boldsymbol{\mu}_{z_{k,t}}^{-},\mathbf{\Sigma}_{z_{k,t}}^{-})$
in closed form, where

$\displaystyle\boldsymbol{\mu}_{z_{k,t}}^{-}$
$\displaystyle=\mathbf{A}\boldsymbol{\mu}_{z_{k,t-1}}^{-}+\mathbf{B}\boldsymbol{a}_{k,t-1}+\mathbf{C}\boldsymbol{\mu}_{l_{k}}^{-},$
(14) $\displaystyle\mathbf{\Sigma}_{k,t}^{-}$
$\displaystyle=\mathbf{A}\mathbf{\Sigma}_{k,t-1}^{+}\mathbf{A}^{T}+\mathbf{C}\boldsymbol{\Sigma}_{l_{k}}^{-}\mathbf{C}^{T}+\mathbf{Q}.$

####  A.2.2 Observation Update (fts-SSM)

In this stage, we compute the posterior belief
$p(\boldsymbol{z}_{k,t}|\boldsymbol{w}^{k}_{1:t},\boldsymbol{a}^{k}_{1:t},\boldsymbol{\beta}_{1:k},\boldsymbol{\alpha}_{1:k-1})=\mathcal{N}(\boldsymbol{\mu}_{z_{k,t}}^{-},\mathbf{\Sigma}_{z_{k,t}}^{-})$.
using the same closed-form update as in [1]. 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}(\boldsymbol{\mu}_{z_{k,t}}^{-},\mathbf{\Sigma}_{z_{k,t}}^{-})=\mathcal{N}(\left[\begin{array}[]{c}\boldsymbol{\mu}_{t}^{u+}\\\
\boldsymbol{\mu}_{t}^{l+}\end{array}\right],\left[\begin{array}[]{cc}\boldsymbol{\Sigma}_{t}^{u}&\boldsymbol{\Sigma}^{s}_{t}\\\
\boldsymbol{\Sigma}^{s}_{t}&\boldsymbol{\Sigma}^{l}_{t}\end{array}\right]^{+})$,
with
$\boldsymbol{\Sigma}_{t}^{u}=\textrm{diag}(\boldsymbol{\sigma}_{t}^{s}),\;\boldsymbol{\Sigma}_{t}^{l}=\textrm{diag}(\boldsymbol{\sigma}_{t}^{l})$
and $\boldsymbol{\Sigma}_{t}^{s}=\textrm{diag}(\boldsymbol{\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,

$\displaystyle\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],$

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

$\displaystyle\boldsymbol{\sigma}^{\mathrm{u},+}_{t}$
$\displaystyle=\boldsymbol{\sigma}^{\mathrm{u},-}_{t}\odot\boldsymbol{\sigma}^{\mathrm{u},-}_{t}\oslash\left(\boldsymbol{\sigma}_{t}^{\mathrm{u},-}+\boldsymbol{\sigma}_{t}^{\mathrm{obs}}\right),$
$\displaystyle\boldsymbol{\sigma}^{\mathrm{s},+}_{t}$
$\displaystyle=\boldsymbol{\sigma}^{\mathrm{u},-}_{t}\odot\boldsymbol{\sigma}^{\mathrm{s},-}_{t}\oslash\left(\boldsymbol{\sigma}_{t}^{\mathrm{u},-}+\boldsymbol{\sigma}_{t}^{\mathrm{obs}}\right),$
$\displaystyle\boldsymbol{\sigma}^{\mathrm{l},+}_{t}$
$\displaystyle=\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),$

, where $\odot$ denotes the elementwise vector product and $\oslash$ denotes
an elementwise vector division.

###  A.3 Modelling Assumptions

####  A.3.1 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 [25].

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.

####  A.3.2 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.

###  A.4 Training

####  A.4.1 Training Objective Derivation

We further expand on the training objective in Section 4.2 here. The training
objective for the MTS3 involves maximizing the posterior predictive log-
likelihood which for a single trajectory, can be derived as,

$\displaystyle L$ $\displaystyle=\sum_{k=1}^{N}\sum_{t=1}^{H}\log
p(\boldsymbol{o}_{k,t+1}|\boldsymbol{\beta}_{1:k-1},\boldsymbol{\alpha}_{1:k-1},\boldsymbol{w}_{k,1:t},\boldsymbol{a}_{k,1:t})$
$\displaystyle=\sum_{k=1}^{N}\sum_{t=1}^{H}\log\iint
p(\boldsymbol{o}_{k,t+1}|\boldsymbol{z}_{k,t+1})p(\boldsymbol{z}_{k,t+1}|\boldsymbol{w}_{k,1:t},\boldsymbol{a}_{k,1:t},\boldsymbol{l}_{k})p(\boldsymbol{l}_{k}|\boldsymbol{\beta}_{1:k-1},\boldsymbol{\alpha}_{1:k-1})d\boldsymbol{z}_{k,t+1}d\boldsymbol{l}_{k}$
$\displaystyle=\sum_{k=1}^{N}\sum_{t=1}^{H}\log\int
p(\boldsymbol{o}_{k,t+1}|\boldsymbol{z}_{k,t+1})p_{\boldsymbol{l}_{k}}(\boldsymbol{z}_{k,t+1}|\boldsymbol{w}_{k,1:t},\boldsymbol{a}_{k,1:t})d\boldsymbol{z}_{k,t+1}.$
(15)

The extension to multiple trajectories is straightforward. The approximation
to the objective is done based on a moment-matching perspective as discussed
in Section 4.2 of the main paper.

####  A.4.2 Initialization

We initialize the states $\boldsymbol{l}_{1}$ and $\boldsymbol{z}_{1,1}$ at
both timescales for the first-time window $k=1$ with an all zeros vector and
corresponding covariance matrices as
$\mathbf{\Sigma}_{l_{1}}=\mathbf{\Sigma}_{z_{1,1}}=10\cdot\mathbf{I}$. For
subsequent windows, the prior belief $p(\boldsymbol{z}_{k,1})$ for the first
time step of time window $k$, is initialized using the posterior belief
$p_{\boldsymbol{l}_{k-1}}(\boldsymbol{z}_{k-1,H}|\boldsymbol{w}_{k-1,1:H},\boldsymbol{a}_{k-1,1:H})$
of the last time step of time window $k-1$.

It is also crucial to correctly initialize the transition matrix at both time
scales so that the transition does not yield an unstable system. Initially,
the transition model should focus on copying the encoder output so that the
encoder can learn how to extract good features if observations are available
and useful. We initialize the diagonal elements of the transition matrix at
both timescales with 1 and the off-diagonal elements with 0.2, while the rest
of the elements are set to 0, a choice inspired from [1].

####  A.4.3 Learnable Parameters

The learnable parameters in the computation graph are as follows:

##### Fast Time Scale SSM:

The linear transition model A, the non-linear control factor b, the linear
latent task transformation model C, the transition noise Q, along with the
observation encoder and the output decoder.

##### Slow Time Scale SSM:

The linear transition model X, the linear control model Y, the transition
noise S, along with the observation set encoder and the action set encoder.

$r_{i}$$\ell$$N$

Figure 8: Graphical Model For Bayesian conditioning with $N$ observations.

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

###  A.5 Bayesian Conditioning As Permutation Invariant Set Operations

######  Gaussian Update Rule 1 (Bayesian Conditioning)

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

$\displaystyle\mathbf{\Lambda}_{l}^{+}$
$\displaystyle=\mathbf{\Lambda}_{l}^{-}+\left[\begin{array}[]{ll}\textrm{diag}(\sum_{i=1}^{n}\frac{1}{\boldsymbol{\sigma}_{i}^{obs}}),&\boldsymbol{0}\\\
\boldsymbol{0},&\boldsymbol{0}\\\ \end{array}\right]$ (16)
$\displaystyle\boldsymbol{\mu}_{l}^{+}$
$\displaystyle=\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(\boldsymbol{r}_{i}-\boldsymbol{\mu}^{\mathrm{u},-}_{l}\right)\odot\frac{1}{\boldsymbol{\sigma}_{i}^{obs}}\\\
\sum_{i=1}^{N}\left(\boldsymbol{r}_{i}-\boldsymbol{\mu}^{\mathrm{u},-}_{l}\right)\odot\frac{1}{\boldsymbol{\sigma}_{i}^{obs}}\\\
\end{array}\right]\hskip 42.67912pt$

Note that $\boldsymbol{\Sigma}_{l}$ is the covariance matrix which is the
inverse of the precision matrix $\boldsymbol{\Lambda}_{l}$. Due to the
observation model assumption $\mathbf{H}=[\mathbf{I},\mathbf{0}]$, they take
block diagonal form,

$\boldsymbol{\Sigma}_{l}=\left[\begin{array}[]{cc}\boldsymbol{\Sigma}_{l}^{u}&\boldsymbol{\Sigma}^{s}_{l}\\\
\boldsymbol{\Sigma}^{s}_{l}&\boldsymbol{\Sigma}^{l}_{l}\end{array}\right],\textrm{
with
}\boldsymbol{\Sigma}_{u}=\textrm{diag}(\boldsymbol{\sigma}_{l}^{u}),\;\boldsymbol{\Sigma}_{l}=\textrm{diag}(\boldsymbol{\sigma}_{l}^{l})\textrm{
and }\boldsymbol{\Sigma}_{s}=\textrm{diag}(\boldsymbol{\sigma}_{l}^{s}).$

##### Proof:

##### 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 $\boldsymbol{r}$.
If the marginal Gaussian distribution for the latent variable $\boldsymbol{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 he single observation
$\boldsymbol{r}$ given $\boldsymbol{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,

$\displaystyle p(\mathbf{l}\mid\mathbf{r})$
$\displaystyle=\mathcal{N}\left(\mathbf{l}\mid\mathbf{\Sigma}\left\\{\mathbf{H}^{\mathrm{T}}\mathbf{L}\mathbf{r}+\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).$
(17)

We refer to [2] 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 $\mathbf{H}=[\mathbf{I},\mathbf{0}]$. In the
derivation, we represent the set of N observations as a random vector

$\bar{r}=\left[\begin{array}[]{ll}\boldsymbol{r_{1}}\\\ \boldsymbol{r_{2}}\\\
\boldsymbol{.}\\\ \boldsymbol{.}\\\
\boldsymbol{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},\mathbf{\Sigma}_{r}\right)$,
where the diagonal covariance matrix has the following form:

$\mathbf{\Sigma}_{r}=\left[\begin{array}[]{llllll}\textrm{diag}(\boldsymbol{\sigma}_{r_{1}}),&0,&0,&..,&0\\\
0,&\textrm{diag}(\boldsymbol{\sigma}_{r_{2}}),&0,&..,&0\\\
\boldsymbol{.},&.,&.,&..,&.\\\ \boldsymbol{.},&.,&.,&..,&.\\\
0,&0,&0,&..,&\textrm{diag}(\boldsymbol{\sigma}_{r_{N}})\end{array}\right]_{Nd\times
Nd}.$

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

$\bar{\mathbf{H}}=\left[\begin{array}[]{ll}\mathbf{H}\\\ \mathbf{H}\\\
\mathbf{.}\\\ \mathbf{.}\\\ \mathbf{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 (17) as follows.

$\displaystyle\Lambda_{l}^{+}$ $\displaystyle=(\Sigma_{l}^{+})^{-1}$ (18)
$\displaystyle=\boldsymbol{\Sigma}_{l}^{-1}+\bar{\mathbf{H}}^{T}\mathbf{\Sigma}_{r}\bar{\mathbf{H}}$
$\displaystyle=\boldsymbol{\Sigma}_{l}^{-1}+\left[\begin{array}[]{ll}\textrm{diag}(\boldsymbol{\sigma}_{r_{1}}),\textrm{diag}(\boldsymbol{\sigma}_{r_{2}}),\textrm{diag}(\boldsymbol{\sigma}_{r_{3}}),.,.,\textrm{diag}(\boldsymbol{\sigma}_{r_{N}})\\\
\boldsymbol{0},\quad\quad\boldsymbol{0},\quad\quad\boldsymbol{0},\quad.,.,\quad\boldsymbol{0}\\\
\end{array}\right]_{2d\times nd}\bar{\mathbf{H}}$
$\displaystyle=\boldsymbol{\lambda}_{l}^{-}+\left[\begin{array}[]{ll}\mathrm{diag}(\sum_{i=1}^{n}\frac{1}{\boldsymbol{\sigma}_{r_{i}}}),&\boldsymbol{0}\\\
\boldsymbol{0},&\boldsymbol{0}\\\ \end{array}\right]_{2d\times 2d}$
$\displaystyle\mu_{l}^{+}$
$\displaystyle=\boldsymbol{\mu}_{l}^{-}+(\Lambda^{+})^{-1}\bar{\mathbf{H}}^{T}\left(\sigma_{\boldsymbol{r}}^{-2}\boldsymbol{I}\right)\left(\boldsymbol{y}-\bar{\mathbf{H}}\boldsymbol{\mu}_{\boldsymbol{x}}\right)$
$\displaystyle=\boldsymbol{\mu}_{l}^{-}+\Sigma^{+}\bar{\mathbf{H}}\left(\sigma_{\boldsymbol{r}}^{-2}\boldsymbol{I}\right)\left(\boldsymbol{y}-\bar{\mathbf{H}}\boldsymbol{\mu}_{\boldsymbol{x}}\right)$
$\displaystyle=\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{\mathbf{H}}\boldsymbol{\mu}_{\boldsymbol{x}}\right)$
$\displaystyle=\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]$
$\displaystyle=\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}{\boldsymbol{\sigma}_{r_{i}}}\\\
\sum_{i=1}^{N}\left(\mathbf{r_{n}}-\mathbf{\mu}^{\mathrm{u},-}_{l}\right)\odot\frac{1}{\boldsymbol{\sigma}_{r_{i}}}\\\
\end{array}\right]$

Here $\mu_{l}^{+}$ is the posterior mean and $\boldsymbol{\Lambda}_{l}^{+}$ is
the posterior precision matrix.

###### Corollary 1

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

###  A.6 Derivation For Matrix Inversions as Scalar Operations

###### Inversion Of Block Diagonal Matrix 1

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

$\displaystyle\begin{split}\boldsymbol{b}^{u}&=\boldsymbol{a}_{l}\oslash{(\boldsymbol{a}_{u}\odot\boldsymbol{a}_{l}-\boldsymbol{a}_{s}\odot\boldsymbol{a}_{s})}\\\
\boldsymbol{b}^{s}&=-\boldsymbol{a}_{s}\oslash{(\boldsymbol{a}_{u}\odot\boldsymbol{a}_{l}-\boldsymbol{a}_{s}\odot\boldsymbol{a}_{s})}\\\
\boldsymbol{b}^{l}&=\boldsymbol{a}_{u}\oslash{(\boldsymbol{a}_{u}\odot\boldsymbol{a}_{l}-\boldsymbol{a}_{s}\odot\boldsymbol{a}_{s})}\end{split}$
(19)

.

##### Proof:

To prove this we will use the following matrix identity of a partitioned
matrix from [2], which states

$\displaystyle\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)$
(24)

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 24. The algebraic manipulations to arrive at scalar operations in
Equation 19 are straightforward.

##  Appendix B Additional Experiments and Plots

###  B.1 Additional results on ablation with discretization step $H.\Delta t$

Figure 9: Ablation on discretization step $H.\Delta t$. The long-term
prediction results in terms of RMSE, with different $H$ on the mobile dataset.

In addition to the Hydraulics Dataset discussed in Section 6.4, we report the
results of the ablation study with different values of $H.\Delta t$, for the
mobile robot dataset. The higher the value of H, the slower the timescale of
the task dynamics relative to the state dynamics. As seen in Figure 9, 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.

###  B.2 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.

####  B.2.1 Franka Kitchen

Figure 10: 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.

####  B.2.2 Hydraulic Excavator

Figure 11: 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.

####  B.2.3 Mobile Robot

Figure 12: 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.

##  Appendix C Robots and Data

In all datasets, we only use information about agent/object positions and we
mask out velocities to create a partially observable setting. All datasets are
subjected to a mean zero, unit variance normalization during training. During
testing, they are denormalized after predictions. The details of the different
datasets used are explained below:

###  C.1 D4RL Datasets

Details: We use a set of 3 different environments/agents from D4RL dataset
[7], which includes the HalfCheetah, Franka Kitchen and Maze2D (medium)
environment. (a) HalfCheetah: We used 1000 suboptimal trajectories collected
from a policy trained to approximately 1/3 the performance of the expert. The
observation space consists of 8 joint positions and the action space consists
of 6 joint torques collected at 50 Hz frequency. 800 trajectories were used
for training and 200 for testing. For the long horizon task, we used 1.2
seconds (60 timesteps) as context and tasked the model to predict 6 seconds
(300 timesteps) into the future. (b) Franka Kitchen: The goal of the Franka
Kitchen environment is to interact with the various objects to reach a desired
state configuration. The objects you can interact with include the position of
the kettle, flipping the light switch, opening and closing the microwave and
cabinet doors, or sliding the other cabinet door. We used the "complete"
version of the dataset and collected 1000 trajectories where all four tasks
are performed in order. The observation space consists of 30 dimensions (9
joint positions of the robot and 21 object positions). The action space
consists of 9 joint velocities clipped between -1 and 1 rad/s. The data was
collected at a 50 Hz frequency. 800 trajectories were used for training and
200 for testing. For the long horizon task, we used 0.6 seconds (30 timesteps)
as context and tasked the model to predict 2.7 seconds (135 timesteps) into
the future. The dataset is complex due to multi-task, multi-object
interactions in a single trajectory. (c) Medium Maze: We used 20000
trajectories from a 2D Maze environment, where each trajectory consists of a
force-actuated ball (along the X and Y axis) moving to a fixed target
location. The observation consists of as the (x, y) locations and a 2D action
space. The data is collected at 100 Hz frequency. 16000 trajectories were used
for training and 4000 for testing. For the long horizon task, we used 0.6
seconds (60 timesteps) as context and tasked the model to predict 3.9 seconds
(390 timesteps) into the future. Rendering of the three environments is shown
in Figure 13.

Figure 13:  D4RL Environments: (left) HalfCheetah (middle) Franka Kitchen
(right) Maze2D-Medium

###  C.2 Hydraulic Excavator

Details: We collected the data from a wheeled excavator JCB Hydradig 110W show
in Figure 14. The data was collected by actuating the boom and arm of the
excavator using Multisine and Amplitude-Modulated Pseudo-Random Binary
Sequence (APRBS) joystick signals with safety mechanisms in place. A total of
150 mins of data was collected at a frequency of 100 Hz. of which was used as
a training dataset and the rest as testing. The observation space consists of
the boom and arm positions, while the joystick signals are chosen as actions.
For the long horizon task we used 1.5 seconds (150 timesteps) as context and
tasked the model to predict 12 seconds (1200 timesteps) into the future.

Figure 14: (left) JCB Hydradig 110W Excavator (right) Franka Emika Panda Robot

###  C.3 Panda Robot With Varying Payloads

Details: We collected the data from a 7 DoF Franka Emika Panda manipulator
during free motion and while manipulating loads with weights 0kg (free
motion), 0.5 kg, 1 kg, 1.5 kg, 2 kg and 2.5 kg. The robot used is shown in
Figure 14. Data is sampled at a frequency of 100 Hz. The training trajectories
were motions with loads of 0kg(free motion), 1kg, 1.5kg, and 2.5 kgs, while
the testing trajectories contained motions with loads of 0.5 kg and 2 kg. The
observations for the forward model consist of the seven joint angles in
radians, and the corresponding actions were joint Torques in Nm. For the long
horizon task we used 0.6 seconds (60 timesteps) as context and tasked the
model to predict 1.8 seconds (180 timesteps) into the future.

###  C.4 Wheeled Mobile Robot

Figure 15: Wheeled Mobile Robot traversing terrain with complex variations in
slopes induced by a mix of sine functions.

Observation and Data Set: We collected 50 random trajectories from a Pybullet
simulator a wheeled mobile robot traversing terrain with slopes generated by a
mix of sine waves as shown in Figure 15. Data is sampled at high frequencies
(500Hz). 40 out of the 50 trajectories were used for training and the rest 10
for testing. The observations consist of parameters which completely describe
the location and orientation of the robot. The observation of the robot at any
time instance $t$ consists of the following features:

$\begin{array}[]{r}o_{t}=[x,y,z,\cos(\alpha),\sin(\alpha),\cos(\beta)\\\
\sin(\beta),\cos(\gamma),\sin(\gamma)]\end{array}$

where, $x,y,z$ \- denote the global position of the Center of Mass of the
robot, $\alpha,\beta,\gamma-$ Roll, pitch and yaw angles of the robot
respectively, in the global frame of reference [28]. For the long horizon task
we used 0.6 seconds (150 timesteps) as context and tasked the model to predict
3 seconds (750 timesteps) into the future.

##  Appendix D Hyperparameters and Compute Resources

##### Compute Resources

For training MTS3, LSTM, GRU and Transformer models we used compute nodes with
(i) Nvidia 3090 and (ii) Nvidia 2080 RTX GPUs. For training more
computationally expensive locally linear models like RKN, HiP-RSSM we used
compute nodes with NVIDIA A100-40 GPUs.

##### Hyperparameters

Hyperparameters were selected via grid search. In general, the performance of
MTS3 is not very sensitive to hyperparameters. Among all the baselines,
Transformer models were most sensitive to hyperparameters (see Appendix E.5
for details of Transformer architecture).

##### Discretization Step:

For MTS3, the discretization step for the slow time scale SSM as discussed in
Section 3.1 for all datasets was fixed as $\mathbf{H}\cdot\Delta t=0.3$
seconds. In our experiments, we found that discretization values between
$0.2\leq\mathbf{H}\cdot\Delta t\leq 0.5$ seconds give similar performance.

##### Rule Of thumb for choosing discretization step in MTS3:

For any N-level MTS3 as defined in Section 3.4, we recommend searching for
discretization factor $H_{i}$ as a hyperparameter. However, as a general rule
of thumb, it can be chosen as $H_{i}=(\sqrt[N]{T})^{i}$, where $T$ is the
maximum prediction horizon required / episode length. This ensures that very
long recurrences are divided between smaller equal-length task-reconfigurable
local SSM windows (of length $\sqrt[N]{T}$) spread across several hierarchies.

##### Encoder Decoder Architecture:

For all recurrent models (MTS3, HiP-RSSM, RKN, LSTM and GRU) we use a similar
encoder-decoder architecture across datasets. Small variations from these
encoder-decoder architecture hyperparameters can still lead to similar
prediction performance as reported in the paper.

Observation Set Encoder (MTS3): 1 fully connected + linear output:

  * •

Fully Connected 1: 240, ReLU

Action Set Encoder (MTS3): 1 fully connected + linear output:

  * •

Fully Connected 1: 240, ReLU

Observation Encoder (MTS3, HiP-RSSM, RKN, LSTM, GRU): 1 fully connected +
linear output:

  * •

Fully Connected 1: 120, ReLU

Observation Decoder (MTS3, HiP-RSSM, RKN, LSTM, GRU): 1 fully connected +
linear output:

  * •

Fully Connected 1: 120, ReLU

Control Model (Primitive Action Encoder) (MTS3, HiP-RSSM, RKN): 1 fully
connected + linear output:

  * •

Fully Connected 1: 120, ReLU

The rest of the hyperparameters are described below:

###  D.1 D4RL Datasets

####  D.1.1 Half Cheetah

Recurrent Models

Hyperparameters | MTS3 | HiP-RSSM | RKN | LSTM | GRU
---|---|---|---|---|---
Learning Rate | 3e-3 | 1e-3 | 1e-3 | 1e-3 | 1e-3
Latent Observation Dimension | 15 | 15 | 15 | 15 | 15
Observation Set Latent Dimension (sts-SSM) | 15 | - | - | - | -
Latent State Dimension | 30 | 30 | 30 | 45 | 45
Latent Task Dimension | 30 | 30 | - | - | -
Latent Abstract Action Dimension (sts-SSM) | 30 | - | - | - | -

Transition Model (HiP-RSSM, RKN): number of basis: 32

  * •

$\alpha(\boldsymbol{z}_{t})$: No hidden layers - softmax output

Autoregressive Transformer Baseline
Learning Rate: 1e-5
Optimizer Used: Adam Optimizer
Embedding size: 96
Number of Decoder Layers: 4
Number Of Attention Heads: 4

Multistep Transformer Baseline
Learning Rate: 1e-5
Optimizer Used: Adam Optimizer
Embedding size: 128
Number Of Encoder Layers: 2
Number of Decoder Layers: 1
Number Of Attention Heads: 4

####  D.1.2 Franka Kitchen

Recurrent Models

Hyperparameters | MTS3 | HiP-RSSM | RKN | LSTM | GRU
---|---|---|---|---|---
Learning Rate | 3e-3 | 9e-4 | 9e-4 | 1e-3 | 1e-3
Latent Observation Dimension | 30 | 30 | 30 | 30 | 30
Observation Set Latent Dimension (sts-SSM) | 30 | - | - | - | -
Latent State Dimension | 60 | 60 | 60 | 90 | 90
Latent Task Dimension | 60 | 60 | - | - | -
Latent Abstract Action Dimension (sts-SSM) | 60 | - | - | - | -

Transition Model (HiP-RSSM, RKN): number of basis: 15

  * •

$\alpha(\boldsymbol{z}_{t})$: No hidden layers - softmax output

Autoregressive Transformer Baseline
Learning Rate: 5e-5
Optimizer Used: Adam Optimizer
Embedding size: 64
Number of Decoder Layers: 4
Number Of Attention Heads: 4

Multistep Transformer Baseline
Learning Rate: 1e-5
Optimizer Used: Adam Optimizer
Embedding size: 64
Number Of Encoder Layers: 2
Number of Decoder Layers: 1
Number Of Attention Heads: 4

####  D.1.3 Maze 2D

Recurrent Models

Hyperparameters | MTS3 | HiP-RSSM | RKN | LSTM | GRU
---|---|---|---|---|---
Learning Rate | 3e-3 | 9e-4 | 9e-4 | 1e-3 | 1e-3
Latent Observation Dimension | 30 | 30 | 30 | 30 | 30
Observation Set Latent Dimension (sts-SSM) | 30 | - | - | - | -
Latent State Dimension | 60 | 60 | 60 | 90 | 90
Latent Task Dimension | 60 | 60 | - | - | -
Latent Abstract Action Dimension (sts-SSM) | 60 | - | - | - | -

Transition Model (HiP-RSSM, RKN): number of basis: 15

  * •

$\alpha(\boldsymbol{z}_{t})$: No hidden layers - softmax output

Autoregressive Transformer Baseline
Learning Rate: 5e-5
Optimizer Used: Adam Optimizer
Embedding size: 96
Number of Decoder Layers: 4
Number Of Attention Heads: 4

Multistep Transformer Baseline
Learning Rate: 1e-5
Optimizer Used: Adam Optimizer
Embedding size: 128
Number Of Encoder Layers: 2
Number of Decoder Layers: 1
Number Of Attention Heads: 4

###  D.2 Franka Robot Arm With Varying Loads

Recurrent Models

Hyperparameters | MTS3 | HiP-RSSM | RKN | LSTM | GRU
---|---|---|---|---|---
Learning Rate | 3e-3 | 9e-4 | 9e-4 | 3e-3 | 3e-3
Latent Observation Dimension | 15 | 15 | 15 | 15 | 15
Observation Set Latent Dimension (sts-SSM) | 15 | - | - | - | -
Latent State Dimension | 30 | 30 | 30 | 45 | 45
Latent Task Dimension | 30 | 30 | - | - | -
Latent Abstract Action Dimension (sts-SSM) | 30 | - | - | - | -

Transition Model (HiP-RSSM,RKN): number of basis: 32

  * •

$\alpha(\boldsymbol{z}_{t})$: No hidden layers - softmax output

Autoregressive Transformer Baseline
Learning Rate: 5e-5
Optimizer Used: Adam Optimizer
Embedding size: 64
Number of Decoder Layers: 4
Number Of Attention Heads: 4

Multistep Transformer Baseline
Learning Rate: 2e-5
Optimizer Used: Adam Optimizer
Embedding size: 64
Number Of Encoder Layers: 2
Number of Decoder Layers: 1
Number Of Attention Heads: 4

###  D.3 Hydraulic Excavator

Hyperparameters | MTS3 | HiP-RSSM | RKN | LSTM | GRU
---|---|---|---|---|---
Learning Rate | 3e-3 | 8e-4 | 8e-4 | 1e-3 | 1e-3
Latent Observation Dimension | 15 | 15 | 15 | 15 | 15
Observation Set Latent Dimension (sts-SSM) | 15 | - | - | - | -
Latent State Dimension | 30 | 30 | 30 | 45 | 45
Latent Task Dimension | 30 | 30 | - | - | -
Latent Abstract Action Dimension (sts-SSM) | 30 | - | - | - | -

Transition Model (HiP-RSSM,RKN): number of basis: 15

  * •

coefficient net $\alpha(\boldsymbol{z}_{t})$: No hidden layers - softmax
output

Autoregressive Transformer Baseline
Learning Rate: 1e-5
Optimizer Used: Adam Optimizer
Embedding size: 96
Number of Decoder Layers: 4
Number Of Attention Heads: 4

Multistep Transformer Baseline
Learning Rate: 5e-5
Optimizer Used: Adam Optimizer
Embedding size: 64
Number Of Encoder Layers: 2
Number of Decoder Layers: 1
Number Of Attention Heads: 4

###  D.4 Wheeled Robot Traversing Uneven Terrain

Hyperparameters | MTS3 | HiP-RSSM | RKN | LSTM | GRU
---|---|---|---|---|---
Learning Rate | 3e-3 | 8e-4 | 8e-4 | 1e-3 | 1e-3
Latent Observation Dimension | 30 | 30 | 30 | 30 | 30
Observation Set Latent Dimension (sts-SSM) | 30 | - | - | - | -
Latent State Dimension | 60 | 60 | 60 | 90 | 90
Latent Task Dimension | 60 | 60 | - | - | -
Latent Abstract Action Dimension (sts-SSM) | 60 | - | - | - | -

Transition Model (HiP-RSSM,RKN): number of basis: 15

  * •

coefficient net $\alpha(\boldsymbol{z}_{t})$: No hidden layers - softmax
output

Autoregressive Transformer Baseline
Learning Rate: 5e-5
Optimizer Used: Adam Optimizer
Embedding size: 128
Number of Decoder Layers: 4
Number Of Attention Heads: 4

Multistep Transformer Baseline
Learning Rate: 5e-5
Optimizer Used: Adam Optimizer
Embedding size: 64
Number Of Encoder Layers: 4
Number of Decoder Layers: 2
Number Of Attention Heads: 4

###  D.5 Transformer Architecture Details

For the AR-Transformer Baseline, we use a GPT-like autoregressive version of
transformers except that for the autoregressive input we also concatenate the
actions to make action conditional predictions.

For Multi-Transformer we use the same direct multistep prediction and loss as
in recent Transformer time-series forecasting literature [35, 19, 21, 34]. A
description of the action conditional direct multi-step version of the
transformer is given in Algorithm 1.

Require: Input past observations $\mathbf{o_{inp}}\in\mathbb{R}^{S\times C}$;
Input Past Actions $\mathbf{a_{inp}}\in\mathbb{R}^{S\times A}$; Future Actions
$\mathbf{a_{pred}}\in\mathbb{R}^{O\times A}$;Input Length $S$; Predict length
$O$; Observation Dimension $C$; Action Dimension $A$; Feature dimension
$d_{k}$; Encoder layers number $N$; Decoder layers number $M$.

1: $\mathbf{o_{inp}}\in\mathbb{R}^{S\times
C},\mathbf{a}_{inp}\in\mathbb{R}^{S\times
A},\mathbf{a}_{pred}\in\mathbb{R}^{O\times A}$

2:
$\mathbf{X}_{inp}=\mathrm{ConCatFeatureWise}\left(\mathbf{o}_{inp},\mathbf{a}_{inp}\right)$
$\triangleright\mathbf{X}_{inp}\in\mathbb{R}^{S\times(C+A)}$

3:
$\mathbf{X}_{pred}=\mathrm{ConCatFeatureWise}\left(\operatorname{Zeros}(O,C),\mathbf{a}_{pred}\right)$
$\triangleright\mathbf{X}_{pred}\in\mathbb{R}^{O\times(C+A)}$

4:
$\mathbf{X}_{\mathrm{enc}},\mathbf{X}_{\mathrm{dec}}=\mathbf{X}_{\mathrm{inp}},\operatorname{ConCat}\left(\mathbf{X}_{inp},\mathbf{X}_{pred}\right)\quad\quad\triangleright\mathbf{X}_{\mathrm{enc}}\in\mathbb{R}^{S\times(C+A)},\mathbf{X}_{\mathrm{dec}}\in\mathbb{R}^{\left(S+O\right)\times(C+A)}$

5:
$\mathbf{X}_{\mathrm{enc}}^{0}=\operatorname{Embed}\left(\mathrm{X}_{\mathrm{enc}}\right)$
$\triangleright\mathbf{X}_{\mathrm{enc}}^{0}\in\mathbb{R}^{S\times d_{k}}$

6: for  _$l$ in $\\{1,\cdots,N\\}$_ do

7:
$\quad\mathbf{X}_{\mathrm{enc}}^{l-1}=\operatorname{LayerNorm}\left(\mathbf{X}_{\mathrm{enc}}^{l-1}+\operatorname{Attn}\left(\mathbf{X}_{\mathrm{enc}}^{l-1}\right)\right)$
$\triangleright\mathbf{X}_{\mathrm{enc}}^{l-1}\in\mathbb{R}^{S\times d_{k}}$

8: $\quad\mathbf{X}_{\text{enc
}}^{l}=\operatorname{LayerNorm}\left(\mathbf{X}_{\text{enc
}}^{l-1}+\operatorname{FFN}\left(\mathbf{X}_{\text{enc }}^{l-1}\right)\right)$
$\triangleright\mathbf{X}_{\mathrm{enc}}^{l}\in\mathbb{R}^{S\times d_{k}}$

end for

9:
$\mathbf{X}_{\mathrm{dec}}^{0}=\operatorname{Embed}\left(\mathbf{X}_{\mathrm{dec}}\right)$
$\triangleright\mathbf{X}_{\mathrm{dec}}^{0}\in\mathbb{R}^{\left(S+O\right)\times
d_{k}}$

10: for  _for $l$ in $\\{1,\cdots,M\\}$_ do

11:
$\quad\mathbf{X}_{\operatorname{dec}}^{l-1}=\operatorname{LayerNorm}\left(\mathbf{X}_{\operatorname{dec}}^{l-1}+\operatorname{Attn}\left(\mathbf{X}_{\operatorname{dec}}^{l-1}\right)\right)$
$\triangleright$ Decoder

12:
$\quad\mathbf{X}_{\mathrm{dec}}^{l-1}=\operatorname{LayerNorm}\left(\mathbf{X}_{\mathrm{dec}}^{l-1}+\operatorname{Attn}\left(\mathbf{X}_{\mathrm{dec}}^{l-1\prime},\mathbf{X}_{\mathrm{enc}}^{N}\right)\right)$
$\triangleright\mathbf{X}_{\mathrm{dec}}^{l-1}\in\mathbb{R}^{\left(S+O\right)\times
d_{k}}$

13:$\quad\mathbf{X}_{\mathrm{dec}}^{l}=\operatorname{LayerNorm}\left(\mathbf{X}_{\mathrm{dec}}^{l-1}+\operatorname{FFN}\left(\mathbf{X}_{\mathrm{dec}}^{l-1}\right)\right)$
$\triangleright\mathbf{X}_{\mathrm{dec}}^{l}\in\mathbb{R}^{\left(S+O\right)\times
d_{k}}$

end for

14: $\mathbf{y}=\operatorname{MLP}\left(\mathbf{X}_{\mathrm{dec}}^{M}\right)$
$\triangleright\mathbf{y}\in\mathbb{R}^{(S+O)\times C}$

15: Return y $\triangleright$ Return the prediction results

Algorithm 1 MultiStep Transformer

##  Appendix E Limitations

We list some of the limitations of the paper here. (i) We restricted our
definition and experiments to MTS3 with two levels of temporal abstractions,
which was sufficient in many of our tasks. However, for certain tasks like the
Maze2D, we believe more hierarchies can help. As discussed in the main paper
the method and inference scheme allows easy addition of more Feudal [6]
hierarchies with larger discretization steps ($\mathbf{H}\cdot\Delta t$). (ii)
We restrict our application to action conditional long horizon future
predictions and do not use the model for (hierarchical) control. A
probabilistically principled formalism for hierarchical control as an
inference problem, that builds upon MTS3 models is left for future work. (iii)
Finally, we restrict our experiments to proprioceptive sensors from the agent
and objects. The performance of MTS3 which relies on “reconstruction loss” as
the objective is yet to be validated on noisy high dimensional sensor inputs
like Images. Image-based experiments and “non-reconstruction” based losses
[18] can be taken up as future work.

##  Appendix F Broader Impact

While we do not foresee any immediate negative societal impacts of our work,
we do believe that machines that can replicate human intelligence at some
point should be able to reason at multiple levels of temporal abstractions
using internal world models [18]. Having intelligent agents with type 2
reasoning capabilities can have both positive and negative impacts. We believe
identifying and mitigating the potentially harmful effects of such autonomous
systems is the responsibility of sovereign governments.