Kui Jia Team's New Work: Enabling World Action Models to Truly Understand 3D Space

UNITAS unifies observation, action, and scene changes within a single metric 3D coordinate system: it represents the motion of human hands and grippers using point trajectories, describes the displacement of scene points with scene flow, and enables world-action models to no longer rely solely on images to infer environmental evolution. With only 1.7 billion parameters, the model achieves a 99.8% task success rate on the LIBERO benchmark; in action-conditioned scene prediction tasks, its displacement prediction error for moving points is reduced by 35% compared to PointWorld.

On October 8, a joint team from The Chinese University of Hong Kong, Shenzhen, DexForce, Foshan University, Sun Yat-sen University, and South China University of Technology published the paper "Unitas: A 3D-Native World Action Model for Embodied Manipulation." The paper has 17 authors. Ruixiang Wang from The Chinese University of Hong Kong, Shenzhen, and Xiaofen Xing from Foshan University are co-first authors. Guiliang Liu, an assistant professor in the School of Data Science at The Chinese University of Hong Kong, Shenzhen, and a reinforcement learning technical expert at DexForce, participated in the research. Yongyi Su from Foshan University and Kui Jia, a professor in the School of Data Science at The Chinese University of Hong Kong, Shenzhen, and founder of DexForce, served as corresponding authors. Guiliang Liu's main research focuses on embodied AI and reinforcement learning; he has published multiple papers at top conferences such as NeurIPS, ICML, and ICLR, and has served as area chair for NeurIPS and ICLR.

World action models must address a set of coupled questions: given a task instruction, what actions should the robot execute, and how will those actions alter the surrounding environment.

Most prior related studies rely on pre-trained video generators, using future images or visual latent variables to represent environmental changes. However, robot interaction occurs in real three-dimensional space, where images are merely two-dimensional projections of that 3D world. The same action, displacement, or contact event will appear different when viewed from a different camera angle; pixel displacement cannot be directly equated to physical distance. Furthermore, if the camera's intrinsic parameters or the relative pose between the camera and the robot change, this mapping relationship shifts accordingly. This leads models to learn a world that is "visually approximate" rather than "physically measurable." Even if the generated predictive frames are highly realistic, implicit biases may still arise in predictions regarding object displacement, rigid geometric properties, and static regions. UNITAS attempts to replace the underlying representation foundation by incorporating observations, actions, and inference results into a shared metric coordinate system, enabling the world action model to perform prediction and simulation directly within 3D physical space.

01 Three-Dimensional Ruler, Reconstructing Model Representation

The primary modification in UNITAS is its representation system. It redefines how actions are expressed, abandoning a sequence of joint angles in favor of segmented point trajectories formed by the hand and gripper in 3D space; the team calls these "action flows." The displacements generated at various points within the scene under action drive are recorded as "scene flows." Both are unified within the same metric coordinate system, where observations, actions, and scene changes share this coordinate baseline during a single prediction window.

Why change this measurement framework? The paper’s example is quite intuitive. When the same object moves the same physical distance, the pixel displacement differs significantly between a nearby camera and a distant one; once the camera intrinsics or the relative pose between the camera and the robot change, this conversion relationship changes accordingly. World action models based on video routes store environmental changes in pixels or visual latent variables, using precisely this variable scale. Therefore, even if the generated images are realistic enough, implicit deviations may still occur in object displacement, rigid geometric features, and static regions. By anchoring observations, actions, and inference results to a shared metric coordinate system, the measurement standard the model learns shifts from "how many pixels moved on screen" to "how many centimeters the object actually moved."

Coordinates can be bound to visual information through world-aligned 3D positional embeddings. Each visual token is associated with a spatial position in the world coordinate system. When valid depth is available, the team performs back-projection of pixels using depth and then transforms them into world coordinates via camera extrinsics, encoding these positions with a set of fixed Fourier features. In the absence of depth information, 64 candidate positions are uniformly sampled along each camera ray from 0.1 meters to 2 meters; patch features score these candidates, and weighted averaging yields the world coordinates. This design allows the model to accept both depth-enabled and depth-free data, ensuring normal operation during inference even when depth inputs are missing. Relying on weighted averaging of candidate positions along camera rays to solve for world coordinates is the core support enabling the model to achieve 3D alignment under depth-free input conditions.

A depth-less branch does not imply unsupervised learning from scratch. In regions with reliable depth, the team constrains the weighted results of candidate positions along the ray to remain aligned with positions obtained by back-projecting ground-truth depth. Supervisory signals are used to correct routing weights, and only pixels corresponding to valid depth participate in this alignment process. During training, both RGB and RGB-D data can be used simultaneously; during inference, removing the depth input does not significantly degrade model performance. The 3D position embedding does not rely on hardware depth sensors: using the same set of weights, the depth input can be replaced with a ray-routing scheme during testing.

