ReWorld: Learning Better Representations for World Action Models

2026-06-25

AI Summary

Existing World Action Models (WAMs) supervise only the final output, resulting in intermediate representations that lack environmental knowledge and thus limit planning performance. This work proposes ReWorld—the first representation learning framework for WAMs tailored to autonomous driving—which uniquely treats intermediate representations as a direct optimization objective. ReWorld jointly shapes these representations through three complementary objectives: predictive accuracy, cross-modal alignment, and discriminative safety-boundary awareness. Built upon Video DiT and Action DiT architectures, the framework incorporates future prediction supervision, cross-modal alignment losses, and hard negative mining. Evaluated on nuScenes and NAVSIM, ReWorld achieves a 23.9% improvement in FVD and a closed-loop PDMS score of 90.4, while also converging approximately twice as fast from scratch compared to prior approaches.

Abstract

World Action Models (WAMs) model future environment evolution under action conditioning, offering a scalable paradigm for autonomous driving. However, existing approaches focus largely on model architecture design, and how a WAM can efficiently learn better world representations for planning remains underexplored. To address this gap, we propose ReWorld, the first representation learning framework specifically designed for autonomous-driving world action models. In WAMs, standard training supervises only the output ends of the generation and planning modules, leaving the intermediate representations that carry world knowledge to be shaped only indirectly, as byproducts of fitting these outputs. The core idea of ReWorld is to treat intermediate representations as direct targets of optimization, shaping them along three complementary dimensions. On the Video DiT responsible for generation, we impose future-predictive supervision on its intermediate representations. On the Action DiT responsible for planning, we first align its intermediate representations cross-modally with the video world representation, then further shape them to be discriminative around safety-critical boundaries via hard-negative supervision. In addition, we systematically analyze the effectiveness of existing representation learning methods in video generation world models, and discuss why their performance is limited on this task. Experiments on nuScenes and NAVSIM show that ReWorld improves fine-tuned video generation by 23.9% in FVD (81.3 to 61.9), raises closed-loop PDMS from 89.1 to 90.4 without any post-training such as RL or post-processing, and accelerates from-scratch convergence by approximately 2x.

Problem

• World Action Models
• representation learning
• autonomous driving
• intermediate representations
• planning

Innovation

• representation learning
• world action models
• intermediate representation optimization
• cross-modal alignment
• hard-negative supervision