PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Chi Wan, Kangrui Wang, Yuan Si, Pingyue Zhang, Manling Li

arXiv:2609.08162v1cs.AI

TL;DR

VLA models trained on static datasets lack mechanisms to adapt their world-dynamics representations in new environments. WorldAgen jointly learns task-conditioned actions and world modeling in a shared Transformer, then applies lightweight test-time updates from exploratory transitions. The baseline matches or exceeds state-of-the-art performance, while TTT achieves state-of-the-art results on CALVIN and LIBERO.

  • Problem

    Static pretraining leaves VLA models unable to adapt their internal representations of world dynamics to novel environments and distribution shifts.

  • Method

    WorldAgen uses a shared Transformer with policy and world-model heads, then performs lightweight test-time world-model updates from exploratory state transitions.

  • Results

    WorldAgen achieves performance comparable to or better than state-of-the-art methods, and TTT attains state-of-the-art results on CALVIN and LIBERO.

  • Takeaways & Limitations

    World modeling can serve as an active test-time adaptation mechanism that improves VLA generalization across manipulation benchmarks.

  • Takeaways & Limitations

    The study evaluates only one model size and architecture and applies test-time training exclusively to the world-model component.

Abstract

from arXiv · show

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.

Introduction

WorldAgen addresses the limits of static pretraining by unifying task-conditioned action prediction with task-agnostic world modeling and adapting the world model during test-time interaction. Its lightweight TTT strategy improves environment understanding and achieves strong results across CALVIN and LIBERO.

  • Static-dataset training leaves VLA models without mechanisms to adapt world-dynamics representations when object layouts, lighting, or physical properties change.
  • WorldAgen combines a task-conditioned policy head and a task-agnostic world-model head within one shared Transformer backbone.
  • During test time, exploratory actions collect ground-truth transitions, enabling lightweight LoRA updates to refine the world model without changing the agent policy.
  • The framework reframes world modeling from passive pretraining into an active test-time adaptation mechanism for novel environments.
  • With only a small number of test-time samples, WorldAgen reaches state-of-the-art results on CALVIN and LIBERO, while its baseline is comparable to or better than existing methods.

Method

WorldAgen unifies task-conditioned action prediction and task-agnostic world modeling through a shared Transformer, while masking prevents information leakage. Its chunked, teacher-forced training and two-step inference support lightweight test-time world-model adaptation.

  • Trajectory Representation: WorldAgen represents each rollout as trajectory units containing instructions, observation and state chunks, action chunks, and prediction placeholders.Chunks are formed from time-indexed trajectory elements, with observations and states uniformly subsampled and actions grouped into chunks.
  • Joint Prediction: Both predictive components share one Transformer backbone, combining task-aware action prediction with task-agnostic environment-dynamics modeling.The shared representation is trained jointly across the two predictive tasks.
  • Joint Prediction: The agent model predicts action chunks from instructions, current observations, states, and agent history, while the world model predicts future observations from predicted actions and world history.During training, the world model uses ground-truth actions through teacher forcing; during inference, it uses the agent model’s predicted actions.
  • Attention Masking: Mixed unidirectional attention uses local and global masks to block action leakage, separate task instructions from world modeling, and preserve strictly causal prediction.The local mask blocks access to current action chunks, while the global mask keeps the world-model head invisible to task instructions.
  • Inference: At inference, the agent first fills a zero action placeholder with a predicted action chunk, then the world model predicts the next observation chunk from the updated trajectory.These two steps alternate to roll out the trajectory autoregressively.
  • Test-Time Adaptation: Test-Time Training adapts only the world model after distribution shifts, refining environment dynamics without requiring additional task-specific annotations.The adaptation is designed to improve downstream action prediction indirectly while remaining decoupled from the task goal.

Experiments

WorldAgen is evaluated on CALVIN and LIBERO, with ablations examining world modeling, LoRA configuration, test-time data volume, and robustness to noise. Results show consistent or improved performance, while moderate test-time data improves adaptation and excessive data can reduce it.

  • Benchmarks: CALVIN measures long-horizon language-conditioned manipulation, while LIBERO emphasizes knowledge transfer across diverse lifelong-learning tasks.
  • Benchmark Results: WorldAgen achieves consistent improvements across five consecutive CALVIN tasks and best or on-par performance across most LIBERO subsets.These comparisons include recent VLA baselines on both benchmarks.
  • World Modeling Ablation: 30.7%: CALVIN average success length rises from 2.96 to 3.87 when image prediction enables world modeling.On LIBERO, success rate similarly rises from 46.5% to 78.0% (+67.7%).
  • LoRA Ablation: 0.02: TTT performance varies by at most this amount across LoRA ranks under fixed learning rate and data size.The main experiments use LoRA rank 128 because it slightly outperforms alternatives while remaining computationally efficient.
  • TTT Data Volume: 3.871 to 3.928: average success score improves as TTT samples increase from 6 to 204, then falls to 3.917 at 340 samples.The results indicate diminishing returns and mild overfitting to image generation rather than action prediction.
  • Robustness: Gaussian noise with std = 0.1 is used to simulate camera distortion and sensor noise, under which WorldAgen outperforms the baselines.

Related Work

