Time-aware World Model:  Adaptive Learning of Task Dynamics

Authors: Anh N Nhu, Sanghyun Son, Ming Lin
Primary area: reinforcement learning
Keywords: RL, Dynamics, World Model

Abstract:
In this work, we introduce Time-Aware World Model, a model-based approach designed to explicitly incorporate the temporal dynamics of environments. By conditioning on the time step size, $\Delta t$, and training over a diverse range of $\Delta t$ values - rather than relying on a fixed time step size - our model enables learning of both high- and low-frequency task dynamics in real-world control problems. Inspired by the information-theoretic principle that the optimal sampling rate varies depending on the underlying dynamics of different physical systems, our time-aware model enhances both performance and learning efficiency. Empirical evaluations demonstrate that our model consistently outperforms baseline approaches across different observation rates in various control tasks, using the same number of training samples and iterations. We will release our source code on GitHub once the final review decisions are made.

Decision: reject
Reviewer scores: 6, 3, 5, 8, 6

Reviews and discussion:
<thread>
creation date:
1730059712425

thread_type:
official_review

<thread_object>
creator: reviewer_hd4m

<content>
summary:
The paper introduces the Time-Aware World Model (TAWM), which adapts to the temporal dynamics of environments by conditioning on the time step size ($\Delta t$). Different from conventional models that use a fixed time step, TAWM trains across a diverse range of ∆t values, allowing it to capture both high- and low-frequency task dynamics. This approach addresses shortcomings in existing models, such as temporal resolution overfitting and inaccurate system dynamics when applied to real-world scenarios.

strengths:
The time-aware mechanism significantly improves the emprical results compared to the non time-aware methods. The paper writing is improved and easy to understand.

weaknesses:
1. The paper writing clearance and correctness should be improved. In Line 140, what does $\eta(\pi)$ mean? policy or expected return? And why $s_i$ could be sampled from a policy $\pi$?
2. What is the advantage of the proposed method compared to directly setting the simulation frequency 2x bigger than the real-world frequency? It seems that your proposed method and baselines all work under small $\Delta t$.
3. As depicted in Fig.5, your proposed method is not significantly superior to the baselines in empirical performance.

questions:
Could you please provide more results of the comparisons between your proposed method and MTS3?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730670727323

thread_type:
official_review

<thread_object>
creator: reviewer_nmdp

<content>
summary:
The paper proposes time step aware world models for handling real life distribution shifts with lower observation frequencies. It includes the time step as an additional input to the world model and uses log-uniform time step sampling during training to learn the world model for various observation frequencies. It shows impressive results on various control tasks within Meta-World, without any increase in the sample complexity.

strengths:
The proposed strategy shows impressive performance when combined with TD-MPC2 on a variety of control tasks while using the same number of samples.
The proposed method can be combined with any existing MBRL algorithms such as TD-MPC2.

weaknesses:
One of the motivations for the TD-MPC2 work was to create generalist agents, which can be trained in a multi-task setup and then can be easily fine tuned on any new task. However, no comparison with TD-MPC2 has been made for a multi-task setup.
A closely related work (as mentioned by the authors) on multi time scale world models (Shaj et al.) is missing as a comparison baseline in the experimental section.

questions:
A precise description of how 4th order runge kutta method is used for integration is missing from section 4.1.2 or Algorithm 1. Can the authors please state it here as well as in the main paper for clarity?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730699026101

thread_type:
official_review

<thread_object>
creator: reviewer_8dtb

<content>
summary:
This paper emphasizes the importance of incorporating temporal information $\Delta t$ into dynamic modeling and introduces a mixture-of-time-step training framework that learns task dynamics across multiple frequencies. The authors' motivation stems from the existence of multi-scale dynamical systems, where each subsystem may operate at a unique frequency, and from the Nyquist-Shannon sampling theorem, which implies that lower sampling frequencies reduce performance, while higher sampling frequencies improve accuracy but increase sample complexity and reduce learning efficiency. Their approach involves conditioning all components of a world model, except the encoder, on $\Delta t$ and using the Euler or RK4 integration method to reformulate the dynamic model. The authors experimentally validate their method on several control tasks from the Meta-World suite, using TD-MPC2 as the baseline. They demonstrate that their time-aware modifications to TD-MPC2 yield superior performance over the baseline when evaluated across varying frequencies at test time, all while maintaining the same sample efficiency.

strengths:
- The paper is well-written and clear, with strong motivation and thorough explanations of each aspect of their approach.

- The proposed idea is innovative and promising, with potential for significant impact on real-world applications.

