Driving · 2026-10-23

UniFlowWAM: A Unified World-Action Model with Visual Semantics and 4D Geometry

Jialin Tian, Zhuofan Li, Xinkai Wang, Bokai Lin, Chenyi Wang, Xinyu Zhan, Cewu Lu, Lixin Yang

AI-assisted summary based on the paper. Consult the linked source for methods and limitations. Benchmark numbers appear after editor verification.

TL;DR

UniFlowWAM is a unified diffusion transformer that treats action, scene-level 3D flow, and future visual features as co-equal modalities. It uses one expert per modality with joint attention and trains with a ten-objective role grid (target/condition/absent) so the model can act directly without flow at test time. The authors evaluate in simulation on RoboCasa-GR1, LIBERO, LIBERO-Plus and on real bimanual manipulation tasks.

Why it matters

The paper targets how robot action labels omit the intermediate 3D world transition. By integrating scene-level 3D flow as an internal supervised modality (rather than an external module), the authors aim to improve manipulation success and robustness while preserving direct policy deployment that requires no flow input at test time.

Method

Model: a diffusion transformer with one expert per modality (actions, scene-level 3D flow, future visual features) exchanging information via joint attention; experts have separate parameters and meet in joint self-attention.
Training: ten role-conditioned objectives (target/condition/absent) over a horizon, with flow as a supervised target and/or condition; a task-balanced sampler ensures all objectives appear in each batch.
Deployment: under the direct policy objective, the flow stream carries learned placeholders so the policy needs no scene-flow extraction or flow encoding at test time.

Limitation

Flow tracking noise and representation shift: the method uses Track4World to generate offline 3D flow pseudo-labels; these labels can contain temporal inconsistencies (e.g., isolated reversals). The authors also note limited flow pretraining: the flow mid-training uses ~1,200 hours, substantially less than the EI-30k corpus used for LDA pretraining, and larger-scale joint pretraining is left to future work.

Abstract (from arXiv)

Robot demonstrations record what the robot saw and what it did. They do not specify how the 3D world changed in between. Scene-level 3D flow describes it by tracking 3D scene points, providing the world representation with 4D geometry. Prior work has used flow to improve robot policies. Flow served as an intermediate prediction that a separate model converts into actions, as a prior injected at inference time, or as an extra output learned together with the action. In all of these designs, flow stays outside the policy, as a separate module or as an extra prediction head. We introduce UniFlowWAM, a Unified World-Action Model that makes 3D scene flow a third modality on par with vision and action. Its world representation therefore covers both visual semantics, through future visual features, and 4D geometry, through 3D flow, and flow receives its own supervision. UniFlowWAM is a diffusion transformer with one expert per modality, and the three experts exchange information through joint attention. We train UniFlowWAM under a ten-objective role grid, where each objective assigns every modality exactly one of three roles, namely target, condition, or absent. The same model thus serves as a policy, as forward and inverse dynamics, and as a predictor of future visual features or 3D flow. We evaluate UniFlowWAM on RoboCasa-GR1, LIBERO, and LIBERO-Plus in simulation, and on four real-world bimanual manipulation tasks. On RoboCasa-GR1, UniFlowWAM raises the success rate of its base model without flow from 52.1% to 56.4%. It preserves in-distribution success on LIBERO at 98.1% and reaches 87.7% zero-shot on LIBERO-Plus. On the real robot, it achieves 87.0% average success, compared with 68.5% for the strongest baseline.

Publication record

Bibliographic evidence describes where and when a record appeared. It does not validate scientific claims.

First public release: 2026-10-22 · Indexed by Pub-AI: 2026-10-23 UTC

preprint