Time-Aware World Model for Adaptive Prediction and Control

Anh N. Nhu*, Sanghyun Son*, Ming C. Lin
University of Maryland, College Park
*Equal Contribution
Paper arXiv Code Poster

ICML 2025

TL;DR: Time-Aware World Model (TAWM) is a model-agnostic training recipe for improving world dynamics models by explicitly conditioning dynamics prediction on the timestep size \( \Delta t \). By training on a mixture of timestep sizes, TAWM captures both high- and low-frequency dynamics, generalizes across varying control frequencies, and avoids (1) inefficient learning from fixed, single small \( \Delta t \) and (2) the failure of large-\( \Delta t \) training that can skip fast dynamics and fail to converge.
Overview of Time-Aware World Model framework
Overview of TAWM. TAWM explicitly conditions world dynamics models on the timestep size \( \Delta t \). Given latent state \( z_t \) and action \( a_t \), TAWM learns a time-aware latent dynamics function \( d_\theta(z_t, a_t, \Delta t) \) and integrates the dynamics over \( \Delta t\). Our Euler formulation integrates \( z_{t+\Delta t} = z_t + d_\theta(z_t, a_t, \Delta t)\,\tau(\Delta t), \) while RK4 formulation is used for higher-order latent dynamics prediction. During training, TAWM adaptively samples trajectories across varying temporal resolutions, capturing multi-scale behavior and adaptable to changing observation rates. At inference time, control actions are generated through Model Predictive Control, particularly the MPPI planner.

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, ∆t, and training over a diverse range of ∆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.

Lay Summary

“World models” are an emerging class of machine learning algorithms that learn how the world changes over time, essentially modeling how actions lead to different outcomes. They serve as an internal simulation of the world, enabling AI systems to plan and solve complex tasks, much like how humans rely on physical intuition and experience to drive, cook, or navigate.

To learn world models, AI systems collect observations from the surrounding environment, similar to how we humans learn from experience. However, current methods typically train these models using a single, fixed observation rate, learning how things evolve at only a single, small time scale. It's like trying to predict the road ahead every second during a calm drive: when little changes, this becomes redundant and inefficient.

We rethink how to train world models: instead of using one fixed observation rate, we train them across a wide range of observation rates, which we call Time-Aware World Models (TAWM). This simple change allows the model to learn both fast and slow dynamics more effectively in a single training run. With our time-aware approach, AI systems learn how the world evolves across multiple time scales, boosting success rates and performance on various complex control tasks — all without requiring additional data or resources. As a result, our method enables more efficient training, reduces energy and computing costs, and supports greener AI. Our findings can enable more robust, efficient AI systems across different domains, from simulation and physical AI research to autonomous driving and industrial autonomous systems.

Experiments

Testing on Varying
Observation Rates / Control Frequencies

Select a task and inference timestep to compare TAWM against fixed-\(\Delta t\) baselines. Videos are loaded only when selected to keep the page fast.

Same training budget
(steps / samples) Various inference
\(\Delta t\)'s * FPS adjusted for
visualization
Control Task
Inference Timestep / Control Frequency * Effective control frequency: \( f = 1/(10\Delta t) \)
(frame skip = 10; (base env: 5 \(\times\) our implementation: 2))

BibTeX Citation

Thank you for your interest in our work. If you find it useful for your research, we would be grateful if you could consider citing our paper. 
 
@InProceedings{pmlr-v267-nhu25a,
  title = 	 {Time-Aware World Model for Adaptive Prediction and Control},
  author =       {Nhu, Anh N and Son, Sanghyun and Lin, Ming},
  booktitle = 	 {Proceedings of the 42nd International Conference on Machine Learning},
  pages = 	 {46265--46287},
  year = 	 {2025},
  editor = 	 {Singh, Aarti and Fazel, Maryam and Hsu, Daniel and Lacoste-Julien, Simon and Berkenkamp, Felix and Maharaj, Tegan and Wagstaff, Kiri and Zhu, Jerry},
  volume = 	 {267},
  series = 	 {Proceedings of Machine Learning Research},
  month = 	 {13--19 Jul},
  publisher =    {PMLR},
  pdf = 	 {https://raw.githubusercontent.com/mlresearch/v267/main/assets/nhu25a/nhu25a.pdf},
  url = 	 {https://proceedings.mlr.press/v267/nhu25a.html},
  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.}
}