weaknesses:
- Since the core framework relies heavily on pre-existing world models and the main contribution is incremental, I expected to see larger-scale experiments in more complex environments, such as real robots or video games, to better illustrate the approach's applicability. In its current form, the paper lacks sufficient novelty to fully engage readers. I recommend that the authors either incorporate experiments in more complex environments to demonstrate real-world applicability or create more challenging scenarios, such as those requiring adaptation to varying $\Delta t$ within the same episode.


- The authors claim that smaller frequencies increase sample complexity and are therefore lead to inefficient learning. However, the paper lacks theoretical or experimental evidence to explain why this is true. To strengthen this claim and further support the motivation for time-aware world models, I suggest adding smaller frequencies to the experiment in Figure 4 and comparing their sample complexity with that of the time-aware model. This addition would provide valuable insight into the efficiency benefits of the proposed approach.

questions:
- Is there a specific reason why lower frequencies are missing in Figure 4? For instance, I was interested in seeing the baseline's performance with $\Delta t =0.1$ ms, as this frequency is used during the training of the time-aware model.

- Section 3.2.3 appears to be a mitigation for the issue described in 3.2.2. Would it be reasonable to combine these sections, or are they based on distinct motivations?

- Since TD-MPC2 assumes an underlying MDP structure, I suggest replacing $o_t$ with $s_t$ in Section 4.1 to avoid confusion with observations in a POMDP setting.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730701673037

thread_type:
official_review

<thread_object>
creator: reviewer_fxm1

<content>
summary:
For model-based RL in continuous time domains, the paper proposes to learn a time-aware world model whose components depend on the sampling frequency. In the training phase, the time-aware world model is trained with randomly varying time steps. In numerical experiments, the proposed method is evaluated under different time steps and show better performance compared with the baseline trained at the fixed time step.

strengths:
- When approximating a continuous-time dynamical system with a sampled model, having a time-step dependent world model makes a lot of sense to better fit the actual dynamics. The paper proposes such a time-aware world model with the potential to improve model-based RL performance.

weaknesses:
- The motivation referring to the Nyquist-Shannon sampling theorem does not seem to be consistent with the observation of numerical experiments. From the sampling theorem, as long as the sampling frequency is high enough, there won't be any information loss. But the numerical experiments suggest that the non time-aware model performs poorly even when the sampling frequency is the highest. This inconsistency makes the connection to the sampling theorem questionable.

- During the training phase, the sampling step is randomly selected from a log-uniform distribution. It is suggested that this distribution helps stabilize the training process, but no theoretical nor numerical analysis is provided. Some ablation studies using different sampling distribution might provide some insights to the choice.

- In the numerical evaluation, performance is compared across different observation rate. Given any time-step, the time-aware model can provide the appropriate prediction by conditioning on the time-step. But for the model trained with a fixed time-step, the information of the observation time-step does not seem to be adjusted. For example, for a model trained with $\Delta t=1$ms, when evaluating at $\Delta t=2$ms, instead of applying the model once, one might want to apply the model twice to have a more accurate prediction given the knowledge of the doubled time-step. Without doing some kind of adjustments for the baseline models make the fairness of comparisons questionable.

- When trained and evaluated with the same time-step, Figure 5 shows similar performance between the baseline model and the time-aware model. This makes the effectiveness of the proposed methods questionable when the goal is just to achieve good performance. Maybe under some scenarios where the training time-step and evaluation time-step are different, the proposed method might make more sense.

questions:
- It would be great if those points listed in the weaknesses above could be addressed.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731191718102

thread_type:
official_review

<thread_object>
creator: reviewer_mzgr

<content>
summary:
The authors propose an approach to address the issue of differing time step sizes $\Delta t$ between training and test phases, which can lead to a mismatch in dynamics. They propose a training scheme that randomly samples $\Delta t$ as inputs and update the model that explicitly conditions on $\Delta t$. They show that empirically this approach outperforms the baseline method, which was trained with a fixed time step size, when tested with a different $\Delta t$.

strengths:
- This paper considers a novel problem setting which is also practically important, especially for robotics control tasks.

- The paper is well written. The motivation is well-stated and the methodology is presented clear.

- The proposed algorithm is evaluated on multiple tasks. And the empirical results are supportive of the main claim of the paper.

weaknesses:
- I would suggest to add a comparison between the proposed approach with works about Robust RL (e.g. Pinto, Lerrel, et al. "Robust adversarial reinforcement learning." International conference on machine learning. PMLR, 2017). In my understanding, the mismatch in the time step size is one specific instance of the dynamics mismatch. I wonder if the proposed approach can get better performance by explicitly conditioning on $\Delta t$ only.

