WorldAgen: Unified State-Action Prediction with Test-Time World Model Training

2026-03-14

AAAI Conference on Artificial Intelligence

AI Summary

本文提出WorldAgen框架，通过联合学习世界建模与动作预测，并采用测试时训练（TTT）来适应新环境，解决VLA模型在未知场景中泛化能力差的问题。

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 in static datasets, without mechanisms for active adaptation to new environments. 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. 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 to the environments 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.

Problem

• vision-language-action models
• new environments
• world dynamics
• generalization
• adaptation

Innovation

• Test-Time Training
• Unified Framework
• Mixed Unidirectional Attention Mask