Driving · 2026-10-21

Plausible Future, Wrong Choice: Steering World-Action Models toward Intended Future

Zhuohao Li, Jiawei He, Yuxuan Yan, Lang Zhou, Yinghao Li, Yanrui Li, Tianyu Zhang, Yi-Yang Song, Xirui Wu andYi-Lin Wei, Fating Hong, Qing Zhang, Zhilin Zhao, Jingdong Wang, Wei-Shi Zheng

arXiv:2610.28431 PDF

AI-assisted summary based on the paper. Consult the linked source for methods and limitations. Benchmark numbers appear after editor verification.

TL;DR

The authors identify “Stage Confusion” in world-action models (WAMs): locally plausible rollouts can advance to the wrong task stage when similar observation-action segments occur across tasks or stages. They propose TaskWAM, a closed-loop, stage-aware framework using Mode Tokens (task requirements) and State Tokens (execution progress) to steer world-action generation toward intended future. They also introduce StageBench for controlled multi-stage evaluation.

Why it matters

For long-horizon embodied manipulation, success requires stage-correct progression, not just plausible futures. TaskWAM aims to reduce cases where WAMs follow plausible but misdirected rollouts, and provides a benchmark (StageBench) that systematically creates overlapping trajectories to test stage-aware progression.

Method

Proposes TaskWAM: Structured Task-State Grounding with Mode Tokens (instruction-grounded fine-grained requirements via a mode codebook and text grounding) and State Tokens (causal encoding of action-observation transitions for execution progress).
Uses Task-State Conditioned Future Steering: injects the structured task state into world/action generation via attention conditioning (including pooled Mode Tokens with AdaLN modulation) and residual flow correction to steer toward intended stage progression.
Introduces StageBench: a controlled multi-stage manipulation benchmark that constructs overlapping observation-action trajectories across task variants by varying stage compositions/orders, and provides automated stage-aware evaluation using event reconstruction.

Limitation

The resulting progress representation depends on evidence available within a finite window: State Tokens summarize the most recent K interactions.

Abstract (from arXiv)

World-Action Models (WAMs) have emerged as a promising paradigm for embodied generalist policies by jointly modeling future world evolution and robot actions across diverse manipulation tasks. However, representing diverse futures is only the first step: a capable WAM must also generate futures that match the requirements of the current task and execution progress. In long-horizon manipulation where similar observation-action segments occur across different tasks and stages, existing approaches may generate locally plausible yet misdirected rollouts to drift toward wrong task stage, i.e., a phenomenon we term Stage Confusion. To address this problem, we propose TaskWAM, a closed-loop, stage-aware framework that represents fine-grained task requirements with Mode Tokens and tracks execution progress with State Tokens, forming a structured task state to steer world--action generation toward intended future. We further introduce StageBench, a controlled multi-stage manipulation benchmark with parameterized stage variations and automated evaluation.Extensive experiments on both simulation benchmarks and real-world tasks demonstrate that TaskWAM effectively reduces confusion and achieves competitive performance with state-of-the-art methods. Moreover, TaskWAM provides strong plug-and-play transfer across heterogeneous WAM architectures and robot embodiments, including single-arm robots, dual-arm grippers, dexterous hands, and humanoid robots.

Publication record

Bibliographic evidence describes where and when a record appeared. It does not validate scientific claims.

First public release: 2026-10-20 · Indexed by Pub-AI: 2026-10-21 UTC

preprint