- I wonder if the log-uniform distribution is the best way to sample $\Delta t$. A justification for this would be helpful.

questions:
Please check the above Weaknesses section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732565366738

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General responses to common concerns from the reviewers

comment:
We appreciate all the reviewers’ comments. Here we provide answers to the common questions from reviewers.

**Q1) Unclear advantage over the baseline in Figure 5 (Reviewer fXm1, hD4m)**

**A1)** Reviewers questioned the clear time-aware model’s advantage in Figure 5 of our time-aware model over the non time-aware baselines.  As shown in Figure 5, the reward curves of the time-aware models may only appear to be converging marginally faster, when compared to the baselines with the default timestep of Δt=2.5ms. However, we would like to emphasize that **these results actually show the advantage of our model over the baseline in terms of both performance and efficiency**: the time-aware model (trained on varying Δt’s) outperforms the baseline (trained only on Δt=2.5ms) when evaluated on different Δt’s (Fig. 3 and Fig. 4) while still converge noticeably faster than the baseline when evaluated on inference Δt=2.5ms (the exact Δt for which the baseline was specifically trained on, so the best performance of the baseline is expected at this exact timestep).  Yet we observed slightly superior performance of our Time-Aware model over the exact trained case for the baseline, with clearly superior performance on other varying Δt’s.  We have run extra experiments and included additional results **updated Figure 5** and **Appendix C** as additional evidence on the generalization of our claim.

Specifically, we included additional reward curves evaluated on different inference Δt’s. Implementation-wise, we saved intermediate model weights during the training process, and evaluated such intermediate model weights on different inference Δt’s to obtain the reward curves (i.e: different figures show the return curve of the same model trained in the same process but evaluated on different inference Δt’s). As shown in Fig. 5 and Appendix C, we can observe clear performance advantages of the time-aware model over the baselines when evaluated on various inference Δt’s, especially Δt>2.5ms while using the same number of training steps (1.5M). This behavior is observed consistently across different environments. These results show that the **time-aware model is able to outperform the non time-aware model across different inference Δt’s in the same number of training steps (or without requiring additional training data)**.

---

**Q2) No impressive strength and the main contribution is incremental. (Reviewer 8dTB, hD4m)**

**A2)**  First, we would like to emphasize that our **core contribution is the time-aware world model training framework, not the architecture contribution**. Particularly, we propose **a simple and highly efficient training method** by:
1. **Explicitly conditioning the dynamics model on time step size Δt**, one of the most important quantities in any dynamical system yet overlooked by the world model and reinforcement learning community.
2. **Train the world dynamic model on a mixture of Δt’s by randomly sampling Δt during the training process**.

To the best of our knowledge, **despite the simplicity of our method, our work is the first to propose such a time-aware training framework**, which is shown in ***Figure 3, 4, 5, 6, 7*** to consistently improve world model’s performance on varying inference time step sizes across different tasks without requiring additional training data, thus enhancing sample efficiency.

Since our main contribution is a time-aware training framework, our method is model-agnostic, which **can be easily employed to train any world model architecture**. 

One important strength of our work is that the dynamic model is explicitly conditioned on Δt and is trained to directly predict the next state $s_{t+Δt}$ in **one-step prediction**. The main advantages of one-step prediction conditioned on Δt are:
1. Since our model can predict the next state under large Δt in a single step, it does not introduce additional computational overhead caused by multi-step predictions. As a result, it is compatible with real-world constraints of observation rate (e.g: 60fps) and/or control frequency (e.g: 120Hz).
2. Another benefit of single-step prediction is that we can avoid compounding error problems, a well-known problem for long-horizon prediction [1,2,3]. 

**References:**
1. Lambert, Nathan, Kristofer Pister, and Roberto Calandra. "Investigating compounding prediction errors in learned dynamics models." arXiv preprint arXiv:2203.09637 (2022).
2. Clavera, Ignasi, et al. "Model-based reinforcement learning via meta-policy optimization." Conference on Robot Learning. PMLR, 2018.
3. Wang, Tingwu, et al. "Benchmarking model-based reinforcement learning." arXiv preprint arXiv:1907.02057 (2019).
</content>
</thread_object>

</thread>


Metareview:
metareview:
This work proposes a way to learn a time-dependent world model with components depending on the sampling frequency. During training this is achieved by randomly varying time-steps. The same is done at test time, where the proposed method is evaluated under different time-steps and is shown to perform better than the baseline method that uses a single fixed time-step. 

The consensus is that this work shows promise, but that the approach is incremental, the novelty limited, and the experiments too small to ascertain whether the proposed tweaks produce enough gains.