Time-Aware World Model for Adaptive Prediction and Control
Authors: Anh N Nhu
Sanghyun Son
Ming Lin
Venue: ICML 2025 poster (2025)
Keywords: RL
Dynamics
World Model
Primary area: reinforcement_learning

Abstract: In this work, we introduce the Time-Aware World Model (TAWM), a model-based approach that explicitly incorporates temporal dynamics. By conditioning on the time-step size, $\Delta t$, and training over a diverse range of $\Delta t$ values – rather than sampling at a fixed time-step – TAWM learns both high- and low-frequency task dynamics across diverse control problems. Grounded in the information-theoretic insight that the optimal sampling rate depends on a system’s underlying dynamics, this time-aware formulation improves both performance and data efficiency. Empirical evaluations show that TAWM consistently outperforms conventional models across varying observation rates in a variety of control tasks, using the same number of training samples and iterations. Our code can be found online at: github.com/anh-nn01/Time-Aware-World-Model.

Time-Aware World Model for Adaptive Prediction and Control
Anh N. Nhu
Sanghyun Son
Ming Lin
Abstract
In this work, we introduce the Time-Aware World Model (TAWM), model-based approach that explicitly incorporates temporal dynamics. By conditioning the time-step size, AI, and train- ing over a diverse range of At values rather than sampling at a fixed time-step TAWM learns both high- and low-frequency task dynam- ies across diverse control problems. Grounded in the information-theoree insight that the optimal sampling rate depends system's 's underlying dynamics, this time-aware formulation improves both performance and data efficiency. Empirical evaluations show that TAWM consistently outper- forms conventional models across varying obser- vation rates in a variety of control tasks, using the same number of training samples and iterations. Our code be found online at: github.com/anh- nn01/Timen
h
M
aetdtt
Eulen
At
inlegration
Ar
Ot
Encoder
Z
Time-Aware
Models
Figure I. Overall framework of our time-aware world model.
An encoder encodes the given observation into latent vector Sr, which is then fed into various models with action a, and time step size At estimate values for action planning.
interest in model-based RL (MBRL). Unlike model-free approaches, MBRL constructs a model M that captures the underlying task dynamics (Sutton, 1990; Deisenroth Rasmussen, 2011; Parmas et al., 2018; Kaiser al., 2019; Janner et al., 2019), enabling the agent to plan actions using the learned world model M. These model-based methods have gained traction over their model-free counterparts due to their superior sample efficiency and enhanced generaliza- tion capabilities (Ha & Schmidhuber, 2018; Hafner et al., 2019; 2020; 2023; Hansen et al., 2022: 2024).
1. Introduction
Despite the remarkable success of world models in vari- control tasks, critical factor in handling dynamical systems- - the time step size, At - has been largely over- looked in existing research. Specifically, the dynamics (s,,0)) + 8t+1 fundamental component of the world model, governs state transitions from the cur- rent state st to the next state 8t-1. Conventional approaches train D using experience tuples (or, at ot+1, r;) collected from interactions with the environment a fixed time step size. However, this existing practice presents three key limitations (Thodoroff et al., 2022):
Deep reinforcement learning (DRL) has recently demon- strated human-level or even expert-level capabilities on many highly complex and challenging problems, such as Go (Silver et al., 2016; 2017), Chess and Shogi games (Sil- et al., 2018), and StarCraft II video game (Vinyals et al., 2019). DRI. is also effectively adopted in broad range of applications, including robotics (Wu et al., 2023; Koh et al., Johannikk et al., 2019), autonomous vehicles (Kiran ct al., 2021; Guan et al., 2024), and challenging control tasks where classical approaches fail to deliver satisfactory performance (Prasad et al., 2017: Nhu al., 2023; Yang al., 2015). Beyond model-free reinforcement learning (RL) methods, where agent directly maps observation to action a; (Williams & Peng, 1989; Schulman et al., 2015; 2017; Haarnoja et al., 2018), there has been growing
model
1. Temporal resolution overfitting: In current training pipelines, the observation time step At are often fixed (e.g., A/ = 2.5ms or frequency f f= = 400HZ While smaller At values stabilize simulations and prevent system aliasing, world models trained exclusively at a single, fixed At often suffer significant performance degradation when deployed in real-world scenarios with different observation rates (c.g. At = 20ms or f = 50Hz) due to compounding errors (Lambert et al., 2022: Wang et al., 2019) and computational costs of roll-out predictions. This discrepancy presents a
Equual contribution 'Department of Computer Science, Uni- versity of Maryland, College Park, United States. Correspon- dence Anh Nhu Eanhu@umd.cduu> Sanghyun Son <shhh (2952m2d.0dduuurr
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Time-Aware World Model
major challenge in extending the applicability of world models beyond simulated environments.
Although one might expect M to require more train- ing samples than M because of the additional parame- ter At, this is not necessarily the case. According to the Nyquist-Shannon sampling theorem (Shannon, 1949; Jerri, 1977), a signal with highest frequency can be perfectly reconstructed by sampling at any frequency just above 2f. Thus, if the observation rate greatly exceeds 2f, the excess data are redundant- they contribute little to training the world model. A physical environment typically comprises multiple dynamical subsystems operating at different char- acteristic frequencies (Section 3.2. By training with a mixture of time-step sizes, we expose the model to a range of effective sampling frequencies, enabling each subsystem to be learned more efficiently (Section 3.2.3).
2. Inaccurate system dynamics: Training the dynamics model M on fixed time step At can lead to overfit- ting, as the model may not capture the true underlying dynamics of the task without conditioning At. This can result in inaccurate state transitions and diminished generalization capabilities.
3. Inefficient dynamics learning: Although small At values are essential for numerical stability and for cap- turing -igg-frrguency dynamics, relying exclusively on fine-grained time steps results in inefficient sam- pling. Real-world systems often exhibit multiscale be- havior, with some components evolving rapidly while others change more slowly. Sampling solely at high frequency captures redundant information about the slower components, wasting computational resources and training data. This issue is exacerbated in systems with widely varying time scales, where no single sam- pling rate is optimal for every dynamical component.
Inspired by the Nyquist-Shannon sampling theorem (Sec- tion 3.2.2), we prove that our time-aware model trained observation data sampled at a mixture of time steps achieves substantially better performance in learning world dynamics across different time steps at inference time, while using the same training budget as the baseline. We demon- strate these results on a variety of control problems in Meta-World (Yu al., 2020) and PDE-control environ- ments (Zhang et al., 2024). Our contributions are sum- marized follows:
Recent work (Shaj Kumar et al., 2023; Lutter et al., 2021) has explored temporal aspects of world models, but primar- ily to improve long-horizon predictions through roll-outs of fixed-At models. Our approach instead conditions the world model directly on At, enabling single-step predictions across arbitrary temporal intervals (Figure 1). Although pre- vious methods can, in principle, accommodate variable time steps through repeated application, that strategy compounds errors and is computationally inefficient. By sampling strate- gically across multiple time scales, our method lets the model learn dynamics at different frequencies simultane- ously, without increasing the number of training samples. Consequently, it generalizes effectively across varying ob- servation rates- crucial for simulation-to re-aaiity transfer - -wiie maintaining superior sample efficiency compared with fixed-At approaches.
1.
We underscore the importance of a time-aware world model (TAWM) that conditions its dynamics modeling on the time step size At, a fundamental quantity in any dynamical system. By explicitly incorporating At, the model learns to capture the underlying task dynamies across broad spectrum of step sizes. Our model's 's abil- ity to generate one-step prediction of the next state conditioned on At mitigates the well-known problem of compounding errors. This approach is particularly suitable for real-world learning and control, where the observation rate may vary - or be substantially lower than the default observation rate used for training.
2. Motivated by the Nyquist-Shannon sampling theorem and by the fact that real-world dynamics comprise many subsystems operating at different - and often unknown frequencies. we propose a mixture-of-tin step training framework for TAWM that performs well across range of observation rates at inference time, without increasing the number of training steps. This perspective - training a world model with varying sam- pling rates offers a new avenue for developing more efficient training strategies.
In this work. We focus on the following question: How can we efficiently train a world model M to capture the underlying task dynamics across varying time step sizes At without increasing sample complexity?
To answer this, we introduce a time-aware world model, MTA Unlike earlier world models M, our model con- ditions both the dynamics model and the reward function on At (Figure 1). because these depend on the temporal gap between consecutive states. We construct MTA by adapting the TD-MPC2 world model (Hansen et al., 2024) and incorporating the fourth-order Runge-Kutta (RK4) method (Butcher, 1987) to stabilize learning at large At values (Section 4.1). We likewise modify the value model to accept At as an additional input. Both models are trained on At values log-uniformly sampled from a predefined interval.
3. Empirically, we demonstrate that time-aware world model can solve range of control tasks at various ob- servation rates without increasing either the amount of data or the number of training steps. This capabil- ity helps narrow the gap between simulation environ- ments and worl-world control problems.
2

---

Time-Aware World Model
2. Related Work
state-transition function, r S A - R is the reward function, and 7 E [0, 1) is the discount factor. We refer to P and as the ground-truth models, because in Section 4.1 we learn approximations of them within our world model.
Although model-free RL algorithms have gained popularity for their impressive performance, their inherent sample in- efficiency limits their applicability to many real-world prob- lems (Sutton al., 1999; Williams & Peng, 1989; Barto et al., 1983; Schulman et al., 2015; 2017). Several works mitigate this limitation by integrating low-variance analyt- ical gradients into the policy-learning process (Suh et al., 2022; Xu et al., 2022; Son et al., 2024).
The goal of reinforcement learning is to obtain a policy or planner π that maximizes the expected discounted return along a trajectory {so, ou..... 8H OH OH-1 SH) of length H:
TT(T) = Ego 1 Σ r an) En)r(staa))
In contrast, model-based RL (MBRL) tackles sample in- efficiency at its root by training a dynamics model that simulates the environment, enabling agents to predict fu- ture states and action outcomes (Deisenroth et al., 2013). Because prediction is unconstrained by real-time sampling, MBRL is inherently more sample efficient. MBRL meth- ods vary mainly in (1) how they define the world model M and (2) how they exploit M for training or planning. Historically, Gaussian processes (GPs) (Deisenroth & Ras- mussen, 2011; Parmas et al., 2018) were widely used, but modern work favors neural networks (Ha & Schmidhuber, 2018: Hafner et al., 2019; 2020; 2023; Hansen et al., 2022: 2024) for their greater representational power. Among these, Dreamer and its variants (Hafner et al., 2019; 2020; 2023; Wu et al., 2023) train policies directly inside the learned (MPC) methods (Hansen et al., 2022; 2024) rely on action planning.
(1)
In model- model- RL (MBRL), we train world dynamics model consisting of a state-transition function dø : SS A + S and a reward function ra : Sx A - IR. The learned world model can then be exploited in multiple ways to derive a policy, for instance by using a planner such as TD-MPC2. Section 1.2 details model formulation and the training pipeline for the time-aware world model.
3.2. Theoretical Motivations
Here we provide theoretical motivations to explain the sam- ple efficiency of using mixture of time step sizes during the training process of the time-aware world model.
model, whereas model-predictive-conn
3.2.1. MULTI-SCALE DYNAMICAL SYSTEMS
However, these approaches paid little attention to the tempo- ral component when designing the world model M, despite its importance in dynamical systems. To address its limi- World
In many control problems, the environmental dynamics can be decomposed into multiple subsystems, each evolving on different temporal scale (Weinan, 2011). In other words, these subsystems can be described by distinct functions, each with its own highest frequency. Formally, consider general dynamical system å = f(z u, E), where and t denote the state, control input, and time, respectively. Applying the Euler integration method, the next state x' is
tation in long-horizon prediction, Multi-Time-Soee Models (MTS3) (Shaj Kumar et al., 2023) explicitly lever- ages temporal gaps when learning task dynamics. MTS3 captures state transitions over different prediction horizons H but uses a single fixed Al during training; it therefore represents only two discrete time scales - fast and slow dynamics corresponding to At and At. In contrast, we incorporate a continuous-valtee At directly into the model, enabling single-step predictions across a range of inference step sizes. This allows our model to handle large tempo- ral gaps (e.g., At 30ms) that would otherwise require 12-20 smaller steps in fixed-step model. Our approach is architecture agnostic: it builds on the TD-MPC2 framework (Hansen et al., 2024) yet crucially differs by explicitly con- ditioning the world model on At and training on a mixture of step sizes, rather than assuming a fixed step throughout.
a' x + fzz u. t). At
where At is the time-step size. From multi-scale perspec- tive, this update can be rewritten as
a' = x + Σ f((z,u,). At
(3)
where f1 u, t) represents the dynamics of subsystem i. Each subsystem may evolve at its own temporal scale and therefore have its own highest frequency f f
3. Background and Motivations
3.2.2. TYQUIST-SHANNOR SAMPLING THEOREM
3.1. Model-Based Reinforcement Learning
The Nyquist-Shanne sampling theorem states that sig- nal must be sampled at a rate of at least twice its highest frequency to avoid information loss (i.e., to prevent alias- ing): fample > 2f 2fmux (Shannon, 1949). Here, fample is the
We formulate the control problem as a Markov decision process (MDP) (S, A, P,r, 7), where S' is the state space, A
is the action space, P: S x A X S + RR is the (stochastic)
3

