4D-WAM: Infusing Spatiotemporal Awareness into World Action Models through Trajectory Fields
Lishan Yang, Wenxuan Song, Xi Wang, Pingyue Sheng, Zheng Fang, Ziyang Zhou, Junjie He, Haodong Yan, Jiayi Chen, Nan Sun, Qiao Sun, Pengwei Wang, Lingqiao Liu, Yan Wang, Yuxiang Gao, Feras Dayoub, Haoang Li
cs.RO
2026-08-08
Aligning a frozen Trace Anything 4D trajectory teacher into WAM features during training lifts the LIBERO-Plus perturbation average from 62.21% to 71.01%, with the inference architecture unchanged.
World Action Models (WAMs) train video prediction and robot action generation jointly inside one video-generation backbone. The vision side sees 2D pixels; the robot arm lives in 3D. That gap is real: grasping a cup requires knowing where the cup sits, how high, how far from the gripper, and in pixel space all of that is indirect.
Recent work adds 3D signal on top: depth maps or point clouds as extra input, or depth as an auxiliary reconstruction target. The paper argues neither goes far enough. A depth map is a per-frame 3D snapshot with no explicit representation of how structure evolves over time, which is exactly where manipulation dynamics live. And supervising pixel-level depth reconstruction pulls attention into low-level detail instead of high-level 4D understanding (3D structure plus its temporal evolution). What is missing is a compact dynamic representation and a way to inject it without touching the original training pipeline.
The recipe uses a 4D teacher for representation alignment during training only, with zero inference overhead. The teacher is Trace Anything, which represents every pixel in a video as a parametric 3D trajectory over time, encoding spatial structure and temporal evolution together.
A direct approach fails first. Aligning a student's intermediate features to a 3D foundation model works for VLA models, but transferred to WAMs it consistently hurts. The paper's diagnosis: π0.5's intermediate features stay positively similar to VGGT's 3D features across layers, while a WAM's (Lingbot-VA here) features are uncorrelated or negatively correlated with 4D features in most layers, so hard alignment creates optimization conflict. But when each model's adjacent-frame features are differenced and then compared, similarity jumps. That is the clue: absolute features do not match, temporal change does.
Two losses follow, both applied to the DiT backbone's intermediate representation at layer 20:
The total loss is the original WAM loss (latent video denoising plus action prediction) plus the two alignment terms, with weights recommended below 0.1 so original objectives stay untouched. After training, Trace Anything, the projectors, and the alignment losses are all removed; inference architecture and parameter count are identical to the base model.
In-distribution, with FastWAM-Joint as the base:
| Setting | FastWAM-Joint | 4D-WAM |
| RoboTwin 2.0 Clean (%) | 90.84 | 91.92 |
| RoboTwin 2.0 Random (%) | 90.32 | 90.76 |
| LIBERO 4-suite avg (%) | 98.2 | 98.6 |
LIBERO is near saturation; 0.4 points there is incremental. The real gap shows out of distribution.
On LIBERO-Plus (seven perturbations: camera, robot, language, light, background, noise, layout), the average rises from 62.21% to 71.01%; under camera perturbation alone it goes from 27.89% to 45.15%, a 17.26-point jump. Camera shifts change observation geometry, which lands squarely on spatial understanding, so the gain pattern matches the method's design.
On RoboTwin Clean2Rand (clean training data, domain-randomized evaluation) with Lingbot-VA as base, the randomized setting improves from 34.6% to 41.8%, the only model averaging above 60% across both settings; FastWAM in the same comparison collapses to 1.2%.
Video prediction quality (LIBERO-Plus, 100 episodes): PSNR from 19.22 to 21.44, SSIM from 0.75 to 0.83. Ablations: removing either alignment term drops the LIBERO average to 98.0% (full version 98.6%); layer 20 is the best alignment depth among layers 16 through 24.
Real-robot experiments use an ARX LIFT2 bimanual platform, four long-horizon tasks, 550 demonstrations, 50 trials per task. Base Lingbot-VA scores zero success on three of four tasks; 4D-WAM completes all four occasionally, at 5.5% total success. Not pretty, but three-zeros-versus-sometimes-works are different failure modes.
The appendix probing gives mechanistic evidence: with the backbone frozen and only a linear probe trained, layer-20 features predict token positions in the next or final frame. Destination-prediction Top-1 rises from 0.4150 to 0.4487 on LIBERO, mean rank drops from 4.66 to 4.07; the trend holds on LIBERO-Plus.
Insensitivity to 3D is a recognized weakness of VLA-style robot models. The standard fixes change inputs (depth or point-cloud channels) or outputs (geometry prediction heads), which alters the architecture and adds inference overhead. This paper takes the pure training-time knowledge-injection route: architecture untouched, zero inference cost, applicable to any DiT-based WAM. For anyone building robot foundation models it is a gain that stacks on top of other improvements, especially under perturbation and domain shift.
The core observation travels beyond robotics. "If absolute features cannot be aligned, align their differences" applies to any distillation setting with incompatible representation spaces, and the query-key destination constraint shows how to turn long-range correspondence into a supervisable quantity.
In-distribution gains are incremental; 0.4 points on a saturated LIBERO says little, and the method's real selling point is entirely on the OOD side. Real-robot success sits at 5.5%, low in absolute terms, only better than the base's zeros, and 550 demonstrations may be the bottleneck. Layer 20 was selected on this backbone; whether it transfers to deeper or shallower backbones is not established. Trace Anything's own errors flow straight into the supervision signal, and its trajectory quality outside the training distribution is not discussed. In the Clean2Rand experiments 4D-WAM switches base models (FastWAM to Lingbot-VA) because FastWAM collapses in that setting; reasonable choice, but the two headline results use different bases, worth noting when comparing numbers across tables.