Spatial-UMI-Dream: Scaling World Action Models with Metric Geometry for UMI Manipulation

📌 Highlights

Spatial-UMI-Dream combines a 5B video-generation backbone with metric depth and action prediction for UMI manipulation. Its central claim is that explicit geometry improves control when wrist-mounted fisheye cameras move with both grippers.

The strongest reported evidence is 79.10% success on RoboTwin-UMI in-domain, 72.15% under combined embodiment and randomized-lighting shifts, and 85.0% average success across six real-world tasks. The paper's matched ablations show gains from both RGB-D inputs and progressive UMI pretraining.

• 605 hours of RGB-D UMI pretraining.
• Twenty-task RoboTwin-UMI benchmark with controlled environment and embodiment shifts.
• RGB-D improves in-domain simulation success by 5.75 percentage points over RGB-only pretraining.
• World-only adaptation raises in-domain simulation success from 39.50% to 61.25%; joint world-action pretraining reaches 79.10%.
• The main unresolved deployment questions concern calibration sensitivity, broader robot transfer, and compute or latency requirements.

🎯 Introduction

UMI decouples demonstration collection from a specific robot by recording close-range video and embodiment-agnostic end-effector trajectories with portable handheld grippers. The resulting wrist-mounted fisheye views are useful for scalable data collection, but they do not provide a stable external camera: image changes combine camera motion, object motion, and scene geometry, and the two grippers move independently in a shared workspace.

The paper argues that RGB-only World Action Models leave the metric geometry needed for precise manipulation implicit in appearance and image-space motion. Spatial-UMI-Dream therefore makes depth an explicit world-state component and evaluates whether robot-independent demonstrations transfer across embodiments and environmental conditions rather than relying only on in-domain task success.

• Prior action-only policies do not explicitly model future scene evolution.
• RGB-only WAMs must infer scale and relative spatial relationships from moving wrist observations.
• RoboTwin-UMI is introduced to measure environment robustness, embodiment transfer, and their combination under a common UMI interface.

🔬 Methodology

The input at timestep t contains RGB and aligned metric depth for each of two grippers. Depth is recovered from the middle and right cameras on each gripper using stereo rectification, Fast FoundationStereo disparity estimation, stereo calibration, and mapping back into the original middle-camera fisheye coordinates. The normalized depth is encoded with a Turbo color mapping and placed beside its corresponding RGB view in a fixed dual-gripper canvas.

A pretrained Wan2.2-TI2V-5B VAE encodes the complete canvas as an ordinary video frame. The shared transformer uses block-causal attention: visual and action tokens interact within the current temporal block, while the block attends to preceding visual history. Language enters through cross-attention, and robot state and noisy actions are added in the joint training stage.

Training is progressive. World-only adaptation predicts future RGB-D observations with the action branch bypassed. Joint world-action training restores state and action encoders, initializes from the adapted checkpoint, and optimizes the combined world and action objective. Before training, demonstrations are filtered using inverse-kinematic replay in ManiSkill and dataset-wide motion-percentile checks. At deployment, gripper commands are compensated for measured response delay.

• Each RGB or depth panel is resized to 160 × 160, producing a 320 × 320 canvas.
• The downstream action is 20-dimensional, with relative translation, 6D rotation, and absolute gripper opening for both grippers.
• Both pretraining stages use the same 605-hour RGB-D UMI corpus.

fig:method-overview

Caption: Model architecture and two-stage training of . In both stages, dual-gripper RGB and metric depth (Fig.~fig:umi-data) form a unified RGB-D canvas. The Wan VAE encodes observation history and current RGB-D inputs into visual tokens, while a text encoder supplies language tokens through cross-attention. The shared world-action transformer uses block-causal attention to condition prediction on preceding visual history. Stage 1: world-only adaptation predicts future RGB-D observations, decoded through the Wan VAE and supervised by $ L_world$. Stage 2: the transformer backbone is initialized from the Stage 1 world-only adaptation checkpoint. Joint world-action training additionally encodes UMI st...

Why It Matters: The architecture figure shows how RGB-D world modeling, state, language, and action prediction are combined across the two training stages.

