Time-Aware World Model for Adaptive Prediction and Control

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

ICML 2025

TL;DR: Time-Aware World Model (TAWM) is a model-agnostic training method that improves dynamics learning by explicitly incorporating time step size Δt and sampling observations at varying frequencies. This addresses the real-world constraint of varying observation rates, enabling efficient learning across temporal scales and outperforming fixed-Δt baselines under the same training budget.

Paper arXiv Code
Model Diagram
TAWM explicitly incorporates time step size Δt into dynamics modeling (Time-Aware) and learns temporal dynamics at various scales by adaptively sampling Δt during training. Control actions are generated using Model Predictive Control (MPC), specifically 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 on Varying Observation/Control Frequencies

TASK: MW-Assembly
* the FPS is adjusted to highlight observation rate differences
TAWM successfully learn to solve the task across different inference Δt, without needing additional training samples nor training data.
The non-time-aware models, even when trained specifically on Δt = 10 ms or Δt = 50 ms, fails to solve the tasks at low observation rates.
We hypothesize this is due to missing important transitions needed to learn the environment dynamics, leading to a collapse in world model learning.
Note: Click the left or right button to view evaluations for more observation rates.
TASK: MW-Basketball
* the FPS is adjusted to highlight observation rate differences
TAWM successfully learn to solve the task across different inference Δt, without needing additional training samples nor training data.
The non-time-aware models, even when trained specifically on Δt = 10 ms or Δt = 50 ms, fails to solve the tasks at low observation rates.
We hypothesize this is due to missing important transitions needed to learn the environment dynamics, leading to a collapse in world model learning.
Note: Click the left or right button to view evaluations for more observation rates.
TASK: MW-Box-Close
* the FPS is adjusted to highlight observation rate differences
TAWM successfully learn to solve the task across different inference Δt, without needing additional training samples nor training data.
The non-time-aware models, even when trained specifically on Δt = 10 ms or Δt = 50 ms, fails to solve the tasks at low observation rates.
We hypothesize this is due to missing important transitions needed to learn the environment dynamics, leading to a collapse in world model learning.
Note: Click the left or right button to view evaluations for more observation rates.
Experiments on Additional Tasks
* the FPS is adjusted to highlight observation rate differences
Note: Click the left or right button to view evaluations for more observation rates.
* the FPS is adjusted to highlight observation rate differences
Note: Click the left or right button to view evaluations for more observation rates.

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{
    nhu2025timeaware,
    title={Time-Aware World Model for Adaptive Prediction and Control},
    author={Anh N Nhu and Sanghyun Son and Ming Lin},
    booktitle={Forty-second International Conference on Machine Learning},
    year={2025},
    url={https://openreview.net/forum?id=gZ5N3TLjwv}
}