---

Time-Aware World Model
observation rate, and max is the highest frequency of the environment's dynamics in the MBRL context. If the obser- vation rate is too low such that fsamll < 2fmax we lose important dynamic details due to the large temporal gap A.. This loss causes high-frequency components to fold back into lower frequencies, leading to inaccurate dynamics learned by the world model (Zeng et al., 2024). Although a higher fsample enables more accurate reconstruction of the environment's ' dynamics, oversampling, introducing redundant data that increase sam- ple complexity and reduce learning efficiency. Therefore, choosing an appropriate observation rate sample is crucial for balancing modeling accuracy and sample efficiency.
MPC2 omits decoder that maps latent representations back to raw observations. Recovering latent states in the high- dimensional observation space is computationally expensive and often unnecessary, as many elements are irrelevant to control. Once the latent dynamics are learned, TD-MPC2 rolls out latent predictions for local trajectory optimization with planning algorithms such as MPC (Hansen et al., 2022; 2024). The baseline TD-MPC2 model comprises five key
excessively high rate leads to
components:
Encoder:
21 h(oe)
Latent dynamic:
2 = D(2,, at)
Reward:
F. = B(z,, az)
Terminal value:
I. = Q an)
Policy prior: à. = p(z))
3.2.3. SIMULTANEOUSLY TRAINING ON MULTIPLE TEMPORAL RESOLUTIONS
where is the raw observation at time step i, 2; is its latent encoding, 2 is the predicted next latent state, fr is the predicted immediate reward, q. is the estimated Q- Q-value of the current statc-action pair, and â is the action sampled from the policy for the current state. At inference time, the Model Predictive Path Integral (MPPI) planner an instance of Model Predictive Control (MPC) is used for planning and action generation. All five component models are implemented using multilayer perceptrons (MLPs).
To better align observation sampling with task dynamics across subsystems operating at different frequencies, we propose simultaneously training the world model at mul- tiple temporal resolutions by varying the observation rate faamll (i.e., varying At) during training. As shown in Equa- tion 3, the overall dynamics comprise several subsystems fi()). each with its own maximum frequency f'max. Ac- cording to the Nyquist-Shannon sampling theorem, each subsystem is learned most efficiently at distinct fsample By randomly varying faamll during training, we avoid un- dersampling high-frequency components and oversampling low-frequency ones, thereby training the subsystems fi(-) more efficiently. Consequently, our time-aware world model learns the underlying task dynamics at multiple temporal resolutions without needing additional data (Section 5).
To train TD-MPC2, we maintain a replay buffer B that stores trajectories (or: Gt. O1+1, Hf ) collected from the environ- ment after each episode of length H. At the end of every episode, the model parameters are updated with batches ran- domly sampled from B. The encoder h. dynamics model D, reward model R. and terminal-value model Q are optimized jointly using a self- sulf-supervised consistency loss, supervised reward loss, and a supervised cemporal-differeme the terminal value. See Hansen et al. (2024) for details. The agent then resumes interaction with the environment, gath- cring additional data augment for subsequent training.
loss for
4. Methodology
In this section, we present our time-aware model formula- tion and training pipeline, designed to learn a world model that performs well across a range of observation rates. We introduce novel time-aware training method that can be seamlessly integrated into any existing world-model archi- tecture, enhancing robustness to observation-rate variations at inference time. We adopt TD-MPC2 (Hansen et al., 2024) as our baseline and adapt its architecture to train time-aware world models for various control tasks. We begin with a high-level overview of TD-MPC2 and its key architectural components. The overall framework depicted in Figure
4.1.2. TIME-AWARE WORLD MODEL
One notable limitation of TD-MPC2 and other otate-of- the-art world models is that they ignore the effect of the observation rate Al at inference time. To introduce time awareness, we condition every component of the world model on At, as described below:
Encoder:
21 h(or)
Latent dynamics:
++ = + d(2,, at, T(t))
where T T(x) = max(0, logio x + 5)
4.1. Model Architecture
Reward model:
f, R[ az, At)
4.1.1. BASELINE WORLD MODEL
Value model:
q: = Q( at At)
Policy prior:
ā. p(z Al)
TD-MPC2 (Hansen et al., 2024) a model-based RL al- gorithm that captures task dynamics in a latent space an "implicit" world model. Unlike reconstruction-based archi- tectures (Ha& & Schmidhuber, 2018; Hafner et al., 2023), TD-
Our time-aware model formulation is architecture-agnonn and can be readily integrated into any sfate-of-the-ar world model. Because the observation encoder merely maps raw
4

---

Time-Aware World Model
observations to the latent space, without modeling dynamics, it is not conditioned on At. All other components that
Algorithm Twae-Aware World Model Training Paradigm
1: Initialize task environment & Initialize time-aware world model MTA
depend on the system's dynamics are conditioned the
time step size by taking At as an explicit input.
2: Set experience buffer B t- 0
While the baseline latent dynamics model D maps a state-action pair (ze, ae) directly to the next latent state 2141, we reformulate D using Euler integration:
3: repeat
4: for each episode do
5: Set At
Amiin Atmax) D> Meta-World: Al 0.001, Almar - 0.05
-: Di2,, oc: At) = + d(ze, Ce At) r(at),
(Aldefault 0.0025)
D Can be either
or Uniform
where d(-) is implemented as an MLP. This formulation enforces the intrinsic property
6:
Set step <-
7:
while slep < Horizon H do
8:
ag + +M'A .act (Ot, At) Execute or in get back (ot+ V2) Add transition (o,, a2, O1+A0, re. At) to B {(o,, at, Fr, Ot+ At)1: B: update time- aware world model 2 slep 4- slep + end while
2440/AA=0 24, V22 Cg.
9:
10:
Thus, instead of learning the transition function D directly, TAWM learns the latent-state derivative (gradient) function d, which is itself conditioned on At to capture higher-order effects, and then advances the state by a single Euler step.
11:
12:
13:
14: end for
State Transition Rather than integrating over the time
15:
until reach N training steps
step At, our dynamics model integrates using
16:
return M
7 ((a)) = max(0, ( + 5).
The possible values of At span several orders of magnitude from minimum of about 10-3 s to maximum of roughly 5 x 10-2, s creating numerical challenges. Because latent vectors change only slightly between steps, the dynamics model d(=,, dz, At) must scale appropriately across this wide range. Empirically, we observed convergence failures in some tasks (e.g., mw-assembly y) when employing lin- ear At integration. By allowing the latent state evolve with respect to the logarithm of the time step via r(Al), we effectively normalize At to narrower range, mitigating numerical issues and enabling more stable learning.
4.2. Training Procedure Using a Mixture of A
Algorithm I summarizes our training procedure, in which we vary the observation rate to encourage the model to lear the underlying dynamics multiple temporal resolutions. The world model acquires spatiotemporal representation of the environment by sampling observations from the various dynamical subsystems at specific points in time and space. According to the Nyquist-Shannon sampling theorem, a signal must be sampled at rate of at least 1/(2faax) where fmax is the highest frequency present in the band-limited signal. Because there is no systematic way to determine fuax for every subsystem, we instead sample observations at several temporal rates to capture the underlying dynamics more effectively.
Before proceeding, we note that we experimented with two popular integration methods for our dynamics model: (1) Euler and (2) fourth-order Kunge-Kutte (RK4), even though only Euler integration was introduced earlier. See Appendix A for RK4 implementation details. We include RK4 because it is broadly applicable to both simple linear and complex nonlinear systems and is standard in physical simulation. The integration method thus serves as a tunable hyperparameter for maximizing TAWM's performance and learning efficiency across control tasks.
At the start of each episode (Algorithm 1), we sample AL from a log-uniform distribution and set the observa- tion/control rate to 1/2t. A log-uniform distribution as- signs equal probability mass every octave of time scales, yielding good coverage across multiple orders of magnitude. This is especially useful when Atmax 2 2maaz where f.ax is the highest frequency of the task dynamics. For tasks with dynamics slow enough to be captured by Almax. however, uniform sampling may outperform log-unifora sampling, particularly at larger At values (Figure 7).
1
For most Meta-World robot-manipulation tasks, the dynam- ics are simple enough to be well approximated by Euler integration. Our ablation study shows that TAWM with Euler integration outperforms its RK4 counterpart on the majority of Meta-World tasks, indicating that Euler suffices in these settings (Figure 3). By contrast, RK4's benefits become more pronounced for tasks with complex dynamics, such those in PDE-control environments (Figure 3).
4.3. Theoretical Analysis of TAWM's Sample Efficiency
We present theoretical insights into TAWM's sample effi- ciency, showing that it can be trained capture task dynam-
5

---

Time-Aware World Model
ics across temporal scales without increasing the number of samples training steps and without sacrificing perfor- mance. We begin with two core assumptions:
1. Model capability: Because our time-aware training framework model-agnostic, it inherits the limitations of the underlying world-model architecture. We as- sume that this architecture has enough representation power to capture the system's 's complex dynamics.
Figure 2. Visualizations of the Meta-Worle control tasks. From left to right: Assembly, Basketball, Box Close, Faucer Open. The default time step size of all environments is At 2.5ms.
Lemma 4.2.
When the environment's dynamics can be fully captured with At, reducing the modeling error at At lowers the upper bound of the error for every At < At.
2. Sampling frequency: We focus our analysis on those subsystems and time intervals for which Afma can effectively capture the underlying dynamics. We as- sume that the empirically determined range [At = 0.001s, mas = 1s] adequately captures the behavior of most subsystems.
Proof:
Please see Appendix H.2.
Building this lemma, we conjecture that improvements gained at one temporal scale transfer to all smaller scales, thereby enhancing TAWM's sample efficiency. In the next section, we empirically confirm this conjecture by demon- strating TAWM's high sample efficiency.
In the following discussion, we denote by f the ground-truth dynamics function, which satisfies the equation:
Z++ = 24 + f(ze, a,, At) At.
5. Experiments
Also, we say that an environment's dynamics can be fully captured with a time step At if its dynamics function f satisfies:
We address three key questions in our experiments: (1) Un- der the same planner, does TAWM match the baseline's performance at the default observation rate while avoiding degradation at lower rates? (2) At which observation rates does TAWM outperform the baseline?
ae....tt Af) - f(z,, ar, At)ll < Vát < At, [lf(ze, ai
(3) Does TAWM require more training data than the baseline? We primar- ily investigate these questions in Section 5.1 and present ablation studies sampling strategies in Section 5.2.
where < 1 is sufficiently small constant.
Then, the following lemma shows that the optimal time- aware dynamics function d" (Section 4. 4.1) can approximate dynamics for all At < At through simple interpolation.
5.1. Evaluations Control Problems
When the environment's 's dynamics can be fully captured with At, the optimal dynamics function d' satisfies the following relation for any smaller time step At < At:
Lemma 4.1.
To answer our main questions and evaluate the performance and learning efficiency of our time-aware world model, we conducted experiments on several control tasks from the Meta- World simulation suite, which uses a default time step of A = 2.5ms (0.0025s) (Yu et al., 2020). We tested nine diverse tasks, each with distinct goals and motion charac- teristics:
At
T(Ar)
||d (zt, at, At) - d' d' AF)
|| <.
T(At)
At
Assembly, Basketball, Box Close, Faucet Open, Hammer, Handle Pull, Lever Pull, Pick Out of Hole, and Sweep Into. Representative environments shown in Fig- ure 2, and Appendix B provides renderings of all nine tasks. Following Hansen et al. (2024), we report the success rate (%) as the primary metric for comparing the time-aware and baseline models on these Meta- Mool-Woolddaoo control tasks.
Proof. Please see Appendix H.1.
Note that the interpolation factor =(S)) (A))) is common to every training sample; regardless of 2, or oz, this rela- tionship holds. Based this observation, introduce our third core assumption follows.
At
On top of that, our experiments also include three non-linear, one-dimensional PDE-control tasks from control-gym (Zhang et al., 2024): Burgers, and Wave. For these tasks, we use the (continuous) total episode reward as the primary metric. Because of space constraints, we defer detailed problem descriptions, visual- izations, and additional analysis of these environments to Appendix G.
3. Interpolation Learning: ics function d effectively learns the relationship de- seribed in Lemma 4.1 as the interpolation factor is shared across all samples.
During training, the dynam-
With this assumption, we can prove the following lemma, which states that error reductions at larger temporal scales contribute to reductions at smaller scales during training.
6

