PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Ying Li, Xiaobao Wei, Jiajun Cao, Hao Wang, Xiaowei Chi, Chengyu Bai, Qianpu Sun, Jiajun Li, Xiaojie Zhang, Peidong Jia, Jian Tang, Sirui Han, Shanghang Zhang

arXiv:2606.14048v3cs.CVcs.RO

TL;DR

Most world action models use 2D video or latent future states, leaving out geometric detail relevant to precise manipulation, while dense 4D prediction can add inference cost. WAM4D transfers geometric priors through training-time spatial registers and causal mixture attention, then removes the geometry readout for deployment. Experiments report improved spatial consistency and action prediction across simulation and real-world long-horizon tasks.

  • Problem

    Most WAMs model futures in 2D video or latent spaces, while explicit dense 4D prediction can increase deployment cost and latency.

    Source: §1 · p. 2

  • Method

    WAM4D uses spatial registers to transfer geometric priors through training-time depth supervision and causal mixture attention, then removes the geometry readout for deployment.

    Source: §1 · p. 2

  • Results

    WAM4D improves spatial consistency and action prediction across simulation and real-world long-horizon tasks.

    Source: §1 · p. 3, §5 · p. 15

  • Takeaways & Limitations

    WAM4D improves geometry-aware video-action representations without requiring dense geometry during policy deployment.

    Source: §5 · p. 15

    WAM4D lacks explicit long-term object memory during autoregressive rollout, which can yield identity-inconsistent completions after severe occlusion in generated futures.

    Source: §5 · p. 15

Abstract

from arXiv · show

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.

1 Introduction

WAM4D addresses the limits of 2D or latent future-state modeling for precise manipulation by transferring geometric priors into a causal video-action model through spatial register tokens. It uses geometry supervision during training and removes the geometry readout for deployment; experiments report improved spatial consistency and success rates with efficient causal action inference.

1 Introduction

  • WAMs jointly predict future observations and executable robot actions, coupling scene imagination and control in one pipeline.Source: §1 · pp. 1–2
  • Unlike designs that project geometry into a 4D scene, concatenate geometric and visual tokens, or synthesize future RGB-D, WAM4D distills geometric priors through spatial registers and removes the readout for action inference.Source: Figure 1 · p. 2, §1 · p. 2
  • Most WAMs predict future states in 2D video or latent spaces, which may omit the object extent, occluded surfaces, free space, and contact geometry needed for precise manipulation.Dense 4D prediction can increase deployment cost and latency through geometry decoding or optimization.Source: §1 · p. 2
  • Spatial registers query history video tokens and use a pretrained geometric head to decode depth, backpropagating depth loss into video features used for action prediction.The geometry head and depth readout are removed at deployment, preserving the lightweight 2D observation-to-action interface.Source: §1 · p. 2
  • Causal mixture attention defines visibility across video, action, and spatial-register tokens to support geometry supervision during training and lightweight inference at deployment.Source: §1 · p. 2
  • WAM4D improves spatial consistency and success rates on RoboTwin 2.0 and real-world long-horizon tasks while preserving efficient causal action inference.Source: §1 · p. 3

2 Related Work

Related work spans video-based embodied world models and WAMs, geometric foundation models, and direct manipulation policies. WAM4D differs by using future depth as a training signal through spatial registers rather than requiring geometry as an inference-time input, output, or reconstructed scene state.

Embodied World Models

  • Embodied world models use video generation to model interactive environments, while WAMs jointly model future observations and executable actions (Wang et al., 2026; Bi et al., 2025; Team et al., 2026; Li et al., 2026; Yuan et al., 2026; Guo et al., 2026).WAM4D transfers 4D geometric priors into video-action representations with lightweight spatial registers.Source: §2 · p. 3

Geometric Foundation Models

Robotic Manipulation Policy

  • Direct manipulation policies map observations to actions, and spatially grounded VLAs inject geometric priors, but typically learn geometry, contacts, and dynamics through action labels (Qu et al., 2025; Li et al., 2026; Sun et al., 2025; Li et al., 2025; Wang et al., 2026).WAMs instead jointly model future observations and executable actions.Source: §2 · p. 3

3 Method

WAM4D keeps a causal video-action backbone and adds a training-time spatial-register depth readout to transfer geometric priors into shared video features. Modality-specific visibility keeps auxiliary geometry outside the action-generation path.

3.1 WAM4D Backbone

  • The causal backbone conditions on language, multi-view RGB history, and historical actions while jointly modeling future video latents and action chunks as prediction targets.Future video and action tokens are flow-matching states, not additional causal inputs.Source: §3.1 · pp. 3–4

