WAM4D: Fast 4D World Action Model via Spatial Register Tokens

Abstract: World action models (WAMs) have recently shown promise in jointly modeling future observations and executable robot actions. However, most existing WAMs still operate in 2D video or latent spaces, where visually plausible rollouts miss the 3D spatial constraints and occluded contact geometry required for precise manipulation. While geometric foundation models offer strong priors for recovering dense 3D structure and motion from visual observations, forcing WAMs to predict the dense 4D representation introduces costly geometric decoding and slows down causal action generation. To address the trade-off, we present WAM4D, a fast 4D world action model that uses lightweight spatial register tokens as training-time future-depth readouts to transfer pretrained geometric priors into a causal video-action transformer, then removes the register branch for lightweight action inference. To prevent non-causal shortcuts, we further design causal mixture attention for the Mixture-of-Transformers (MoT) WAM backbone, defining modality-specific visibility among video, action, and geometry tokens. Comprehensive experiments on RoboTwin 2.0 and challenging real-world manipulation tasks show that WAM4D improves spatial consistency and achieves competitive action prediction while maintaining efficient inference.

• The paper introduces WAM4D, which distills future-depth supervision into spatial register tokens during training so geometric reasoning improves causal video-action features without adding a geometry decoder at deployment.
• WAM4D achieves 93.8% clean and 89.9% randomized success on RoboTwin 2.0, plus 0.90 average sub-action success on four real-world AstriBot S1 tasks, while using 9.71 GiB of memory and 525 ms chunk latency.
• Ablations show that middle-layer, unidirectional registers and a trainable pretrained Depth Anything 3 head provide the strongest control and geometry gains, while random depth-head initialization underperforms the no-depth baseline.

WAM4D addresses a specific inefficiency in 4D world action modeling: existing approaches such as TesserAct, Kinema4D, and X-WAM treat dense geometry as an explicit inference-time output, which forces costly geometric decoding during action generation and risks shifting the training objective toward geometric reconstruction rather than control. WAM4D instead uses future depth purely as a training-time supervision signal, distilled through lightweight spatial register tokens into a causal video-action transformer, and removes the entire geometry path at deployment (2606.14048).

Motivation and positioning

The paper argues that most world action models (WAMs) operate in 2D video or latent spaces, where visually plausible rollouts conceal errors in contact geometry, occluded surfaces, and object extent that matter for precise manipulation. Prior 4D embodied world models inject depth, normals, point maps, or point clouds as explicit inputs or outputs; this improves physical fidelity but adds dense decoding at inference and does not guarantee that geometric priors strengthen the causal coupling between video prediction and action generation. WAM4D's design principle is that geometry should shape the history video features used for action prediction, without ever appearing on the inference path.

Architecture

WAM4D builds on the causal video-action backbone of LingBot-VA, initialized from its pretrained weights with the Wan2.2 video VAE. At decision step $t$, the model conditions on a language instruction, multi-view RGB history (a three-view mosaic of one head camera and two wrist cameras), and historical actions, jointly denoising future video latents and a 32-step action chunk via flow matching.

Spatial register distillation: learnable register tokens are replicated over future depth timesteps and aligned to mosaic pixel positions (each register covers a $32\times32$ input cell; the three-view model uses 960 registers). At selected transformer layers (default: 12, 14, 16, 18), registers act as queries attending to themselves and valid history video tokens through depth extraction blocks with RoPE encoding. Their features are projected into a pretrained geometric head — a DA3-GIANT-1.1 any-view DualDPT head from Depth Anything 3 — which decodes future depth supervised by SmoothL1 loss. The depth loss backpropagates into the shared video-action backbone, distilling geometric priors into the causal video features that feed action prediction.

Causal mixture attention: the Mixture-of-Transformers visibility mask is designed to prevent non-causal shortcuts. Future action tokens may attend only to history video, history actions, and their own noised future-action tokens; they are masked from future video tokens and spatial registers. Registers attend only to themselves and history video. This keeps the auxiliary geometry branch strictly off the policy path while still allowing it to regularize shared features.