---

Time-Aware World Model
imw-leverppull
mw-pick-out-of-hne
mw mw-sweep-into
1.0
Bzseline
Baseline (Repeat)
0.8
0.8
0.8
TAWM (RK4)
0.6
TAWM (Fuler)
0.6
Reseline (Repear)
Roseline
0.6
Raseline (Ropeit)
Baseline
TAWM
TAWM (RK4)
0.4
0.4
TAWM (Euler)
TAWM (Euler)
02
0.0
0.0
0.0
50
10
pde-burgers
pde-allen cahn
pde-wave
-220
-20
2000
-40
-400
-240
Baseline
Baseline
Baseline
Baseline (Repeat)
Reseline (Ropeat)
Baselinn: (Repent)
-600
TAWM (RK4)
250
TAWM (RK4)
TAWM (RK4)
TAWM (Ealer)
TAWM (Euler)
-100
TAWM (Buler)
200
400
600
800 1000
200
400
600
300
1000
200
400
600
800
1000
Figure 3. Performance comparisons for tasks from Meta-Worke (Up), and tasks from PDE-Control (Down). The x-axis corresponds to the evaluation At (in ms), and the y-axis corresponds to success rate (Meta-world) and episode reward (PDE-Control). We trained our TAWM model using either RK4 or Euler integration method with log-uniform sampling strategy. The baseline method was trained using only default time step (At 2.5 ms for Meta- World tasks; At = 50 ma// 110 m0/m// Allen-Cahn Cahn, PDE- Wave, respectively), which is signified with the black dashed line in each environment. For fair comparisons, we extended the baseline method by repeating the same actions for the larger evaluation time steps than the default one. For instance, when the evaluation At- = 5ms, we repeat the same action for 5/2.5 2 times. Plots show mean and 95% confidence intervals over seeds, with 10 evaluation episodes per seed.
ms/100 110 ms for PDE-Burgers, PDE-
Training Setup. We built our time-aware model the base TD-MPC2 architecture, using the same default training hyperparameters (e.g.. model size, learning rate, and hori- zon). As shown in Algorithm 1, we randomly varied At during training to create a mixture of training observation rates. Because there no systematic way to identify the highest frequency in each task's dynamics, we set At to the range [0.001, 0.05] for Meta-Works tasks and [0.01, 1.0] for PDE-control tasks. For Meta-World tasks, each TAWM was trained for 1.5 million steps, which required roughly 40-45 hours on a single NVIDIA RTX 4000 GPU (16 GB VRAM) and 32 CPU cores. For PDE-Allen-Cahn and PDE- Wave, the TAWM and baseline models were trained for IM steps. For PDE-Burgers, all models were trained for 750k steps.
namics of some WeradWorld tasks are sufficiently simple to be well approximated by TAWM with Euler integration, leading to better performance than with RK4 integration.
Effects of using Mixtures of At. To demonstrate the ef- fectiveness of training the world model at multiple temporal
resolutions At, we compare our approach with baselines trained only at fixed At values different from the default At 2.5ms.. Figure 4 shows that our time-aware model outperforms all baselines across three Meta- World tasks. Most notably, when baselines are trained solely at low ob- servation rates (c.g., At Z 10ms), they fail to converge and achieve zero success on every task at every evaluation rate. In contrast, by training with mixture of time steps, our time-aware model consistently surpasses any baseline trained at single step size, regardless of that fixed At. These results suggest that environmental dynamics comprise multiple subsystems, each describable as a time-dependent, spatially parameterized function with its own highest fre- quency. Varying the observation rate (i.e., varying Al), therefore enables the world model to learn these underlying subsystems more effectively.
Performance Comparisons Across Different At. To assess the time-aware model's performance at different inference-time observation rates, we evaluated it on multiple tasks with varying At. As shown in Figures and 11, both TAWM variants one using RK4 integration and the other using Euler integration outperform the baseline (trained at fixed AL = 2.5 ms) across all tasks under identical hyperparameters. We observed similar gains in every other Meta- World task (Appendix E), indicating that the time- aware model effectively learns both fast and slow dynamics without increasing sample complexity. Empirically, the dy-
Convergence Rate on Various At. To evaluate the sam- ple efficiency of TAWM, we compare the episode suc- cess rate curves across different control tasks between our
7

---

Time-Aware World Model
mw-assembly
mw-baskethall
mv-box-closc
Baselino
(1me)
0.6
Beaiine (2 San4)
0.6
Reecling (2 3.5m))
0.6
Bowelin (2.5mg
(10me)
(10mg
Honelino (10ms)
0.4
Beatine (50ma)
Reecling (50md)
0.4
Boweling (50mms)
TAWM
0.2
0.2
0.0
0.0
0.0
10
40 30
10 200
30
40
10
340
50
Figure 4. Performance comparisons to baseline methods trained on different At values for Meta- Meta-World tasks. The x-axis corresponds to the evaluation At, and the y-axis corresponds to the success rate. The non-time-awan baseline models were trained with different fixed At values (shown in the legend). Our TAWM model employs RK4 integration and log-uniform Át sampling. TAWM ousperform all baseline models trained with fixed time-step sizes. When the baselines trained only low observation rates (Ár 10ms). they fail on all tasks. The dashed black line indicates the default time-step size (At = 2.5ms). Means and 95% confidence intervals are computed over 3 seeds, each evaluated for 10 episodes.
0.6
0.6
0.6
0.4
0.4
0.4
0.2
(RK4)
TXWM (Bule)
Eval. A 2.5 me
Eval. 2. 10.0 ma
Eval. 1- 30,0
Eval. 2 50.0
1.0
1.0
1.0
1.0
0.8
0.5
1.5
01.
0.2
022
0.2
012
0.0
0.0
0.0
01.0
0.5M
IM
1.5M
0.5M
IM
1.5M
0.5M
IM
1.5M
0.5M
1M
1.5M
Eval,
2.5 ms
Evall 2i 10.0 ms
Eval. à 30.0 It
Evall
50.0 175
Figure 5. Learning curve comparison Meta-World Basketball (Up) and to the evaluation At (in ms), and the y-axis corresponds to the success rate. We trained our TAWM using either RK4 or Euler integration method with log-uniform sampling strategy. The baseline method was trained using only the default time step (At 2.5mm). Plots show mean and 95% confidence intervals over 3 seeds, with 10 evaluation episodes per seed.
Lever-Pull (Down). The x-axis corresponds
TAWM models and non-time-aware baseline models trained with fixed default time step of At = 2.5 ms. We assess the success rate curves of TAWM and baseline models under different evaluation At values across various control tasks throughout training. In Figure 5, we present the learning curves for the Meta-World Basketball and Lever-Pull tasks. Despite having to learn state transitions under varying time step sizes, our TAWM with RK4 integration method con- verges at least as quickly as the bascline when evaluated at the default At = 2.5 ms - the exact time step on which the baseline was specifically trained. Compared to the non- time- aware baseline specialized at only AL 2.5 ms, TAWM (RK4) still converges to higher success rate at At 2.5 ms after 1.5M training steps. Both TAWM variants (RK4 and Euler) reach 100% success the Meta- World Basketball task at At = 2.5 ms, outperforming the baseline.
fault (At > 2.5 ms), both TAWM (RK4) and TAWM (Euler) significantly outperform the non- non-time-and aware baselines. We refer the readers to Figure 12 and Figure 13 in Appendix F, as well as Figure 17 Appendix G, for additional learning curves on other control tasks. These results demonstrate that, despite having learn task dynamics across differ- ent temporal scales, time-aware model does not re- quire additional training steps samples to converge to sufficiently accurate model capable of effectively solving control tasks at varying observation rates.
Comparison to MTS3 Finally, we compare our model's performance with the Multi Time Scale World Model (MTS3) (Shaj Kumar et al., 2023), closely related ap- proach that shares the high-level motivation of modeling world dynamics across multiple temporal scales, as intro- duced in Section 2. We would like to refer the readers to Appendix D for detailed descriptions of MTS3's experi-
When evaluated at inference A/ values greater than the de-
8

---

Time-Aware World Model
mental settings. In Figure 6, we observe that our TAWM consistently outperforms MTS3 across various settings of MTS3's slow dynamics hyperparameter H. Notably, MTS3 exhibits rapid performance degradation as the evaluation time step increases, suggesting that it suffers from com- pounding errors due to long-horizon predictions. In con- trast, our TAWM shows greater robustness at lower obser- vation rates, especially on the Meta-World Faucet-Open task in which TAWM's success rate remains at 90% at At = 50 ms s0.
mww-basketblll
mww mww-fucct-open
10
MIST
0.5
MTST (1-1)
06-
MIST (1-11))
MESS
MISS (1-15)
MTS3 (1-11
MESI (l-Se
04
TAAM
07
00
Figure 6. Performance comparisons to MTS3 across different evaluation At's on Meta-World tasks. The x-axis corresponds to the evaluation At, and the y-axis corresponds to the success rate. The MTS3 models were trained under different slow dynamics set- tings, indicated by the H values shown in the legend, using offine data with 4M transitions. Our TAWM employs RK4 integration method with log-uniform sampling strategy and was trained for 1.5M steps. The dashed black line indicates the default time-step size (At = 2.5ms). Means and 95% confidence intervals are computed over seeds for TAWM, each evaluated for episodes. TAWM outperforms MTS3 across different H settings, especially at large At values.
5.2. Ablation studies on sampling strategies
The At sampling strategy is tunable hyperparameter to optimize TAWM's performance and efficiency. In this section, we conduct an ablation study on the im- pact of the At sampling strategy during training on the performance and sample efficiency of TAWM. In Meta-World tasks, we trained the same TAWM archi- tecture in the same At range using two different sam- pling strategies: (1) Log- Sogg-nnfform(00000 8, 0.05 s) and (2) Uniform(0.000 s, 0.05 s). In some tasks, such as Meta- World Assembly and Basketball, uniform sampling may outperform log-uniform sampling- particularly when the task dynamics are sufficiently slow to be captured by Ammx-
mw-assembly
nnw wew-baskelball
08
0.1
Badis (2.3m0)
Bevd ine Guaheo) TAWM Dugh TAWM (unierno
TAWN (he)
Figure shows that while TAWM trained with the uniform sampling strategy generally performs better in most envi- ronments and achieves higher performance at low sampling rates (At 2 30 ms), it exhibits lower success rates at smaller inference At in some environments, such as Meta-World Assembly. Additional results in other environments are pre- sented in Figure 8, which show a similar pattern. Regardless of the At sampling strategy, time-aware models consistently outperform con-time-awand baselines. Therefore, depend- ing on the task dynamics, TAWM can be effectively and ef- ficiently trained with any reasonable At sampling strategy and is not restricted to log-uniform or uniform schemes.
TAWK (aeifoe)
0.4
022
00
Figure 7. Ablation Study on At Sampling Strategy. We trained our TAWM models with log-uniform and uniform sampling strate- gies, both of which use the RK4 integration method. The baseline method was trained using only default time step (At = 2.5ms), which is signified with the black dashed lines. The x-axis corre- sponds to the evaluation At, and the y-axis corresponds to the success rate. Mcans and 95% confidence intervals are computed over 3 seeds, each evaluated for 10 episodes.
6. Conclusion
Limitations and Future Work. We have yet to develop a systematic methodology analytically compute the highest frequency of the underlying task dynamics; therefore, the upper bound for the training time step size Atmax is deter- mined empirically. This limitation necessitates a search for A Afmux when applying TAWM to new environments (e.g.. autonomous driving or PDE controllers) to improve training efficiency and convergence. An interesting direction for fu- ture work is to develop an automatic method to identify the highest frequency of task dynamics, which would determine the lower bound for sampling frequency (or the upper bound for Atmax) in training TAWM. We also plan to extend our current deterministic dynamics model to a probabilistic one to better capture world state transitions, particularly at larger At values and/or over longer horizons.
In this paper, we introduce a novel time-aware world model (TAWM) that adaptively learns task dynamics across differ- ent temporal scales. By explicitly conditioning the dynamic model on the time step size At and training it with a mix- ture of temporal scales (via At sampling), TAWM achieves robust performance across varying observation rates on di- verse control tasks without requiring additional training steps or samples. Empirical results show that our model consistently outperforms the non-time-aware baseline mod- els, which are trained only a fixed time step size At for different training At values. We hope that the insights and results in this paper offer new perspective on world model training, thereby contributing to the community a new, efficient, yet simple world model training method.
9

---

Time-Aware World Model
Acknowledgments
learning with stochastic actor. In International Con- ference Machine Learning (ICML), pp- 1861-1870. PMLR, 2018.
This research is supported in part by the National Science Foundation and the U.S.Army Research Labs Cooperative Agreement on " "I and Autonomy for Multi-Agent Systems".
Hafner, D., Lillicrap, T., Ba, J., and Norouzi, M. Dream to control: Learning behaviors by latent imagination. arXiv preprint 1rXiv:19/..01000 2019.
Impact Statement
Hafner, D., Lillicrap, T., Norouzi, M., and Ba, J. Mas-
This paper presents a novel, model-agnostic approach to ef- ficiently train Time-Aware World Models (TAWM) capable of handling varying observation rates and control frequen- cies without increasing sample complexity. Our method enables: (1) training single world model usable across mul- tiple observation rates and control frequencies, (2) avoiding the need to retrain a different model for each frequency, and (3) significantly reducing computational and energy consumption. This training method is model-agnostic, is compatible with any existing world model architecture, and demonstrates the feasibility of time-aware training without increasing training time or data. This work promotes effi- cient, green computing in academic and industrial rescarch of world models. Another potential societal impact includes applications in autonomous military systems. The authors do not foresee any other immediate consequences resulting from this work.
tering atari with discrete world models. arXiv preprint
arXiv:2010.02/93 2020.
Hafner, D., Pasukonis, J., Ba, J., and Lillicrap. T. Mastering diverse domains through world models. arXiv preprint urXiv:2301.cm 104, 2023.
Hansen, N., Su, H., and Wang, X. TD-MPC2: Scalable, ro- bust world models for continuous control. In The Twelfth International Conference on Learning Representations (ICLR), 2024.
Hansen, N. A., Su, H., and Wang, X. Temporal differ- ence learning for model predictive control. In Chaudhuri, K., Jegelka, S., Song, I.., Szepesvari, C., Niu, G., and Sabato, S. (eds.), Proceedings of the 39th International Conference on Muchine Learning (ICML), volume 162, pp. 8387-8406. PMLR, 17-23 Jul 2022.
References
Janner, M., Fu, J., Zhang, M., and Levine, S. When to trust your model: Model-based policy optimization. Advances in Neural Information Processing Systems (NeurIPS), 32, 2019.
Barto, A. G., Sutton, R. S., and Anderson, C. W. Neuronlike adaptive elements that can solve difficult learning control problems. IEEE Transactions on Systems, Man, and Cybernetics, (5):834-846 1983.
Jerri, A. J. The shannon sampling theorem itss various ex- tensions and applications: A tutorial review. Proceedings of the IEEE, 65(11):1565-1556, 1977.
Butcher, J. C. The numerical analysis of ordinary differen- lial equations: Runge-Kullt Kulla and generul linear methods. Wiley-Interscience, 1987.
Johanninn, T., Bahl, S., Nair, A., Luo, J., Kumar, A., Loskyll, M., Ojea, J. A., Solowjow, E., and Levine, S. Residual reinforcement learning for robot control. In 2019 International Conference on Robotics and Automa- tion (ICRA), pp. 6023 0029.602..22 IEEE, 2019.
Deisenroth, M. and Rasmussen, C. E. Pilco: A model-based and data-efficient approach to policy search. In Proceed- ings of the 28th International Conference on Machine Learning (ICML), pp. 465 472, 2011.
Kaiser, L.., Babaeizadeh, M., Milos, P., Osinski, B., Camp- bell, R. H., Czechowski, K., Erhan, D., Finn, C., Koza- kowski, P., Levine, S., et al. Model-based reinforce- ment learning for atari. arXiv preprint arXiv: ma0:::903000330 2019.
Deisenroth, M. P., Neumann, G., Peters, J., et al. survey on policy search for robotics. Foundations and Trendy@ in Robotics, 2(1-2))1-143 2013.
Guan, Y., Liao, H., Li, Z., Hu, J., Yuan, R., Li, Y., Zhang, G., and Xu, C. World models for autonomous driving: An initial survey. IEEE Transactions on Intelligent Vehicles, 2024.
Kiran, B. R., Sobh, I., Talpaert, V., Mannion, P., Al Sal- lab, A. A., Yogamani, S., and Pérez, P. Deep reinforce- ment learning for autonomous driving: A survey. IEEE Transactions on Intelligent Transportation Systems, 23 46)4499-4992, 2021.
Ha, D. and Schmidhuber, J. World models. arXiv preprint arXiv: 1803. 2018.
Koh, J. Y., Lee, H., Yang, Y., Baldridge, J., and Anderson, P. Pathdreamer: A world model for indoor navigation, 2021. URL wwt.h:tt//////..t. //w/xvvvvv///////////////////xvv www.ww.ww......t.. 20....///11.00... 08756.
Haarnoja, T., Zhou, A., Abbeel, P., and Levine, S. Soft actor- critie: Off-policy maximum entropy deep reinforcement
10

---

Time-Aware World Model
Lambert, N., Pister, K., and Calandra, R. Investigating compounding prediction emrors in learned dynamics mod- els, 2022. URL https: / /taaaaiiiiiv////// org.abs//rg//////////// 09637.
Bolton, A., al. Mastering the game of go without human knowledge. Nature, 550(7676):354-3335 2017.
Silver, D., Hubert, T., Schrittwieser, J., Antonoglou, I.,
Lai, M., Guez, A., Lanctot, M., Sifre, I.,, Kumaran, D., Graepel, T., Lillicrap, T., Simonyan, K., and Has- sabis, D. A general reinforcement learning algorithm that masters chess, shogi. and go through self-play Sci- ence, (62((666)))))11 1atp11140-11444 2018. doi: 10.11201100000 1166sciience aar6404, URL htwp:: /www http:::ss///wwsssienneeeem/e abs 1126 /ww//cciinccc.c.cc..........cccc aar 6404.
Langley, P. Crafting papers on machine learning. Langley, P. (ed.), Proceedings the 17th International Conference Machine Learning (ICML 2000), pp. 1207-1216, Stan- ford, CA, 2000. Morgan Kaufmann.
science. org doi/
Lutter, M., Hasenclever, L., Byravan, A., Dulac-Arold, G., Trochim, P., Heess, N., Merel, J., and Tassa, Y. Learning dynamics models for model predictive agents, 2021. URL httpps://at /arxivv.// org/ abs///////////// 2209.. 14311.
Son, S., Zheng, L., Sullivan, R., Qiao, Y.-L., and Lin, M. Gradient informed proximal policy optimization. Advances in Neural Information Processing Systems (NeurIPS), 36, 2024.
Nhu, A. N., Le, N.-A., Li, S., and Truong, guided reinforcement learning system for realistic vehicle active suspension control. In 2023 International Confer- ence on Machine Learning and Applications (ICMLA), IEEE, 2023.
Physics-
Suh, J., Simchowitz, M., Zhang, K., and Tedrake, R. Do differentiable simulators give better policy gradients? In International Conference on Machine Learning (ICML), pp. 10668-206966 PMLR, 2022.
pp.
Parmas, P., Rasmussen, C. E., Peters, J., and Doya, K. Pipps: Flexible model-based policy search robust to the curse of chaos. In International Conference Machine Learning PMLR, 2018.
Sutton, R. S. Integrated architectures for learning, plan- ning, and reacting based on approximating dynamic pro- gramming. In Machine learning proceedings 1990, pp. 216-224. Elsevier, 1990.
(ICML), pp.
Prasad, N., Cheng, L.-F., Chivers, C.,, Draugelis. M., and Engelhardt, B. E. A reinforcement learning approach to weaning of mechanical ventilation in intensive care units. arXiv preprint arXiv: 7704006000 2017.
Sutton, R. S., McAllester, D., Singh, S., and Mansour, Y. Policy gradient methods for reinforcement learning with function approximation. In Solla, S., Leen, T., and Müller, K. (eds.), Advances in Neural Information Processing Systems, volume 12. MIT Press, 1999.
Schulman, J., Levine, S., Abbeel, P., Jordan, M., and Moritz, P. Trust region policy optimization. In International Con- ference Machine Learning (ICML), pp. 1889-1897. PMLR, 2015.
Thodoroff, P., Li, W., and Lawrence, N. D. Bench- marking real-time reinforcement learning. In Al- banie, S., Henriques, J. F., Bertinetto, L., Hernández- Garcia, A., Doughty, H., and Varol, G. (eds.), NeurIPS 2021 Workshop on Pre-registration in Ma- chine Learning, volume 181, pp. 26-41, PMLR, 13 Dec 2022. URL httptttss///ppoceeeiii v M181/thodoroff22a.. html..
Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. Proximal policy optimization algorithms. arXiv preprint arXi:: 0177700770077777 2017.
Shaj Kumar, V., Gholam Zadeh, S., Demir, O., Douat, L., and Neumann, G. Multi time scale world mod- els. Advances Neural Information Processing Systems (NeurIPS), 01-26664-267755 2023.
mlr www.mpp.tt.... ..rrsss///
Vinyals, O., Babuschkin,, I., Czarnecki, W. M., Mathieu, M., Dudzik, A., Chung, J., Choi, D. H., Powell, R., Ewalds, T., Georgiev, P., al. Grandmaster level in starcraft using multi-agent reinforcement learning. Nature, 575 (7782)350-335 2019.
Shannon, C. Communication in the presence of noise. Pro- ceedings of the IRE, 3701):10-21, jan 1949, doi: 10. 1109/jrproc. 222999.. URL https:tpps...mmc ptattts:://o..... 10. 10.1109/ wtt00//rrrooo.o 1949.999949 232969.
Silver, D., Huang, A., Maddison, C. J., Guez, A., Sifre, L., Van Den Driessche, G., Schritwieserr J., Antonoglou, I., Panneershelvam, V., Lanctot, M., al. Mastering the game of go with deep neural networks and tree search. Nature, 529(7587):484 2016.
Wang, T., Bao, X., Clavera, I., Hoang, J., Wen, Y., Langlois, E.. Zhang, S., Zhang, G., Abbeel, P.. and Ba, J. Bench- marking model-based reinforcement learning, 2019. URL https://arrxiv.orr/bbbb///
Silver, D.. Schrittwieser, J., Simonyan, K., Antonoglou, I,, Huang, A., Guez, A., Hubert, T., Baker, L., Lai, M.,
Weinan, E. Principles of multiscale modeling. Cambridge University Press, 2011.