3.2 Spatial Register Distillation

  • Learnable spatial registers query history-video features at selected layers, then a pretrained geometric head decodes future depth whose loss distills spatial priors into shared video-action features.The register branch is used only during training.Source: §3.2 · p. 5

3.3 Causal Mixture Attention

  • Causal mixture attention lets future actions use history video, history actions, and their own noised tokens, while masking future video and registers from action prediction.Registers attend only to themselves and valid history-video tokens; registers, depth blocks, and the geometric head are removed at inference.Source: §3.3 · p. 5

3.4 Training Objective

  • The video-action stream uses conditional flow matching, while future depth uses SmoothL1 supervision whose loss updates the backbone and geometry-readout components.The stated loss weights are λact = 1 and λdepth = 1.Source: §3.4 · p. 6

3.5 Implementation Details

  • The default setup samples up to 17 frames, predicts eight future video and depth frames, and uses 32-step action chunks with a three-view mosaic and 960 spatial registers.Actions are 16-dimensional absolute end-effector pose and gripper vectors; spatial registers correspond to 32 × 32 image cells.Source: §3.5 · p. 6

4 Experiments

Experiments evaluate WAM4D on RoboTwin 2.0 and AstriBot S1 real-world manipulation, measuring policy success alongside video, depth, point-cloud, and inference performance.

4.1 Experimental Setup

  • RoboTwin 2.0 evaluation covers the complete task suite in clean and randomized settings, with depth-supervised training using demonstrations annotated with depth.The randomized setting varies clutter, illumination, background, tabletop height, object placement, and language instruction; each task has 50 clean and 500 randomized trajectories.Source: §4.1 · p. 7
  • Real-world evaluation uses four AstriBot S1 manipulation tasks, with 100 demonstrations per task and 10 physical rollouts per task.Depth supervision is obtained using Depth Anything 3 through an offline pseudo-depth pipeline.Source: §4.1 · p. 7
  • Comparisons include VLA and video-action baselines, with reproduced baselines using matched splits, camera configuration, and training budget when possible.Latency comparisons use 10 action-denoising steps and, when applicable, 5 video-denoising steps; WAM4D deployment removes the geometry path and predicts actions only.Source: §4.1 · pp. 7–8
  • Evaluation measures policy success, video-generation quality, and depth accuracy using success rates, FVD, PSNR, SSIM, LPIPS, AbsRel, δ1, and δ2.For depth metrics, AbsRel is lower-is-better, while δ1 and δ2 measure the fraction of pixels within fixed multiplicative error thresholds and are higher-is-better.Source: §4.1 · p. 8
  • Point-cloud evaluation back-projects predicted and target depth with fixed dataset-level camera intrinsics and compares geometry and temporal overlap.CD1 and CD2 are L1 and L2 Chamfer Distance variants; F-score measures thresholded overlap and F-score-T its temporal consistency across future frames.Source: §4.1 · p. 8

4.2 Simulation Results on RoboTwin 2.0

On RoboTwin 2.0, WAM4D is reported to achieve competitive performance while maintaining low inference cost; results cover aggregate and per-task success across clean and randomized settings.

4.2 Simulation Results on RoboTwin 2.0

  • The aggregate success rate averages results over 50 tasks under both clean and randomized settings.The overall success rate is reported in Table 1, while full per-task rates appear in Table 3.Source: §4.2 · p. 8
  • WAM4D achieves competitive RoboTwin 2.0 performance while maintaining low inference cost.Table 1 reports full-task success rate, action-generation latency, and VRAM.Source: §4.2 · p. 8
  • Table 3 reports per-task success rates for 50 tasks under Easy and Hard settings, corresponding to clean and randomized evaluation protocols.The full per-task results complement the aggregate comparison in Table 1.Source: Table 3 · p. 9, §4.2 · p. 9
  • Figure 5 visualizes spatial-register attention on randomized RoboTwin samples.The accompanying text describes the visualization as register-to-history attention.Source: Figure 5 · p. 10

4.3 Real-World Results

Across four AstriBot S1 real-robot evaluations, WAM4D achieves the best overall performance; the section also defines sequential sub-action scoring for these tasks.