At deployment, registers, depth blocks, and the geometric head are removed entirely, leaving a pure observation-to-action generation path with KV-cache-based autoregressive inference.

Main results

On the full RoboTwin 2.0 suite (50 tasks), WAM4D achieves 93.8% clean / 89.9% randomized average success, competitive with LingBot-VA (92.9/91.6) and Fast-WAM (91.9/91.8), and well above $\pi_0$ (65.9/58.4) and $\pi_{0.5}$ (82.7/76.8). Inference latency is 525 ms per chunk with 9.71 GiB peak memory — slower than VLAs ($\pi_0$: 64 ms) but faster than LingBot-VA (844 ms) and Motus (1516 ms), and with lower VRAM than all other WAM baselines.

Real-world evaluation on AstriBot S1 across four contact-rich tasks yields an average sub-action success of 0.90, versus 0.84 for LingBot-VA, 0.80 for Fast-WAM, and 0.74 for $\pi_{0.5}$. The largest gains appear in long-horizon LEGO sorting (0.8 vs. 0.4–0.5 for baselines on the third sub-step), supporting the claim that geometry distillation benefits precision over extended horizons.

Ablations

Ablations on a ten-task RoboTwin split isolate three design choices:

• Depth readout interface: attaching a depth head directly to future VAE latents improves RGB metrics but yields weak geometry (AbsRel 0.081, F-score 0.519), because RGB VAE latents are not suited to metric depth decoding. Spatial registers querying only history tokens give substantially better geometry (middle-layer: AbsRel 0.053, F-score 0.685).
• Register placement: shallow layers (2/4/6/8) yield the best RGB synthesis (FVD 168.8) but weaker control; middle layers (12/14/16/18) give the best unidirectional control (75.2% clean SR) and best geometry quality. A bidirectional variant reaches the highest clean SR (76.6%) but degrades most geometry metrics and adds computation, so unidirectional middle placement is chosen.
• Geometric head initialization: this is the most striking ablation result. A trainable pretrained DA3 head achieves 80.1% clean SR with the best video and geometry metrics across the board, whereas random initialization drops below even the no-depth baseline (70.0% vs. 71.7%), and a fixed pretrained head lands in between (75.2%). The authors conclude that depth supervision alone is insufficient — the pretrained geometric prior is essential, and light adaptation improves on freezing it.

Attention visualizations show registers attending consistently to objects across views, static background regions, and the initial gripper pose, indicating interpretable geometric query behavior. A qualitative mode retaining the depth branch enables autoregressive RGB-D rollout back-projected into point clouds from a single initial frame.

Limitations

The paper is candid about several constraints. First, WAMs remain slower than direct VLAs at inference (525 ms vs. 64 ms for $\pi_0$), which the authors explicitly flag as unresolved. Second, the model lacks explicit long-term object memory: in long autoregressive rollouts, occluded objects can be completed as visually plausible but incorrect identities. This does not affect closed-loop policy success because fresh observations are continuously ingested, but it limits open-loop generated futures. Third, real-world depth supervision relies on offline pseudo-depth from Depth Anything 3 rather than sensor ground truth, so the reported geometry gains inherit the teacher's biases. Fourth, the bidirectional register variant outperforming the default on clean success rate suggests the optimal information flow between geometry and policy streams is not fully settled.

Conclusion

WAM4D demonstrates that geometric foundation priors can be transferred into causal video-action models through a compact spatial-register interface used only at training time, yielding improved spatial consistency and manipulation success at deployment cost comparable to a 2D WAM. The central empirical finding — that a trainable pretrained geometric teacher is necessary and sufficient for these gains, while random-init depth heads underperform no supervision at all — clarifies when geometry distillation helps. Open questions include closing the latency gap with VLAs, persistent object memory for long rollouts, and whether the bidirectional variant's higher clean success rate can be retained without its geometry-quality degradation.