Related work spans VLA policies, predictive world models, and test-time training. WorldAgen targets the gap between these areas by adapting the world model within a unified language-grounded action framework.

  • Vision-Language-Action Models: VLA models unify perception, language, and control, but most directly predict actions without explicitly modeling environmental dynamics.
  • World Models: World models support predictive dynamics for planning and decision-making, yet robotic applications remain largely decoupled from language-grounded action generation.
  • Test-Time Training: Test-time training adapts models to distribution shifts during inference, but robotic applications have mainly focused on perception modules or low-level policies.WorldAgen instead incorporates TTT into a unified VLA framework by adapting the world model during deployment.

Conclusion

WorldAgen jointly learns future-state prediction and task-conditioned action prediction through a shared Transformer, then uses exploratory rollouts for lightweight test-time adaptation. This produces consistent CALVIN and LIBERO gains and supports continual interaction as an adaptation signal.

  • WorldAgen uses a shared Transformer backbone with world-model and policy heads for joint state prediction and task-conditioned action prediction.
  • Short exploratory rollouts provide an online TTT signal that adapts the world model during deployment.
  • Consistent CALVIN and LIBERO gains demonstrate benefits from continual test-time interaction in adaptive VLA systems.

Limitations and Future Work

The study evaluates one model size and architecture and applies test-time training only to the world-model component. Future work will explore broader architectures and joint adaptation of action prediction and world modeling.

  • Limitations: The evaluation covers only a single model size and architecture, while TTT is restricted to the world-model component.
  • Future Work: Future work includes larger, more diverse architectures and joint test-time training for action prediction and world modeling.

Appendix

The appendix documents implementation details and supplementary experiments for WorldAgen. It covers benchmarks, architecture, trajectory processing, chunk configurations, and TTT fine-tuning comparisons.

  • Scope: The appendix provides supplementary materials for the WorldAgen framework.It accompanies the main framework description with additional technical and experimental information.
  • Implementation: It specifies benchmark descriptions, model architecture components, and trajectory-processing methods for variable-length chunks.
  • Experiments: It reports experiments on chunk configurations and compares LoRA with full fine-tuning for test-time training.

Benchmark

The paper evaluates WorldAgen on CALVIN and LIBERO, which test long-horizon manipulation and knowledge transfer across diverse tasks. CALVIN emphasizes sequential language-conditioned control, while LIBERO provides multiple task suites with varying transfer demands.

  • LIBERO: LIBERO evaluates lifelong robot-manipulation learning through knowledge transfer across four distinct task suites.The suites target spatial, object-centric, procedural, and highly entangled knowledge transfer.
  • LIBERO: LIBERO-100 contains 100 tasks divided into LIBERO-90 for policy pretraining and LIBERO-10 for downstream long-horizon evaluation.
  • CALVIN: CALVIN is a simulated benchmark for long-horizon language-conditioned robot manipulation.It contains 34 complex manipulation tasks and evaluates sequences of up to five consecutive instructions.
  • CALVIN: CALVIN varies desk colors and object configurations across four environments to test generalization across visual contexts.

Model Architecture

WorldAgen uses a Transformer-based multimodal architecture that processes visual, language, action, and robot-state inputs with flexible chunking. Its encoder compresses visual tokens, while separate decoding pathways predict images and actions.

  • Architecture: WorldAgen supports variable-length image and action chunks across multimodal inputs and outputs.
  • Visual encoder: A MAE-pretrained ViT-B encodes dual-view RGB observations, while a Perceiver Resampler compresses image tokens for efficiency.
  • Input modalities: Language uses CLIP ViT-B/32, while robot states and actions are projected into a shared embedding space from 7-dimensional vectors.The vectors encode six arm-state dimensions and one gripper-state dimension.
  • Decoder: The decoder uses a Vision Transformer pathway for image prediction and separate linear or MLP layers for action and gripper-state outputs.
  • Backbone: Qwen3 is used as the backbone, with dense and mixture-of-experts variants designed for parameter and computation efficiency.

Baseline Results

Baseline experiments show that chunk configuration strongly affects performance, with intermediate action and chunk lengths preferred over excessively long or short settings. Across CALVIN and LIBERO, lightweight LoRA-based TTT consistently improves results, whereas excessive adaptation or full fine-tuning can hurt performance.

  • Chunk configuration: Longer action chunks degrade performance through accumulated prediction errors, while excessively short chunks provide too little training information.
  • Chunk configuration: The authors use action chunk lengths of 5 for CALVIN and 3 for LIBERO based on benchmark experiments.
  • Chunk configuration: Three chunks achieve the best reported performance, and this configuration transfers to another chunk setting that reaches state-of-the-art results.
  • Test-time training: Full fine-tuning decreases performance on CALVIN, whereas LoRA balances scenario adaptation with retention of pretrained scene perception.
  • Test-time training: TTT improves performance across settings using around 200 samples, supporting adaptation across architectural configurations.
  • Test-time training: TTT performance improves with small data amounts but degrades when excessive samples cause LoRA to overfit world modeling.The paper therefore emphasizes balancing adaptation data against action execution.
  • Additional results: Qwen3 outperforms GPT-2 across language-backbone comparisons, while visualization shows TTT enabling completion of a longer task sequence.
Loading 2609.08162v1…