Trace-WAM: Learning to Act via 3D Motion Prediction for World Action Models

Xu Hu; Changjie Wu; Linzhe Shi; Hang Zhang; Chen Yang; Lingjun Zhang; Guowen Zhang; Mu Xu

cs.ROOct 27, 2026

Trace-WAM: Learning to Act via 3D Motion Prediction for World Action Models

Authors: Xu Hu, Changjie Wu, Linzhe Shi, Hang Zhang, Chen Yang, Lingjun Zhang, Guowen Zhang, Mu Xu,

Organizations: The Hong Kong Polytechnic University · Amap, Alibaba Group · Tsinghua University · University of Science and Technology Beijing ·

Abstract

World Action Models (WAMs) learn representations for robot control by coupling future visual prediction with action generation. Successful manipulation depends primarily on how the robot and task-relevant objects evolve, whereas future-video prediction must also reconstruct the appearance and background of the entire scene. We therefore propose Trace-WAM, which augments video-action learning with future 3D motion prediction as an auxiliary objective. This objective supervises temporal point correspondence together with image-plane displacement and depth change, directly describing the interaction dynamics and spatial changes involved in manipulation. To this end, we adopt a three-branch Mixture-of-Transformer (Tri-MoT) architecture to jointly learn video prediction, motion prediction, and action generation. The motion expert predicts future trajectories from current visual features, allowing motion supervision to shape the representations used for action generation. Because the action expert does not consume predicted trajectories, inference retains the current-observation policy path: a single visual forward pass followed by action denoising. Experiments on RoboCasa-GR1, LIBERO, and LIBERO-Plus demonstrate strong in-distribution performance and out-of-distribution generalization, with particularly clear gains over Fast-WAM under limited demonstrations and distribution shifts.

Figures & tables

Figure 2: Depth complements image-plane motion. (a) Depth-change survival curves over 32-step point tracks whose maximum 2D displacement is at least 2% of the image diagonal. Depth change is

\max_{k}|d_{t+k}-d_{t}|/d_{t}

. LIBERO uses simulator depth; Real uses estimated relative depth. (b) Retaining depth in motion supervision improves LIBERO-Plus success ( Tab. 4 ).

Figure 3: Overview of Trace-WAM. Left: three Transformer experts jointly learn future video prediction, 3D motion prediction, and action generation. Right: colored cells allow attention among the illustrated video, motion, and action streams. At inference, future-video and motion prediction are removed, and action denoising uses cached visual keys and values.

Method	P.T.	Few-shot	Full
GR00T N1.6 [ 2 ]	$\checkmark$	–	47.6
StarVLA-OFT [ 9 ]	$\times$	–	48.8
StarVLA- $\alpha$ [ 45 ]	$\times$	–	53.8
ABot-M0 [ 43 ]	$\checkmark$	–	58.3
RLDX-1 [ 15 ]	$\checkmark$	–	58.7
JoyAI-RA [ 50 ]	$\checkmark$	–	63.2

Table 1: In-distribution performance on RoboCasa-GR1. Success rates (%) in the few-shot and full-data settings. P.T. indicates additional robotic pretraining. DIAL is trained without human data; Best and second-best results are shown in bold and underlined, respectively.

Figure 4: Few-shot OOD generalization on RoboCasa-GR1.

Method	P.T.	Task Suite				Avg.
Method	P.T.	Spatial	Object	Goal	Long	Avg.
OpenVLA [ 18 ]	$\checkmark$	84.7	88.4	79.2	53.7	76.5
OpenVLA-OFT [ 16 ]	$\checkmark$	97.6	98.4	97.9	94.5	97.1
UniVLA [ 5 ]	$\checkmark$	96.5	96.8	95.6	92.0	95.2
$\pi_{0}$ [ 3 ]	$\checkmark$	96.8	98.8	95.8	85.2	94.2
$\pi_{0}$ -FAST [ 33 ]	$\checkmark$	96.4	96.8	88.6	60.2	85.5

Table 2: Quantitative comparison on LIBERO [ 27 ] . Success rates (%) across four task suites.