fig:umi-data

Caption: Metric depth construction for UMI observations. Each UMI gripper has three cameras. We undistort and rectify the middle and right camera views, then use Fast FoundationStereo to recover metric depth. We map the depth back to the original middle-camera fisheye view to align it with the RGB observation.

Why It Matters: The depth pipeline is the key mechanism that turns moving fisheye observations into a metric RGB-D world representation.

fig:trajectory-filtering

Caption: UMI trajectory filtering. Wrist observations (top) and simulation replay (bottom) of a folding demonstration. The third frame violates robot kinematic limits, so the entire episode is discarded.

Why It Matters: The example clarifies how infeasible handheld demonstrations are removed before policy training.

fig:deployment

Caption: Gripper compensation at deployment. Desired actions (blue), measured states (purple), and compensated actions (orange) of one gripper (a) without and (b) with compensation under the same continuous target trajectory. Both panels show the first 6 seconds without smoothing. Mean absolute tracking error over this interval decreases from 12.49 to 2.85\,mm.

Why It Matters: The figure connects the learned action trajectory to physical gripper execution and quantifies the benefit of response compensation.

📊 Experiments

The evaluation uses the twenty-task RoboTwin-UMI simulation benchmark and six real-world bimanual tasks on the Dexmate Vega1: Pick-and-Place, Bimanual Handover, Scan-and-Place, Pack-and-Close, Handover-and-Pour, and Table Wiping. RoboTwin-UMI uses Aloha-AgileX and ARX-X5 with common UMI grippers and wrist-mounted fisheye cameras. Clean Aloha measures in-domain performance, randomized-lighting Aloha tests environmental robustness, clean ARX X5 tests embodiment transfer, and randomized-lighting ARX X5 tests combined shifts. Each simulation task-setting uses 100 evaluation episodes, while each real-world task uses 20 episodes.

On RoboTwin-UMI, Spatial-UMI-Dream reaches 79.10% mean success on clean Aloha, 73.35% under randomized lighting on Aloha, and 72.15% under randomized lighting on ARX X5. These results exceed π0.5 by 4.65, 7.20, and 10.30 percentage points in the respective settings. On the real-world suite, it reaches 85.0% average full-episode success and 93.2% average progress score, compared with 60.8% and 73.3% for π0.5.

The ablations support two mechanisms. Replacing depth with duplicated RGB panels lowers in-domain simulation success by 5.75 percentage points and joint-shift success by 8.25 points; real-world mean success falls from 85.0% to 80.8%. World-only adaptation raises in-domain simulation success from 39.50% for Wan2.2 to 61.25%, while adding joint world-action pretraining reaches 79.10%. The paper reports that the full model also improves over video-only pretraining by 15.35 points under embodiment transfer and 13.40 points under joint shifts.

• Simulation metric: success rate, with higher values preferred.
• Real-world metrics: full-episode success rate and average progress score, both reported in percent.
• Baselines: π0.5, RDT2-FM, Fast-WAM, and LingBot-VA.
• Generalization tests use trained policies without additional training data or fine-tuning.

fig:benchmark-setup

Caption: benchmark construction. We retain the same RoboTwin tasks while adapting the grippers and observations for UMI. From left to right: the original setup; UMI grippers and wrist-mounted fisheye cameras on Aloha-AgileX and ARX-X5; and randomized lighting. These settings support evaluation of cross-embodiment and environmental generalization, separately and jointly.

Why It Matters: It shows how the controlled benchmark preserves tasks while varying robot embodiment and lighting.

fig:real-world-tasks

Caption: Real-world task suite. We show egocentric keyframes from six real-world tasks, including unimanual manipulation and bimanual coordination. These views are for visualization only and are not used as policy inputs.

Why It Matters: It contextualizes the range of real-world manipulation behaviors used for evaluation.

fig:sim-results-mean

Caption: Average success across embodiments.

Why It Matters: It summarizes the cross-embodiment simulation comparison across the evaluated policies.

🔮 Conclusion

The paper presents a geometry-aware World Action Model that treats aligned metric depth as part of the predicted world state rather than as an auxiliary representation. Combined with two-stage UMI pretraining, this design produces strong closed-loop manipulation results in simulation and on six real-world tasks.

