MT-WAM: Reorienting the One-Pass Predictive Representation Toward Action Generation

Yiguang Yang; Jiankun Peng; Xiaoming Wang; Yiran Zhang; Zhibo Fang

cs.CVSep 18, 2026

MT-WAM: Reorienting the One-Pass Predictive Representation Toward Action Generation

Authors: Yiguang Yang, Jiankun Peng, Xiaoming Wang, Yiran Zhang, Zhibo Fang

Organizations: Aerospace Information Research Institute, Chinese Academy of Sciences · University of Chinese Academy of Sciences

Abstract

Fast-WAM shows that video-action co-training improves control without generating future video at inference, making the representation from a single video diffusion Transformer forward central to action generation. However, future-observation prediction does not explicitly prioritize the future dynamics and visual structure needed for control. We present MT-WAM, which retains the original training objectives and adds complementary supervision for future two-dimensional point trajectories and visual features. A lightweight dual-stream branch copied from the video backbone's final blocks provides target-specific processing, while a structured attention mask prevents cross-stream attention. Motion-stream tokens supply additional dynamics conditions to the action expert. Future visual-feature prediction provides supervision in a feature space that captures object and spatial structure. This supervision trains the video backbone to provide more informative visual context for action generation under changing visual conditions, without adding visual-feature-stream tokens to action conditioning. At inference, MT-WAM uses video and motion caches computed once per replan and skips future-video prediction. Without additional embodied policy pretraining, MT-WAM achieves 98.2% success on LIBERO and 73.7% on LIBERO-Plus, exceeding Fast-WAM by 23.8 percentage points on the latter. On RoboTwin 2.0 Clean2Rand, Random success increases from 6.30% to 19.40%; across four real-world tasks, average success increases from 67.0% to 77.8%.

Figures & tables

Fig. 1: MT-WAM overview. MT-WAM augments future-observation prediction with 2D point-trajectory and future visual-feature targets through a dual-stream dynamic branch. Motion-stream tokens supply additional dynamics conditions, while future visual-feature supervision shapes the scene context that the video backbone provides to the action expert.

Fig. 2: MT-WAM training, architecture, and inference. (a) Training uses future-observation, action, 2D point-trajectory, and future visual-feature objectives. (b) The video path retains future-observation prediction, the action expert generates actions, and an $M$ -block dynamic branch predicts the two added targets. (c) The streams use separate FFNs and a structured attention mask; action queries attend to motion-stream tokens but not visual-feature-stream tokens. (d) The video backbone and dynamic branch populate the video and motion KV caches once per replan, and the action expert reuses both throughout denoising. The future-video, motion, and visual-feature heads are training-only; Fig. 3 shows the complete masks.

Fig. 3: Structured attention masks used by MT-WAM. Rows are queries and columns are keys; colored cells denote allowed attention.

f_{0}^{\pi}

and

f_{0}^{\beta}

are the reference copies available to action queries and the two branch streams;

f_{1:T}

a_{1:H}

m

, and

s

denote the future video, action, motion, and visual-feature groups. Inference omits

f_{1:T}

Fig. 4: Real-world tasks and success rates (%). Representative initial, key-interaction, and completion frames for T1–T4 are aligned with task-level SRs. All methods use the same 400 demonstrations and 10-epoch training budget, and are evaluated over the same initial configurations with 100 trials per task; Average is the arithmetic mean over T1–T4.

Method	PT	Spatial	Object	Goal	Long	Average
OpenVLA [ 13 ]	✓	84.7	88.4	79.2	53.7	76.5
$\pi_{0}$ [ 16 ]	✓	96.8	98.8	95.8	85.2	94.1
$\pi_{0.5}$ [ 36 ]	✓	98.8	98.2	98.0	92.4	96.9
LingBot-VA [ 37 ]	✓	98.5	99.6	97.2	98.5	98.5
Motus [ 38 ]	✓	96.8	99.8	96.6	97.6	97.7
Fast-WAM [ 1 ]	✗	98.2	100.0	97.0	95.2	97.6

TABLE I: Success rates (%) on LIBERO

Method	PT	Cam.	Rob.	Lang.	Light	Bg.	Noise	Lay.	Ovr.
UniVLA [ 39 ]	✓	1.80	46.20	69.60	69.00	81.00	21.20	31.90	42.87
OpenVLA-OFT [ 40 ]	✓	56.40	31.90	79.50	88.70	93.30	75.80	74.20	69.59
$\pi_{0}$ [ 16 ]	✓	13.80	6.00	58.80	85.00	81.40	79.00	68.90	53.63
$\pi_{0}$ -FAST [ 41 ]	✓	65.10	21.60	61.00	73.20	73.20	74.40	68.80	61.59
WorldVLA [ 42 ]	✓	0.10	27.90	41.60	43.70	17.10	10.90	38.00	25.03
Fast-WAM [ 1 ]	✗	16.26	44.13	66.82	79.77	52.60	38.54	61.38	49.86

TABLE II: Success rates (%) on LIBERO-Plus

Method	PT	Clean	Rand	Average
$\pi_{0}$ [ 16 ]	✓	65.92	58.40	62.16
$\pi_{0.5}$ [ 36 ]	✓	82.74	76.76	79.75
ABot-M0 [ 46 ]	✗	81.20	80.40	80.80
Motus [ 38 ]	✓	88.66	87.02	87.84
Motus from Wan 2.2 [ 38 ]	✗	77.56	77.00	77.28
LingBot-VA [ 37 ]	✓	92.90	91.50	92.20

