4D-WAM: Infusing Spatiotemporal Awareness into World Action Models through Trajectory Fields

1 Introduction

• Existing 3D approaches add geometric inputs or reconstruction targets but represent frames separately and may emphasize low-level details over high-level 4D understanding.The introduction calls for a compact representation and a scalable training recipe.Source: §1 · p. 1
• 4D-WAM aligns WAM intermediate visual features with 4D features encoded by a foundation model, using 3D trajectory fields that associate pixels across frames with parametric 3D trajectories.These trajectories represent 3D positions over time, coupling spatial structure with temporal evolution.Source: §1 · p. 1
• Motion alignment matches frame-to-frame feature changes, while destination alignment matches source-to-final-frame similarity distributions to guide the model toward manipulation targets.The two objectives provide local motion supervision and target-oriented guidance beyond local temporal changes.Source: §1 · pp. 1–2
• 4D-WAM improves spatial understanding, execution accuracy, robustness, and generalization across in-distribution and OOD evaluations, with particularly large gains under OOD shifts.Its effectiveness is also validated with FastWAM and Lingbot-VA, and video predictions show improved structural consistency.Source: §1 · p. 2, Figure 1 · p. 2
• The paper presents 4D-WAM as model-agnostic and reports improved precision, robustness, and generalization across different WAM backbones.Figure 1 also reports consistent baseline outperformance across settings, particularly under OOD shifts.Source: §1 · p. 2, Figure 1 · p. 2

2.1 World Action Models

• WAM approaches use video-model generative priors with action experts, spanning visual planning followed by action inference, inverse-dynamics methods, joint video-action prediction, and autoregressive prediction.Examples include UniPi, LAPA, Mimic-Video, UWM, DreamZero, Motus, WorldVLA, and Lingbot-VA.Source: §2.1 · p. 2

2.2 Spatiotemporal Understanding for WAMs

• Spatiotemporal approaches introduce 3D scene outputs, point trajectories, action videos, or other 4D priors, often with added modalities or specialized objectives; 4D-WAM instead aligns intermediate features with motion and source-to-destination trajectory cues.Examples include TesserAct, PointWorld, Kinema4D, Action Images, WAM4D, X-WAM, and µ0.Source: §2.2 · pp. 2–3

3.1 Limitations of Typical Representation Alignment

• Directly aligning WAM representations with 4D reconstruction features performs worse, while their temporal feature differences show higher similarity and offer a more suitable alignment signal.The passage contrasts WAM representations, which are uncorrelated or negatively correlated with 4D features across most layers, with the positive similarity observed for VLA and 3D features.Source: §3.1 · p. 3

3.2 Motion Alignment via Token-wise Frame Differences

• Motion Alignment matches temporal variations across consecutive frames rather than absolute features, using token-wise cosine distance between representations projected into a shared space.The method targets consistent motion directions while avoiding direct absolute-feature matching, which the authors say alleviates optimization conflict and preserves the DiT model’s representation capacity.Source: §3.2 · pp. 3–4

3.3 Source-Conditioned Destination Constraint

• Source-conditioned destination alignment uses source and destination token similarity distributions, minimizing their KL divergence to transfer the teacher’s correspondence structure.The implementation uses the first frame as the source and the final frame as the destination, treating the latter as a completion signal for the manipulation trajectory.Source: §3.3 · pp. 4–5

3.4 Overall Training Objective

• The overall objective retains latent reconstruction and action prediction losses, adding motion and destination alignment as auxiliary supervision with balancing coefficients recommended below 0.1.The alignment losses are intended to improve temporal consistency and goal awareness without overly affecting video generation.Source: §3.4 · p. 5

4 Experiments

• The experiments cover LIBERO and RoboTwin 2.0 in-distribution, LIBERO-Plus and RoboTwin Clean2Rand out-of-distribution, plus video prediction and real-world tasks.The setup uses FastWAM-Joint and Lingbot-VA as base models in the reported settings.Source: §4 · p. 5, §4.1 · p. 5, §4.2 · p. 6

4.1 In-distribution Experiments

• 1.08% clean and 0.44% randomized absolute success-rate gains over FastWAM-Joint on RoboTwin 2.0 accompany LIBERO’s 98.2%-to-98.6% average success-rate rise.LIBERO improvements occur on the Object, Goal, and Long suites.Source: §4.1 · p. 5

4.2 Out-of-Distribution (OOD) Experiments

• 8.8 percentage points higher average success rate and wins across all seven LIBERO-Plus perturbation categories are reported for 4D-WAM over FastWAM-Joint.The largest reported gain is 17.26 points under camera perturbations.Source: Table 2 · p. 6, §4.2 · p. 6
• 4D-WAM is the only model with an average success rate above 60% across clean and randomized RoboTwin settings, substantially improving randomized performance over Lingbot-VA.The evaluation uses Lingbot-VA as the base model and trains on the clean split.Source: Table 3 · p. 6, §4.2 · p. 6

