SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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  $\sigma_{l_k}^{u-}$ ,  $\sigma_{l_k}^{l-}$  and  $\sigma_{l_k}^{s-}$  to its corresponding precision representation  $\lambda_{l_k}^{u-}$ ,  $\lambda_{l_k}^{l-}$  and  $\lambda_{l_k}^{s-}$  which can be performed using block-wise matrix inversions of  $\Sigma_{l_k}^-$ . Due to the factorization of the covariance matrix, this operation can be performed solely by scalar inversions. As the update equations are rather lengthy, they are given in Appendix A, B. Subsequently, we compute the posterior precision, where only  $\lambda_{l_k}^{u}$  is changed by
$$
\boldsymbol {\lambda} _ {l _ {k}} ^ {u +} = \boldsymbol {\lambda} _ {l _ {k}} ^ {u -} + \sum_ {t = 1} ^ {H} \mathbf {1} \oslash \boldsymbol {\nu} _ {k, t} \tag {8}
$$
while  $\lambda_{l_k}^{l+} = \lambda_{l_k}^{l-}$  and  $\lambda_{l_k}^{s+} = \lambda_{l_k}^{s-}$  remain constant. The operator  $\oslash$  denotes the element-wise division. From the posterior precision, we can again obtain the posterior covariance vectors  $\sigma_{l_k}^{u+}$ ,  $\sigma_{l_k}^{l+}$  and  $\sigma_{l_k}^{s+}$  using only scalar inversions, see Appendix A, B. The posterior mean  $\mu_{l,k}^{+}$  can now be obtained from the prior mean  $\mu_{l,k}^{-}$  as
$$
\boldsymbol {\mu} _ {l, k} ^ {+} = \boldsymbol {\mu} _ {l, k} ^ {-} + \left[ \begin{array}{c} \boldsymbol {\sigma} _ {l _ {k}} ^ {u +} \\ \boldsymbol {\sigma} _ {l _ {k}} ^ {s +} \end{array} \right] \odot \left[ \begin{array}{c} \sum_ {t = 1} ^ {H} \left(\boldsymbol {\beta} _ {k, t} - \boldsymbol {\mu} _ {l _ {k}} ^ {\mathrm {u}, -}\right) \oslash \boldsymbol {\nu} _ {k, t} \\ \sum_ {t = 1} ^ {H} \left(\boldsymbol {\beta} _ {k, t} - \boldsymbol {\mu} _ {l _ {k}} ^ {\mathrm {u}, -}\right) \oslash \boldsymbol {\nu} _ {k, t} \end{array} \right]. \tag {9}
$$
Note that for  $H = 1$ , i.e a single observation, the given equation is equivalent to the RKN updates. Moreover, the given rule constitutes a unification of the batch update rule for Bayesian aggregation [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,\dots,S_{N - 1}\}$ . Each of the state space model  $S_{i}$  is given by  $S_{i} = (Z_{i},A_{i},O_{i},f_{i},h_{i},H_{i}\Delta t,L_{i})$ , where  $Z_{i}$  is the state space,  $A_{i}$  the action space, and  $O_{i}$  the observation space of the SSM. The parameter  $H_{i}\Delta t$  denotes the discretization time-step and  $f_{i}$  and  $h_i$  the dynamics and observation models, respectively. Here,  $l_{i}\in L_{i}$  is a task descriptor that parametrizes the dynamics model of the SSM and is held constant for a local window of  $H_{i + 1}$  steps.  $l_{i}$  is a function of the latent state of SSM one level above it, i.e.,  $S_{i + 1}$ . The boundary cases can be defined as follows: for  $i = 0$ ,  $H_0 = 1$ . Similarly, for  $i = N - 1$ , the latent task descriptor  $L_{i}$  is an empty set. For all  $i$ ,  $H_{i} < H_{i + 1}$ .
Even though our experiments focus on MTS3 models with 2 hierarchies, extensive experimentation with more hierarchies can be taken as future work.
4 MTS3 as a Hierarchical World Model
MTS3 allows for a natural way to build world models that can deal with partial observability, 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 sts-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(\pmb{z}_{k,1})$  for the first time step of the time window  $k$ , we use the prior belief  $p(\pmb{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.,
$$
\begin{array}{l} L = \sum_ {k = 1} ^ {N} \sum_ {t = 1} ^ {H} \log p \left(\boldsymbol {o} _ {k, t + 1} \mid \boldsymbol {\beta} _ {1: k - 1}, \boldsymbol {\alpha} _ {1: k}, \boldsymbol {w} _ {k, 1: t}, \boldsymbol {a} _ {k, 1: t}\right) \\ = \sum_ {k = 1} ^ {N} \sum_ {t = 1} ^ {H} \log \iint p (\boldsymbol {o} _ {k, t + 1} | \boldsymbol {z} _ {k, t + 1}) p (\boldsymbol {z} _ {k, t + 1} | \boldsymbol {w} _ {k, 1: t}, \boldsymbol {a} _ {k, 1: t}, \boldsymbol {l} _ {k}) p (\boldsymbol {l} _ {k} | \boldsymbol {\beta} _ {1: k - 1}, \boldsymbol {\alpha} _ {1: k}) d \boldsymbol {z} _ {k, t + 1} d \boldsymbol {l} _ {k} \\ = \sum_ {k = 1} ^ {N} \sum_ {t = 1} ^ {H} \log \int p \left(\boldsymbol {o} _ {k, t + 1} \mid \boldsymbol {z} _ {k, t + 1}\right) p _ {\boldsymbol {l} _ {k}} \left(\boldsymbol {z} _ {k, t + 1} \mid \boldsymbol {w} _ {k, 1: t}, \boldsymbol {a} _ {k, 1: t}\right) d \boldsymbol {z} _ {k, t + 1}. \tag {10} \\ \end{array}
$$
The extension to multiple trajectories is straightforward and omitted to keep the notation uncluttered. Here,  $\pmb{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}\left(\pmb{z}_{k,t + 1}|\pmb{w}_{k,1:t},\pmb{a}_{k,1:t}\right)$  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(\pmb{o}_{k,t+1}|\pmb{\beta}_{1:k-1}, \pmb{\alpha}_{1:k}, \pmb{w}_{k,1:t}, \pmb{a}_{k,1:t}) \approx \mathcal{N}(\pmb{\mu}_{\pmb{o}_{k,t+1}}, \mathrm{diag}(\pmb{\sigma}_{\pmb{o}_{k,t+1}}))$  where we use the mean of the prior belief  $\pmb{\mu}_{z_{k,t+1}}^{-}$  to decode the predictive mean, i.e.,  $\pmb{\mu}_{\pmb{o}_{k,t+1}} = \mathrm{dec}_{\pmb{\mu}}(\pmb{\mu}_{z_{k,t+1}}^{-})$  and the variance estimate of the prior belief to decode the observation variance, i.e.,  $\pmb{\sigma}_{\pmb{o}_{k,t+1}} = \mathrm{dec}_{\pmb{\sigma}}(\pmb{\Sigma}_{z_{k,t+1}}^{-})$ . This approximation can be motivated by a moment-matching perspective and allows for end-to-end optimization of the log-likelihood without using auxiliary objectives such as the ELBO [1].
Gradients are computed using (truncated) backpropagation through time (BPTT) [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 sts-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.
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.
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
(c)
Figure 3: (a) and (b) Comparison with baselines in terms of RMSE for long horizon predictions (in seconds) as discussed in Section 6.3. (c) Ablation with different MTS3 variants is discussed in 6.6.
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. AR-Transformers also fail possibly due to error accumulation caused by the autoregression. However, Multi-Transformers are a strong baseline that outperforms MTS3 in the Medium Maze and Panda dataset by a small margin. However, on more complex tasks like the Kitchen task, which requires modelling multi-object, multi-task interactions [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.
Prediction Horizon	Algorithm
MTS3	Multi-Trans	AR-Trans	LSTM	GRU	RKN	HiP-RSSM
Half Cheetah	6 s	-2.80 ± 0.30	0.25 ± 0.05	X	7.34 ± 0.06	7.49 ± 0.04	X	X
Kitchen	2.5 s	-25.74 ± 0.12	-7.3 ± 0.2	X	32.45 ± 1.64	32.72 ± 0.65	X	X
Medium Maze	4 s	-0.21 ± 0.022	-0.88 ± 0.02	X	4.03 ± 0.32	7.76 ± 0.07	X	X
Panda	1.8 s	2.79 ± 0.32	3.77 ± 0.33	X	7.94 ± 0.39	7.91 ± 0.23	X	X
Hydraulic	12 s	-2.64 ± 0.12	-2.46 ± 0.03	X	7.35 ± 0.061	7.35 ± 0.06	X	X
Mobile Robot	3 s	-6.47 ± 0.71	-5.17 ± 0.23	X	11.27 ± 2.3	14.55 ± 5.6	X	X
Table 1: Comparison in terms of Negative Log Likelihood (NLL) for long horizon predictions (in seconds). Here bold numbers indicate the top methods and  $\pmb{x}$  denotes very high/nan values resulting from the highly divergent mean/variance long-term predictions.
6.5 How important are the modelling assumptions and training scheme?
Now, we look at three important modelling and training design choices: (i) splitting the latent states to include an unobservable "memory" part using observation model  $h^{sts} = h^{fts} = \mathbf{H} = [\mathbf{I},\mathbf{0}]$  as discussed in Sections 3.3 and 3.2, (ii) action abstractions in Section 3.1.2, (iii) training by imputation.
Figure 4: Ablation on discretization step  $H. \Delta t$  (a) The long-term prediction results in terms of RMSE, with different  $H$  values as discussed in Section 6.6 on the hydraulics dataset. (b) The predictions by MTS3 variants with different values of timescale parameter  $H. \Delta t$  on a trajectory picked from the hydraulics excavator dataset. The top figures are for  $H = 3$  and  $H = 10$ . Bottom figures are for  $H = 30$  and  $H = 75$ . Note that the results reported in the paper are with  $H = 30$ .
To analyze the importance of the memory component, we derived and implemented an MTS3 variant with an observation model of  $h^{sts} = h^{fts} = \mathbf{I}$  and a pure diagonal matrix representation for the covariance matrices. As seen in Figure 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.\Delta t$ ?
Finally, we perform ablation for different values of  $H$ .  $\Delta t$ , which controls the time scale of the task dynamics. The results reported are for the hydraulics dataset. The higher the value of  $H$ , the slower the timescale of the task dynamics relative to the state dynamics. As seen in Figure 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.
In this work, we introduce MTS3, a probabilistic formalism for learning the dynamics of complex environments at multiple time scales. By modelling the dynamics of the world at multiple levels of temporal abstraction we capture both the slow-changing long-term trends and fast-changing short-term trends in data, leading to highly accurate predictions spanning several seconds into the future. Our experiments demonstrate that simple linear models with principled modelling assumptions can compete with large transformer model variants that require several times more parameters. Furthermore, our inference scheme also allows for principled uncertainty propagation over long horizons across multiple time scales which capture the stochastic nature of environments. We believe our formalism can benefit multiple future applications including hierarchical planning/control. We discuss the limitations and broader impacts of our work in Appendix G and H.
We thank the anonymous reviewers for the valuable remarks and discussions which greatly improved the quality of this paper. This work was supported by funding from the pilot program Core Informatics of the Helmholtz Association (HGF). The authors acknowledge support by the state of Baden-Württemberg through bwHPC, as well as the HoreKa supercomputer funded by the Ministry of Science, Research and the Arts Baden-Württemberg and by the German Federal Ministry of Education and Research.