UniFlowWAM

A Unified World–Action Model with Visual Semantics and 4D Geometry

Jialin Tian1,2,* Zhuofan Li1,3,* Xinkai Wang2,3 Bokai Lin1,2,3 Chenyi Wang2 Xinyu Zhan1 Cewu Lu1,2,3 Lixin Yang1,2,3,✉

1MVIG@SJTU 2Shanghai Innovation Institute 3Noematrix *Equal contribution · ✉Corresponding author

MVIG@SJTU Shanghai Innovation Institute Noematrix
RoboCasa-GR1
58.9%best single evaluation; 56.4% mean over 11 seeds, from 52.1% for LDA-1B
LIBERO-Plus, zero-shot
87.7%seven perturbation types, trained on LIBERO only; 98.1% in-distribution
Real robot, four bimanual tasks
87.0%average success; π0.5, the strongest baseline, 68.5%

Actions say what the robot did. Flow shows how the scene moved.

Scene-level 3D flow tracks where every point of the scene goes in 3D. UniFlowWAM learns it as a third modality next to vision and action, yet deploys as a direct policy that needs no flow.

Read the abstract

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.

Scene flow trajectories drawn over a robot gripper moving above two pens next to an open drawer
Scene-level 3D flow on a pen-and-drawer demonstration, blue to red over the action horizon. Extracted with Track4World from the same RGB videos used for training, no depth sensor needed.

Three modalities. One diffusion transformer.

A diffusion transformer with one expert per modality; the three experts meet only in joint self-attention. The instruction and the current frame enter through Qwen3-VL-4B, which every expert reads by cross-attention.

UniFlowWAM architecture: current and future observations, point map and scene flow, and proprioception with action chunks are encoded and fed to a three-expert MM-DiT with joint self-attention
Select the figure to enlarge.

Mid-trained on 214,443 episodes from ten embodiments, then post-trained per benchmark.

Ten objectives. What makes the model unified.

In each objective, every modality is either a target, a condition, or an absent placeholder. Each modality is a target in some objectives and a condition in others. Select an objective to see what it asks.

Scroll sideways to see all ten objectives.

Simulation. Gains where the base is weakest.

RoboCasa-GR1, 24 humanoid tabletop tasks

Success rate (%). The gain concentrates on the hard tasks: on the 10 tasks where released LDA-1B scores below 50%, UniFlowWAM improves all 10, by 6.5 points on average.

Published resultsSR
UWM14.2
GR00T-N1.647.6
StarVLA47.8
Being-H0.749.2
FLARE55.0
LDA-1B55.4
UniFlowWAM†58.9
OpenWAM-α60.5
Repeated evaluation, seeds 0–10Mean ± SD
LDA-1B, released52.11 ± 2.24
LDA-1B, reproduced52.76 ± 1.16
UniFlowWAM without mid-training54.84 ± 1.00
UniFlowWAM with mid-training56.38 ± 1.47

†Best of seeds 0–10; published results are single evaluations. Bold is best and underline second best within each block.

LIBERO and zero-shot LIBERO-Plus

Success rate (%). Post-trained on LIBERO only, then evaluated on LIBERO-Plus without adaptation. The widest margin is under camera perturbations.

ModelLIBEROLIBERO-Plus
Cam.RobotLang.LightBg.NoiseLayoutTotal
OpenVLA-OFT97.156.431.979.588.793.375.874.269.6
π0.596.972.276.885.196.295.587.086.984.7
StarVLA-α–48.763.486.895.894.675.080.277.8
Fast-WAM97.640.065.368.891.458.263.875.565.3
4D-WAM98.645.264.390.694.357.769.179.271.0
FLEX-π98.573.753.184.391.993.893.482.980.9
OpenWAM-α99.333.876.188.097.087.139.877.569.2
LaMP98.364.569.688.295.397.476.973.879.3
Track4Action97.076.660.787.196.396.188.478.782.3
LDA-1B92.833.171.276.089.991.172.066.269.5
UniFlowWAM (ours)98.196.966.087.498.087.398.181.787.7

Bold is best and underline second best in each column. Leading methods score 97–99% on LIBERO, so it is reported only to verify that in-distribution success is preserved.

Four bimanual tasks. Strongest policy on every one.

An ALOHA-style platform with two 6-DoF arms. Four tasks that ask different things of the two arms; 50 demonstrations each, 50 trials per policy.

Success rate and process score

Success rate (%)Process score: share of milestones completed

Cabbage: grasp the cabbage, place it into the pot, grasp the lid by its knob, cover the pot
Chili: grasp the green chili and place it on the plate, then the red chili
Pen: pick up and place two pens into the drawer, hold the drawer unit steady, push the drawer closed
Towel: grasp opposite edge midpoints, lift, lower to fold in half, grasp a corner, fold into a quarter-sized towel

Instructions and scored milestones of the four tasks.

Show as a table

Rollouts

Trained on one green towel. Folds what it has never seen.

Out-of-distribution test. All 50 towel demonstrations use one green towel, laid flat at varying positions across the table, long edge toward the robot, give or take a small rotation. We step outside that range in two directions: how the towel is placed, and the towel itself. Success is less frequent than in distribution, but the policy still completes the folds.

Placement

Vertical. The training towel turned 90°, short edge toward the robot, so the grasp points and the fold direction are new. Two episodes, both folded without human help.

A different towel

Color. The same type of towel in a different color, beige.
Texture. A thinner towel, also in a different color, yellow.
Size. A small white cloth on a white tablecloth.
Softness. Much softer plush, also smaller and purple.

Flow objectives pay off where demonstrations are sparse.

Nested objective sets, everything else fixed. Every variant deploys as the same direct policy with no flow input.

Real-robot success rate (%) by objective set

On chili transfer, 4 of the 6 failures of policy-only training fall in the sideways 90° group. In that group, policy-only training succeeds in 20% of trials, while the full ten-objective grid succeeds in 80%.

The gap is rooted in the training data, which is not balanced: most chili grasps cluster near the lengthwise placement, and few pens start near the edge of their region. Therefore, the full ten-objective grid, with 3D flow as both a target and a condition, holds up better where demonstrations are sparse.

Histogram of demonstrated chili orientations: 6, 16, 11, 35, 22 and 10 percent from minus 90 to 60 degrees
Share of chili demonstrations at each angle: 0° is lengthwise, ±90° sideways.
Top camera view with the region of initial pen positions in red and its sparse boundary band in blue
Initial pen positions in the top camera. The red box spans all 49 demonstrations; the blue band is its boundary, where demonstrations are sparse.

Cite

The arXiv identifier will be added when the paper is released.

@article{tian2026uniflowwam,
  title   = {UniFlowWAM: A Unified World--Action Model with Visual Semantics and 4D Geometry},
  author  = {Tian, Jialin and Li, Zhuofan and Wang, Xinkai and Lin, Bokai and
             Wang, Chenyi and Zhan, Xinyu and Lu, Cewu and Yang, Lixin},
  year    = {2026}
}