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

cs.ROOct 20, 2026

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

Authors: Zhuohao Li, Jiawei He, Yuxuan Yan, Lang Zhou, Yinghao Li, Yanrui Li, Tianyu Zhang, Yi-Yang Song,

Organizations: 1Sun Yat-sen University · 2Shenzhen Loop Area Institute · 3The Hong Kong University of Science and Technology (Guangzhou) · 4Xi’an Jiaotong University ·

Abstract

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.

Figures & tables

Figure 1: From plausible futures to reliable stage progression. Top (left): Overlapping task trajectories can induce Stage Confusion: WAMs may follow plausible yet incorrect execution routes. Top (right): StageBench systematically evaluates this challenge through controlled overlapping multi-stage trajectories across eight task families and more than 2.2 million demonstration frames. Bottom: Visualization of TaskWAM performing reliable long-horizon manipulation across diverse robot embodiments and stage progressions. Visually similar stages are highlighted with matching color borders to indicate potential stage confusion.

Figure 2: Overall framework. TaskWAM augments a World-Action DiT through explicit task-state grounding and conditioning. (a) Mode Tokens retrieve mode representations from a learnable codebook and ground them with text context to encode fine-grained task requirements. (b) State Tokens aggregate action-observation interactions into a causal interaction history that tracks execution progress. (c) Task-state-conditioned steering conditions the world and action streams through attention and modulation, while residual flow guidance steers generation toward intended futures.

Figure 3: The visualization of task families in StageBench (see Appendix B.1 for more details).

Method	Kitchen Manipulation				Bimanual Manipulation				Avg. Suc.
Method	Cooking & Cheese	Cooking & Juice	Cooking & Cleanup	Cooking & Breakfast	Size-Order Cubes	Color-Order Cups	Arm-Assigned Hammering	Repeated Bread Cutting	Avg. Suc.
GR00T N1.7	43	33	9	28	60	28	55	43	37.4
$\pi_{0.5}$	52	54	43	40	69	46	75	57	54.5
Motus	27	28	27	12	15	10	47	20	23.3
Cosmos Policy	35	53	35	55	18	7	45	13	32.6
w/ TaskWAM	57	61	45	63	29	22	56	31	45.5

Table 1: Performance comparison on diverse long-horizon manipulation tasks of StageBench. Our method outperforms all baselines in Success Rate ( Suc. , %

\uparrow

). Best results are in bold .

Task	Motus		$\pi_{0.5}$		LingBot-VA		LaWAM		AHA-WAM		Fast-WAM		w/ TaskWAM
Task	Clean	Rand.	Clean	Rand.	Clean	Rand.	Clean	Rand.	Clean	Rand.	Clean	Rand.	Clean ( $\Delta$ )	Rand.( $\Delta$ )
Open Laptop	95	91	90	96	92	94	100	100	98	97	98	100	100	100
Place Bread Basket	91	94	77	64	97	95	92	85	89	94	91	93	98	97
Place Bread Skillet	86	83	85	66	95	90	90	83	91	87	90	93	94	98
Place Object Scale	88	85	86	80	96	95	95	88	96	87	90	97	96	99
$\cdots$	$\cdots$

Table 2: Performance comparison on RoboTwin 2.0 Benchmark (Clean vs. Randomized, 50 tasks). We show the representative tasks and full comparison results is in Appendix B.2 .

Method	Mode	State	Residual Flow	Suc. $\uparrow$
Baseline				40.9
+ Mode Tokens	✓			56.4
+ State Tokens	✓	✓		61.6
Full TaskWAM	✓	✓	✓	69.4

Table 3: Component Ablation.

Method	Suc. $\uparrow$	SPR $\uparrow$
FastWAM Baseline	40.9	68.4
w/ TaskWAM	69.4	87.1

Table 4: Stage Progression Rate ( SPR , %

\uparrow

) on StageBench.

Figure 4: Stage-wise completion over long-horizon execution. TaskWAM maintains higher completion across all stage transitions.

Figure 5: Real-world robot platform. (a) Franka Research 3 with third-person and wrist-view RealSense D435 cameras; (b-c) dual-arm AgileX Piper (6-DoF) and Nero (7-DoF) equipped with two wrist cameras and one head-mounted camera, respectively; (c) Unitree G1 humanoid robot equipped with BrainCo Revo2 dexterous hands and operated via PICO VR for whole-body teleoperation.