The reported evidence supports complementary roles for the components: RGB-D inputs improve spatially sensitive control, world-only adaptation transfers the pretrained video model to moving UMI observations, and joint world-action learning connects those representations to executable behavior. RoboTwin-UMI further shows that in-domain success and transfer performance can differ substantially.

• Best reported real-world average success: 85.0%.
• Best reported combined-shift simulation success: 72.15%.
• Metric depth and UMI-specific action pretraining both contribute measurable gains.

🛠️ Future Research Improvements

The next research step is to test whether the reported geometry and pretraining gains survive beyond the two simulated embodiments, the randomized-lighting conditions, and the Vega1 real-world platform used here. Additional robots, camera placements, object categories, and physical environmental shifts would clarify the scope of cross-embodiment transfer.

A second direction is robustness and efficiency analysis. The method relies on calibrated stereo depth and a 5B backbone, but the main-body evidence does not report depth-failure sensitivity, inference latency, memory requirements, or wall-clock deployment cost. These measurements would be important for practical robot integration.

Further ablations could isolate the effects of depth quality, depth encoding, canvas layout, filtering, and command compensation rather than comparing only RGB versus RGB-D and video-only versus joint pretraining.

• Evaluate calibration drift and missing or noisy stereo depth.
• Measure latency and compute separately for pretraining, fine-tuning, and inference.
• Test transfer to additional robot kinematics and camera configurations.
• Report standalone world-prediction and depth-quality metrics alongside task success.
• Expand environmental shifts beyond lighting.

🏭 Potential Industry Use Scenarios

The task suite suggests potential use in bimanual assembly and packing, where gripper-to-object distance and relative placement are critical. The reported Pick-and-Place, Pack-and-Close, and Handover results are relevant to workflows that require coordinated manipulation rather than isolated single-arm motions.

The model may also be applicable to robot-assisted inspection and tool use, liquid transfer, and sustained-contact operations. Scan-and-Place, Handover-and-Pour, and Table Wiping provide evidence for these behavior classes, although the paper does not establish performance in production environments.

The UMI data-collection interface could support organizations that want to gather demonstrations independently of the eventual robot platform. RoboTwin-UMI indicates how such systems might be screened for embodiment and environment transfer before physical deployment, but broader validation is still required.

• Bimanual assembly, packing, and transfer.
• Inspection or scan-guided placement.
• Liquid handling and handover workflows.
• Surface cleaning and sustained-contact manipulation.
• Pre-deployment evaluation of robot-independent demonstration policies.

💬 Critical Analysis

The strongest aspect of the paper is the alignment between its problem diagnosis and evaluation design. Moving wrist cameras make geometry difficult to infer from RGB, and the method adds aligned metric depth directly to the world state. The matched RGB-only comparison, the staged pretraining ablation, and the controlled RoboTwin-UMI shifts provide more informative evidence than a single in-domain success number.

The reported gains are substantial, but the comparisons should be interpreted within the stated setup. Spatial-UMI-Dream uses a spatial world-action pretrained checkpoint, while baselines use their released base checkpoints, and all methods are then fine-tuned on the same demonstrations. The paper also evaluates only a narrow set of embodiments and environmental changes, so the generality of the transfer result is not yet established.

Reproducibility depends on several system details: calibrated three-camera grippers, stereo depth reconstruction, fisheye rendering, trajectory filtering, action representation, and gripper-response compensation. The absence of reported latency, compute cost, standalone perceptual metrics, and robustness to depth failure leaves important deployment questions unanswered. Overall, the evidence supports geometry-aware UMI pretraining as a promising control strategy, while broader hardware and environmental validation is needed before treating it as a general solution.

• Evidence strength: matched modality and pretraining ablations plus closed-loop simulation and real-world evaluation.
• Main caveat: transfer is demonstrated on two simulated embodiments and one real-world robot platform.
• Main engineering risk: dependence on calibrated stereo depth and fisheye alignment.
• Main evaluation gap: no reported latency, inference-cost, or standalone depth/world-prediction analysis.