TABLE III: Success rates (%) on RoboTwin 2.0 Clean/Rand

Method	PT	Clean	Random	Absolute drop (pp)	Random retention (%)
$\pi_{0}$ [ 16 ]	✓	46.40	16.30	30.10	35.1
RDT [ 47 ]	✓	34.50	13.70	20.80	39.7
UP-VLA [ 48 ]	✓	52.90	15.20	37.70	28.7
Diffusion Policy (DP) [ 15 ]	✗	28.00	0.60	27.40	2.1
HALO [ 49 ]	✓	80.50	26.40	54.10	32.8
BagelVLA [ 50 ]	✓	75.30	20.50	54.80	27.2

TABLE IV: Success rates (%) on RoboTwin 2.0 Clean2Rand

Method	Future video	Overall SR (%)	Latency (ms)	TFLOPs
Fast-WAM [ 1 ]	✗	49.86	296.325	3.736949
Fast-WAM-Joint [ 1 ]	✓	68.41	503.729	30.672534
MT-WAM	✗	73.66	313.480	5.142984

TABLE V: Inference efficiency on LIBERO-Plus using one NVIDIA RTX PRO 6000 96 GB GPU

Fig. 5: LIBERO-Plus ablation effects relative to MT-WAM . Each cell reports the change in SR in percentage points; the Overall SR entries additionally report the absolute SR. MT-WAM achieves an Overall SR of 73.7%.

Fig. 6: Success rates (%) across branch depths on LIBERO-Plus .

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Configuration	LIBERO	RoboTwin 2.0
		Clean/Rand	C2R
Camera views	2	3	3
Camera layout	Horizontal	Wrist- horizontal + vertical	Wrist- horizontal + vertical
Input resolution	$224\times 448$	$384\times 320$	$384\times 320$
Observation window	33	33	33
Action chunk length	32	32	32

Appendix

TABLE VI: Dataset-specific training configurations

Protocol	Time	Batch size
LIBERO	16 h	96
RoboTwin 2.0	4 days	48
RoboTwin 2.0 C2R Stage 1	4 days	48
RoboTwin 2.0 C2R Stage 2	18 h	48
Real-world multi-task	17 h	48

Appendix

TABLE VII: Training cost and global batch size

World Action Models

World Action Models connect predictions of how scenes change under action to executable policies for robots, vehicles, and other embodied agents.

See all 210 papers

MoWAM: Explicit Future Motion Prediction for Efficient World Action Models
Sep 17, 2026Jiayu Wang, Bin Zhu, Yue Yu +1World Models for Robotics World Model-Based Planning
College of Computer Science and Artificial Intelligence, Fudan University, Shanghai, China · Singapore Management University, Singapore · Institute of Trustworthy Embodied AI, Fudan University, Shanghai, China
Fast-WAM: Do World Action Models Need Test-time Future Imagination?
Mar 17, 2026Tianyuan Yuan, Zibin Dong, Yicheng Liu +1World Action Models Efficient Inference for World Action Models
IIIS, Tsinghua University · Galaxea AI
Foresight Without Seeing: Latent Futures for World Action Models
Aug 12, 2026Jiakai Huang, Zhongbo Wu, Siyu Xu +5World Models for Robotics Visuomotor Policy Learning
Shanghai Jiao Tong University · ACE Robotics · Nanyang Technological University

Sep 17, 2026J/K move · Enter open · S save

MoWAM: Explicit Future Motion Prediction for Efficient World Action Models

Jiayu Wang, Bin Zhu, Yue Yu, Jingjing Chen

College of Computer Science and Artificial Intelligence, Fudan University, Shanghai, China · Singapore Management University, Singapore · Institute of Trustworthy Embodied AI, Fudan University, Shanghai, China

World Models for Robotics World Model-Based Planning Efficient Inference for World Action Models Robot Motion Prediction Robotic Manipulation World Action Models+2

Open paper AskPDF arXiv

World Action Models (WAMs) improve robot policy learning by incorporating future dynamics, yet explicitly generating future videos at inference introduces substantial computational overhead. Removing future generation improves efficiency, but leaves future dynamics only implicitly encoded in observation features, which can limit robustness under distribution shifts. We propose MoWAM, an efficient WAM that replaces future video generation with explicit future motion prediction. Instead of reconstructing the complete future scene, MoWAM models structured robot motion as a compact abstraction of the future, capturing how the robot is expected to evolve under the current scene and interaction constraints. A Mixture-of-Transformer architecture learns future visual dynamics during training while jointly predicting motion and action, allowing video generation to be removed entirely at inference while retaining an explicit representation of the future. The compact motion representation further enables efficient inference-time scaling by sampling multiple candidates of motion and action pairs and selecting among them with a motion-aware task-progress verifier. Experiments on LIBERO, LIBERO-Plus, and real-world manipulation tasks demonstrate that MoWAM achieves strong in-distribution performance, improved out-of-distribution robustness, and higher average real-world success than representative WAM baselines. In addition, performance improves as more candidates are explored, demonstrating that explicit future motion provides an effective and efficient basis for inference-time scaling.

Abstract

World Action Models

Explore similar work

MoWAM: Explicit Future Motion Prediction for Efficient World Action Models

Fast-WAM: Do World Action Models Need Test-time Future Imagination?

Foresight Without Seeing: Latent Futures for World Action Models