---

Time-Aware World Model
Williams, J. and Peng, Reinforcement learning algo- rithms function optimizers. In Proceedings of the In- ternational Joint Conference on Neural Networks, Wash- ington DC, volume 2, pp. 89-95, 1989.
Wu, Escontrela, A., Hafner, D., Abbeel, and Gold- berg, K. Daydreamer: World models for physical robot learning. In Conference Robot Learning (CoRL), pp. 2226-2240. PMLR, 2023.
Xu, Makoviychuk, V., Narang, Y., Ramos, F., Matusik, W., Garg, A., and Macklin, M. Accelerated policy learn- ing with parallel differentiable simulation. In Interna- tional Conference on Learning Representations (ICLR), 2022.
Yang, L., Nagy, Z., Goffin, P., and Schlueter, A. Reinforce- ment learning for optimal control of low exergy buildings. Applied Energy, 156:577-586, 20155
Yu, T., Quillen, D., He, Z., Julian, R., Hausman, K., Finn, C., and Levine, S. Meta-world: benchmark and evalua- tion for multi-task and meta reinforcement learning. In Kaelbling, L. P., Kragic, D., and Sugiura, K. (eds.), Pro- ceedings of the Conference on Robot Learning (CoRL), volume 100 of Proceedings of Machine Learning Re- search, pp. 1094-1100. PMLR, 30 Oct-01 Nov 2020.
Zeng, Z., Yue, Z., Mauroy, A., Gonçalves, J., and Yuan, Y. A sampling theorem for exact identification of continuous- time nonlinear dynamical systems. IEEE Transactions on Automatic Control, pp. 1-16, 2024. doi: 10.1109/TAC. 2024.3409639.
Zhang. X., Mao, W., Mowlavi, S., Benosman, M., and Başar, T. Controlgym: Large-scale control environments for benchmarking reinforcement learning algorithms. In 6th Annual Learning for Dynamics & Control Conference, pp. 181-196. PMLR, 2024.
12