Method	P.T.	Perturbation Type							Avg.
Method	P.T.	Camera	Robot	Language	Light	Background	Noise	Layout	Avg.
OpenVLA [ 18 ]	$\checkmark$	0.8	3.5	23.0	8.1	34.8	15.2	28.5	15.6
OpenVLA-OFT [ 16 ]	$\checkmark$	56.4	31.9	79.5	88.7	93.3	75.8	74.2	69.6
UniVLA [ 5 ]	$\checkmark$	1.8	46.2	69.6	69.0	81.0	21.2	31.9	42.9
$\pi_{0}$ [ 3 ]	$\checkmark$	13.8	6.0	58.8	85.0	81.4	79.0	68.8	53.6
$\pi_{0}$ -FAST [ 33 ]	$\checkmark$	65.1	21.6	61.0	73.2	73.2	74.4	68.8	61.6

Table 3: Zero-shot transfer from LIBERO to LIBERO-Plus [ 11 ] . Success rates (%) across seven perturbation types.

1) Motion Supervision Ablation
Method	Camera	Robot	Language	Light	Background	Noise	Layout	Avg.
w/o Motion FM-split	55.7	84.5	89.2	96.1	70.5	85.1	83.7	80.4
w/o Static-Motion Penalty	60.9	86.9	93.4	97.3	75.3	88.1	85	83.6
2) Motion Conditioning Ablation
Motion Inject	61.7	82.8	93.4	96.3	73.1	81	84.9	81.6
Video–Motion Bid-Attention	54.4	83.5	94	96.2	68.3	82.3	84.8	80.3

Table 4: Ablation Studies on LIBERO-Plus. Success rates (%) across seven perturbation types. (1) Motion supervision ablates moving/static balancing in

\mathcal{L}_{M}

and the static-motion penalty

\mathcal{L}_{\mathrm{zero}}

. (2) Motion conditioning compares auxiliary supervision with Motion Inject and bidirectional video–motion attention. (3) Motion representation compares 2D/3D motion and raw-motion/B-spline targets. Best and second-best results are shown in bold and underlined, respectively.

Motion Representation	8 Steps	16 Steps	32 Steps
Image-Plane ADE (pixels, $\downarrow$ )
Raw Motion (Direct32)	0.92	1.25	1.87
B-spline Motion	0.67	0.95	1.43
Depth MAE (mm, $\downarrow$ )
Raw Motion (Direct32)	1.634	2.297	3.427
B-spline Motion	1.277	1.808	2.768

Table 5: Motion Prediction Accuracy on LIBERO-Plus. We aggregate 140 rollout episodes at each horizon; lower is better.

Figure 5: Real-world Experiments. Left: a Piper arm with fixed-view and wrist-view RealSense D435i cameras. Center: four observed frames from a successful Trace-WAM execution of each task. Right: task success rates over 20 trials per task.

Figure 6: Qualitative future 3D motion predictions on LIBERO-Plus. The overlays visualize B-spline motion predictions over a 32-step horizon.

World Model Learning

World model learning builds predictive models of how environments change under actions to support planning, simulation, and decision-making.

See all 414 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
TrajWAM: An Efficient Trajectory-Guided World Action Model for Dynamic Robotic Manipulation
Oct 19, 2026Haohan Geng, Sha Zhang, Weishuo Sun +7Dynamic Robotic Manipulation World Action Models
University of Science and Technology of China. · The Chinese University of Hong Kong. · University of Macau. +1
De-Diffusing WAMs: Action-Sufficient Direct Prediction for Robot Control
Oct 20, 2026Shaoyi Zheng, Zhongtao Jiang, Yue Peng +4Efficient Inference for World Action Models World Models for VLA Models
ByteDance. · New York University (NYU).

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 Model Learning

Explore similar work

MoWAM: Explicit Future Motion Prediction for Efficient World Action Models

TrajWAM: An Efficient Trajectory-Guided World Action Model for Dynamic Robotic Manipulation

De-Diffusing WAMs: Action-Sufficient Direct Prediction for Robot Control