Multi Time Scale World Models

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.

Strengths

* The paper addresses an important problem that is under-explored. * For the most part the writing is quite clear. * 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. * 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. * 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. *Minor comments* - Covariance matrices Q and R do not appear to be defined. - Line 31: “under non-stationary” words appear to be missing from this sentence

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? * 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. * 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. * Equation 5 (fully linear) appears to contradict Appendix A.3.1 (non-linear action model). Can you please explain this? * 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). I generally like the paper, and I am open to increasing my score if I feel these questions are adequately addressed during the rebuttal period.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

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.

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.

Strengths

- Formal derivation of the method, including the derivation of the uncertainty bounds. - Comparison against large number of valid baselines in multiple different settings, and all results indicate the proposed method is better - Ablation studies on which components of the proposed method are important - Experiments on both simulators and robots

Weaknesses

- Some weaknesses in the baselines: - 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 - 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. - 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. - 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. - 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. - 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? - 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. - Given the complexity of the algorithm, having at least a pseudocode to refer to is crucial for correct implementation in the future. - 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.

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. 2) 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.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

2 fair

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). ## Rebuttal acknowledgement I 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).

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.

Strengths

- Multi-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. - Efficient 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. - I 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. - The 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. - Robust Performance: The model has been shown to outperform recent methods on several system benchmarks for long horizon predictions as noted in the figure 3. - Captures 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)

Weaknesses

- In 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? - The results have presented do not show how well the prediction settings vary across different tasks. - The results haven’t shown how well large deviances in the dynamics would be encoded? Example some large spike. - Higher 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$?

Questions

Conceptual: - Both 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? - It 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? - How 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. - Is there a reason why the observation abstraction was chosen over the state space? Formulation, results and presentation: - How 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. - The posterior calculation $P(z_t|o,a)$ involves a factorised over the covariance matrix. Where are $s, l$ defined? - In eq1. How are Q and R parameterised; are these the $diag \sigma_{z;o}? - What happens when H is manipulated? - Can tasks from different environments be learnt? Or would this be non-trivial? - What 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? - Would it be possible to clarify what $m_t$ action entails; is this the derivative or something else? - The 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? Minor: - Is there a reason why some equations are numbered and others not?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

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.

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.

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 - 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. - Spacing between Table 1 and the caption is too tight. - Figure 1 should be raised so that the caption does not overlap with 3.1 (try to align with line 79 paragraph).

Questions

- How were the hyperparameters for MTS3 tuned? How were the baselines tuned? - 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.?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

- The modeling assumptions constitute most of the limitations, but these are also what enable the convenient closed-form updates. - 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. - 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. - 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. - 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.