World Tokens: Enhancing Embodied Policies with Training-Time World Modeling

2026-08-10

AI Summary

This work addresses the limitation of existing vision-language-action (VLA) models in explicitly modeling physical dynamics, while avoiding the high inference overhead typically incurred by integrating video world models. The authors propose the World Tokens architecture, which during training employs a World Adapter to map vision-language features into shared world tokens, enabling joint optimization of future video denoising and action prediction. At deployment, the computationally intensive world model branch is discarded, retaining only a lightweight inference path. This approach achieves strong performance on the LIBERO benchmark, sets a new state-of-the-art average score on SIMPLER, and significantly improves task success rates on the real-world R1 Pro robot—all while maintaining action generation latency comparable to standard VLA models.

Abstract

Vision-language-action (VLA) models are a widely adopted paradigm for embodied policies. They excel at efficient closed-loop control but do not explicitly model how physical scenes evolve as a task unfolds. Recently emerging world-action models (WAMs) leverage pretrained video world models to capture spatiotemporal evolution, yet retaining future generation or a large video backbone in the control loop substantially increases inference cost. We introduce World Tokens, an embodied policy architecture built around a World Adapter that bridges visual-language understanding, world-dynamics modeling, and action generation. It uses world modeling during training to enhance the action policy while preserving efficient deployment. Specifically, the World Adapter transforms VLM features into a fixed set of world tokens, which condition a jointly fine-tuned future-video denoiser and simultaneously serve as the action expert's sole visual-language context. This shared conditioning allows gradients from future-video denoising to directly shape the representation used for action prediction, while exclusive routing prevents the policy from bypassing that representation. At deployment, the world-model branch is removed, leaving only the VLM, World Adapter, and action expert, with no online video-model inference. With a 2B backbone and no embodied action pretraining, World Tokens is highly competitive on LIBERO, attains the best reported averages on SIMPLER, substantially improves real-world R1 Pro success over a matched action-only baseline, and generates each action chunk at VLA-level latency.

Problem

• embodied policies
• world modeling
• vision-language-action models
• inference cost
• spatiotemporal evolution

Innovation

• World Tokens
• World Adapter
• embodied policy
• world modeling
• vision-language-action