UNITAS: A 3D-Native World Action Model for Embodied Manipulation

AI summary

Overview

Research area: Robotics — world action models (WAMs), vision-language-action (VLA) policy learning, 3D scene/flow prediction for embodied manipulation.

Technical level: Advanced. The paper assumes familiarity with diffusion/rectified-flow transformers, Mixture-of-Transformers architectures, camera intrinsics/extrinsics, point-flow representations, and standard manipulation benchmarks (LIBERO, RoboTwin, VLABench).

Scope in one sentence: The paper proposes Unitas, a 1.7B-parameter model that represents observations, robot/human motion, and resulting scene change as point trajectories in a single metric 3D frame, and evaluates it both as a manipulation policy and as an action-conditioned scene predictor.

What This Paper Is About

Most world action models predict the future in image space: they generate future video frames or visual latents and tie action learning to those pixels. But a robot acts in metric 3D space, and pixels are a view-dependent projection where equal pixel motion does not mean equal physical displacement — the same motion looks different across cameras, intrinsics, and robot poses. Unitas instead asks both halves of the question — "what motion should the robot execute?" and "how will that motion change the world?" — inside one shared 3D coordinate frame, using point trajectories as a common language for robot grippers, human hands, and scene points.

Key Contributions

1. A 3D-native world action model. Unitas is presented as the first WAM that aligns observations, actions, and scene dynamics in a shared metric frame within each prediction window, using one representation across robot embodiments and human hands.
2. Action flow and scene flow as a shared interface. Action flow encodes human hands and robot grippers as 3D point trajectories; scene flow encodes scene-point displacements conditioned on those trajectories. The same interface serves both a policy mode (direct executable actions) and a simulator mode (action-conditioned scene prediction).
3. World-aligned 3D positional embeddings and a physical-time trajectory tokenizer. Visual tokens are grounded in world coordinates with or without depth input (depth-derived embeddings, or learned weights over 64 candidate depths sampled from 0.1 to 2.0 m along calibrated camera rays). The tokenizer encodes an entire point trajectory as a single token keyed to elapsed seconds rather than frame indices, anchored at the point's current 3D position.
4. A coupled two-expert architecture with efficient inference. A layer-coupled Mixture-of-Transformers pairs a point dynamics expert (action and scene flow) with a smaller action expert (end-effector action chunks), with an attention mask that lets policy-only inference skip both future-flow branches, and causal key–value caching across denoising steps.

Main Findings

