Title: Multi Time Scale World Models
Venue: NeurIPS 2023
Accepted: yes
Keywords: Hierarchical Models; Multi Time Scale Learning; World Models
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

Abstract

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

For efficient long-horizon prediction and planning, the model needs to predict at multiple levels of temporal abstractions [27, 22]. 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 [26, 17].

3

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 [15, 12] 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 ( Z , A , O , f, h, ∆ t ) , where Z is the state space, A the action space and O the observation space of the SSM. The parameter ∆ 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

Here z t ∈ Z , a t ∈ A and o t ∈ O are the latent states, actions and observations at time t. The vectors ϵ t and 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 [9, 1, 25] 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 o with a latent observation w . This latent observation is obtained by an encoder w o t = enc w ( o t ) along with the uncertainty of this observation, i.e., σ o t = enc σ ( 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

where a simple observation matrix of H = [ I , 0 ] is used. The underlying assumption behind this observation model is that the latent state z t = [ p T t , d T t ] T has twice the dimensionality of the latent observation w t and only the first half of the latent state, i.e., p t , can be observed. The second half of the latent state, i.e., 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 ( z t | o 1: t , a 1: t -1 ) = N ( µ t , Σ t ) such that only the diagonal and one off-diagonal vector of the covariance need to be computed, i.e.

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, [29] introduce Bayesian aggregation in the context of Meta-Learning. They again use an encoder to obtain a latent observation vector r o t and its uncertainty vector σ o t . Given the observation model p ( r o t | z ) = N ( Hz , diag ( σ o t )) with H = I and a prior p ( z ) = N ( µ 0 , diag ( σ 0 )) , the posterior p ( z | 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, 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 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 ∆ 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 ∆ 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, 9]. 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 [5] 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 S fast = ( Z , A , O , f fts l , h fts , ∆ t, L ) . Here, l ∈ 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 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

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 l k , we have to consider a distribution over task-descriptors p ( 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

Latent observations. Following [1], we replace the observations by latent observations and their uncertainty, i.e., we use latent observation encoders to obtain w k,t = enc w ( o k,t ) and an uncertainty encoder σ k,t = enc σ ( o k,t ) . The observation model is hence given by p ( w k,t | z k,t ) = N ( h fts ( z k,t ) , diag ( σ k,t )) .

Latent Task

Dynamics

Latent State

Dynamics

Time Scale 2 (Slow)

Time Scale 1 (Fast)

3.1.2 Slow time-scale SSM

The slow time-scale (sts) SSM only updates every H time step and uses the task parameter l as latent state representation. Formally, the SSM is defined as S slow = ( L , E , T , f sts , h sts , H ∆ t ) . It uses an abstract observation β ∈ B and abstract action α ∈ A that summarize the observations and actions respectively throughout the current time window. The general dynamics model is hence given by

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

β

1

,t

t

= 1

..H

α

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

β

2

,t

t

= 1

..H

α

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

β

3

,t

α

3

t

= 1

..H

Observation abstraction. In terms of the abstract observation model, we choose to model H observations β k,t , t ∈ [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,

and p ( β k,t | l k ) = N ( h sts ( l k ) , diag ( ν k,t )) . Hence, the abstract observation β k,t contains the actual observation 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 α k causes the transitions to the latent task l k from l k -1 . It should contain the relevant information of all primitive actions a k,t , t ∈ [1 , H ] executed in the time window k . To do so, we again use Bayesian conditioning and latent action encoding. Each control action a k,t and the encoding of time-step t is encoded into its latent representation and its uncertainty estimate, i.e.,

The single latent actions α k,t can be aggregated into a consistent representation α k using Bayesian aggregation [29]. To do so, we use the likelihood p ( α k,t | α k ) = N ( α k , diag ( ρ k,t )) and obtain the posterior p ( α k | α k, 1: H ) = N ( µ α k , Σ α 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

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 ( l k | β 1: k -1 , α 1: k ) and posterior p ( l k | β 1: k , α 1: k ) belief state for the sts-SSM as well as the belief states for the fts-SSM. The prior belief p ( l k | β 1: k -1 , α 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,

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 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 z t = [ p t , m t ] T into its observable part p t and a part m t that needs to be observed over time. We also use a linear observation model p ( w t | z t ) = N ( Hz t , diag ( σ t )) with H = [ I , 0 ] .

We will assume that the distribution over the task descriptor is also given by a Gaussian distribution, i.e., p ( l k ) = N ( µ l k , Σ 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

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 ( l k | l k -1 , α k ) = N ( Xl k -1 + Y α k , S ) , where X , Y and S are learnable parameters. The marginalized transition model for the abstract actions is then given by

We can directly use this transition model to obtain the Kalman prediction update which computes the prior belief p α 1: k ( l k | β 1: k -1 ) = N ( µ -l k , Σ -l k ) from the posterior belief p α 1: k -1 ( l k -1 | β 1: k -1 ) = N ( µ + l k -1 , Σ + 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 ( β k,t | l k ) = N ( H l k , diag ( ν k,t )) with H = [ I , 0 ] . 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., l k = [ u k , v k ] . In difference to the standard factorized Kalman observation update given in Appendix A, we have to infer with a set of observations ⃗ β k,t with t = 1 . . . 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 α 1: k ( l k | β 1: k ) = N ( µ + l k , Σ + 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 [31]. We found that this update rule is easier to formalize using precision matrices. Hence, we first transform the prior covariance vectors σ u -l k , σ l -l k and σ s -l k to its corresponding precision representation λ u -l k , λ l -l k and λ s -l k which can be performed using block-wise matrix inversions of Σ -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 λ u l k is changed by

while λ l + l k = λ l -l k and λ s + l k = λ s -l k remain constant. The operator ⊘ denotes the element-wise division. From the posterior precision, we can again obtain the posterior covariance vectors σ u + l k , σ l + l k and σ s + l k using only scalar inversions, see Appendix A, B. The posterior mean µ + l,k can now be obtained from the prior mean µ -l,k as

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 [29] 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 ∆ 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 ∆ t denotes the discretization time-step and f i and h i the dynamics and observation models, respectively. Here, l i ∈ 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, nonstationarity and uncertainty in long-term predictions, properties which are critical for model-based control and planning. Furthermore, introducing several levels of latent variables, each working at a different time scale allows us to learn world models that can make action conditional predictions/'dreams' at multiple time scales and multiple levels of state and action abstractions.

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 ( z k, 1 ) for the first time step of the time window k , we use the prior belief p ( 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.,

The extension to multiple trajectories is straightforward and omitted to keep the notation uncluttered. Here, 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 l k ( z k,t +1 | w k, 1: t , 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 ( o k,t +1 | β 1: k -1 , α 1: k , w k, 1: t , a k, 1: t ) ≈ N ( µ o k,t +1 , diag ( σ o k,t +1 )) where we use the mean of the prior belief µ -z k,t +1 to decode the predictive mean, i.e, µ o k,t +1 = dec µ ( µ -z k,t +1 ) and the variance estimate of the prior belief to decode the observation variance, i.e., σ o k,t +1 = dec σ ( Σ -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) [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 p ( z k, 1 ) with the prior belief p l k -1 ( z k -1 ,H +1 | w k -1 , 1: H , 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 [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 [17]. 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. [9, 1, 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.

[16, 13, 10] 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 [28, 23, 2], 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 [33, 18] and learning world models [19]. [33, 18, 20] use transformer architectures based on a direct multistep loss [32] and show promising results for long-term 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 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 F, a brief description of them is given here: (a) RNNs - We compare our method to two widely used recurrent neural network architectures, LSTMs [11] and GRUs [4]. (b) RSSMs - Among several RSSMs from the literature, we chose RKN [1] and HiP-RSSM [25] 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 [28] variants. The first variant (AR-Transformer) relies on a GPT-like autoregressive prediction [23, 2]. Whereas the second variant (Multi-Transformer) uses direct multi-step loss [32] from recent literature on long horizon time-series forecasting [33, 18, 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 use 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-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 nonstationary 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.

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 3. We see that MTS3 gives consistently good long-term action conditional future predictions on all 6 datasets. Deep Kalman

0

1

2

3

4

5

6

Seconds

0.2

0.3

0.4

0.5

HalfCheetah

0

1

2

3

4

Seconds

0.2

0.4

0.6

0.8

1.0

1.2

1.4

Medium Maze

0.0

0.5

1.0

1.5

2.0

2.5

Seconds

0.04

0.06

0.08

0.10

0.12

Kitchen

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

Seconds

0.05

0.10

0.15

0.20

0.25

0.30

Panda

RMSE

MTS3

LSTM

GRU

RKN

HiP-RSSM

AR-Transformer

Multi-Transformer

(a)

0

2

4

6

8

10

12

Seconds

0.005

0.010

0.015

0.020

0.025

0.030

Excavator

(b)

(c)

models [1, 25] which operate on a single time scale fail to give meaningful mean estimates beyond a few milliseconds. Similarly, widely used RNN baselines [11, 4] which form the backbone of several world models [8, 10] give poor action conditional predictions over long horizons. ARTransformers 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 [7], MTS3 is the only model that gives meaningful long horizon predictions. A detailed description of 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 of whether the principled probabilistic inference translates to accurate uncertainty quantification during long-horizon predictions. We trained all the baselines with a negative log-likelihood loss and used the same as a metric to quantify the quality of uncertainty estimates. During the evaluation we rely on a sliding window approach (see Appendix C) and report the results for the last timestep in Table 1. 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.

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

0

2

4

6

8

10

12

Seconds (Each Second = 100 timesteps)

0.01

0.02

0.03

0.04

0.05

0.06

Flat MTS3

2

3

5

10

30 (chosen in paper)

75

RMSE

(a)

0

200

400

600

800

1000

1200

0.3

0.4

0.5

0.6

0.7

0.8

MTS3 (H=3)

Ground Truth

Predictions

Masked

0

200

400

600

800

1000

1200

0.35

0.40

0.45

0.50

0.55

MTS3 (H=10)

Ground Truth

Predictions

Masked

0

200

400

600

800

1000

1200

0.35

0.40

0.45

0.50

0.55

MTS3 (H=30)

Ground Truth

Predictions

Masked

0

200

400

600

800

1000

0.35

0.40

0.45

0.50

0.55

Ground Truth

Predictions

Masked

(b)

To analyze the importance of the memory component, we derived and implemented an MTS3 variant with an observation model of h sts = h fts = I 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 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 3c 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. ∆ t ?

Finally, we perform ablation for different values of H. ∆ 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 4a, 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 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 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.

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

References

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.

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.

Eduardo F Camacho and Carlos Bordons Alba. Model predictive control . Springer science & business media, 2013.

Kyunghyun Cho, Bart van Merriënboer, Dzmitry Bahdanau, and Yoshua Bengio. On the properties of neural machine translation: Encoder-decoder approaches. In Proceedings of SSST-8, Eighth Workshop on Syntax, Semantics and Structure in Statistical Translation , pages 103-111. Association for Computational Linguistics, October 2014.

Peter Dayan and Geoffrey E Hinton. Feudal reinforcement learning. Advances in neural information processing systems , 5, 1992.

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.

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.

David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv:1803.10122 , 2018.

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.

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.

Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Comput. , 9(8): 1735-1780, November 1997. ISSN 0899-7667.

Michael I Jordan. Graphical models. Statistical science , 19(1):140-155, 2004.

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.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

Daphne Koller and Nir Friedman. Probabilistic graphical models: principles and techniques . MIT press, 2009.

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.

Yann LeCun. A path towards autonomous machine intelligence version 0.9. 2, 2022-06-27. Open Review , 62, 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 .

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 .

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 .

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.

Doina Precup and Richard S Sutton. Multi-time models for temporally abstract planning. Advances in neural information processing systems , 10, 1997.

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.

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.

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.

Murray Shanahan and Melanie Mitchell. Abstraction for deep reinforcement learning. arXiv preprint arXiv:2202.05839 , 2022.

Richard S Sutton. Td models: Modeling the world at a mixture of time scales. In Machine Learning Proceedings 1995 , pages 531-539. Elsevier, 1995.

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.

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.

Paul J Werbos. Backpropagation through time: what it does and how to do it. Proceedings of the IEEE , 78(10):1550-1560, 1990.

Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Ruslan Salakhutdinov, and Alexander Smola. Deep sets. arXiv preprint arXiv:1703.06114 , 2017.

Ailing Zeng, Muxi Chen, Lei Zhang, and Qiang Xu. Are transformers effective for time series forecasting? arXiv preprint arXiv:2205.13504 , 2022.

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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\n* 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\n\nQUESTIONS:\n* 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.\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-21"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n1) 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.\n\nLIMITATIONS:\nAuthors 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).", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nLooking 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.\n\nSTRENGTHS:\n-\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\nWEAKNESSES:\n-\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\nQUESTIONS:\nConceptual:\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\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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).\n\nQUESTIONS:\n- 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.?\n\nLIMITATIONS:\n- 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}