The selection of action points and scene points was also carefully designed. Based on the URDF geometric model, each gripper jaw has 20 points selected; following forward kinematics, these points move synchronously with the joints and finger poses. For human hands, 21 MANO keypoints are used per hand, ensuring a unified representation format for both human hands and robotic grippers. Scene points are divided by voxels, with an upper limit of 12,000 points and random downsampling applied; the simulation dataset uses 5-millimeter voxels.

The third core component is the physical trajectory tokenizer. Different data sources have varying frame rates and durations; if each time step were treated as a separate token, identical durations would consume inconsistent model capacity across different data sources. The team instead uses seconds rather than frame numbers to mark time, with the sampling rate denoted as f, so that the k-th moment corresponds to τ = k /f. By leveraging fixed basis functions shared across datasets, the same real elapsed time is mapped to identical features. An input MLP converts each displacement segment into a vector, a time MLP transforms time features into modulation parameters, an encoder pools the entire trajectory into a single latent variable, and a decoder can read and reconstruct the trajectory at specified moments. As a result, regardless of trajectory length or sampling density, each trajectory occupies only one token.

The joint distribution of the three components is decomposed into three factors: an action flow is generated first; action blocks are conditioned on both the action flow and the robot's current state; and a scene flow is conditioned on the action flow. This decomposition enables parallel sampling for the action flow and action blocks. Scene prediction can either use the action flow autonomously generated by the model or, in simulator mode, condition on an externally provided trajectory segment. When running only the policy branch, one can compute solely the action block branch, thereby omitting all computations for the two flow branches entirely. Consequently, the single model possesses dual capabilities: it can directly output robotic action commands, or infer environmental changes given a sequence of actions.

This representation system also natively supports cross-embodiment generalization. Points from human hands and robotic grippers all fall into the same coordinate system, sharing a single tokenizer; observation tokens are arranged according to world coordinates, while action stream and scene stream tokens are anchored to their respective 3D positions, allowing the network to learn motion patterns at fixed spatial anchors. Unifying point trajectories as "displacements generated from the current position" is key to merging data from different sources and perspectives into a single representation space.

The model body consists of a set of coupled expert modules. The point dynamics expert handles action flows and scene flows, while the action expert is responsible for outputting executable end-effector action blocks; both are diffusion Transformers. The action expert is smaller in scale, reading the point dynamics expert's observations and language keys layer by layer, and superimposing its own current state with noisy instructions. During training, scene prediction uses clean action flows from demonstration samples as conditions; during inference, action flows and action blocks denoise in parallel on the same Euler timeline, and the generated action flow then serves as a condition to drive scene denoising. When executing only policy tasks, the two branches for future action flows and scene flows can be skipped to directly output actions.

The visual-language backbone is built on Florence-2-large, while a point-dynamics expert maintains dual branches for action flows and scene flows. The action expert is more lightweight; it reads the point-dynamics expert’s observations and language key-value pairs layer by layer, adds its own current state and noisy instructions, and predicts the relative translation, rotation, and gripper opening for each future time step. Both experts have 28 layers: the point-dynamics branch has a width of 640 with 10 heads; the VLM backbone has 28 layers, a width of 1280, and 20 heads; the tokenizer’s latent variable dimension is 256. This architecture, pairing a large backbone with a lightweight action expert, allows the action branch to retain only shallow coupling, eliminating the need to reconstruct the backbone network from scratch. During inference, context and clean action flow key-values can be cached, allowing computations at every denoising step of the scene to reuse the cache.

02 Nine Data Sources, Two-Stage Training

On the data side, the team integrated nine data sources for pre-training, including public datasets RoboTwin, LIBERO, VLABench, DROID, AgiBot, RoboMIND, and RH20T, as well as the human hand dataset VITRA and the team’s self-collected dataset W1. Differences in motion rates across different datasets are resolved through a physical time-alignment mechanism: AgiBot and RoboMIND’s raw trajectories are at 45 Hz. The team downsampled them by a stride of 3 to 15 Hz, allowing a single action block to cover a longer physical duration so that its spatial span remains comparable to other datasets. Within VITRA, Ego4D operates at 6 Hz, EgoExo4D and Something-Something V2 at 8 Hz, and EPIC-KITCHENS at 10 Hz.

The trajectory tokenizer is not a redundant component. The team conducted point-by-point validation on a reserved set of 600 trajectory windows: encoding an entire point trajectory into a single token and then decoding it back based on the original physical timestamps resulted in reconstruction errors at the millimeter level. This proves that a single token is sufficient to fully carry the motion information of a trajectory, regardless of whether it comes from a robotic arm, a human hand, or scene points.