4.3 Real-World Results

  • Across all four real-robot evaluations, WAM4D achieves the best overall performance, demonstrating robustness in contact-rich, geometry-sensitive, and long-horizon manipulation; LingBot-VA shows lower success rates in these settings.Source: §4.3 · p. 10
  • The four evaluations use contact-rich, geometry-sensitive, and long-horizon manipulation scenarios on the AstriBot S1 platform.Figure 4 illustrates the evaluation setup.Source: §4.3 · p. 10
  • Sub-actions are intermediate goals scored during rollouts; after any failure, later sub-actions are not attempted and are recorded as 0.Table 4 defines the sub-actions used for the real-world tasks.Source: §4.3 · p. 10, Table 4 · p. 10

4.4 Model Analysis and Ablations

The ablations examine how spatial-register attention, register design, and geometric-head initialization affect depth readout, generation quality, and control. Middle-layer unidirectional registers and a trainable pretrained geometric head are selected as the final settings.

4.4 Model Analysis and Ablations

  • Ablations use a fixed ten-task RoboTwin 2.0 subset and controlled comparisons, assessing both generation quality and control because depth loss alone does not establish manipulation gains.The split spans single- and multi-step manipulation tasks; per-task success rates and full ablation metrics are reported in Tables 5 and 6.Source: §4.4 · pp. 10–12, Table 5 · p. 11, Table 6 · p. 11

4.4.2 Spatial Register Design Ablation

  • Middle-layer registers give the strongest unidirectional control performance and overall geometry quality, while bidirectional visibility has the highest success rate but adds complexity and degrades geometry metrics.The default uses unidirectional registers after layers 12, 14, 16, and 18.Source: Table 7 · p. 12, §4.4.2 · p. 12

4.4.1 Spatial Register Attention Analysis

  • Spatial registers attend to object-consistent regions across views, static background regions, and the initial gripper pose when predicting future depth.The attention maps are averaged over heads from selected transformer layers.Source: §4.4.1 · p. 12

4.4.3 Geometric Head Ablation

  • A trainable pretrained geometric head gives the strongest generation quality, while random initialization lowers video and geometry quality and a fixed pretrained head performs between them.The authors use the trainable pretrained head as their final setting.Source: Table 8 · p. 13, §4.4.3 · p. 13

4.5 Qualitative 4D Rollout Visualization

The spatial-register depth branch supports a qualitative 4D rollout visualization, separate from the lightweight deployment path.

4.5 Qualitative 4D Rollout Visualization

  • Starting from one observed frame, WAM4D autoregressively predicts RGB frames and depth maps that can be back-projected into point clouds; deployment removes the depth branch.The rollout path is for qualitative analysis, while default deployment uses lightweight action generation.Source: §4.5 · p. 13

4.6 Failure Analysis

Long autoregressive rollouts can lose object identity after occlusion or departure from visible context because WAM4D has no explicit long-term memory. This limitation affects rollout visualization, while evaluated control receives fresh environment observations.

4.6 Failure Analysis

  • The rollout limitation affects closed-loop visualization, but not evaluated policy success rates, because control receives fresh observations from the environment or simulator.Source: §4.6 · p. 13
  • Without explicit long-term memory, an occluded or out-of-view object may be completed as a different object during long autoregressive rollouts.The figure depicts a representative failure case.Source: §4.6 · p. 13, Figure 7 · p. 15

4.7 Compute Details

Compute measurements are reported on one A800 80GB GPU; training includes the register blocks and geometric head, while default inference removes them.

4.7 Compute Details

  • Latency and peak-memory measurements use a single A800 80GB GPU; training includes register blocks and the pretrained geometric head, which default inference removes.The compute comparison reports latency as mean ± standard deviation in milliseconds when available.Source: §4.7 · p. 15, Table 9 · p. 15

5 Conclusion and Limitation

WAM4D transfers geometric foundation priors into causal video-action representations and improves spatial consistency and action prediction across simulation and real-world long-horizon tasks. It avoids dense geometry at deployment, but lacks explicit long-term object memory, which can produce identity-inconsistent completions after severe occlusion.

5 Conclusion and Limitation

  • WAM4D improves spatial consistency and action prediction across simulation and real-world long-horizon tasks.It transfers geometric foundation priors into causal video-action representations through spatial register distillation.Source: §5 · p. 15
  • WAM4D improves geometry-aware video-action representations without requiring dense geometry during policy deployment.The conclusion frames this as preserving a lightweight observation-to-action deployment path.Source: §5 · p. 15
  • Without explicit long-term object memory during autoregressive rollout, WAM4D can produce identity-inconsistent completions after severe occlusion.The paper suggests persistent object memory or scene-state tracking as future work.Source: §5 · p. 15

Related papers

Loading 2606.14048v3…