Vision language action models improve by training world models at test time
WorldAgen: Unified State-Action Prediction with Test-Time World Model Training
Artificial Intelligence
Summary
Models that see, understand language, and act often struggle when they face new environments they haven't seen before. The authors created WorldAgen, a system that learns both how the world changes and what actions to take based on instructions. When used in a new place, it tries some actions, watches what happens, and quickly updates itself to understand better. This makes the model more accurate in predicting what to do. Tests show it beats or matches current best methods, especially when it adapts during use.
What this means in practice
- •For robotic system developers: Improve robot adaptability in changing environments by enabling on-the-fly updating of environmental models using small samples of interaction data.
- •For autonomous vehicle teams: Enhance autonomous vehicle decision-making by incorporating test-time world model updates to better handle unfamiliar road scenarios and dynamics.
Authors
Chi Wan, Kangrui Wang, Yuan Si, Pingyue Zhang, Manling Li
Abstract
How can vision-language-action (VLA) models adapt to new environments where world dynamics shift? While recent research has combined world modeling and action prediction to improve VLA performance, existing methods largely rely on pretraining on static datasets, without mechanisms for active adaptation at deployment time. As a result, these models often fail to generalize when deployed in unseen scenarios with novel object configurations or dynamics. We present WorldAgen, a unified framework that jointly learns world modeling and action prediction while enabling Test-Time Training (TTT) to adapt to new environments. WorldAgen employs a shared Transformer backbone with two heads: (1) a world model head that predicts future states from past state-action trajectories, and (2) an agent model head that predicts actions conditioned on task instructions. We design a Mixed Unidirectional Attention Mask to separate these two models. During test time, WorldAgen samples exploratory actions, collects ground-truth state transitions, and performs lightweight TTT updates to refine its world model. This adaptation improves the model's understanding of the environment and leads to more accurate action predictions. Experiments on the CALVIN and LIBERO benchmarks demonstrate that our baseline model achieves comparable, and in some cases superior, performance to current state-of-the-art approaches. Moreover, with TTT on a small number of samples, our method surpasses existing state-of-the-art models, highlighting the effectiveness of adapting world models at inference time.