Multi Time Scale World Models
Authors: Vaisakh Shaj
Saleh GHOLAM ZADEH
Ozan Demir
Luiz Ricardo Douat
Gerhard Neumann
Venue: NeurIPS 2023 spotlight (2023)
Keywords: Hierarchical Models; Multi Time Scale Learning; World Models
TL;DR: Propose a principled framework for learing world models at multiple time scales/temporal abstractions.

Abstract: Intelligent agents use internal world models to reason and make predictions about different courses of their actions at many scales. 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. 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.

Multi Time Scale World Models
Vaisakh Shaj"* Saleh Gholam Zadeh 1.2
Ozan Demir?
Luiz Ricardo Douat
Gerhard Neumann'
Karlsruhe Institute Of Technology (KIT), Germany
SAP SE, Germany
3 Robert Bosch GmbH, Germany
Abstract
Intelligent agents use internal world models to reason and make predictions about different courses of their actions at many scales [2I]. 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 [17]. In this work, we propose a probabilistic formalism to learn multi-tima scale world models which we call the Multi Time Scale State Space (MTS3) model. Our model uses computationally efficient inference scheme 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: //tp:/gggtttttttt com/ ALRAub/
Introduction
World models attempt learn compact and expressive representation of the environment dynamics from observed data. These models can predict possible future world states as function of an imagined action sequence and are key ingredient of model-predictive control [3] 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 single level of temporal abstraction, typically at 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 [7].
For efficient long- horizon prediction and planning, the model needs predict multiple levels of temporal abstractions (27 (22]. Intuitively, low-level temporal abstractions should contain precise predict accurately in the short term, while high-level, abstract
details about the input representations should simplify accurate long-term predictions. Both abstractions must also interrelate with each other least in the sense that the higher-level predictions/plans moment-by- moment momoent predictions. For example, in robotic manipulation, the robot must be able to perform precise and coordinated movements to grasp and manipulate the object 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.
be turned into low-level
fast time scale
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 another 266 17
"oorrsponding author. Email <vaisakhs.shaje gmail.com>.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

In this work, we attempt to come up with a principled probabilistic formalism for learning such multi-time scale world models 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 several challenging simulated and real-world prediction tasks such as the D4RL dataset, simulated mobile robot and real manipulators including data from heavy machinery excavators.
Preliminaries
State space models (SSMs) are Bayesian probabilistic graphical models [5 12 that are popular for learning patterns and predicting behaviour in sequential data and dynamical systems. Formally, we define state space model as tuple (Z, A. O, f, h, At), where 2 is the state space, A the action space and O the observation space of the SSM. The parameter At 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
21 -,
Et: Et (0, Ex), and o, = h(zz) + 0t ~ ...0
Here E Z, a, E A and o, O are the latent states, actions and observations at time t. The vectors & and 82 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 [9] 25] where these dlosed-forn 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.
To perform inference in SSMs, we follow [. They use a (locally-)lineal dynamics model. Moreover, they replace the observations with latent observation w. This latent observation is obtained by an encoder w., enc..(o;) along with the uncertainty of this observation, i.e., oo enc, Due the non-linear observation encoder, simplified linear observation model can now be used. Hence, the dynamics and observation models can be described as
(Locally-)Linees Gaussian SSMs.
p(z++1/z,,00 = -N(A,, + + B,a,, diag(o,)), and p(w., = N (o.ag((...);
where a simple observation matrix of H = [I, 0] is used. The underlying assumption behind this observation model is that the latent state z, = Ip. d, has twice the dimensionality of the latent observation w, and only the first half of the latent state, i.e., Pt, be observed. The second half of the latent state, i.e., d,, 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 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 backpropagate for end-to-end learning. Hence, M introduce factorization of the belief p(z, 1011, d1:1-1) [H, Σ,) such that only the diagonal and one off off-oiagonal vector of the covariance need to be computed, i.e.
n Σ
Σ,
=
with Σ. = diag(o)). Σ, = diag(o)) and Σ; = diag(a?))
Using this factorization assumption, closed-form Gaussian inference can be performed using only scalar divisions which are fast and easy to back- --rrppagate.. These factorization assumptions form the basis for the inference update our MTS3 model.
Bayesian Aggregation. To aggregate information from several observations into consistent repre- sentation, [29] introduce Bayesian aggregation in the context Meta-Learning. They again use an encoder to obtain latent observation vector T., and its uncertainty vector o.,- Given the observation model prr,, [2) = 2Hz, diag(a.,)) with H I and a prior p(z) = NM diag(a (r))), the posterior pl 2 :o can again be effectively computed by Gaussian inference that involve only

---

scalar inversions. Note that computing this posterior is simplified case of the Kalman update rule used in Gaussian SSMs M. with no memory units, H 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 [29].
3 Multi Time Scale State Space Models
Our goal is learn a principled sequential latent variable model that can model the dynamics of partially observable robotic systems under multiple levels of temporal abstrac- tions. 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 multi- ple time scales. ii) It allows for single global model be learned that can be shared across changing configurations of the environments. iii) can give accurate long-term pre- dictions and uncertainty estimates. iv) It is probabilistically principled yet scalable during learning and inference.
Latent Task
Dynamics
Scale (Slov)
Latent Dynamic Time Scale (Fact
Figure I: MTS3 captures slow-moving long-term trends the latent task dy- namics and the fast-moving short-term trends the latent state dynamics.
3.1 General Definition
An MTS3 model with 2 timescales is defined by two SSMs a fast and a slow time scale respectively. Both SSMs are coupled via the latent state of the slow time scale SSM, which pwrametrizes/teconiggrre
the system dynamics of the fast time scale SSM. While the fast time scale SSM runs at the original time step At of the dynamical system, the slow time scale SSM is only updated every Il step, i.e., the slow time scale time step is given by HAt. 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 9 The definition with 2-level MTS3 along with the inference and learning schemes that propose directly extendable to arbitrary number of temporal abstractions by introducing additional feudal [ 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 = (2,A, At, L). Here, I L is 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 keth time window of H steps 4. The probabilistic dynamics and observation model of the fast time scale for the /th time step in the kth window can then be described
as
p(zk, Z a0 1,2--1,46)
NNS (2 aR,--, 42), (4),,9. Q), and p(OR.IZZ..) = N(O (NH)) (zx,).
fask-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.c., instead of considering a single task deseriptor 4g, have consider distribution over task-descriptors p(l&) 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
aI.--1 =
an,--1, 1
Pl,
Latent observations. Following M1 we replace the observations by latent observations and their uncertainty, i.e., we use latent observation encoders to obtain Wi,r = (nc,,(Or.r) and an uncer- tainty encoder σk. (oc.) The observation model is hence given by p(wk,llkk.) (h (*x,), di.g(ork.)).
3

---

3.1.2 Slow time-scale SSM
The slow time-scale (sts) SSM only updates every H time step and uses the task parameter as latent state representation. Formally, the SSM is defined as Saoo (C,E, f h tts At). It uses an abstract observation B3 B and abstract action E A that summarize the observations and actions respectively throughout the current time window. The general dynamics model is hence given by
p(l 11110-11000 oy.) (N11 a4),
While there exist many ways to implement the abstraction of observations and actions of the time windows, we choose use consistent formulation by fusing the information from all H time steps of time window & using Gaussian conditioning.
Figure The graphical model corresponding to an MTS3 with 2 timescales. The latent task variable 4 captures the slow-changing dynamics using abstract observation inferred from {/x, ) and abstract action ag as described in section 3.1.2 The inference in the fast time scale uses primitive observations Wk,t, primitive actions as and the latent task descriptor 4 which parameterizes the fast-changing dynamics of k,, for a time window as discussed in the section
Observation abstraction. In terms of the abstract observation model, choose to model H observations BA,,, t E [1, 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 is again obtained by an encoder architecture, i.e,
BE. enc (), Va,t = (on,(o
and p(3 (&) (15*((()))) (v)gg().:)). Hence, the abstract observation 3x, contains the actual observation o at time step & as well as 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 Og causes the transitions to the latent task from 4-11 should contain the relevant information of all primitive actions ag,z. E [1, H] executed in the time window k. To do so, we again use Bayesian conditioning and latent action encoding. Each control action and the encoding of time-step t is encoded into its latent representation and its uncertainty estimate, i.e.,
Okkt = enc, (ak.r, t), Pk.: enc,(ax.t.) t).
The single latent actions Ou can be aggregated into a consistent representation Qn using Bayesian aggregation 29 To do so, we use the likelihood p(cu,,//ak)
A (ax, diag(Pr.)) and obtain the posterior p(0ss) 1:1) NN E.,), 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 express different effects of the actions the uncertainty of the prediction. Due to the Gaussian representations, we can compute the marginal transition model
Pas (4x 14-11 O..,1:H) /
Pax (4,|4-11, on ( hopppoooo (d)
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 33 1 a1:kk
and posterior p(4x] (1:sk belief state for the sts-SSM as well as the belief states for the fts-SSM. The prior belief ( 1.k-1, 011:)) contains all information up to time window k - 1 and serves as 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 probabilisticc graphical model of our MTS3 model 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 whenever the context is clear. We use linear Gaussian task conditional transition model, i.e,
zt a,11 = N (Az,-1 + Ba;- + Cl Q)
where A, B, C and Q are state-independent but learnable parameters. In our formulation, the task descriptor can only linearly modify the dynamics which was sufficient obtain state-of-the-art performance in our experiments, but more complex parametrizations, such as locally linear models, would also be feasible. Following [.]. we split the latent state &; [p.. m,]" into its observable part P, and part m, that needs to be observed over time. We also use linear observation model p(w,/zz) (HHz, z,, diag(or)) ()) with H [I.
We will assume that the distribution over the task descriptor is also given by a Gaussian distribution, ie., p(42) .V (Hr, Σt,): which will be provided by the slow-time scale (sts) SSM, 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
-N(A=-- (Azz-1 + Ba.- + + CEC, C²), CCC
PL,
a.-1) =
(6)
which will be used instead of the standard dynamics equations. We follow the same factorization assumptions in [ 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 the RKN [
3.3 Inference in the Slow-Time Scale SSM
Prediction Update. We follow same Gaussian inference scheme as for the fts-SSM, i.e., we again employ linear dynamics model p(1u/4=-1, ag)) (XX You, S), where X, Y and are learnable parameters. The marginalized transition model for the abstract actions is then given by
St Pas (1 (1x((x-1)= = / p(lx)1 ax. = V (Xlx-1 YM,,, S YEa, Y").
(7)
We directly this transition model to obtain the Kalman prediction update which computes the prior belief Pa... (4 = N( Σī₂) from the posterior belief Part-, N _t_, of the previous time window, see Appendix A.
=
Observation Update. Similarly, will use a linear observation model for the abstract observations p(Br,,14 4x) = A (HLe, diag(uk.))) with H = [I. As can be seen from the definition of the observation matrix H. the latent space is also decomposed into its observable and unobservable part, i.e., 4 = [ux v.). In difference to the standard factorized Kalman observation update given in Appendix A, we have to infer with set of observations 3 with / =1 1... H for single time window While in principle, the Kalman observation update can be applied incrementally H times to obtain the posterior Pa.. 11)) NNtt Et.) such update would be very slow and also
5

---

cause numerical inaccuracies. Hence, we devise new permutation invariant version of the update rule that allows parallel processing with set encoders [3]]. We found that this update rule is easier to formalize using precision matrices. Hence, we first transform the prior covariance vectors o o and is to its corresponding precision representation , N and XIz which can be performed using block-wise matrix inversions of " 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 Appendix A, B. Subsequently, we compute the posterior precision, where only X is changed by
Σ 1QVEL
(8)
while X/ x and X X remain constant. The operator 2 denotes the element-wise division. From the posterior precision, we can again obtain the posterior covariance vectors oi 0 and oft using only scalar inversions, see Appendix A, B. The posterior mean Mi,, can now be obtained from the prior mean Mi.k
Mit ML,
u. == (Bx, DE (3x,t Mi,
Mi
2 VAet
Vk.
Note that for // = 1. i. a single observation, the given equation is equivalent to the RKN updates. Moreover, the given rule constitutes unification of the batch update rule for Bayesian aggregation 229 and the incremental Kalman update for our factorization of the belief state representation ||] 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, {Su, Si,.., -1-1 Each of the state space model S; is given by S; = (Zi, A,, Oi, fi, h4 H, Li), where Z, is the state space, A, the action space, and O the observation space of the SSM. The parameter H;At denotes the discretization time-step and fa and h; the dynamics and observation models, respectively. Here, 4 E Lg is a task descriptor that parametrizes the dynamics model of the SSM and is held constant for a local window of H(+1 steps. S1+1. The boundary cases can be defined as follows: for i = 0, Ho = 1. Similarly, for = N - 1. the latent task descriptor Lę is an empty set. For all i, H, < Hi+1-
a function of the latent state of SSM one level above it, i.e..
Even though our experiments focus on MTS3 models with hierarchies, extensive experimentation with more hierarchies be taken future work.
4 MTS3 as a Hierarchical World Model
MTS3 allows for 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 masel- control and planning. Furthermore, introducing several levels of latent variables, each working different time scale allows us to learn world models that can make action conditional predic- dreamm" 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 sequence of actions. Our principled formalism allows for 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 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( ((xx.) for the first time step of the time window k, we use the prior belief p( p(zx-1,H++) of the last time step of the time window k - 1.
future

---

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 single trajectory, i.e.,
L
log pogp,,,t+1 01:kk Wk,1: (k.1:t)
N
H
log /ʃ plo..:+lERR.HH))(ee..... w1... ax,1a, Ai
N
H
log
plon.tt|EEatt1)a (F WKk 11.12))Z2..+
(10)
k=
The extension to multiple trajectories is straightforward and omitted to keep the notation uncluttered. Here, 0g.t+1 is the ground truth observations at the time step f + and time window & which needs to be predicted from all (latent and abstract) observations up to time step f. The corresponding latent state prior belief Pt, [1b,+1|WWs,.: dx,1:)) has a closed form solution as discussed Section 3.2
We employ a Gaussian approximation of the posterior predictive log- lig-likelihood of the form p(o 104.441/11s-11 01.kk Wk.1:t, aR.1:1) = (N diag o oo....))) where we use the mean of the prior belief = to decode the predictive mean, i.e, Mos.111 = dec,, (M=.))))))) and the variance estimate of the prior belief to decode the observation variance, i.c., do..t+ dec, (E....)) This approximation can be motivated by moment-matching matching perspective and allows for end-to-end optimization of the log-likelihood without using auxiliary objectives such as the ELBO [..
Gradients are computed using (truncated) backpropagation through time (BPTT) (30 and clipped. We optimize the objective using the Adam [14] stochastic gradient descent optimizer with default parameters. We refer to Appendix A for more details. For training, we also initialize the prior belief 1 with the prior belief Pt,-1 (H - dg-1,1.H1) 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- 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 the loss assumes that observations are always available up to time step To increase the long-term prediction performance, we can treat the long-term prediction problem as 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 strong self supervied learning signal for long perg-term prediction with varying prediction horizon length. This imputation scheme is applied at both time scales, masking out single-time steps 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 carly works that enabled environment models at different temporal scales to be intermixed, producing temporally abstract world models was proposed by (27]. 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 [77. A probabilistically principled formalism for these has been lacking in literature and this work is an early attempt address this issue.
Deep State Space Models. Deep SSMs combine the benefits deep neural nets and SSMs by offering tractable probabilistic inference and scalability to high-dimensional and large datasets. [ [ 24 use neural network architectures based on exact inference on SSMs and perform state estimation and dynamics prediction tasks. (25) extend these models to modelling non-stationary dynamics.
7

---

16 13 [0] 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 few milliseconds and model the dynamics at single time scale.
Transformers Recent advancements in Transformers 28 23 2 which rely attention mechanism,
have demonstrated superior performance in capturing Long-range dependency compared to RNN models in several domains including time series forecasting [33 [8] and learning world models [19]
33 18 20 use transformer architectures based on direct multistep loss [2 and show promising results for long-tern snrgttemm forecasting since they avoid error accumulation from autoregression. On the other hand [19] uses a GPT-like autoregressive version of transformers to learn world models. These deterministic models, however, do deal with temporal abstractions and uncertainty estimation in principled manner. Nevertheless, we think Transformers that operate at multiple timescales based on our formalism can be promising alternative research direction.
Experiments
In this section, we evaluate our approach to 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 F, a brief description of them is given here: (a) RNNs We compare method to two widely used recurrent neural network architectures, LSTMs [1 and GRUs 4 (b) RSSMs - Among several RSSMs from the literature, we chose RKN 1 and HiP-RSSM 25 these have shown excellent performance for dynamics learning for short-term predictions and rely on exact inference as our case. (c) Transformers We also
compare with state-of-the-art Transformer 28 variants. The first variant (AR-Transformer) relies
on a GPT-like autoregressive prediction 23 2 Whereas the second variant (Multi-Transforre Transformer) uses direct multi-step loss [2) from recent literature long horizon time-series forecasting kk] [8 20)
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 E, a brief description of them is as follows. (a) D4RL Datasets We a set of 3 different environments/agents from D4RL dataset (6] 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 non-linear non-man 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 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 partially observable setting.
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 "sliding window RMSE" in Figure We see that MTS3 gives consistently good long-term action conditional future predictions all 6 datasets. Deep Kalman
8

---

HallChootohh
Pando
RM455
120
810
1.5
025 1.0
Transforme"
(a)
Excavotor
Mobile Robot
NTS3 Variants on Excavater
NTS3 Variants on Kithenn
0427
6.01
1.14
3.20
(85
5.5
Seconds
Seconds Mulli Transforme?
Scoonds
Sec2nda
MT133
6RU
HIF-RSM
NTS3
No Action Abstraction
1.7M
AR-Tramaformr
No Hemory
No Imputatiin
(b)
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 discussed in 6.6
models 25 which operate on a single time scale fail to give meaningful mean estimates beyond few milliscconds. Similarly, widely used RNN baselines [ which form the backbone of several world models (8] [0]] give poor action conditional predictions over long horizons. AR- Transformers also fail possibly due error accumulation caused by the autoregression. However, Multi-Trann Transformers are 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 (77. MTS3 is the only model that gives meaningful long horizon predictions. detailed description the metric "sliding window RMSE" is given in Appendix C. A visualization of the predicted trajectories vs. ground truth is given in Appendix D.
6.4
Can MTS3 make accurate long-term probabilistic predictions (variance estimates)?
Next, we examine the question whether the principled probabilistic inference translates to accurate uncertainty quantification during long-horizon predictions. We trained all the baselines with negative log-likelihood loss and used the same During the evaluation we rely on a sliding window approach (see Appendix C) and report the results for the last timestep Table I As seen Table I, MTS3 gives the most accurate uncertainty estimates in all datasets except Medium Maze, where it is outperformed by Multi- Teiihrrranffooerr Also, notably, AR- Transformers and deep Kalman models fail learn any meaningful uncertainty representation when it comes to long-term predictions.
metric quantify the quality of uncertainty estimates.
Prediction Horizon
Algorithm
MTS3
Multi Trans
AR- Trans
ISTM
GRU
RKN
HSS-RSSM
Half Cheetah
6s
-2.80 0.30
0.26+0.00
7.34 0.0
7:40
X
Kitchen
2.58
-25.74
0.12
$2.45 # 1.61
32,72 0.66
X
Medium Maze
4s
-0.21 2: 0.022
-0.88 0.02
4.03 1: 0.32
7.76 =: 0.07
Panda
1.8%
2.79 0.32
3.77
0.33
8
7.94 0.39
7.91
0.23
Hydraulie
12s
2.64 0.12
-2.46
0.03
7.35 + 0.06
7.35 = 0.06
Mobile Robot
0.71
-5.17
2.3
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 X denotes very high/nan values resulting from the highly divergent mean/varianc, 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 httt 1 H [I.0] as
discussed in Sections 3.3 and 3.2 (ii) action abstractions in Section 3.1.2 (iii) training by imputation.

---

MTSS (++55
MTS3 (H-10)
66
BD
US0
2.4
04
380
000
031
ATS3 (4--30
MT53 (4-75)
Suceed Tast Secend- 360 01157800
Tv. M753
Showen S1pper
105
(a)
(b)
Figure 4: Ablation on discretization step H.A. (a) The long-term prediction results in terms of RMSE, with different И values as discussed Section 6.6 on the hydraulics dataset. (b) The predictions by MTS3 variants with different values of timescale parameter H. At on trajectory picked from the hydraulics excavator dataset. The top figures are for H = and H = 10. Bottom figures are for H 30 and H 75. Note that the results reported in the paper with H = 30.
To analyze the importance of the memory component, we derived and implemented an MTS3 variant with an observation model of h htt htts and a pure diagonal matrix representation for the covariance matrices. As seen in Figure 3c 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 mon-mankovian dynamical systems. Regarding (ii), we further devised another variant where MTS3 only had access 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 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.AI?
Finally, we perform ablation for different values of H. AI, 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 smaller values of И (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 trajectory handpicked from the hydraulics excavator dataset. As seen in Figure 4b 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 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 accunulation and poor predictions.
Conclusion and Future Work
In this work, we introduce MTS3, probabilistic formalism for learning the dynamics of complex environments at multiple time scales. By modelling the dynamics 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 G and H.
10

---

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 Informaties of the Helmholtz Association (HGF). The authors acknowledge support by the state of Baden- Württemberg through bwHPC, 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
[1] Philipp Becker. Harit Pandya, Gregor Gebhardt, Cheng Zhao, C James Taylor, and Gerhard Neumann. Recurrett kalman networks: Factorized inference in high-dimensional deep feature spaces. In International Conference on Machine Learning, pages 544 552.. PMLR, 2019.
[2] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et Language models are few-shot learners, Advances in neural information processing systems, 33: 33:1877-1901, 2020.
[3] Eduardo F Camacho and Carlos Bordons Alba. Model predictive control. Springer science & business media, 2013.
[4] 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.
[5] Peter Dayan and Geoffrey E Hinton. Feudal reinforcement learning. Advances in neural information processing systems, 1992.
[6] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint 1xXxx:2004002219 2020.
[7] 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. 1956, 2019.
[8] David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv: 10122, 2018.
[9] 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.
[10] 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.
[11] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Comput., 9(8):
1735-1780, November 1997. ISSN 0899-7667.
[12] Michael Jordan. Graphical models. Statistical science, 19(1):10--155, 2004.
[13] 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.
[14] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: arXiv: 6422,690, 2014.
[15] Daphne Koller and Nir Friedman. Probabilistic graphical models: principles and techniques. MIT press, 2009.

---

[16] 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.
[17] Yann LeCun. A path towards autonomous machine intelligence version 0.9. 2, 2022-06-27. Open Review, 62, 2022.
[18] Yong Liu, Haixu Wu, Jianmin Wang, and Mingsheng Long. Non-stationary transformers: Exploring the stationarity in time series forecasting. In Alice H. Oh, Alckh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, wet. wetpr//oeeee/ooouo/dddddddDDDDDd
2022. httpttttt///ppereeee
[19] Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample-efficient world models. In The Eleventh International Conference on Learning Representations, 2023. URL
attps://openrevill
[20] 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 httpss/:::// http:::///ppeeevieee wwtprrvviwwwwwweeffreeeeeeee net/forum? id=JbdcOvTOcol
[21] Giovanni Pezzulo and Paul Cisek. Navigating the affordance landscape: feedback control as process model of behavior and cognition. Trends in cognitive sciences, 20(6):414-424, 2016.
[22] Doina Precup and Richard S Sutton. Multi-time models for temporally abstract planning. Advances in neural information processing systems, 10, 1997.
[23] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, al. Language models are unsupervised multitask learners. OpenAl blog, 188);9, 2019.
[24] 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.
[25] Vaisakh Shaj. Dieter Buchler, Rohit Sonker, Philipp Becker, and Gerhard Neumann. Hidden pa- rameter recurrent state space models for changing dynamics scenarios. International Conference On Learning Representations, 2022.
[26] Murray Shanahan and Melanie Mitchell. Abstraction for deep reinforcement learning. arXiv preprint wwxXi::2022.08839, 2022.
[27] Richard S Sutton. Td models: Modeling the world at a mixture of time scales. In Machine Learning Proceedings 1995, pages 531-539. Elsevier, 1995.
[28] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[29] Michael Volpp, Fabian Flürenbrock, Lukas Grossherger, Christian Daniel, and Gerhard Neu- mann. Bayesian context aggregation for neural processes. In International Conference on Learning Representations, 2020.
[30] Paul Werbos. Backpropagation through time: what it does and how to do it. Proceedings of the IEEE, 18000) 1990.
[31] Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Ruslan Salakhutdinov, and Alexander Smola. Deep sets. arXiv preprint arXiv: 14, 2017.
[32] Ailing Zeng, Muxi Chen, Lei Zhang, and Qiang Xu. Are transformers effective for time series forecasting? arXiv preprint 2355: 2022.
[33] 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.
12


Reviews:
[{"review_id": "l0oHbDv2wv", "reviewer": "Reviewer_tb9F", "summary": "The authors introduce the Multi Time Scale State Space (MTS3) model in this work. The model uses closed-form equations derived using exact inference, spread across two time-scales, to produce long-horizon predictions and uncertainty estimates. They demonstrate the superiority/competitiveness of their inference approach across a number of offline datasets, both in terms of long-term deterministic predictions, and long-term uncertainty quantification.", "questions": "- How were the hyperparameters for MTS3 tuned? How were the baselines tuned? \n- Transformers are achieving impressive results in control as of late. Where might they be integrated into this work in the future, to improve scalability, generalization, etc.?", "limitations": "- The modeling assumptions constitute most of the limitations, but these are also what enable the convenient closed-form updates.\n- The authors note that the their model is limited to two levels of temporal abstraction, and that for certain tasks (e.g. Maze2D) more hierarchies may help. The method allows for addition of more complex abstractions such as Feudal hierarchies.\n- The authors are restricting their application to action conditional long-horizon future predictions. Although they note that future work can use these predictions for hierarchical control, they leave this for future work—it would have been nice to demonstrate this possibility with real control results.\n- Their method relies on reconstruction loss, which may have limited direct application to image-based domains. However, as the authors note, non-reconstruction based losses can be integrated. Again, it would have been nice to see this integration of different losses directly. \n- Overall, while I believe the authors have extensively noted the major limitations, especially for listed limitations (ii) and (iii), it would have been nice to see results in this paper to demonstrate that these limitations can indeed be relieved with simple substitution/introduction of new components.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- To the best of my knowledge, this closed-form multi time-scale inference approach using SSMs is novel, and has not be explored in previous works. The produced inference model is a principled (but non-trivial) integration of pre-existing components, and the results demonstrate the promise of such an approach\n- Generally, the work is of high quality. The writing is clear, and the paper is well-organized. The related work section is brief, but appears to adequately address prior work relevant to the aims of this paper.", "weaknesses": "- In the Figure 3 ablation plots, MTS3 should remain as Red in (c) to improve clarity. It would also be much clearer if the ablations were displayed in their own figure. It also does not make much sense, beyond organization of the plots onto separate lines, why Figure 3a and 3b are separated, as they are displaying the same findings across different environments. The a/b grouping contains no semantic difference. I understand the need for conserving space, but the way the figures are grouped and color-coded, it is not clear at a first glance what the relationship between a, b, and c are.\n- Spacing between Table 1 and the caption is too tight.\n- Figure 1 should be raised so that the caption does not overlap with 3.1 (try to align with line 79 paragraph).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for looking at our submission in a positive light. We would like to address a few questions raised by the reviewer below:\n\n**Question: How were the hyperparameters for MTS3 tuned? How were the baselines tuned?**\n\nAll hyperparameters for MTS3 and baselines including Transformers were tuned using grid search. For all recurrent models (MTS3, HiP-RSSM, RKN, LSTM, and GRU), we use a similar encoder-decoder architecture across datasets to ensure a fair comparison but allow the latent state dimensions to vary based on hyperparameter search. Small variations from these encoder-decoder architecture hyperparameters can still lead to similar prediction performance as reported in the paper.\n\n**Question: Transformers are achieving impressive results in control as of late. Where might they be integrated into this work in the future, to improve scalability, generalization, etc.?**\n\nWe do not see a way as of now to integrate transformers into the \"exact inference\" scheme for MTS3. However, a variational formulation of MTS3 (on the same graphical model that we propose) would allow for integrating transformers, especially to replace the set encoders used for Bayesian aggregation with transformer-style (self-attention-based) task encoders. Here the task posterior would be parameterized with transformers operating on sets of observations. Deriving an ELBO for variational approximations on the MTS3 graphical model is straightforward and would result in mathematically convenient decomposable objectives and can be a future research direction.\n\n**Weakness: Confusing figure placement and formatting issues**\n\nThank you for the suggestions. These would be duly addressed in the final version.\n"}, {"review_id": "BpV6GVqJIA", "reviewer": "Reviewer_BLVQ", "summary": "Looking to tackle the lack of temporal granularity in existing world models, the paper proposes a multi-time scale linear Gaussian state space model (MTS3). The model uses an efficient closed-form inference scheme on multiple time scales for highly accurate long-horizon predictions and uncertainty estimates over longer horizons. The experiments focus on action conditional long horizon future predictions, showing that MTS3 outperforms recent methods on several system identification benchmarks, including the D4RL dataset, a simulated mobile robot, and real manipulators including data from heavy machinery excavators.", "questions": "Conceptual:\n\n-\tBoth abstractions interrelate with each other in the sense that the higher-level predictions can be turned into low-level moment-by-moment predictions? How closely aligned where the encoded spaces across the levels for k-1, k when considering transition from t at k-1 to the next task? \n-\tIt is stated that it is trivial to extend to arbitrary number of temporal abstractions – would it be possible how this would work when there aren’t explicit factors like task that can determine the timestep H?\n-\tHow similar do the tasks, $l$ have to be for the results to hold? Change in environment dynamics? It was not clear to me how well the dynamics are encoded as the model transitions predictions from l to l+n, etc.\n-\tIs there a reason why the observation abstraction was chosen over the state space? \n\nFormulation, results and presentation:\n\n-\tHow is the second half of the latent state $d_t$ defined as the derivative that the model uses to the estimate the change of the observable part? Would this work for image-based observations formulations of the same model. \n-\tThe posterior calculation $P(z_t|o,a)$ involves a factorised over the covariance matrix. Where are $s, l$ defined? \n-\tIn eq1. How are Q and R parameterised; are these the $diag \\sigma_{z;o}? \n-\tWhat happens when H is manipulated? \n-\tCan tasks from different environments be learnt? Or would this be non-trivial?\n-\tWhat does the temporal encoding for both encoding of abstracted observation and action look like? Is this the number of steps taken or latent encoding of it?\n-\tWould it be possible to clarify what $m_t$ action entails; is this the derivative or something else?  \n-\tThe prior belief for the first time step of time window is initialised using the posterior belief of the last time step of time window; does this work for instances there is a massive shift in dynamics?\n\nMinor:\n\n-\tIs there a reason why some equations are numbered and others not?\n", "limitations": "The authors have appropriately addressed the limitations of MTS3, specifically regarding the consideration of only two timescales and the model's exclusive evaluation for predictions. Typically, it is customary to validate methods in a controlled setting when assessing the performance of the world model. However, I appreciate the learning the policy from such formulations is tricky and requires further consideration. ", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "-\tMulti-Time Scale Predictions: Training scalable hierarchical world models that operate at multiple timescales is an open challenge in multiple fields including compressed video prediction, reinforcement learning, control theory, etc. The paper presents one way to interleave between different levels conditioned on a task descriptor. The low-level learns the dynamics conditioned on a particular task and the higher-level is trained to predict the next task. \n-\tEfficient Inference Scheme via factorised formulation: Following Becker, et al, 2019, and Volpp, et al 2020, the paper proposes multi-scale inference via closed-form solutions using simplifying locally linear assumptions.\n-\tI like the way the two levels are connected: through the specification of the prior belief $p(l_k| B_{1:k-1}, \\alpha_{1:k-1})$ that defines the $p(l_k)$ in the fast-time scale.\n-\tThe use of a probabilistic formalism allows the model to handle uncertainty in predictions, particularly in prediction of changes to change, which is a common challenge across continual learning task settings. \n-\tRobust Performance: The model has been shown to outperform recent methods on several system benchmarks for long horizon predictions as noted in the figure 3. \n-\tCaptures Complex Dynamics: The model can better capture the complex, non-linear dynamics of a system in a more efficient and robust way than models that learn on a single time scale (Fig. 5)\n", "weaknesses": "-\tIn the formulation, the slow time scale SSM is only updated every step, i.e., the slow time $H$ scale time step is given by $H \\triangle t$. Therefore, the results are contingent on a design choice. I am cognizant that the this was evaluated for the range $0.2 – 0.5 $. However, how dramatically could results / predictions change if the wrong discretization step was chosen. Is there a way to systematically infer this through the system? \n-\tThe results have presented do not show how well the prediction settings vary across different tasks. \n-\tThe results haven’t shown how well large deviances in the dynamics would be encoded? Example some large spike. \n-\tHigher level encoded space is based on the observation space. Therefore, the higher-level is inherently conditioned on the first level. So, how would the scaling up work for multiple levels? Would $o \\to B$ change to $o \\to z \\to B$?\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive feedback and valuable suggestions. We write our replies to the weaknesses/questions listed by the reviewer below.\n\n**(Weakness) \"However, how dramatically could results ... wrong discretization step chosen. Is there a way to systematically infer this through the system?\"**\n\nThank you for your insightful comment. The discretization step is an important hyperparameter. We chose $0.2s \\leq H \\Delta t \\leq 0.5s$, based on related works in meta-learning[3], HiP-MDP / HiP-SSM [4], where for most physical systems like robots, dynamics can be assumed to be fixed for such a short duration.\n\nHowever, we further performed an ablation and some intuitive prediction plots for different $H$ on the mobile robot and hydraulic excavator dataset. The results of this can be found in **the rebuttal pdf**. We mention a rule of thumb on how to choose $H$ in different levels in the next section.\n\n**(Question) It is stated that it is trivial to extend to an arbitrary number of temporal abstractions – would it be possible how this would work when there aren’t explicit factors like tasks that can determine the timestep $H$?**\n\nWe had to limit our definition to 2 timescales because of space constraints. But a definition and details of a generic N-level MTS3 is defined in the \"Global Comment\" section. We also give a general rule of thumb on how the discretization step can be chosen (without hyperparameter search) there. We kindly request the reviewer to read the details in the global comments section.\n\n**(Weakness) Higher level encoded space is based on the observation space...  how would the scaling up work for multiple levels? Would it be o -> B or o -> z -> B ?**\n**(Question) Is there a reason why the observation abstraction was chosen over the state space?**\n\nWe hope the general definition of MTS3 and the computational/implementation details answer the question on how scaling up would work. We chose this approach to maintain a Feudal [5] (Top to Bottom) hierarchy for the latent variables. Ideally, the model should first make the easier higher-level task predictions (more abstract information without worrying about too many details as in lower levels), independent of the lower-level tasks and use the higher-level predictions to condition predictions on lower-level windows. Using observation abstractions allows for this feudal causal relationship.\n\n**(Question) How is the second half of the latent state $d_t$... Would this work for image-based observations .. ?**\n\nChoosing the observation model as $H=[I,0]$ allows for this division of the latent state vector into two distinct components. The upper part utilizes the identity matrix $I$ in $H$ to directly extract information from the observations. Meanwhile, the second lower part remains unobservable and is meant to hold information inferred over time, such as velocities in ordinary dynamical systems or images. The key aspect that contributes to the effectiveness of this choice is the selection of the covariance matrix structure (as mentioned in line 65 of the paper). The covariance matrix is designed to incorporate both diagonal and off-diagonal elements, ensuring that the correlation between the memory and the observation parts is effectively learned in the off-diagonal part. On the contrary, if we were to utilize a pure diagonal covariance structure, it would not update the memory units (the later half) or their variance adequately during the Observation/Task/Kalman Update step.\n\nYes, this assumption can effectively deal with image-based observations too as demonstrated by [6].\n\n**(Question) How are Q and R parameterized?**\nWe apologize for not making this clear in the paper. The transition noise $Q$ is assumed to be diagonal. It is learned and is independent of the state following [6]. The observation noise $R$ is again diagonal and is the output of the observation encoder, which is used in the Kalman Update / Bayesian Conditioning stage. This would be clarified and made more evident in the final version.\n\n**(Question) What happens when $H$ is manipulated?**\nAnswered above.\n\n**(Question) Can tasks from different environments ... be non-trivial?**\nThe challenge here would be dealing with varying dimensional observation and action space as the agent/environment changes. If this challenge can be addressed, MTS3 should scale to different environments in our opinion by learning environment-specific temporal abstractions for adaptation.\n\n**(Question) What does the temporal encoding ...?**\nIt is the normalized value of the absolute number of steps taken in a particular window (the absolute number of steps in a window is always $1 \\leq t \\leq H$). We didn't use any encoding.\n\n**(Question) Would it be possible to clarify what $m_t$ action entails; is this the derivative or something else?**\nCould the reviewer clarify which notation in the manuscript they are referring to? We would be happy to clarify this.\n\n**(Weakness / Question): How well large/massive shift in dynamics would be encoded? For example some large spike.**\nWe understand that since we have Gaussian assumptions, it's natural to pose this question about discrete changes. We can't learn/predict changes in dynamics caused by external rare events that are not predictable, like a spike (up and down suddenly) due to a human pushing a robot for a brief duration unexpectedly. However, we do think it can handle dynamics changes in the form of step functions (e.g., the load of the robot changes at set durations etc) as long as these changes have a pattern. The Franka Kitchen environment and the Panda real robot encounter such scenarios in our experiments.\n\nWe hope to have addressed most of the questions that the reviewer pointed out here and in global comments. We thank the reviewer again for their insights and remain hopeful that they would consider increasing the score.\n\nPlease refer to the global comment section for references.\n"}, {"review_id": "RGwerSU4AI", "reviewer": "Reviewer_HTHM", "summary": "The paper proposes a formalized multi-scale world model, which works at two timescales: a fast-timestep module that predicts individual timesteps, and a slower one that is only updated over fixed number of steps. The slower module defines a \"task\" that controls how the fast module functions, and the slower module reads abstract observations/actions over the time window. The paper derives Bayesian updates and ways to optimize the models. Experiments on three different datasets indicate that the proposed method is significantly better at predicting steps long into the future over baselines, and is also able to better predict the uncertainty.", "questions": "1) Transformer models are known improve in many different tasks as you scale them up in terms of parameters (or: number of layers, dimension size and number of heads per layer). The transformer models used in this paper were rather small (e.g., ~100 dimensionality, few layers). Did you experiment with larger model scales? However, even if transformer models get better at scale, they become slower to infer with, which can be an argument against them while employed in robotics with tight time and compute constraints. \n2) How long did it took to train each of the baselines, and how long does it take to run the model (e.g.,, infer steps 5 seconds in to the future)? Knowing the time-to-train and time-to-infer for each method, even if different code was ran on different pieces of hardware. This could strengthen the proposed method's case if it was faster than transformers.", "limitations": "Authors acknowledge the limitations of the work, and propose good paths to extend the work. Authors also discuss the broader impact (no immediate societal impact from the work).\n\n\n## Rebuttal acknowledgement\n\nI have read authors' rebuttal which addressed my concerns, and raised my score from 4 to 7 and confidence from 2 to 3 (before discussion period closed).", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- Formal derivation of the method, including the derivation of the uncertainty bounds.\n- Comparison against large number of valid baselines in multiple different settings, and all results indicate the proposed method is better\n- Ablation studies on which components of the proposed method are important\n- Experiments on both simulators and robots", "weaknesses": "- Some weaknesses in the baselines:\n\t- Only the proposed method had \"two-layer\" approach timewise, while others methods were \"single-layer\". While the formal \"two-layer\" approach was the main selling point of the proposed method, one can also naively create similar effect with RNNs or transformers\n\t\t- E.g., see this paper for two-layer RNN in Section 2.1: Jaderberg, Max, Wojciech M. Czarnecki, Iain Dunning, Luke Marris, Guy Lever, Antonio Garcia Castaneda, Charles Beattie et al. \"Human-level performance in 3D multiplayer games with population-based reinforcement learning.\" Science 364, no. 6443 (2019): 859-865.\n\t\t- Alternatively, a \"single-layer\" approach of the proposed method could also be included to better demonstrate the value of \"two-layer\" approach. However, if two-layer version outperforms single-layer approach, then two-layer results of the other baselines is even more warranted to ensure the benefit is not alone from two-layer approach.\n\t- Transformers used were rather small (4 layers, ~100 dimension), but many of the previous work has shown that larger transformers perform better (e.g., + 8 layers, 500 dimensions). While larger models are hard to use in robotics (not much compute / tight time constraints for control), testing the scalability of the different methods in terms of parameters to train would provide a better picture of the methods.\n\t\t- If larger transformer turns out to be better at world modelling without big impact in inference time, this method still has the benefit of capturing uncertainty, and can potentially be a better fit in planning.\n- While paper does mention the potential applications of the world modelling approach, there are no experiments to demonstrate this usefulness. The results seem positive for the proposed method (better at modelling the world), but it is unclear how useful this is down the line. For example, how accurate do you have to be to perform good planning for control?\n- No code available. The paper does detail the algorithm and setup used to great detail, but the code might contain details that researchers would need to replicate experiments. Small changes to the underlying libraries or how data is pre-processed may have big effects, or there might be parts in the code that were not reported in the paper. Any code, even if messy, is better than no code. \n\t- Given the complexity of the algorithm, having at least a pseudocode to refer to is crucial for correct implementation in the future.\n\t- Example of how code-level implementation details matter: Engstrom, Logan, Andrew Ilyas, Shibani Santurkar, Dimitris Tsipras, Firdaus Janoos, Larry Rudolph, and Aleksander Madry. \"Implementation matters in deep rl: A case study on ppo and trpo.\" In International conference on learning representations. 2019.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the comments and suggestions. Here are our replies to some of the weaknesses and questions listed.\n\n**Weakness:** Transformers used were rather small (4 layers, ~100 dimensions), testing the scalability of the different methods in terms of parameters to train. **Questions:** The transformer models used in this paper were rather small (e.g., ~100 dimensionality, few layers). Did you experiment with larger model scales? ... inference time/parameters in robotics with tight time and compute constraints.\n\nLarge transformers perform well and generalize when there is no scarcity of data, but it is unfortunately not the case with Robotics. We extensively tested different hyperparameter settings before arriving at this architecture. Larger models resulted in overfitting.\n\nAs the reviewer suggested, we are providing an analysis of the number of parameters and inference time for Transformer Vs MTS3 across two datasets in Table 1 in the **rebuttal pdf**. As seen in the table, MTS3 uses a fraction of parameters as transformers.\n\n**Questions:** How long did it take to train each of the baselines, and how long does it take to run the model? ...\n\nWe report a comparison for MTS3 vs Transformer baselines for training and inference time in Table 1.  As far as training time is considered, transformers scale better, due to the high parallelization possible at this stage. However, GPT-like Autoregressive transformers used widely in literature take several times more inference time due to autoregression (something not desirable in online robotic deployment). The non-autoregressive Multi-step transformers are faster but have the drawback that it's not flexible as the number of timesteps it can predict ahead is fixed and is hardcoded into the output decoder. Also, note that our code is not optimized as the standard transformer libraries. The performance was evaluated on similar hardware for all models.\n\n**Weakness:** Only the proposed method had a \"two-layer\" approach timewise, while others methods were \"single-layer\"... Alternatively, a \"single-layer\" approach of the proposed method could also be included to better demonstrate the value of the \"two-layer\" approach.\n\nWe include an extensive ablation in the rebuttal pdf, with varying discretization parameters (H) and a \"Flat-MTS3\" which works at a single time scale as the reviewer suggested. All other aspects of the model were fixed including training regime with imputation. It clearly shows the advantage of learning the dynamics with observation/action abstractions at multiple time scales.\n\nAdditionally, we referred to the paper suggested by the reviewer, Jaderberg et al. \"Human-level performance in 3D multiplayer games with population-based reinforcement learning.\" Science 364, no. 6443 (2019): 859-865. Unfortunately, since there is no public codebase available for the paper, we modified our existing LSTM/GRU baselines, with an additional slow-moving LSTM/GRU Cell as shown in Figure 2 of the reference. The paper shows a \"bottom to top\" architecture, where the low-level states of the fast RNN update the higher-level slow RNN. The hidden states of both RNNs were decoded to make predictions.\n\nOur experiments on both Mobile and Hydraulic robots resulted in worse performance than the single-layer LSTM for this particular architecture. However, we think it would be a fairer comparison only if we have access to their actual codebase or more architectural details.\n\nAlso, the 2-layer bottom-up RNN in the said reference was designed to capture the non-markovian state representation in the context of model-free RL policies, and we do not think this was meant for multi-step ahead predictions. We think that a Top to Bottom hierarchy for the latent variables as in our case is essential for making long-term predictions. Ideally, the model should first make the easier higher-level task predictions (more abstract information without worrying about too many details as in lower levels), independent of the lower-level tasks and use the higher-level predictions to condition predictions on lower-level windows. MTS3 maintains this top-to-bottom causal relationship between latent variables at multiple temporal abstractions in a principled way.\n\nWe would also request the reviewer to take a look at the global comment section, where we place our work in a different light. We think principled formalisms (e.g., SSMs, MDPs, Semi-MDPs in hierarchical RL context) are important for algorithmic advancements, and we lay a foundation for such a formalism for learning models of the world at multiple temporal abstractions.\n\n**Weakness:** No Code Available\n\nWe submit a link to an unofficial version of our codebase [here](https://drive.google.com/file/d/103fbLz1ahrfbSCGN0Ko58A7v76EVLPlu/).\n\nWe hope we were able to provide sufficient answers to the questions posted by the reviewer and hope that the reviewer would look at the work in a positive light. Please let us know if you need any further clarifications.\n"}, {"review_id": "t3hMLAYsUW", "reviewer": "Reviewer_hxxZ", "summary": "The paper addresses learning predictive world models that operate at multiple (i.e. 2) time scales. At the slower time scale, the belief over the “task” (i.e. the high-level state) is updated at every H time steps by aggregating the influence of the low-level observations and actions received over that period. The transition dynamics for the fast time scale model resemble those for a standard state space model, with the exception that the they are conditioned upon the high-level state. Throughout, the authors utilise linear transition models in latent space, with Gaussian process and observation noise.", "questions": "* How do the results support the case that it is the multi- time scale modelling of your approach that leads to accurate long-term predictions, rather than other factors such as the imputation-based training? From the results, it appears that the imputation-based training scheme (which is a fairly common technique) is crucial for obtaining strong performance. Were the baselines also trained in this manner, or were they only trained to make one-step predictions?\n* Can you include an ablation where *H* is set to 1 in your approach (i.e. there is only a single time scale in the model)? This would help to strengthen the case that this work demonstrates that the multiple time scales lead to improved long-term predictions.\n* Can the authors explain how they achieve strong results with an entirely linear model as stated in Equation 5? The authors do not appear to use the state-conditional linear model from [1], as they state the linear model is state independent. In principle, we might expect that there exists a high-dimensional latent space where the dynamics are linear, but the authors use fairly low-dimensional latent spaces (e.g. 30 or 60 dimensions). Thus, I find it difficult to understand how the authors approach can outperform methods that allow for non-linear latent dynamics.\n* Equation 5 (fully linear) appears to contradict Appendix A.3.1 (non-linear action model). Can you please explain this?\n* How should the results in Section 6.4 be interpreted? It is unclear to me how the log-likelihood are supposed to indicate good uncertainty estimates. Are these values the negative log-likelihood of real trajectories evaluated under each model? In which case, wouldn’t the higher likelihood values of MTS3 just indicate that the model is more accurate (not that the uncertainty estimates are better).\n\nI generally like the paper, and I am open to increasing my score if I feel these questions are adequately addressed during the rebuttal period.", "limitations": "The limitations section in the appendix is appreciated. A discussion of any limitations associated with the assumption of state-independent linear Gaussian latent dynamics would be useful.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* The paper addresses an important problem that is under-explored.\n* For the most part the writing is quite clear.\n* The experimental results appear strong in comparison to a number of baselines.", "weaknesses": "* It is unclear whether the experimental results support the authors’ claim that their model can more accurately capture system dynamics because of the two-time scale approach, or whether their method outperforms the baselines to due to other factors (such as the imputation-based training). An obvious baseline of updating both the low-level and high-level states at the same frequency (i.e. setting H = 1, and therefore removing the multi time scale aspect) is missing. Please see the questions.\n* In the preliminaries, the authors imply that they are building upon the locally-linear state space model from [1], that learns linear dynamics that are *conditional upon the state*. However, the authors model the dynamics in latent space to be linear in the latent state, action, and task context (Equation 5), and this linear transition model is *state-independent* (Line 150, i.e. it is the same linear dynamics for all states). It is unclear how the approach can achieve such strong performance with a completely linear model in a fairly low-dimensional latent space. This is further confused by the fact that Section 3.2 is contradicted by Appendix A.3 which instead tells us that the dependence on the actions is in fact non-linear, and learnt by a multi-layer MLP. Please see the questions.\n* I think the paper could be stronger if it presented a more general framework (e.g. arbitrary number of hierarchy levels, arbitrary non-linear dynamics (non-Gaussian)), and then presented their 2-timescale, linear Gaussian model as one example of that framework that is computationally simple. Presenting a more general framework would help to set the groundwork for future works that consider the problem of multi time scale world models.\n\n*Minor comments*\n- Covariance matrices Q and R do not appear to be defined.\n- Line 31: “under non-stationary” words appear to be missing from this sentence", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for posing these critical questions about our model and the suggestions given. We would attempt to answer these with the following paragraphs and global comment (and attached rebuttal pdf document).\n\n**Weakness:** It is unclear whether the experimental results support the authors’ claim that their model can more accurately capture system dynamics because of the two-time scale approach, or other reasons like imputation. **Question:** Can you include an ablation where $H$ is set to 1 in your approach?\n\nWe perform ablation for different values of $H$ (everything else remains fixed) in the rebuttal pdf. A value of $H = 1$ resulted in training instabilities and NaN values. The lowest $H$ value that resulted in stable training was 2. As seen in Figure 1 of the rebuttal pdf, smaller $H$ values (relatively fast time scale) for the task levels give significantly worse performance. Furthermore, we devised a \"Flat-MTS3\" variant with only 1 hierarchy, but all other parameterization and training schemes including imputation kept fixed. This again results in significantly worse performance indicating that it is, in fact, the multi-time scale approach that results in strong empirical results over the baselines.\n\nIn the section below we suggest a more systematic way of choosing the $H$ value for an MTS3 with an arbitrary number of hierarchies.\n\n**Weakness** I think the paper could be stronger if it presented a more general framework (e.g. arbitrary number of hierarchy levels, arbitrary...)\n\nWe had to limit our definition to 2 timescales because of space constraints. But a definition and details of a generic $N$-level MTS3 is defined in the \"Global Comment\" section. We also give a general rule of thumb on how the discretization step can be chosen by the rule of thumb (without hyperparameter search) there. We kindly request the reviewer to read the details in the global comments section.\n\nA section for a generic definition of an $N$ Level MTS3 can be included in the final script.\n\n**Question** Were the baselines also trained in this manner or were they only trained to make one-step predictions?\n\nThe baselines GRU, LSTM, RKN, HiP-RSSM were in fact trained with imputation schemes as in previous works. Moreover, we also tried a variety of imputation strategies that were not used by the original authors for these models to make them as fair a comparison as possible. However, none of those made much difference for “flat” / single time scale models (deterministic or stochastic). So this also provides evidence that it is, in fact, this multi-time scale formulation that results in a huge improvement in performance.\n\nThe ablations with different $H$ in Figure 1 of the rebuttal document and also the prediction plots in Figure 2 further strengthen our claim and we hope this addresses the reviewer's concerns.\n\n**Question** Equation 5 (fully linear) appears to contradict Appendix A.3.1 (non-linear action model). Can you please explain this?\n\nWe apologize for causing this confusion and the reviewer is right regarding this. The control model $B$ for the fast ($fts$) SSM is assumed to be non-linear, everything else is fully linear based on insights provided in [7] which studies the effects of action conditioning in great detail. Since the primitive action is known and not uncertain (zero variance), unlike the state transition model $A$ or the task transformation model $C$, we do not need to constrain the model $B$ to be linear as it neither affects the Kalman gain nor the closed-form **covariance** updates in the Kalman predict step.\n\nAlso, we think this is a matter of interpretation. The non-linear action model $b$ (mentioned in Appendix), can also be thought of as an encoder which projects the actions to a linear latent space and the matrix $B$ can be interpreted as the identity in this case ($B=I$). Thus, the dynamics still remain fully linear in the latent space.\n\nWe will clarify and make a discussion on this in the final version of the paper.\n\n**Weakness:** In the preliminaries, it is unclear ..  strong performance with a completely linear model. **Question:** Can the authors explain how  ... entirely linear model  .. in Equation 5?\n\nWe would also like to refer to works [11,12] that marry deep learning and Koopman Theory where finite-dimensional linear embeddings are learned with deep encoders for modeling non-linear dynamics.\n\nHowever, the real workhorse behind MTS3's surprisingly good results is its capability to model dynamics at multiple temporal abstractions in a top-down fashion with upper levels reconfiguring the lower levels. The multi-time idea has been in discussion in the machine learning community for some time. For eg, the idea of the Hierarchical JEPA architecture proposed (but not implemented) by LeCun in [2] is very similar to what we have in MTS3. We essentially confirmed the hypothesis proposed by LeCun et al. by formalizing and implementing it.\n\n**Question:** How should the results in Section 6.4 be interpreted? ... (not that the uncertainty estimates are better).\n\nWe would like to refer to some highly cited literature [8] on how marginal (Negative) Log Likelihood is a combined metric that measures both accuracy and uncertainty estimation of regression tasks. Similar metrics as ours were also used in related literature [9,10] to quantify the uncertainty of probabilistic models. We further plot the uncertainty in predictions (shaded regions) in Appendix C and also the rebuttal pdf (Section 2) to further get an intuitive feel. But we agree that there is still room for improvement in improving the quality of uncertainty estimates and also using better calibration metrics.\n\nWe hope the responses here and in the global comment section were able to address most of the questions/concerns raised by the reviewer and we remain hopeful that they would consider raising the score. References are listed in the global comments section.\n"}]

Author rebuttal:
We thank all the reviewers for their valuable suggestions / insightful questions and comments. We would like to post answers here to some common questions/weaknesses raised by multiple reviewers.

**1. Questions on whether the strong experimental results are out of learning at multiple time scales or some training strategies like imputation**

We hope the extensive ablations that we report in the **rebuttal pdf**, with different values of the discretization step (which changes the magnitude of time scale) show that it's in fact learning temporally abstract hierarchical world models that result in improved performance. We also compare with a "Flat"  single time scale MTS3. Note that all ablations/baselines were, trained with the best possible imputation schemes to ensure a fair comparison. Moreover, we provide intuitive prediction plots that explain to some degree how MTS3 behaves when the discretization step $H$ is changed.

**2. A general definition and implementation sketch for MTS3 with arbitrary hierarchies. How to choose H.**

**Definition:** An $N$-level MTS3 is defined as a family of $N$-state space models, $\{S_0, S_1, ..., S_{N-1}\}$. Each of the state space models $S_i$ is given by $S_i = (Z_i, A_i, O_i, f_i, h_i, H_i \Delta t, L_i)$, where $Z_i$ is the state space, $A_i$ the action space, and $O_i$ the observation space of the SSM. The parameter $H_i \Delta t$ denotes the discretization time-step and $f_i$ and $h_i$ the dynamics and observation models, respectively. Here, $l_i \in L_i$ is a task descriptor that 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}$.

**Choosing Discretization Step:** Though we recommend searching for $H_i$ as a hyperparameter, 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.

**Computation and Implementation:** From a computational standpoint, the higher-level temporal abstractions are inferred by the aggregation of $H_i$ observations/actions at each level $i$. However, since we derive this aggregation as a permutation-invariant set operation, this can be efficiently parallelized. The aggregation rules derived can be thought of as simple "probabilistic attention," where attention weights are given by the learned variances in the set encoder and has a linear computational complexity of $O(H_i)$, when similar deep-set operations (self-attention) in transformers have a complexity of $O(H_i^2)$.

**3. The advantage over transformers / Can transformers be integrated into MTS3 formalism?**

The goal of this paper is to come up with a formal probabilistic framework to learn world models leveraging hierarchical temporal abstractions. This also lays a foundation for designing more principled hierarchical planning methods, that can leverage control as an inference framework. Though we outperform Transformers on several benchmarks, we think Transformers (or any stateful RNN) that operate at multiple timescales taking inspiration from this formalism can be a promising alternative research direction as pointed out in line 260 in the main paper. But coming up with such an architecture in a non-ad-hoc way is not trivial in our opinion. 

There are also debates within the community as to whether LLMs/ Transformers capture causal relationships [1] and can autoregressive models like GPTs be effective for planning and reasoning [2].

We think our model combines the benefit of both worlds, set-based processing of temporal information as in Transformers via principled 
aggregation schemes (a sort of probabilistic attention), yet maintaining the sequential dependency of RNNs which enforces causality by design. They also use a fraction of parameters as Transformer baselines and are more computationally efficient during inference as discussed in the previous section.

**Additional Points**

We are attaching a [zip file](https://drive.google.com/file/d/103fbLz1ahrfbSCGN0Ko58A7v76EVLPlu/) to our unofficial code. We once again thank all the reviewers for their insightful comments and questions. We hope we were able to answer most of the questions raised by the reviewers. Please let us know if you need more clarification.

**Reference**

[1] https://www.cclear.cc/2023/CLeaR23_roundtable_discussion.pdf

[2] LeCun, Yann. "A path towards autonomous machine intelligence version 0.9. 2, 2022-06-27." Open Review 62 (2022).

[3] Nagabandi et al. "Learning to adapt in dynamic, real-world environments through meta-reinforcement learning." ICLR (2018).

[4] Shaj  et al. "Hidden parameter recurrent state space models for changing dynamics scenarios." ICLR (2022).

[5] Dayan and Hinton. "Feudal reinforcement learning." NIPS (1992).

[6] Becker et al. "Recurrent Kalman networks: Factorized inference in high-dimensional deep feature spaces." ICML, 2019.

[7] Shaj et al. "Action-conditional recurrent Kalman networks for forward and inverse dynamics learning." CoRL, 2021.

[8] Wilson and Izmailov. "Bayesian deep learning and a probabilistic perspective of generalization." Neurips (2020):

[9] Volpp, Michael, et al. "Bayesian context aggregation for neural processes." ICLR. 2020.

[10] Singh, Gautam, et al. "Sequential neural processes." Neurips (2019). 

[11] Lusch et al. "Deep learning for universal linear embeddings of nonlinear dynamics." Nature Communications (2018): 4950.

[12] Weissenbacher et al. "Koopman q-learning: Offline reinforcement learning via symmetries of dynamics." ICML, 2022.




Meta-review: The paper presents a novel probabilistic formalism called Multi Time Scale State Space (MTS3) to allow intelligent agents to learn world models that operate at multiple levels of temporal abstraction. Through a computationally efficient inference scheme, the proposed model achieves highly accurate long-horizon predictions and uncertainty estimates in both simulated and real-world dynamical systems.

Reviewers have highlighted the paper's strengths in offering a novel closed-form multi-time scale inference approach and robust performance on several benchmarks. The formal derivation of the method and comparison against a wide range of valid baselines were also commended, as were the experiments that span both simulators and robots. However, the paper could further benefit from addressing limitations such as the two-level temporal abstraction and the model's exclusive focus on predictions. Some reviewers also pointed out the potential for a more general framework and noted some weaknesses in the baselines used for comparison. Despite these areas for improvement, some of which can be addressed in revisions, the paper stands as a high-quality contribution to the field and is strongly recommended for acceptance.
Decision: Accept (spotlight)