4.3 Experiements on Video Prediction

• 21.44 PSNR versus 19.22 and 0.83 SSIM versus 0.75 show higher video-prediction quality than FastWAM-Joint on LIBERO-Plus.PSNR measures pixel-level reconstruction fidelity; SSIM measures structural similarity.Source: Table 4 · p. 6, §4.3 · p. 6
• 4D-WAM better preserves object motion, target regions, and temporal consistency than FastWAM-Joint in visualized future frames, including under occlusion.Examples include consistent tracking of a microwave door despite occlusion and preserving a book’s position through grasping.Source: Figure 4 · p. 6, §4.3 · pp. 6–7

4.4 Ablation Study

• Above 98.5% accuracy is achieved when alignment is at or beyond layer 20; both objectives at layer 20 perform best, while either ablation reaches 98.0%.Shallower-layer alignment yields accuracy below 98%.Source: Table 5 · p. 7, §4.4 · p. 7

4.5 Real-world Experiments

• More than 10% average subtask-progress gains accompany 4D-WAM’s higher performance than Lingbot-VA across four real-world tasks.The tasks use a bimanual manipulator and are evaluated in 50 trials per task.Source: Table 6 · p. 7, §4.5 · p. 7

• 4D-WAM is a model-agnostic alignment strategy that guides WAMs to capture temporal dynamics and identify episode-level target regions from 3D trajectory fields.Source: §5 · p. 7
• Across multiple benchmarks, 4D-WAM achieves competitive manipulation performance with clear gains in in-distribution, challenging OOD, and real-world settings.Source: §5 · p. 7
• Video-prediction analyses show more accurate object motion and stronger spatiotemporal consistency, with reported improvements in downstream performance and generalization.Source: §5 · p. 8

A Auxiliary Training Cost

• Pre-extracting and caching trajectory representations makes trajectory feature extraction the main added training cost, while alignment operations remain lightweight.The online setting extracts trajectory features during each training iteration.Source: A Auxiliary Training Cost · p. 8
• Cached 4D-WAM raises iteration time from 7.3 s to 7.5 s and total training cost from 51 to 52 eight GPU hours versus the baseline.Its peak GPU memory usage increases by 5.4 percentage points.Source: A Auxiliary Training Cost · p. 8
• Online trajectory extraction requires 59 GPU hours and raises peak GPU memory usage to 96.0%, while caching substantially reduces the overhead.The alignment projectors add 56 MiB of trainable parameter storage.Source: A Auxiliary Training Cost · p. 8
• After training, Trace Anything, the alignment projectors, and alignment objectives are removed, leaving the original WAM inference architecture and parameter count.4D-WAM introduces trajectory supervision only during training, with additional trainable parameters coming only from lightweight projectors.Source: A Auxiliary Training Cost · p. 8

B Real-world Task and Dataset

• The ARX LIFT2 platform has two 6-DoF arms with 1-DoF grippers and three RGB cameras positioned between and beside the arms.The cameras capture the scene from complementary viewpoints.Source: B Real-world Task and Dataset · p. 8
• The dataset contains 550 demonstrations: 100 for cylinder insertion, 150 for towel flattening, 200 for wrench retrieval, and 100 for grasping a rotating-turntable light.Two tasks are bimanual and two are single-arm tasks.Source: B Real-world Task and Dataset · p. 8
• The four manipulation tasks are shown as complete sequential processes from left to right.They involve inserting four size-matched cylinders in ascending order, flattening a towel, retrieving a wrench into a tray, and grasping a red light on a running turntable.Source: B Real-world Task and Dataset · p. 8, Figure 7 · p. 9

C Probing Experiments: Mechanism of 4D Awareness

• The frozen WAM backbone’s visual tokens from the 20th DiT block are fed to two independent lightweight MLP probes.The selected layer is the same one used for trajectory field alignment.Source: C Probing Experiments: Mechanism of 4D… · p. 9
• Frame-to-frame probes predict a token’s position in the next frame, while destination probes predict a first-frame token’s position in the final frame.Probe predictions are compared with Trace Anything target distributions using KL divergence, cross-entropy, Top-k accuracy, and average rank.Source: C Probing Experiments: Mechanism of 4D… · p. 9
• On LIBERO destination prediction, Top-1 accuracy rises from 0.4150 to 0.4487 and mean rank falls from 4.66 to 4.07; the same metric pattern holds across both datasets.Across LIBERO and LIBERO-Plus, 4D-WAM also achieves lower KL and CE and higher Top-k accuracy than FastWAM-Joint variants.Source: C Probing Experiments: Mechanism of 4D… · p. 9
• The probing results suggest that 4D-WAM captures local dynamics and precise destinations, capabilities described as essential for coherent rollouts and contributing to downstream performance.Source: C Probing Experiments: Mechanism of 4D… · p. 10