TaskWAM
STAGE-AWARE WORLD–ACTION MODELS

TaskWAM

Plausible Future,
Wrong Choice.

Steering World-Action Models
toward Intended Future

Grounding generation in task requirements
and execution progress.

Explore the robots
Zhuohao Li1,2Jiawei He3Yuxuan Yan4,2Lang Zhou1,2Yinghao Li1,2
Yanrui Li5Tianyu Zhang6,2Yi-Yang Song1,7Xirui-Wu8Yi-Lin Wei1,9Fating Hong10
Qing Zhang1Zhilin Zhao1,2Jingdong WangWei-Shi Zheng1,2,7,†
  1. 1 Sun Yat-sen University
  2. 2 Shenzhen Loop Area Institute
  3. 3 The Hong Kong University of Science and Technology (Guangzhou)
  4. 4 Xi'an Jiaotong University
  5. 5 Southern University of Science and Technology
  6. 6 Beijing Institute of Technology
  7. 7 Peng Cheng Laboratory
  8. 8 The Chinese University of Hong Kong Shenzhen
  9. 9 Ainoro
  10. 10 King Abdullah University of Science and Technology

lizhh268@mail2.sysu.edu.cn† Corresponding authors

Sun Yat-sen University
Shenzhen Loop Area Institute
The Hong Kong University of Science and Technology (Guangzhou)
Xi’an Jiaotong University
Southern University of Science and TechnologySouthern University
of Science and Technology
Beijing Institute of TechnologyBeijing Institute
of Technology
Peng Cheng LaboratoryPeng Cheng
Laboratory
The Chinese University of Hong Kong Shenzhen
King Abdullah University of Science and TechnologyKAUST
01 / TASKWAM IN THE PHYSICAL WORLD

REAL-WORLD
DEMONSTRATIONS

Four embodiments.
A shared task-state framework.Single arm · Dual arms · Dexterous hands

01 AGILEX PIPER

Bread toasting
& spreading

From loading and heating to retrieval and spreading, follow the intended sequence through a multi-stage breakfast routine.

Embodiment
Dual 6-DoF arms
Task requirement
Timing & stage progression
2.5×
A CLOSER LOOK

Retrieval & spreading

Coordinated bimanual execution during the final stages.

4×
02 FRANKA RESEARCH 3

Sequential
cup nesting

Arrange and nest colored cups in the instructed order, while tracking which steps have already been completed.

Embodiment
Single arm · Parallel gripper
Task requirement
Instruction-defined ordering
2×
03 AGILEX NERO

Laboratory
fluid manipulation

Coordinate two arms to handle laboratory containers and pour fluid into the receiving dish.

Embodiment
Dual 7-DoF arms
Task requirement
Coordinated bimanual execution
2×
04 UNITREE G1

Repeated
bread cutting

Carry out repeated cuts with dexterous hands while tracking execution progress.

Embodiment
Humanoid · Revo2 hands
Task requirement
Repetition & execution progress
02 / THE CHALLENGE

What should happen next?

A plausible continuation can still miss the task.
Both the instruction and the execution history matter.

Same scene. Different instructions.

The timing of object collection changes the required sequence of actions.

INTERLEAVED EXECUTION

Collect during heating

Object collection happens while the microwave is heating.

SEQUENTIAL EXECUTION

Collect after heating

Object collection follows the heating and reopening steps.

03 / THE FRAMEWORK

A structured task state
for world and action.

A plausible future can still follow the wrong task or execution stage. TaskWAM grounds world–action generation in fine-grained task requirements and execution progress. Mode Tokens represent the instructed requirements, while State Tokens track progress through closed-loop interaction. Together, they steer generation toward the intended future.

TaskWAM framework: Mode Tokens ground task requirements, State Tokens track closed-loop execution progress, and structured task state conditions world and action generation.
04 / OUR BENCHMARK

StageBench

We introduce StageBench to evaluate stage confusion arising from task requirements and execution progress in multi-stage manipulation.

StageBench task families, each shown from wrist and scene views: Cooking and Cheese, Juice, Cleanup, Breakfast; Size-Order Cubes; Color-Order Cups; Arm-Assigned Hammering; Repeated Bread Cutting.
StageBench spans kitchen and bimanual manipulation, with instruction variants that change the required execution.
TASK REQUIREMENTS

Same scene, different instructions.

Vary order, timing, arm assignment, and repetition within matched scenes.

EXECUTION PROGRESS

Similar observations, different stages.

Distinguish repeated interaction stages through the preceding execution history.

EXPLORE A TASK

Color-Order Cups

Place the colored cups in the instructed sequence. Choose an order to see how the required execution changes.

Recorded task examples

05 / RESULTS & ANALYSIS

Better completion.
More reliable progression.

Paper-reported results on StageBench.
Average performance across eight task families.

SUCCESS RATE · FASTWAM
40.9→69.4%

+28.5 percentage points with TaskWAM

STAGE PROGRESSION · FASTWAM
68.4→87.1%

Correct next-stage completion

SUCCESS RATE · COSMOS POLICY
32.6→45.5%

+12.9 percentage points with TaskWAM

TaskWAM framework

Enlarged TaskWAM framework

StageBench task families

Full-resolution StageBench overview showing kitchen and bimanual task families from two camera views.