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

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.

Link To Code: https://github.com/anh-nn01/Time-Aware-World-Model

Primary Area: Reinforcement Learning

Keywords: RL, Dynamics, World Model

Submission Number: 6714