Task	Setting	Cosmos Policy	FastWAM	TaskWAM	$\bm{\Delta}$
Task	Setting	$S_{1}\rightarrow S_{2}\rightarrow S_{3}$	$S_{1}\rightarrow S_{2}\rightarrow S_{3}$	$S_{1}\rightarrow S_{2}\rightarrow S_{3}$	$\bm{\Delta}$
Sequential Cup Nesting	Single-Arm Franka	45.0 $\rightarrow$ 40.0 $\rightarrow$ 30.0	55.0 $\rightarrow$ 45.0 $\rightarrow$ 35.0	70.0 $\rightarrow$ 65.0 $\rightarrow$ 55.0	+20.0
Bread Cutting	Dexterous Hands	95.0 $\rightarrow$ 80.0 $\rightarrow$ 20.0	100.0 $\rightarrow$ 80.0 $\rightarrow$ 35.0	100.0 $\rightarrow$ 95.0 $\rightarrow$ 70.0	+35.0
Lab. Fluid Manipulation	Bi-Gripper Nero	65.0 $\rightarrow$ 40.0 $\rightarrow$ 40.0	80.0 $\rightarrow$ 55.0 $\rightarrow$ 50.0	90.0 $\rightarrow$ 70.0 $\rightarrow$ 70.0	+20.0
Bread Toasting & Spreading	Bi-Gripper Piper	$S_{1}\rightarrow S_{2}\rightarrow S_{3}\rightarrow S_{4}\rightarrow S_{5}$	$S_{1}\rightarrow S_{2}\rightarrow S_{3}\rightarrow S_{4}\rightarrow S_{5}$	$S_{1}\rightarrow S_{2}\rightarrow S_{3}\rightarrow S_{4}\rightarrow S_{5}$
Bread Toasting & Spreading	Bi-Gripper Piper	90.0 $\rightarrow$ 20.0 $\rightarrow$ 20.0 $\rightarrow$ 0.0 $\rightarrow$ 0.0	95.0 $\rightarrow$ 25.0 $\rightarrow$ 25.0 $\rightarrow$ 10.0 $\rightarrow$ 5.0	95.0 $\rightarrow$ 45.0 $\rightarrow$ 45.0 $\rightarrow$ 30.0 $\rightarrow$ 20.0	+15.0

Table 5: Results of real-world experiments. We report the stage-wise success rate for each task,

S_{k}

denotes cumulative success through stage

k

. Best results are bolded .

Method	Suc. $\uparrow$	SPR. $\uparrow$	Latency (s) $\downarrow$
FastWAM	40.9	68.4	0.190
TaskWAM	69.4	87.1	0.209

Table 6: Performance-efficiency trade-off. TaskWAM achieves significant gains in multi-stage reliability with minimal inference overhead.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Hyperparameter	Setting
Backbone	Wan2.2-TI2V-5B + FastWAM Action DiT
Training GPUs	$16\times$ NVIDIA H100 80GB
Effective batch size	128
Optimizer updates	58,720
Optimizer	AdamW
Learning rate	$1\times 10^{-4}$

Appendix

Table 7: Training configuration on kitchen manipulation tasks of StageBench.

ID	Task family	Execution requirements
Kitchen Manipulation
K1	Cooking & Cheese	Interleaved; Sequential
K2	Cooking & Juice	Interleaved; Sequential
K3	Cooking & Cleanup	Interleaved; Sequential
K4	Cooking & Breakfast	Interleaved; Sequential
Bimanual Manipulation

Appendix

Table 8: Task families and execution requirements in StageBench.

Figure 6: StageBench benchmark visualization. Representative task variants and their multi-stage execution processes are illustrated, highlighting diverse stage structures and instruction- dependent task requirements.

Figure 7: Execution orders in a LIBERO household task. Both variants share the mug-placement and heating prefix. The Interleaved route collects the bowl and butter during heating and then reopens the microwave. The Sequential route waits for heating to finish and reopens the microwave before collecting the objects.

Figure 8: Instruction-dependent cup placement in RoboTwin. From a shared initial state, the robot places the red (R), green (G), and blue (B) cups in four prescribed orders: RGB, RBG, BGR, and GRB. Stage labels identify cup-specific operations, whose execution order varies with the instruction.

Figure 9: Visualization of Sequential Cup Nesting.

Figure 10: Visualization of Bread Toasting & Spreading.

Figure 11: Visualization of Laboratory Fluid Manipulation.

Figure 12: Visualization of Bread Cutting.