Choosing seconds rather than frame numbers as the time base involves clear considerations. The frame rate range across the nine data sources is 6 Hz to 45 Hz; the same action occupies several times more frames in different recordings. If time were measured in frames, identical lengths of physical time would allocate unequal model capacity across different data sources. After aligning to real seconds, actions with varying speeds can learn under a unified metric scale.

Each of the nine data sources serves a specific role. RoboTwin, LIBERO, and VLABench provide single-arm and dual-arm tasks in simulation environments; DROID offers large-scale demonstration data in real-world scenarios; AgiBot, RoboMIND, and RH20T come from real-robot collection on different robot bodies; VITRA provides first-person human-hand action data; and W1 is the team’s self-collected dataset. The robot branch outputs gripper point trajectories, the human-hand branch outputs keypoint trajectories, and simulation datasets can provide ground truth for scene points. Missing annotations are directly masked and excluded from loss calculations.

The training process is divided into two stages. In the first stage, the trajectory tokenizer is trained separately to reconstruct historical and future action flows, as well as available scene trajectories. Once training is complete, the tokenizer weights are frozen. In the second stage, two groups of expert modules are jointly trained based on corrected flows. The total loss consists of four additive components: mean squared error for the action flow, scene flow, and individual action blocks, plus position embedding supervision loss. The weights for the first three losses are set to 1, while the position embedding loss weight is set to 0.1; samples with missing annotations are directly masked. The visual-language backbone is initialized with Florence-2-large; the VLM has 28 layers, a width of 1280, and 20 heads; both the point-dynamics expert and the action expert have 28 layers, a width of 640, and 10 heads; the tokenizer’s latent variable dimension is 256. Pre-training totaled 200,000 steps with a total batch size of 1,024, using the AdamW optimizer. The learning rate for the expert modules was set to 1e-4, and for the VLM to 1e-5, with weight decay at 1e-8. After 1,000 steps of linear warmup, the learning rate remained constant. Training utilized 64 BW1000_H GPUs.

By pre-training and freezing the tokenizer first, then training expert modules subsequently, this two-stage architecture separates motion representation learning from action generation into independent phases, preventing interference between the two tasks.

Inference efficiency is another advantage of this design. When performing only action output, UNITAS does not need to generate action streams or scene streams. On a single RTX 5090 D, the inference time for a single action block is 55.77 milliseconds, a 2.20x speedup compared to the version that generates both actions and action streams simultaneously; at the 1.7 B parameter scale, it achieves a 2.81x faster inference speed than Fast-WAM, which lacks video generation capabilities. When used as a simulator, the number of scene query points can be dynamically adjusted during inference, dropping from 3,000 to 1,024, reducing joint prediction latency from 561.60 milliseconds to 305.49 milliseconds, allowing trade-offs between scene coverage density and computational cost. Context and clean action stream key-value pairs are cached and reused in every step of scene denoising, reducing redundant calculations.

03 Policy Performance: Small Model, Hard Metrics

At the policy level, UNITAS achieves leading results across multiple standard benchmarks while maintaining a lightweight model size. The average task success rate on the LIBERO benchmark is 99.8%; even without mixed-data pre-training, it reaches 99.5%. In RoboTwin 2.0’s 50 bimanual tasks, the success rate under clean conditions is 92.94%, and 93.04% under random perturbation conditions. These results are achieved by a 1.7 B parameter model, with fewer parameters than Fast-WAM’s 6 B and Motus’s 8 B. On VLABench, the success rate is 56.2%, higher than π0.5’s 48.1%, with superior progress scores as well. The performance of the 1.7 B small model surpasses that of 6 B and 8 B-level video world-action models, proving that incorporating interactive behavior into 3D coordinate representation inherently offers parameter efficiency advantages.

Resisting distribution drift is another prominent advantage of the model. In the LIBERO-Plus suite, which includes seven types of observation, instruction, and initial state perturbations, UNITAS ranks first with a total score of 89.1%. In camera perturbation scenarios, it achieves 82.1%, outperforming Cosmos Policy’s 75.8% and LaMP’s 64.5%. Mixed-data pre-training can further improve the total score by 2.6 percentage points, with the most significant gains coming from changes in the robot’s initial state, rising from 84.8% to 94.2%; language rewriting task scores increase from 77.6% to 85.0%. The team also conducted control validations: keeping the number of RoboTwin samples constant but introducing additional data sources, the end-effector error on the RoboTwin test set actually decreased, demonstrating that cross-dataset transfer effects provide positive gains.

