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

• URL: http://arxiv.org/abs/2609.08162v1
• Date: Tue, 08 Sep 2026 02:46:13 GMT
• Title: WorldAgen: Unified State-Action Prediction with Test-Time World Model Training
• Abstract summary: We present WorldAgen, a unified framework that jointly learns world modeling and action prediction.<n>WorldAgen samples exploratory actions, collects ground-truth state transitions, and performs lightweight TTT updates to refine its world model.<n>Our method surpasses existing state-of-the-art models, highlighting the effectiveness of adapting world models at inference time.
• Score: 19.89058127984308
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Reinforcing VLAs in Task-Agnostic World Models [15.056103368344596]
  Post-training Vision-Language-Action (VLA) models via reinforcement learning (RL) have emerged as an effective strategy to adapt to new tasks without costly real-world interactions.<n>We argue that world and reward models should capture transferable physical priors that enable zero-shot inference.<n>We propose RAW-Dream, a new paradigm that completely disentangles world model learning from downstream task dependencies.
  arXiv  Detail & Related papers  (2026-05-12T16:16:15Z)
• HarmoWAM: Harmonizing Generalizable and Precise Manipulation via Adaptive World Action Models [58.191567345416836]
  World Action Models (WAMs) have emerged as a promising paradigm for robot control by modeling physical dynamics.<n>HarmoWAM fully leverages a world model to unify predictive and reactive control, enabling both general transit and precise manipulation.<n>We construct three training-unseen test environments across six real-world robotic tasks, covering variations in background, position, and object semantics.
  arXiv  Detail & Related papers  (2026-05-11T17:59:56Z)
• Test-Time Mixture of World Models for Embodied Agents in Dynamic Environments [29.514831254621438]
  Test-time Mixture of World Models (TMoW) is a framework that enhances adaptability to unseen and evolving domains.<n>TMoW updates its routing function over world models at test time, unlike conventional MoE where the function remains fixed.<n>We evaluate TMoW on VirtualHome, ALFWorld, and RLBench benchmarks, demonstrating strong performance in both zero-shot adaptation and few-shot expansion scenarios.
  arXiv  Detail & Related papers  (2026-01-30T07:06:40Z)
• Agent2World: Learning to Generate Symbolic World Models via Adaptive Multi-Agent Feedback [51.22403664895878]
  Agent2World is a tool-augmented multi-agent framework that achieves strong inference-time world-model generation.<n>It also serves as a data engine for supervised fine-tuning, by grounding generation in multi-agent feedback.
  arXiv  Detail & Related papers  (2025-12-26T18:54:14Z)
• Co-Evolving Latent Action World Models [57.48921576959243]
  Adapting pre-trained video models into controllable world models via latent actions is a promising step towards creating generalist world models.<n>We propose CoLA-World, which for the first time successfully realizes this synergistic paradigm.<n>This unlocks a co-evolution cycle: the world model acts as a knowledgeable tutor, providing gradients to shape a high-quality LAM.
  arXiv  Detail & Related papers  (2025-10-30T12:28:40Z)
• Transformer World Model for Sample Efficient Multi-Agent Reinforcement Learning [2.3964255330849356]
  We present the Multi-Agent Transformer World Model (MATWM), a novel transformer-based world model for reinforcement learning.<n>MATWM combines a decentralized imagination framework with a semi-centralized critic and a teammate prediction module.<n>We evaluate MATWM on a broad suite of benchmarks, including the StarCraft Multi-Agent Challenge, PettingZoo, and MeltingPot.
  arXiv  Detail & Related papers  (2025-06-23T11:47:17Z)
• Learning Transformer-based World Models with Contrastive Predictive Coding [58.0159270859475]
  We show that the next state prediction objective is insufficient to fully exploit the representation capabilities of Transformers.<n>We propose to extend world model predictions to longer time horizons by introducing TWISTER, a world model using action-conditioned Contrastive Predictive Coding.<n>TWISTER achieves a human-normalized mean score of 162% on the Atari 100k benchmark, setting a new record among state-of-the-art methods that do not employ look-ahead search.
  arXiv  Detail & Related papers  (2025-03-06T13:18:37Z)
• Pre-Trained Video Generative Models as World Simulators [59.546627730477454]
  We propose Dynamic World Simulation (DWS) to transform pre-trained video generative models into controllable world simulators.<n>To achieve precise alignment between conditioned actions and generated visual changes, we introduce a lightweight, universal action-conditioned module.<n> Experiments demonstrate that DWS can be versatilely applied to both diffusion and autoregressive transformer models.
  arXiv  Detail & Related papers  (2025-02-10T14:49:09Z)
• An Efficient Occupancy World Model via Decoupled Dynamic Flow and Image-assisted Training [50.71892161377806]
  DFIT-OccWorld is an efficient 3D occupancy world model that leverages decoupled dynamic flow and image-assisted training strategy.<n>Our model forecasts future dynamic voxels by warping existing observations using voxel flow, whereas static voxels are easily obtained through pose transformation.
  arXiv  Detail & Related papers  (2024-12-18T12:10:33Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.