Time-Aware World Model for Adaptive Prediction and Control

TL;DR

The paper introduces TAWM, a time-aware world model that explicitly incorporates temporal dynamics, improving control performance and data efficiency across diverse tasks by conditioning on variable time-step sizes.

Contribution

It presents a novel time-aware modeling approach that explicitly conditions on time-step size, enabling better handling of diverse temporal dynamics in control problems.

Findings

TAWM outperforms conventional models across various control tasks.

It improves data efficiency and performance by learning from diverse {}t values.

Empirical results demonstrate consistent gains in control accuracy.

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.