• Policy results on LIBERO: 99.8% average success (Spatial 99.8, Object 100.0, Goal 100.0, LIBERO-10 99.4), the best reported in Table 1. Without mixed-dataset pretraining, 99.5%.
• RoboTwin 2.0: 92.94% (Clean) and 93.04% (Randomized) with pretraining; 91.08% and 90.74% without. The 1.7B model outperforms the compared methods while using fewer parameters than Fast-WAM (6B) and Motus (8B).
• Robustness on LIBERO-Plus: 89.1% overall success across 10,030 evaluation episodes, versus 86.5% without pretraining. Under camera perturbations it reaches 82.1%, compared with 75.8% for Cosmos Policy and 64.5% for LaMP.
• VLABench: 56.2% success rate and 67.6% progress score, versus 48.1% success rate for π0.5. The no-pretraining variant scores lower on success rate (50.3%) but higher on intention score (71.2%) than the pretrained model (69.3%).
• Action-conditioned scene prediction on RoboTwin (simulator mode): best on all four metrics — Moving ADE 0.427 cm vs 0.656 for PointWorld, Static ADE 0.043 vs 0.084, Moving FDE 0.825 vs 1.050, and Moving Acc@1cm 92.31% vs 88.22%. Reported reductions are 35% (Moving ADE), 21% (Moving FDE), and 49% (Static ADE) relative to PointWorld.
• Video-based baselines lifted to 3D perform worse on scene prediction: Ctrl-World + Depth Anything 3 + SpatialTrackerV2 scores Moving ADE 3.310 cm and Acc@1cm 41.54%; X-WAM + SpatialTrackerV2 scores 2.286 cm and 49.69%.
• Real-robot tasks (20 trials per method per task): Organize Books 18/20 (90%) versus π0.5 15/20 (75%) and Fast-WAM 13/20 (65%); Pack the Ball 16/20 (80%) versus π0.5 11/20 (55%) and Fast-WAM 12/20 (60%). Note an internal inconsistency: the introduction states 65% for Pack the Ball, while Table 2 and Section 4.3 report 16/20 (80%). The abstract's "average of 85% across real-world tasks" corresponds to the 90% and 80% figures.
• Ablations on LIBERO-Plus (no pretraining, overall success): full model 86.47%; reference model without scene-flow supervision 85.74%; removing 3D positional encoding 81.35% (camera perturbation drops from 79.11% to 66.73%); removing action-flow supervision 78.89%; removing action-flow history tokens 75.92% (robot-initial-state perturbation drops 19.87 percentage points); removing all three 67.58%. Adding scene-flow supervision to the reference improves 85.74% to 86.47%.
• Pretraining transfer: mixed-dataset pretraining adds 2.6 percentage points overall on LIBERO-Plus, with the largest gains under robot initial-state changes (84.8% to 94.2%) and language perturbations (77.6% to 85.0%). Adding data sources while holding RoboTwin samples per batch fixed lowers RoboTwin test end-effector XYZ error.
• Action-scene consistency: on three RoboTwin tasks, Unitas reportedly gives the best agreement between predicted scene evolution and executed actions among the compared WAMs, with details deferred to the appendix (not included in the provided content).

Methodology in Plain English

The model works in windows. At each step it takes multi-view images (depth optional), camera calibration, a task instruction, and the recent motion history, and fixes one 3D coordinate frame for the whole window. Robot demonstrations use a robot-base or shared embodiment frame; egocentric human data uses the fixed camera frame of the window anchor.

Every visual token is given a position in that world frame. If depth is available, each pixel is back-projected and transformed into world coordinates, then encoded with a Fourier map. If depth is missing, the model looks along the camera ray: it samples 64 candidate depths between 0.1 and 2.0 m, scores each candidate using the patch feature, and takes a weighted sum of their world-frame positional features. On data where depth is reliable, these routing weights are trained to match depth-derived features.

Motion is described as displacement of specific points over real time. For a robot, 20 points are sampled from the URDF gripper geometry of each arm and tracked through forward kinematics, preserving correspondence as fingers and arms move. For human data, 21 MANO keypoints per hand are used. Scene points come from back-projecting RGB-D pixels after masking out the robot using its URDF and joint state, following PointWorld.

The key representational trick is the trajectory tokenizer: instead of spending one token per timestep, a whole trajectory is compressed into one latent token, using time in seconds (with a shared 3-second scale) rather than frame indices, so data recorded at 6, 8, 10, 15, or 20 Hz can be mixed coherently. The token carries a spatial anchor at the point's current 3D position, so the network reasons about motion through features attached to fixed locations.

Two transformer experts are coupled. The point dynamics expert handles action flow and scene flow; the smaller action expert reads the dynamics expert's observation/language keys and values plus the current end-effector state, and predicts relative translation, rotation, and gripper openness. Scene tokens attend to clean action flow but not to noisy action-flow latents, so scene prediction is conditioned on motion without contaminating motion denoising. Action chunks deliberately do not attend to future flow tokens, which is what makes policy-only deployment cheap — those branches are simply skipped.

Training proceeds in two stages: first the tokenizer is trained (and frozen) to reconstruct historical and future action flow plus available scene trajectories; then both experts are trained with rectified flow, using a loss that sums action-flow, scene-flow, and control terms at equal weight plus the positional-embedding term at weight 0.1. The vision-language backbone is initialized from Florence-2 and pretrained on nine data sources (RoboTwin, LIBERO, VLABench, DROID, AgiBot, RoboMIND, RH20T, VITRA, and self-collected W1 data).