---

Time-Aware World Model
Appendix
A. Descriptions of 4th-order Runge-Kutte integration.
In this section, we extend our description of the 4th-order Runge-Kutta (RK4) integration mentioned in Section 4.1.2. The detailed RK4 integration is as follows:
kj = d(21; Ot, At)
21 = 21 + (2,, at At/2) - (At/22)
k2 d (21, at, At)
2 21 + (2), at. A1/2) - (A1/2)
33 d ($2, ae, At)
(4)
23 + d (22, At) 7 (Att)
Ka = d (23, at, At)
2+t = 2t 6
(ki 2k + 2h + k4) ((tt)
Consistent with the notations in Section 4.1.2, 2t. ag denotes the latent state-action pairs at time t, d(.) denotes our dynamic model parameterized by a neural network, and &, (i E 2, 3) are the intermediate middle points. The final prediction of next latent state under time step size At 244
B. Meta-World Task Dynamics Visualizations.
To illustrate the diversity and key differences in dynamics across control tasks in our experiments, we present additional visualizations of different Meta- World tasks. Figure 9 and Figure 10 show sequences renderings from the task initialization to the task completion (from left to right) of different control tasks.
C. Additional results for ablation study of At sampling strategies.
In this section, we provide additional results for the ablation study on the impact of At sampling strategy the TAWM's performance and efficiency. Specifically, we trained the same TAWM over the same range of Á/ values using two different sampling strategies: (1) Log-Uniform(1 ms, 50 ms) and (2) Uniform( ms, 50 ms). The performance of TAWMs trained with different At sampling strategies, along with the non-time-aware baselines, is shown in Figure 8.
Figure 8 shows that while our time-aware models trained with the uniform sampling strategy generally perform better across most environments and significantly outperform at low sampling rates (inference At 2 30 ms) - they exhibit lower success rates at smaller inference At in certain environments, such Meta-Worte Lever-Pull. Regardless of the At sampling strategy. TAWMs Ms consistently outperform the non-time-aware baselines. Therefore, our time-aware model can be efficiently and effectively trained with any reasonable sampling strategy and is not restricted to either log-uniform or uniform sampling.
mww mww-box men-bbx--llese
mIW wpw.fancce-opeen
ITW hw--hhmmerr
IIW lever-pall
1.5
1.5
1.5
(og)
TAWM toatoons
0.4
TAWM Dogi
05
0.5
Ea Seo
Baceiaer" Smo
TAWM (og)
TAWM (ogi
0.
03
DXRM (aniker))
0.8 LAWI (andeeE))
0.2
10
15
93
10
10
10
15
10
10
33
Figure 8. Ablation Study At Sampling Strategy. We trained TAWM models with log-uniform and uniform sampling strategies, both of which use the RK4 integration method. The bascline method was trained using only default time step (At = 2.5mss, which is signified with the black dashed lines. The x-axis corresponds the evaluation At, and the y-axis corresponds the success rate, Mcans and 95% confidence intervals are computed over 3 seeds, each evaluated for 10 episodes.
13

---

Time-Aware World Model
Figure 9. Visualizations of the Meta- World control tasks. From top to bottom: (1) Assembly, (2) Basketball, (3) Box Close, (4) Faucet
Open, (5) Hammer. From left right: sequences of renderings from task initialization completion.
14

---

Time-Aware World Model
Figure Visualizations of the Meta- World control tasks. From top to bottom: (6) Handle Pull, (7) Lever Pull, (8) Pick Out Of Hole (9) Sweep Into. From left to right: sequences of renderings from task initialization to completion.
15

---

Time-Aware World Model
Developing a more adaptive and optimal Al sampling strategy tailored to specific objectives and scenarios is promising direction for future work particularly to achieve strong performance across both small and large At values and to accelerate convergence rates. In the meantime, TAWM with either log-uniform or uniform sampling strategy has demonstrated practical effectiveness, as shown by our experimental results.
D. Comparisons with Multi Time Scale World Model (MTS3).
As mentioned in Section 2. Multi Time Scale World Model (MTS3) (Shaj Kumar et al., 2023) a closely related work sharing the high-level motivation with our work: to model the world dynamics at multiple temporal levels. Specifically, MTS3 proposes probabilistic approach to jointly learn the world dynamics at two temporal abstractions: task level (slow dynamics/timescale) and state level (fast dynamies/tmescale). These 2 timescales are separately learned by two state space models (SSMs): SSM fast and SSM where SSM fast learns the dynamics evolving at original small timestep At of the dynamical systems and SSMfast learn the slow dynamics evolving at H At. Although this approach also explicitly considered different temporal abstraction levels in learning the world dynamics, there are several critical differences between MTS3 compared to our work:
1. Models vs Training method: (Shaj Kumar et al., 2023) proposes model architecture to learn a world model with several discrete temporal abstraction levels. On the other hand, we proposed a simple yet effective and efficient time-aware training method that can be employed to train any world model architecture
2. Discrete vs continous timescales: The original MTS3 currently only handle only 2 timescales: At and H where both At and H is fixed in the training process. Although the MTS3 can be adapted to learn multiple timescales, the number of timescales is limited to a discrete value. Furthermore, the SS (slow dynamic model) does not directly model state transition under large temporal gap At (or low observation rate) but rather learns the task latents to guide M/s/ to long-horizon predictions. On the other hand, our time-aware approach can directly predict the future states St+Ar under large At.
3. Multi-step vs one-step prediction: MTS3 considers the future prediction under large At as a long-horizon prediction problem. Specifically, to predict St MTS3 discretizes the long temporal gap into several smaller timesteps:
At = fust, where fast is the original timestep size SaMf is trained with and M E N MTS3 then iteratively applies the model M times to predict 81+ Ar. This timestep discretization approach has 2 critical limitations: (1) MTS3 cannot model state transitions under At that is not divisible by At and (2) multi-step predictions are vuinerable to compounding errors, a well-known problem in long-horizon modeling. On the other hand, our model can directly predict the next state with a one-step prediction, effectively alleviating the compounding error problem.
4. Inference efficiency: Another disadvantage of multi-step prediction is inference inefficiency. In contrast, our time- aware model can efficiently predict long-term future states without sacrificing computational efficiency by using one-step prediction.
5. Prediction-only vs Control: As acknowledged by (Shaj Kumar et al., 2023), the original MTS3 is strictly a prediction model. On the contrary, our time-aware model can be used efficiently with a planner to solve control problems.
We conduct empirical comparisons between MTS3 and our proposed time-aware model the control problem, extending beyond the prediction-only scope in MTS3. First, MTS3 is trained with offline data consisting of 4 x 10 (4M transitions) collected from random trajectories (10%), partially-train trainned policy's trajectories (20%), and fully-trained expert policy trajectories (80%). Since MTS3 is strietly prediction-focused and is not designed for controls, we combined MTS3 with MPPI plamners and our world model's s trained value and reward functions. Implementation-wise we replaced our dynamic model with MTS3 and kept all other components unchanged, including the planner (MPPI), learned value function, and learned reward function. This design ensures fair comparison between the models, as any performance gap is attributed solely to the difference between MTS3 and our dynamic model.
We kept the default hyperparameter settings as in the original MTS3 paper and codebase while varying the slow dynamics H to investigate the impacts of H on MTS3's performance. The authors of MTS3 suggested that using H = VT, which is / /99 xe 10 in our experiments on Meta- World environments. We chose H = 11 to divide the episodes into lguthh local SSM windows. The MTS3 inference stepping are adjusted such that when evaluated on Atreal > Attrain, the model applied Atceat/ Attrain times (Atrain is the fast time step between
observations).
16

---

Time-Aware World Model
E. Additional performance comparisons between TAWM and baseline models
In this section, we provide additional performance comparisons between TAWM and baseline models across varying inference At values, as shown in Figure 11. On each task, both TAWM variants (using RK4 and Euler integration methods) were trained for 1.5 million steps, with At values sampled from Log-Uniforn I ms, 50 ms) distribution during training. In contrast, the baseline models were trained solely fixed default time step of At 2.5 ms. For these baselines, the pretrained weights provided by the original TD-MPC2 paper (Hansen et al., 2024) for each corresponding task.
mw-assembly
mw-baskethall
mw-box-close
1.0
1.0
1.0
0.8
0.8
0.8
Baseline
Hoseline
Haselane
0.6
Baseline (Rapeat)
0.5
Reseline (Rapeat)
0.6
Raselate (Ropezt)
TAWM (RK4)
EAWM (RK4)
TAWM (RK4)
0.4
TAWM (Euler)
0.4
0.4
10
20
30
40
50
10
20
30
40
50
10
20
30
40
50
owffaucetoopen
mw-hammer
mw-handle-pull
1.0
1.0
Buscline
1.0
Buseling (Rapeat)
0.8
0.8
0.8
0.5
TAWM (Fuler)
Baseline
0.4
0.4
Buseling
0.4
TAWM (RK4)
(Repeat)
0.2
TAWM (RK4)
(Repett)
TAWM (Huler)
LAWM (buler)
10
20
30
40
50
10 20 30
40
50
10
20
30
40
50
Figure 11. Performance comparisons for 6 additional tasks from Meta- Meta-World World tasks. The x-axis corresponds to the evaluation At (in ms), and the y-axis corresoonds to success rate. We trained our TAWM model using either RK4 or Euler integration method with uni-uniform sampling strategy. The baseline method trained using only the default time step (At 2.5ms), which is signified with the black dashed lines. For fair comparisons, we extended the baseline method by repeating the same actions for the larger evaluation time steps than the default one. For instance, when the evaluation At 5mss, repeat the same action for 5/2.5 2 times. Plots show mean and 95% confidence intervals over 3 seeds, with evaluation episodes per seed.
F. A Addiionnal learning curve comparisons between TAWM and baseline models
In this section, we present additional learning curve results for TAWM and baseline models various Meta-World control tasks. Specifically, Figure 12 and Figure 13 show the learning curves of our time-aware models (TAWMs) and the non- time-aware baseline models (original TD-MPC2), evaluated at differrnnt inference At values. Each curve is generated by evaluating the models at various training steps using fixed inference At [1,50] ms.
17