Task	Motus		$\pi_{0.5}$		LingBot-VA		LaWAM		AHA-WAM		Fast-WAM		w/ TaskWAM (ours)
Task	Clean	Rand.	Clean	Rand.	Clean	Rand.	Clean	Rand.	Clean	Rand.	Clean	Rand.	Clean / $\Delta$	Rand. / $\Delta$
Adjust Bottle	89	93	100	99	90	94	100	100	100	100	100	100	100	100
Beat Block Hammer	95	88	96	93	96	98	90	93	100	94	99	97	100	98
Blocks Ranking RGB	99	97	92	85	99	98	97	100	100	98	100	100	100	99
Blocks Ranking Size	75	63	49	26	94	96	93	89	98	98	94	98	93	93
Click Alarmclock	100	100	98	89	99	100	100	100	100	100	100	100	100	100

Appendix

Table 9: Complete RoboTwin 2.0 task-wise comparison.

Setting	Suc. $\uparrow$
TaskWAM	69.4
w/ Shuffled Temporal Order	50.3
w/ Repeated Latest State	58.1

Appendix

Table 10: Analysis of State Tokens.

Figure 13: Training convergence. TaskWAM converges faster and reaches a lower action loss than FastWAM under the same training setup.

Figure 14: Real-world stage confusion in Bread Toasting & Spreading. Stage 1 (loading bread) and Stage 4 (retrieving toasted bread) exhibit similar visual configurations, while Stage 3 introduces a long waiting period with redundant observations. Baseline WAMs may confuse the current progress with earlier stages, resulting in repeated or incorrect execution.

World Models for Robotics

World models for robotics predict how scenes and robot states change under actions to support planning, control, and learning.

See all 380 papers

StageWAM: Joint-Embedding Stage Prediction for World-Action Models in Robot Manipulation
Aug 11, 2026Xiao Liu, Yuguang Yang, Xi Wang +6World Models for Robotics Joint-Embedding Predictive Architecture
Institute for AI Industry Research (AIR), Tsinghua University · School of Electronic Information Engineering, Beihang University · AIR Wuxi Innovation Center, Tsinghua University +3
What Limits Cross-Task Generalization in World Action Models?
Oct 12, 2026Chen Zhao, Zipeng Ye, Yihan Lin +2World Action Models World Models for Robotics
Renmin University of China, Beijing 100872, China · Key Laboratory of Data Engineering and KnowledgeEngineering, MOE, China · Meituan, Beijing 100102, China
When to Trust Imagination: Adaptive Action Execution for World Action Models
May 7, 2026Rui Wang, Yue Zhang, Canyang Chen +3Long-Horizon Robotic Manipulation World Models for Robotics
Southern University of Science and Technology, Shenzhen, China · The University of Hong Kong, Hong Kong, China

Aug 11, 2026J/K move · Enter open · S save

StageWAM: Joint-Embedding Stage Prediction for World-Action Models in Robot Manipulation

Xiao Liu, Yuguang Yang, Xi Wang, Kai Jiang+5

Institute for AI Industry Research (AIR), Tsinghua University · School of Electronic Information Engineering, Beihang University · AIR Wuxi Innovation Center, Tsinghua University+3

World Models for Robotics Joint-Embedding Predictive Architecture Robotic Manipulation

Open paper AskPDF arXiv

Generalist robot policies aim to map multimodal observations and linguistic task instructions to actions across diverse tasks. However, existing methods typically represent the future as a fixed, short video-action chunk. This short-term future captures local scene evolution for action execution, but it does not explicitly describe the stage-level future that specifies how a task should progress from its current stage to the next. We therefore distinguish two complementary futures for robot manipulation: a short-term physical future to capture local scene evolution and a stage-level semantic future to represent task progress. We introduce StageWAM, which augments a Motus-based World Action Model (WAM) with Stage-JEPA, a goal-conditioned Joint-Embedding Predictive Architecture (JEPA) predictor. Given the current observation and task instruction, Stage-JEPA uses a frozen V-JEPA2 encoder to extract the current-state representation and predicts the latent target of the next inferred stage. Across 50 RoboTwin 2.0 tasks in clean and randomized environments, StageWAM achieves 90.25% overall success and reduces the mean number of execution steps in successful rollouts by 5.97% relative to the strongest baseline.

Abstract

World Models for Robotics

Explore similar work

StageWAM: Joint-Embedding Stage Prediction for World-Action Models in Robot Manipulation

What Limits Cross-Task Generalization in World Action Models?

When to Trust Imagination: Adaptive Action Execution for World Action Models