First, clarify the evaluation focus of each benchmark. LIBERO focuses on single-arm manipulation tasks, divided into four subsets: Spatial, Object, Goal, and long-horizon tasks (LIBERO-10); RoboTwin 2.0 includes 50 bimanual tasks, setting up two conditions: clean environment and randomized perturbation; VLABench tests object and instruction generalization capabilities, divided into five tracks, evaluated using three metrics: success rate, progress score, and intent score; LIBERO-Plus specifically constructs seven types of perturbations covering camera, lighting, background, noise, layout, robot state, and language instructions, measuring the model’s robustness in out-of-distribution scenarios. All models run under the same evaluation protocol, with success rates calculated using a unified standard to support horizontal comparison.

In real-robot experiments, the team repeated tests for two tasks 20 times each. For the bookshelf organization task, which requires aligning books into shelf compartments, UNITAS succeeded 18 times, achieving a 90% success rate; compared to π0.5’s 15 successes and Fast-WAM’s 13 successes. For the ball packing task, which involves placing a sphere into a bag and zipping it up, UNITAS succeeded 16 times, achieving an 80% success rate; compared to π0.5’s 11 successes and Fast-WAM’s 12 successes.

Ablation experiments decompose the contribution of each module. Removing 3D positional embeddings caused the LIBERO-Plus total score to drop from 85.74% to 81.35%, with the camera perturbation metric showing the most significant decline, falling from 79.11% to 66.73%. Removing action flow supervision reduced the total score to 78.89%. Removing action flow history tokens lowered the total score to 75.92%, with the robot initial state perturbation metric dropping by 19.87 percentage points. When all three modules were removed simultaneously, the model degraded into a standard VLA, achieving a score of only 67.58%. In contrast, adding scene flow supervision to the baseline model increased the total score from 85.74% to 86.47%, indicating that the ability to predict scene changes can reverse-optimize action learning performance.

04 Action-Driven 3D Scene Prediction

Another core capability of UNITAS is its simulator mode: given current observations and a sequence of actions, the model predicts the subsequent displacement changes of 3D points. The team evaluated it on RoboTwin’s reserved trajectory dataset, comparing UNITAS with PointWorld, which directly predicts 3D trajectories, and video-based approaches like Ctrl-World and X-WAM; for the latter, predictions are first lifted into 3D space and then calibrated using ground-truth depth from the initial frame. UNITAS leads across all four metrics: Average Displacement Error (ADE) for moving points drops from 0.656 cm in PointWorld to 0.427 cm, drift in static regions falls from 0.084 cm to 0.043 cm, Final Displacement Error (FDE) decreases from 1.050 cm to 0.825 cm, and accuracy within a 1 cm landing radius improves from 88.22% to 92.31%. UNITAS dominates all four metrics, with drift in static regions nearly halved—a clear dividing line between solutions that “measure accurately in 3D” and those that “look convincing in images.”

The training supervision signals for scene prediction come solely from simulation data, yet the model can generalize to unseen real-world scenarios. The team used the DROID and RH20T datasets, inputting only ground-truth gripper trajectories as conditions to have the model predict three-dimensional changes in real scenes. Even without supervised training on streaming data from real environments, the model still accurately aligns with the displacement and spatial configuration of objects within the frame, producing reasonable motion extrapolation results.

The evaluation phase standardized the comparison methodology. The team used RoboTwin's reserved trajectories, employing the same set of initial observations and a single segment of robot actions as conditions, with simulator states serving as the 3D ground truth; prediction results from video-route models were first mapped into 3D space before comparison, with the calibration process using only the depth of the initial frame's ground truth. Four metrics each serve their respective roles: Moving Point ADE represents average displacement error, FDE represents endpoint error at the final step, Static ADE measures drift in background static points, and Acc@1cm counts the proportion of moving point predictions falling within 1 centimeter of the ground truth.

This interface natively supports cross-embodiment generalization. Human hands and robotic grippers share the same action flow representation, while action flows, scene flows, and historical trajectories utilize a single tokenizer. Data from different embodiments and observation perspectives can be integrated into the same representation space for computation. The three-dimensional dynamics of the real world are thus bridged from simulation to real-world scenarios through this shared cross-embodiment representation.

Integrating policy generation and scene prediction into a single interface is another key design in this work. A single model outputs action commands when running only the policy branch; switching to simulator mode allows it to predict scene evolution. Both modes share the same 3D representation. Action flows and scene flows are two sides of the same coin in physical interaction, rather than being handled by two separate models. The team has named this technical route 'Physical Space World Action Modeling': interaction geometry shifts from an implicit background to an explicit prediction target, aligning action learning and world modeling along a single logical axis.