---

Time-Aware World Model
mw-assee nww-assembly
mww mw-box-elose
mw mw-faucet-open
1.00
1.00
1,00
0.75
Baseline
TAWM (RK4)
TAWM (RK())
TAWM (RK()
TAWM (Fuler)
TAWM (Fuler)
TAWM (Fuler)
0.00
0.00
0.00
0.5M
1M
1.5M
0.5M
1M
1.5M
0.5M
IM
1.5M
Eval. &
Eval. & 1- 1.0
Eval. A 1.0
1.0
1.0
M
0.8
0.6
0.6
0.5
0.2
0.2
0.2
0.5M
IM
1.5M
0.5M
1M
1.5M
0
0.5M
IM
1.5M
Eval, 1 2.5
Eval. A 2.5 ms
Eval. A 2,5 ms
0.8
0.8
0.8
0.4
0.4
0.2
0.2
0.2
0.5M
IM
1.5M
0.5M
IM
1.5M
0.5M
IM
1.5M
Eval. == 10.0 ms
Eval. A 10.0 ms
Eval. A 10.0 ms
0.8
0.8
0.8
0.4
0.4
0.4
0.0
0.0
0.5M
Eval. 1- 20.0 mss
Eval. A 20.0 ms
Eval. A 20.0 mas
1.0
1.0
0.8
0.8
0.4
0.4
0.8
0.0
0.0
0.0
1.0
1.0
0.8
0.8
0.6
0.6
0.4
0.4
0.0
0.0
0.0
0.5M
U.5M
1.5M
1M
Eval.
Figure 12. Additional Meta-World tasks: Success Rate Curve under different evaluation time step sizes.
18

---

Time-Aware World Model
mww hammer
mw-handle-pul.
mw-lever-pull
0.75
0.75
0.75
Taselline TAWM (KK4)
0.50
0.50
0.50
Daseline
Daseline
LAWM (Eauler)
0.25
(R344)
0.25
(RK4)
0.25
TANM (Fulec)
TAWM (Eules)
0.00
0.00
0.00
0.5M
1.5M
1.5M
1M
1.5M
1.0
1.0
1.0
0.6
0.6
0.5
0.4
0.0
0.0
0.0
0.5M
1.5M
1M
2.5
2,5
1.0
1.0
1.0
0.6
0.6
0..
0.2
0.2
0.2
0.5M
IM
1.5M
0.5M
1M
1.5M
0.5M
1M
1.5M
Eval. 1- 10.0
Eval. A 10.0
Eval. A 10.0
1.0
1.0
0.6
0.6
0.5
0.2
0.2
0.2
0.5M
IM
1.5M
0.5M
JM
1.5M
0.5M
1M
1.501
Eval. A 20.0 ms
Eval. A 20.0 ms
Eval. A 20.0 mss
1.0
1.0
0.8
0.8
0.6
0.6
0.2
0.2
0.2
0.5M
IM
1.5M
0.5M
IM
1.5M
0
0.55M
1M
1.5M
Eval. A 1 30.0 mss
Eval. A t- 30.0 ms
Eval. A 30.0 mas
0.8
0.8
0.8
0.1
0.4
0.2
0.5M
1.5M
0.5M
IM
1.5M
0.5M
IM
1.5M
Eval. t- 50.0 ms
Eval. A - 50.0 ms
Eval. A 50.0
Figure 13. Additional Meta-Worle tasks: Success Rate Curve under different evaluation time step sizes.
19

---

Time-Aware World Model
G. Additional experiments on PDE control problems.
In this section, we provide detailed descriptions of the PDE control tasks and present additional experiments and analysis of our TAWM on various PDE control problems in control cerol-gyye -gmm environments (Zhang al., 2024).
G.1. PDE Control Problems.
The PDE problems are one-dimensional PDE control problems featuring periodic boundary conditions and spatially distributed control inputs. The spatial domain is defined as Ω [0, L] C R. The continuous field of the PDE defined as t( Q x IR+ + R. where z, represent the spatial coordinates in the field and time, respectively. Generally, the PDE in each control task is defined as:
du F (Br Du du ôt Oz dx2:..) a(x,t)
In the equation above, Fis a differential operator (linear or non-linear) defined differently for each control task. The specific formulations of FF are carefully described in the next subsection for each environment. a(x, t) the distributed control force over the PDE field. a1x, () is defined as:
Ra - a(x, = Σ 1z(2)a,(t) 1-0
The control force is composed of n. scalar control inputs ay (t), each influencing a specific subset of the domain through its corresponding foreing support function (x). Generally, each action introduces external forces/energy the PDE fields to control their dynamics and steer them toward the target state. More comprehensive details of the PDE environments are
described in (Zhang et al., 2024). In our experiments, the target state of all PDE control tasks is Starget = Ô. At each step 4 the reward is computed as the LQ-error between the current state 9t and the target state Stargeti
r2 [(a) -E [(se Starget) TQ Starget) ag Roe
G.2. Burgers' Equation.
1.0
st5e
0.92
domain coordinates
0.8
0.82
0.8
0.5
0.72
2000 lime
0.6
0.6
0.52
0.6
0.4
0.42
0.32
0.12
0.2
1.0
Ammittion time
0.0
0.02
0.0
domain coordmates
0.00
0.25
0.50
0.75
1.00
62
simulation time
domain coordinates
Figure 14. Visualizations of systems dynamics of uncontrolled Burgers PDE. The spatial domain has length I, 1 with the diffusivity (viscosity) parameter = 10-3 The initial state u(x, f 0) = sechhh1100 10z 5).
Burgers' equation is a simplified PDE that captures the essential dynamics of fluid waves and gas dynamics. The visualization of the Burgers' Equation is shown in Figure 14. The field velocity u(x, 1) in Burgers' equation is defined as:
Ou Ou
Pu dx
a(x,t) t)
Øt
dr
where v > 0 is the diffusivity (viscosity) parameter, and a(x, t) is a source term modeling an external force acting on the PDE system at coordinate and time
20

---

Time-Aware World Model
0.80
stae
0.8
0.40
0.5
simulal dimulation time
025
050
0.20
0.8
0.00
state
state
0.0
-0.20
0.4
-3.7
-0.40
0.5
-0.5
-0.80
condationn time
-1.00
-1.0
14
0.0
0.0
denain coordnatees
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.5
1.0
1.5
2.0
0.5
simulation time
domainn coordinates
Figure 15. Visualizations of systems dynamics of uncontrolled Allen-Cath PDE. The spatial domain has length L = with the diffusivity (viscosity) parameter 10-* and potential constant 5.0. The initial state is u(2, t = (x 1)a cos c 1)).
G.3. Allen-Cahn Equation.
In material sciences, the Allen-Cahn PDE is non-linear PDE used to model the binary alloy systems' phase separation. The visualization of the Allen-Cahn equation is shown in Figure 15. The temporal dynamics of field w(x; t) is:
Ou
Ou Ox
12
+V
- a(x,t)
where u = #1 indicates the presence of each phase, v/ > is the diffusivity/viscosity parameter, V is the potential constant, and (xx, z) a source term modeling an external force acting on the PDE system at coordinate x and time t.
G.4. Wave Equation.
20.0
1.0
state (first half)
0.92
domain coordinates
0.8
0.72 c
state (fisst halft
0.8
15.0
0.6
0.62
12.5
0.5
10.0
simuaation
0.4
0.42
0.2
0.32
0.2
0.12
0.2
2.5
5.0
time
20
0.0
10
15
0.0
0.00
0.25
0.50
0.75
1.00
0.0
68 domger coordiratos
0
022
simulation time
domain coordinates
0.0
20.0
0.40
state (second half)
domain coordinates
0.8
0.30
state (second haif)
15.0
0.2
0.6
0.10
12.5
0.00
0.0
10.0
simulation
0.2
0.10
-0.20
-0.2
0.44
0.2
-0.30
5.0
-0.40
-0.4
2.5
oonaaation time
25
0.0
10
15
0.00
0.25
0.50
0.75
1.00
0.0
donano coordinates
08
simulation time
domain coordinates
0.0
Figure 16. Visualizations of systems dynamics of uncontrolled Wave PDE. The spatial domain has length I, 1 with 0.1. The initial state u(x,L = 0) = sech(10x 5) and (æx t 0) = O.
The wave equation is second-order linear PDE describing the spatial propagation of waves in homogeneous mediums. Such wave propagation PDE has many important applications in physics and engineering problems. The visualization of Wave PDE is shown Figure 16. The scalar quantity u(x, t) has wave-propagation dynamics defined as:
pu
pu
alz,t)
012
where is spatial wave speed, and a(x, t) is source term of a force acting on the system coordinate and time t.
21

---

Time-Aware World Model
G.S. Additional learning curve comparisons on PDE control tasks.
We present additional learning curves on PDE control tasks to provide insights into TAWM training efficiency on these tasks. Due to the small simulation time step required for stable simulation in PDE- PDE-Burgers (At 10 seconds), training this task takes longer in wall-clock time. Therefore, curves of TAWM and the baseline under different observation At values are shown in Figure 17. For each model and environment, the same set of intermediate weights is used for evaluation at each training step.
trained TAWM for only 750k steps PDE-Burgers The learning
pde-burgers
pde-allen cahn
pde-wave
-220
-200
8
-24
-10
-8
0.5M
0.5M
0.5M
IM
Eval, A - 10.0 ms
Pval, A t= 10.0 ms
Evall À 10.0 ms
-1
-25
-250
-8
-30
-10
-235
-35
12
-240
-40
-245
-45
0.5M
0.5M
0.5M
IM
Eval. à - 100.0 ms
Eval. A t- 100.0 ms
Eval. A 100.0 ms
230
-30
-235
-15
-20
-245
-250
-50
0.5M
0.5M Eval, A 200.0 ms
0,5M
IM
Eval. A 200.0 ms
Eval. A 200.0 ms
-10
-20
-25
-20
-235
-30
-35
-40
-50
-40
-250
-60
Eval. A 500.0 ms
Eval. A 500.0 ms
Eval. À 500.0 ms
-25
30
200
-35
400
-600
-20
50
Eval. 1000.0
Eval. A 1000.0
Eval. 1000.0
Figure 17. PDE control tasks: Average Negative LQ Error learning curves under different evaluation time step sizes. The default time step size are At = 50 ms/ ms/1000000000m0000 mss for PDE-Burgers, PDE-Allen-Cahn, PDE-Wave, respectively.
22

---

Time-Aware World Model
H. Proofs
In this section, we present the proofs of Lemma 4. and Lemma 4.2 from the theoretical analysis of TAWM's sample efficiency in Section 4.3.
H.1. Proof of Lemma 4.1
Proof. According to the definitions of the true dynamics function and the optimal learned dynamics function d', we have
24+ 2t f(2,, ae, At) At d' (2t, at, At) T(At)
(5)
+ I(z, on Al) = d' (z,, ae At) T(at) At
Since we are considering scenarios where the environment dynamics can captured with AT, the following holds:
||d" ae, At) :(AI) d' (2, AE)
(AA)
<E
At
At
(6)
Al
((A))
A T(At)
= ||d2 (2t, 0t, At) - df (2t, 0t, Al)
<
T(At)
Af
Given the sampling-frequenyy frequency assumption in Section 4.3 (Af [0.001, 1.0) and the definition of T()), we observe that
At < 1. Consequently, Lemma 4.1 holds. 7(A))
H.2. Proof of Lemma 4.2
Proof. Under Assumption 3 in Section 4.3, the following holds for current dynamics model d:
A
T(AI)
[[d(z,, at, At) d(21, az AT)
<
(7)
((tt) Ai
we denote I(At)= 7(AI]
r(A()
At then the following holds:
[ld(z), ar, At) (24, ae, At)ll
||d(zz 0r, At) d(zr, 0t, A0) I(AI) + d(z), ae, AT) I(A1)
d' (2t, (ts At) d' (zt, a2 At) I(At) d' de, At) [(a\)))
[|d(z), 0z, At) d(2:, 0t, Af) I(At)]l
(8)
||ds (2,, 0₁ At) - d' (2,, o,, At) T(At)]
+ ||d(z,, 0t At) d' At)l [(At)
= 2e + |(2,, a, At) d' (Zz, a, At]l I(At).
Likewise, the modeling error at time step At is upper bounded by the error at AŁ. Therefore, by minimizing the error at Ar, we expect that the error the smaller time step would safely decrease as well.
23


Reviews:
[{"review_id": "ACPqdWPTPJ", "reviewer": "Reviewer_AknR", "summary": "This work introduces Time-Aware World Model (TAWM), a model-based approach designed to explicitly incorporate the temporal dynamics of environments. By conditioning on the time step size, ∆t, and training over a diverse range of values ∆t – rather than relying on a fixed time step size – TAWM enables learning of both high- and low-frequency task dynamics in diverse control problems.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Please see additional experimental results [here](https://sites.google.com/view/anonymous-site-rebuttal-6714).\n\n**Q1. Lack of theoretical claims**\n\nA1. To corroborate our empirical results, we offer additional theoretical analysis on the sample efficiency of our proposed time-aware world model. Here we provide only a brief overview of the analysis in this rebuttal; the full analysis will be included in the final version of our paper (we will provide proofs if requested in the next rebuttal).\n\nWe start with following definitions:\n\n* $z(t)$: Function that gives state vector of the world at time step $t$. We assume this function is a locally Lipschitz function with a constant $L_1$. That is, for a sufficiently small $\\Delta t > 0$, we assume $| z(t + \\Delta t) - z(t) | \\le L_1 \\cdot \\Delta t.$ This is a mild assumption, as we are only concerned about physical environments, where every physical property (e.g. position, velocity) changes continuously. Also, note that $L_1$ is proportional to the frequency of the underlying dynamics -- if the environment is almost static, $L_1$ would be near zero. \n\n* $a(t)$: Function that gives an action at time step $t$.\n\n* $f(z(t), a(t), \\Delta t)$: Ground truth dynamics function that we want to approximate. This function satisfies the following equation for any pair of $z(t), a(t), $ and $\\Delta t$: $z(t + \\Delta t) = z(t) + f(z(t), a(t), \\Delta t) \\cdot \\Delta t$.\n\n* $d(z(t), a(t), \\Delta t)$: Dynamics function that we optimize to approximate $f$. We denote the predicted next state vector using this dynamics function as $\\hat{z}(t + \\Delta t) = z(t) + d(z(t), a(t), \\Delta t) \\cdot \\Delta t$.\n\nThen, the following Lemma 1 holds.\n\n**Lemma 1**. For sufficiently small $0 < \\Delta t_1 < \\Delta t_2$, $| f(z(t), a(t), \\Delta t_1) - f(z(t), a(t), \\Delta t_2) \\cdot \\frac{\\Delta t_2}{\\Delta t_1} | < \\frac{\\Delta t_2 - \\Delta t_1}{\\Delta t_1} \\cdot L_1$.\n\nNote that this relationship holds for every training data in our buffer -- therefore, we assume our dynamics function $d$ captures this relationship easily during training.\n\n**Assumption 1**. For sufficiently small $0 < \\Delta t_1 < \\Delta t_2$, $| d(z(t), a(t), \\Delta t_1) - d(z(t), a(t), \\Delta t_2) \\cdot \\frac{\\Delta t_2}{\\Delta t_1} | < \\frac{\\Delta t_2 - \\Delta t_1}{\\Delta t_1} \\cdot L_2$.\n\nWe can expect that $L_1$ would converge to $L_2$ during training. Based on these relationships, we can prove the following lemma.\n\n**Lemma 2**. For sufficiently small $0 < \\Delta t_1 < \\Delta t_2$, if $| f(z(t), a(t), \\Delta t_2) - d(z(t), a(t), \\Delta t_2) | = \\epsilon$, we can compute the approximation error of the state vector as $| z(t + \\Delta t_2) - \\hat{z}(t + \\Delta t_2) | = \\epsilon \\cdot \\Delta t_2$. Then, for $\\Delta t_1$, following holds: $| z(t + \\Delta t_1) - \\hat{z}(t + \\Delta t_1) | \\le \\epsilon \\cdot \\Delta t_2 + (\\Delta t_2 - \\Delta t_1) \\cdot (L_1 + L_2)$.\n\nThis lemma tells us that when we decrease the approximation error $\\epsilon$, it not only reduces the state approximation error at $(t + \\Delta t_2)$, but also that of $(t + \\Delta t_1)$. That is, when we optimize our model for one time step, it transfers to another time step. Also, it is more effective for the systems whose dynamics has lower frequency, and thus lower $L_1$. **Likewise, this lemma shows why our TAWM shows superior, or at least similar sample efficiency than the baseline, even though it has to learn additional temporal element.** \n\n--------------------------\n\n**Q2. Relationship to the video generative-model based world model (e.g. VideoAgent).**\n\nA2. Thank you for highlighting these works, which provide valuable insights that help connect our research with VLM Q&A. Both approaches share a common philosophy regarding the importance of temporal information. While VLM Q&A searches for the most relevant frames from pre-existing data to enhance video understanding, TAWM adaptively samples in time during the **data generation** process to capture the dynamics of underlying subsystems. In essence, VLM Q&A focuses on identifying key frames, whereas TAWM emphasizes aligning data generation with the system’s temporal dynamics. We will elaborate on this connection in the revised version.\n\n--------------------------\n\n**Q3. Possible application of TAWM for reducing sim2real gap.**\n\nA3. Although we don't yet have results on TAWM's learning transferability, we conjecture that TAWM can reduce the sim2real gap by adaptively sampling the frequency space to better synchronize **temporal effects** between simulated and real-world dynamics. This adaptive approach could enhance the robustness of the learning process by mitigating unexpected temporal noise from capturing devices or environmental factors—issues that are often absent in simulated environments. By training our world model across a wide range of frequencies, we expect it to become more resilient against these disturbances and improve its ability to generalize from simulation to real-world applications.", "claims_and_evidence": "The authors claim TAWM efficiently trains the world model M to accurately capture the underlying task dynamics across varying time step size ∆t’s without increasing sample complexity.\n\nThe authors demonstrate the results on diverse control problems in MetaWorld environments.", "methods_and_evaluation_criteria": "TAWM conditions estimation of the next state and reward on ∆t, as they depend on the temporal gap between the current and next state. The authors formulate M by modifying the world model of TD-MPC2 using 4-th order Runge-Kutta (RK4) method to enforce certain dynamical properties. Additionally, they modify the value model to take ∆t as an extra input. The authors train these models using various values of ∆t, which are log-uniformly sampled from a predefined interval.", "theoretical_claims": "No significant theorectial claims in my oppinion.", "experimental_designs_or_analyses": "Empirically, the authors show that our time-aware world model can effectively solve various control tasks under different observation rates without increasing data nor training steps.", "relation_to_broader_scientific_literature": "This work belongs to the familiar of Model-based RL", "essential_references_not_discussed": "Most references have been included in my opinion.", "other_strengths_and_weaknesses": "More theoretical insights (e.g., theorems) are appreciated.", "other_comments_or_suggestions": "More discussions about Generative World Model, especially video generative model-based (e.g., VideoAgent), are appreciated. Since Video generation directly learns the underlying dynamics. Note I am not requiring the authors to compare with VideoAgent, and I understand it’s different task setting.", "questions_for_authors": "How can TAWM help reduce the sim2real gap? Is there any empirical or theoretical evidence?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dimSjlcFkC", "reviewer": "Reviewer_V9z9", "summary": "This paper proposes Time-Aware World Models (TAWM) to enhance the robustness of world models in various frequency control tasks. During training, TAWM takes the randomly sampled transition time interval ($\\Delta t$) as an additional condition, enabling it to adapt to the control frequency of the test environment during evaluation. The authors validate the effectiveness and robustness of TAWM through experiments on the MetaWorld benchmark.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Please see additional experimental results [here](https://sites.google.com/view/anonymous-site-rebuttal-6714).\n\n**Q1. Gap between experimental design and motivation (data collection frequency matches the testing frequency in most manipulation tasks).**\n\nA1. We appreciate the reviewer’s concerns but believe this is a misunderstanding of our work. While data collection frequency can match testing frequency in laboratory settings, **this is not always feasible or effective in practice**.\n\nPrior works have acknowledged such limitations [1]. Consider training a model at $\\Delta t=$ 2.5 ms and deploying it on a physical robot with sensors operate at 100 fps ($\\Delta t=$10 ms). One could either (1) use the trained model directly or (2) repeat the same action for 4 substeps to compensate for the mismatch. However, as Fig. 2 shows, TAWM outperforms these strategies across various frequencies, making it flexible for any $\\Delta t$ encountered in practice w/o additional data or training.\n\nOne might argue that we could train the model at 10 ms in simulation to match the real-world frequency.  Figure 4 shows that this solution is not robust when $\\Delta t$ grows larger, as larger $\\Delta t$ introduces instability due to missing important high-frequency dynamics (see Introduction).  TAWM overcomes this issue by concurrently sampling multiple frequencies with better performance.\n\nWhile approaches like DayDreamer [2] can train models for manipulation tasks at low sampling frequencies, their robot motions between time steps are very slow, as shown in their demo. Such slow motions prevent the loss of important dynamics and stabilize training, but result in slow robot actions in practice. In contrast, our TAWM can effectively learn to solve tasks at large $\\Delta t$ without limiting the robot's motion speed, effectively learning both fast and slow task dynamics simultaneously.\n\nWe further conducted experiments on PDE-control environments, whose dynamics fundamentally differ from manipulation tasks (see **A6 of Reviewer 3**). Our results indicate that TAWM is effective beyond manipulation tasks and generalizes well across different classes of control problems.\n\n----\n**Q2. Lack of novelty and unfair comparison**\n\nA2. We respectfully disagree with the reviewer’s comment that the comparison is unfair. **Our training strategy—sampling multiple frequencies—is precisely the contribution of our work**. The baseline TD-MPC2 model doesn't incorporate the temporal element, so it cannot be trained the same way.  A comparison with TD-MPC2 is either (1) use it as-is or (2) manually apply substeps and repeat the same action (see A1).  TAWM outperforms both baselines by a considerable margin in as fair a comparison as possible, given the fundamental difference in approaches.\n\nRegarding novelty,  one might view our primary modification as simply adding adaptive time intervals to the model input. However, we do not believe this should be dismissed as a lack of novelty for the following reasons:\n\n* **Simplicity with Clear Benefits**:  Often simpler approaches are preferable when they offer clear advantages. Even if the idea may appear 'simple' at first, our extensive experiments demonstrate that it yields substantial performance gains over the baseline without increasing model size or the required training samples. This simplicity also means that other world-model learning methods can readily adopt our approach.\n\n* **Integration of Physical Dynamics and Temporal Elements**: Prior work has generally neglected explicit temporal modeling and adherence to basic physics principles in world models. Simply adding a time interval input, $z_{t+Δt}​=d(z_t​,a_t​,Δt)$, does not ensure compliance with these principles. In contrast, TAWM incorporates two key, model-agnostic components from physics simulation—time stepping and an integration method—that have been overlooked by previous methods. By integrating these components, we efficiently train our dynamics model by reducing the optimization space (e.g., the state remains unchanged when \n$\\Delta t = 0$). Our main contribution lies in embedding physical dynamics into world models—a concept applicable across various architectural designs. We emphasize that our method is model-agnostic, and our contribution is in the novel consideration of BOTH *temporal* and *physical dynamics* in the world model -- a significantly new concept that none of the prior works has explored.   \n\nWe also support our experimental results by theoretical proofs, as requested by the reviewers.\n\n----\n**Q3. Lack of supplementary material.**\n\nA3. We already provided an appendix with additional experimental results across all environments.  We refer this reviewer to the supplementary material and hope that it addresses any concerns regarding the breadth of our extensive evaluation.\n\n----\n\n[1] Thodoroff et al., \"Benchmarking real-time reinforcement learning.\"\n\n[2] Wu, Philipp, et al. \"Daydreamer: World models for physical robot learning.\"", "claims_and_evidence": "I am confused about the motivation of this paper. The authors claim that learning a time-aware model is essential to address issues like temporal resolution overfitting and inaccurate system dynamics caused by differences in observation frequencies during training and testing. However, the manipulation tasks in the experiments do not seem to encounter these problems. Therefore, the authors should provide more examples of tasks that face these challenges and validate their approach on a wider range of benchmarks.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "Yes.", "experimental_designs_or_analyses": "Yes.", "relation_to_broader_scientific_literature": "This paper presents an efficient, model-agnostic approach to training Time-Aware World Models (TAWM) that adapts to varying control frequencies without increasing sample complexity, reducing the need for retraining and lowering computational costs.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "***Strengths***\n1. This paper is clearly written and easy to follow.\n2. The proposed method is reasonable, and the experimental results demonstrate TAWM's robustness to control frequency and the efficiency of policy learning.\n\n***Weaknesses***\n1. This method lacks novelty, as the model learning in TAWM is largely based on TDMPC-2, with the main difference being the addition of time interval as an input.\n2. There is a gap between the experimental design and the motivation of the method. In most manipulation tasks, the data collection frequency matches the testing frequency. The authors need to provide a clearer explanation of which tasks experience discrepancies between training and testing frequencies and conduct experiments on the relevant benchmarks.\n3. The experimental comparison is unfair, as the baseline, TDMPC-2, is trained using a single control frequency, while TAWM is trained with data from various frequencies.", "other_comments_or_suggestions": "No.", "questions_for_authors": "See weaknesses.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "D97qzzKMvI", "reviewer": "Reviewer_ojxn", "summary": "The work presents a world model conditioned on time step size, showing that training with a sampling of different time step sizes can improve long horizon prediction stability. The algorithm conditions the world model on the time step size and uses 4th order Runge-Kutta to integrate the dynamical model. Experiments demonstrate performance improvements across a variety of simulated robotics tasks when using the world model for model-predictive control with different observation rates without increasing training budget (in data or steps). The new model out-performs an existing world model that models two different time scales.\n\n## update after rebuttal\n\nIncreased to accept based on the greater statistical rigor with some promising results, combined with clear text to address my comments.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your comprehensive reviews and suggestions!\n\n**Weakness 1.** The Nyquist sampling motivation is weak given the lack of a theoretical or empirical method to establish the connection to the theoretical signal optimization needs.\n\nA1. We have added the theoretical analysis of the sample efficiency and effectiveness of our proposed time-aware world model. Due to the word limit, we would like to refer to **our answer A1 to reviewer AknR** above for the theorem, lemmas, and their high-level idea for theoretical analysis. We would be happy to provide the proof details if the reviewer requests it in the next rebuttal.\n\n----\n**Weakness 2.** A few points were not clearly addressed, including the benefit of using 4th-order Runge-Kutta instead of Euler and the lack of substantial differences between log-uniform and uniform sampling.\n\nA2. Since this concern overlaps with the reviewer's **Q1** and **Q3**, we would like to refer to our answers A5 and A6 in this rebuttal.\n\n----\n**Suggestion 1.** Additional references.\n\nA3. We agree that AlphaGo and its subsequent works as well as DayDreamer are indeed relevant to this paper. We have included these works in our references and will be available in the final version.\n\n----\n**Suggestion 2.**  This (Fig. 12) should be Fig. 2 (in the body).\n\nA3. We have updated the main text accordingly to refer to both Figure 2 and Figure 12.\n\n----\n**Q1.** How much of the performance is due to the use of RK4 instead of Euler's method?\n\nA6. Thank you for raising your concern regarding the choice integration method. We have included additional ablation studies on the RK4 vs Euler integration method in our anonymous website: https://sites.google.com/view/anonymous-site-rebuttal-6714. We also include additional experiments on `PDE-Control` tasks from the `control-gym` envs: arxiv.org/abs/2311.18736.\n\nWe emphasize that our proposed method’s central contribution is the adaptive time stepping, which is **integration-agnostic** and can be used with any integration, depending on the nature of task dynamics. \n\nWe employed the RK4 integration method because it is generalizable to both systems with simple, linear dynamics and highly complex, nonlinear dynamics and is a standard method for simulations of physical systems. For most robot manipulation tasks (Meta-World), the dynamics are sufficiently simple to be approximated with the Euler integration method. Empirically, our ablation shows that TAWM Euler performed better than RK4-based TAWM in most simple Meta-World tasks, suggesting the underlying dynamics of Meta-World tasks are sufficiently simple to be captured by the Euler integration method. The advantages of RK4 are more apparent in tasks with complex, non-linear dynamics, such as the PDE-Control tasks.\n\nThese results reinforce the merit of our core contribution of adaptive time stepping for training world models. The integration method and $\\Delta t$-sampling method are two parameters we can adjust to maximize TAWM's performance and efficiency --both of which outperform the baseline.\n\n----\n**Q2.** What evidence is that log-uniform sampling is superior to uniform scaling? The claims ultimately state that any sampling can be used. If this is the claim being made that should be explicit at the start of the paper.\n\nA7. Thank you for your comments! As explained in A6, the integration method and $\\Delta t$-sampling method are two parameters that we can adjust to maximize TAWM's performance and efficiency. Therefore, uniform sampling and/or the Euler method can be better choices depending on the tasks. For example, uniform sampling performs better than log-uniform sampling in tasks with sufficiently slow dynamics to be captured by $\\Delta t_{max}$. For tasks demanding fast inference time with sufficiently simple dynamics, Euler method is preferred. \n\nWe will be sure to make this point much more clear at the beginning of our paper in the final revision.  Thank you for the suggesiton.\n\n----\n**Q3.** How fast is inference/planning in TAWM?\n\nA8. We provide additional details of the inference time below. Inference time is averaged over 1000 planning steps for each model.\n* Baseline: $\\mu=$0.027 s; Q1 = 0.026 s; Q3 = 0.027 s\n* Euler: $\\mu=$0.028 s; Q1 = 0.028 s; Q3 = 0.028 s\n* RK4: $\\mu=$0.048 s; Q1 = 0.048 s;  Q3 = 0.05 s\n\nThe tradeoff between Euler and RK4 is well-known in simulation: Euler is faster but can be unstable/less accurate than RK4 depending on the underlying dynamics. As mentioned in our answers A6 and A7, the integration method is one of the adjustable parameters. If underlying dynamics are sufficiently simple and slow to be approximated by the Euler method, the Euler method is preferred. Otherwise, RK4 is a more generalizable method.\n\nWe have incorporated your comments and new results into our paper. We thank you again for your comprehensive reviews and valuable suggestions, and we hope our answers sufficiently address your concerns and questions.", "claims_and_evidence": "The claims and their evidence:\n- claim: Training on variable time step sizes (TAWM) addresses compounding errors with one-step prediction.\n\t- evidence: The MetaWorld experiments examine error when sampling inference with long time steps, showing success when using long time steps.\n- claim: TAWM handles variable frame rate data without altering training budget.\n\t- evidence: The MetaWorld experiments include many comparisons of the TAWM to fixed time step models on varying scales. In most cases TAWM does better, primarily when facing longer time steps (lower frequency). The main exception is the hammer environment.\n- claim: TAWM works in control tasks with different observation rates without increasing training budget (data or steps)\n\t- evidence: Same evidence as above.\n- claim: (implicit) Euler method fails where 4th order Runge-Kutta succeeds\n\t- evidence: None provided. This merits a separate evaluation or ablation (perhaps in the appendix).\n\nThere is an implied claim that log-uniform sampling is important for the training process, but the evidence shows some advantages to uniform time step sampling at long horizons, with mixed evidence of differences at short horizons (overlapping confidence intervals). See Figures 7 and 8.", "methods_and_evaluation_criteria": "Yes. Robotics tasks are a natural scenario for long-range control tasks (using MPC or otherwise). Testing across different time horizons is the natural task to test.", "theoretical_claims": "No. No proofs were made.", "experimental_designs_or_analyses": "Yes. The Meta-World world evaluations.\n- These only have an ablated version of the TAWM as comparison. They would benefit from alternative methods.\n- No experiments report on inference time requirements.\n- Only one baseline model was compared.", "relation_to_broader_scientific_literature": "The work is contextualized in relation to the model-based reinforcement learning literature. The proposed model performs model-predictive control in a world model, rather than learning a control policy (as in the Dreamer models). \n\nThe key contribution is devising a world model training approach that can train on variable time step durations and perform inference when observations occur at different frame rates. This is contrasted with the existing MTS3 model that is trained for a discrete number of temporal resolutions (two in the existing model).", "essential_references_not_discussed": "It would be reasonable to reference AlphaGo and subsequent work (AlphaGo Zero, AlphaZero, MuZero, Muesli, and so on) in that model family as exemplars of the model-based RL approach. The latter models in this series developed model-based RL algorithms similar to the Dreamer model family with very strong performance across multiple domains.\n\nDayDreamer (Wu et al. in CoRL 2023) may merit consideration as a robotics model that trains in simulation alongside a physical robot (https://proceedings.mlr.press/v205/wu23c.html). This handles the online learning case, which contrasts with the simulation work in this paper.", "other_strengths_and_weaknesses": "# strengths\n- originality: The core idea of conditioning on variable time steps combines simplicity with novelty, showing how to generally improve world modeling tasks.\n- significance: Improving world model temporal behavior is valuable across a number of task domains in robotics and other physical systems. Enabling models to handle variable time steps is also relevant to other architectures like directly learning behavioral policies (instead of planning over a time step aware model).\n\n# weaknesses\n- clarity: A few points were not clearly addressed, including: the benefit from using 4th order Runge-Kutta instead of Euler integration and the lack of substantial differences between log-uniform and uniform sampling.\n\t- The Nyquist sampling motivation is also weak given the lack of a theoretical or empirical method that establishes a connection to the theoretical signal optimization needs.", "other_comments_or_suggestions": "\"To assess our time-aware model’s performance at different inference-time observation rates, we evaluated it on multiple tasks with varying ∆ t. As shown in Figure 12, it outperforms the baseline (trained at a fixed ∆ t = 2.5ms) in across three tasks\"\n- This should be Figure 2 (in the body).", "questions_for_authors": "- [Q1] How much of the performance is due to the use of RK4 instead of Euler's method?\n- [Q2] What evidence is there that log-uniform sampling is superior to uniform scaling?\n\t- Figures 7 and 8 show uniform sampling works better at long time steps. Short time steps overlap in performance with the log-uniform case as well. The claims ultimately state that any sampling can be used. If this is the claim being made that should be explicit at the start of the paper.\n- [Q3] How fast is inference / planning in TAWM?\n\t- Considering the downstream task is robotic manipulation having a sense of how fast model predictions and planning occur is important to assess deployment in the real world (eventually).", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper introduces the Time-Aware World Model (TAWM), a model-based reinforcement learning framework that conditions on variable time step sizes during training to improve generalization across environments with differing temporal dynamics. The reviewers acknowledged the method’s simplicity and empirical strength, highlighting its robustness and sample efficiency without increased training cost. Concerns were raised regarding the novelty of the approach, real-world applicability, and fairness of baseline comparisons. The rebuttal effectively addressed these points by providing theoretical analysis, ablation studies, and clarifications on the experimental setup and motivation. One reviewer raised their score following the response, citing stronger justification and statistical support. Given the clarity, empirical effectiveness, and practical utility of the approach, the overall assessment supports acceptance. Suggested directions for improvement include evaluation on real-world or sim-to-real tasks and broader comparisons with alternative world model approaches.
Decision: Accept (poster)