Why This Matters

Impact on research. The paper argues for treating interaction geometry as an explicit prediction target rather than an auxiliary cue or a byproduct of image generation. It offers a concrete alternative to the video-generator lineage of WAMs and separates action generation from scene prediction at the architectural level, so the same model can act or simulate. Its point-flow interface also gives a common representation for human hand data and multiple robot embodiments, which is directly relevant to cross-embodiment data scaling.

Real-world applications:

• Manipulation policies for双臂 / bimanual robot platforms in structured settings such as shelf organization and packing, where metric precision matters.
• Learned simulators for planning: given a candidate motion, predict where scene points will end up, enabling model-based action selection without executing on hardware.
• Data augmentation and policy evaluation from human demonstration video, since hands and grippers share one representation.
• Robust deployment under camera changes: world-coordinate grounding is directly aimed at the viewpoint and robot-initial-state sensitivity that the paper attributes to video-based WAMs.

Industry relevance. The 1.7B parameter count and the policy-only inference path (skipping future-flow branches, with cached context across denoising steps) target practical latency and compute budgets. Beating larger video-based models such as Fast-WAM (6B) and Motus (8B) on RoboTwin makes the efficiency argument concrete, and the released code lowers the barrier to adoption.

Future Directions

• Longer-horizon prediction. The conclusion explicitly names extending the formulation to reliable prediction over longer interactions; the current design fixes one frame per prediction window with a 3-second trajectory timescale.
• Broader real-world deployment. Real-world evaluation covers only two tasks with 20 trials each, and the paper notes real-world scene prediction is shown on a DROID example without real-world scene-flow supervision — closing that supervision gap is an open step.
• Deeper use of scene-flow supervision. The ablation shows scene-flow supervision gives a modest gain (85.74% to 86.47%) in policy mode; how much more can be extracted from joint action–scene training, and the reported action-scene consistency study in the appendix, remain questions.
• Robustness beyond the tested perturbations. LIBERO-Plus covers seven perturbation types, but real deployment involves lighting, occlusion, and object variation outside that set; the paper's evidence for viewpoint robustness does not establish the rest.

Target Audience

Robotics and embodied-AI researchers working on world models, VLA policies, or 3D perception for manipulation; engineers building manipulation stacks who need metric rather than pixel-space prediction; and graduate students with background in diffusion models and multi-view geometry who want a worked example of unifying action and scene dynamics in a single representation. Readers without familiarity with point flows, rectified flow, and camera geometry will find the methodology sections dense.

Authors’ abstract

World action models (WAMs) aim to answer a coupled physical question: given a task instruction, what motion should the robot execute, and how will that motion change the surrounding world? Most existing WAMs build on pretrained video generators and represent world evolution through images or visual latents. Robotic interaction, however, takes place in metric three-dimensional space, while images are view-dependent projections whose pixel distances do not directly encode physical distances. We introduce UNITAS, to our knowledge the first 3D-native world action model that unifies observations, actions, and scene dynamics in a shared metric 3D frame within each interaction, using a common representation across robot embodiments and human hands. Action flow represents human hands and robot grippers as 3D point trajectories, while scene flow describes scene-point displacements conditioned on these trajectories. World-aligned 3D positional embeddings ground visual tokens with or without depth input, and a physical-time trajectory tokenizer encodes each point trajectory as one token anchored at its current 3D position. This interface supports both direct action execution and action-conditioned scene prediction. With 1.7B parameters, UNITAS achieves the best action-conditioned scene prediction on RoboTwin among the compared methods, with up to 49% lower displacement errors than PointWorld, and state-of-the-art manipulation success, including 99.8% on LIBERO and an average of 85% across real-world tasks. The code is available at https://github.com/DexForce/UNITAS.