Source-linked AI summary
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
TL;DR
Robot actions and future images do not specify how scene points move in 3D, leaving a gap in world-state representation. UniFlowWAM adds supervised 3D flow as a modality alongside actions and future visual features; across evaluated settings, it reports stronger real-world success than the strongest baseline.
Problem
Robot action labels and future visual features do not describe how scene points move in 3D, leaving the transition between observed and resulting states unspecified.
Source: §1 · pp. 1–2Method
UniFlowWAM jointly models actions, future visual features, and 3D flow with ten role-conditioned objectives that assign modalities as targets, conditions, or absent.
Source: §1 · p. 2, §3.4 · p. 6Results
Across simulation and real-world evaluations, UniFlowWAM reports 87.0% average real-robot success versus 68.5% for π0.5, the strongest baseline.
Source: §1 · p. 3Takeaways & Limitations
One jointly trained model serves as a policy, forward and inverse dynamics, and a predictor of future world state.
Source: §1 · p. 2Flow mid-training uses approximately 1,200 hours of demonstrations, substantially less than the EI-30k corpus used for LDA pretraining; larger-scale joint pretraining is left to future work.
Source: §5 · p. 12
Abstract
from arXiv · showhide
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.
1 Introduction
Robot actions and future images leave the 3D path of scene points unspecified, while prior uses of flow place it outside the policy. UniFlowWAM makes 3D flow a co-equal modality in a model that predicts actions and world representations, with reported gains across simulation and real-world manipulation.
1 Introduction
- Prior work used scene-level flow as an intermediate before action generation [38] [42], an inference-time prior [34], and jointly learned sparse 3D keypoint flow with actions [18], showing flow’s utility as a control signal and learnable variable.Scene-level 3D flow describes how each scene point is displaced in 3D over time.Source: §1 · p. 2
- UniFlowWAM joins future visual semantics and 3D flow geometry with actions as co-equal modalities, explicitly supervising scene-point displacement beyond depth-based geometric targets.One jointly trained model serves as a policy, forward and inverse dynamics, and a predictor of future world state.Source: §1 · p. 2
- Ten role-conditioned objectives assign each modality as target, condition, or absent, letting one model cover policy, dynamics, and prediction while retaining direct policy deployment without flow at test time.Under the direct policy objective, the flow stream contains learned placeholders.Source: Figure 1 · p. 2, §1 · p. 3
- 87.0% average real-robot success versus 68.5% for π0.5, while UniFlowWAM preserves 98.1% on LIBERO and reaches 87.7% zero-shot on LIBERO-Plus.On RoboCasa-GR1, adding flow raises the base model's average success rate by 4.3 percentage points.Source: §1 · p. 3
2 Related Work
Prior robot models combine actions with future visual predictions, while other approaches use flow as an interface, inference-time input, or training signal. UniFlowWAM instead places actions, 3D flow, and future visual features in one diffusion transformer.
Unified models for manipulation
- Unified future-image–action models predict visual consequences and infer actions, but their future features do not explicitly describe how the 3D scene evolves.Depth-based approaches provide 3D scene information, while feature-prediction methods use pretrained visual representations as future-state targets.Source: §2 · p. 3
How flow is used in policies
- UniFlowWAM places actions, scene-level 3D flow, and future visual features in one diffusion transformer, with a separate expert per modality and joint self-attention between experts.The objective determines which modalities are noised targets and which are clean conditions; predicted flow latents decode into point maps and scene flow.Source: Figure 2 · p. 4
- Prior policies use flow as an explicit perception-to-control interface, an inference-time condition, or a training signal; concurrent methods also learn 3D keypoint trajectories with actions.Some methods predict flow or point trajectories before a separate action module, while others inject tracked or predicted flow at inference.Source: §2 · pp. 3–4
3 Method
UniFlowWAM models actions, scene-level 3D flow, and future visual features as modalities of one diffusion transformer. Its formulation, flow representation, expert architecture, role-conditioned objectives, and training strategy make the model usable for both action prediction and world-state prediction.
3 Method
- UniFlowWAM adds a dedicated flow expert to the LDA visual–action backbone and trains ten objectives that specify prediction targets and available conditions.The model handles actions, scene-level 3D flow, and future visual features within one multimodal diffusion transformer.Source: §3 · p. 4
3.1 Problem Formulation
- The context comprises the current observation, proprioceptive state, and language instruction; action chunks, future DINOv3 features, and flow jointly define the conditional distribution.Future visual features use the frame at the horizon endpoint, while flow summarizes scene motion over the horizon.Source: §3.1 · p. 4
3.2 Scene-Level 3D Flow Representation
- Scene-level 3D flow combines each point's 3D position with its displacement to the next frame, encoded together by a pretrained 4D VAE.A pretrained dense point tracker supplies these quantities from monocular video, and the VAE decoder recovers both point maps and flow without depth-sensor data collection.Source: §3.2 · p. 5
3.3 Three-Expert Setup
- Each modality has its own expert and prediction head, and the experts exchange information through joint self-attention using private projections and shared spatiotemporal rotary positions.The visual token combines current-frame DINOv3 features with a future-frame slot, while flow enters as an offline-encoded latent.Source: §3.3 · p. 5
3.4 Role-Conditioned Multi-Objective Learning
- The ten objectives assign each modality as target, condition, or absent, covering modality-only prediction, pairwise conditioning, and action prediction conditioned on both flow and vision.Only target slots receive flow-matching loss; condition slots are clean inputs and absent slots receive neither input nor loss.Source: §3.4 · pp. 5–6
3.5 Training Strategy
- The pretrained visual–action model is mid-trained jointly on the ten objectives to add flow as both a prediction target and a conditioning modality, then post-trained on downstream demonstrations.The flow expert is incorporated into the pretrained backbone during this training strategy.Source: §3.5 · p. 6
4 Experiments
Experiments test UniFlowWAM in simulation and on real bimanual manipulation tasks, measuring benchmark performance, robustness, and the contribution of 3D flow. The model improves over baselines on RoboCasa-GR1 and the real robot, while preserving LIBERO performance and achieving strong zero-shot LIBERO-Plus results.
4 Experiments
- Experiments cover RoboCasa-GR1, LIBERO and LIBERO-Plus in simulation, plus four real-world bimanual tasks on an ALOHA-style robot.The four tasks are chili transfer, cabbage transfer, pen manipulation, and towel folding.Source: §4 · p. 6, §4.3 · pp. 8–9, Figure 3 · p. 9
4.1 Experimental Setup
- Mid-training uses 214,443 episodes from three public demonstration datasets, spanning ten embodiments, with offline scene-flow extraction and frozen 4D-VAE encoding.The corpus contains 133.0 million frames and yields around 42 million training windows.Source: §4.1 · p. 7
4.2.2 Post-training and Results
- Across the ten-task subset where released LDA averages below 50% success, UniFlowWAM with mid-training improves every task by 6.47 percentage points on average.On the remaining 14 tasks, the average gain is 2.71 points.Source: §4.2.2 · p. 7
- 98.05% average success on LIBERO and 87.65% on LIBERO-Plus zero-shot; the widest LIBERO-Plus margin is under camera perturbations, at 96.9% versus 76.6%.LIBERO is reported as an in-distribution performance check because leading methods already score 97–99%.Source: §4.2.2 · p. 8, Table 3 · p. 8
- 87.0% average success on the real robot, versus 68.5% for π0.5; on RoboCasa-GR1, UniFlowWAM with mid-training reaches 56.38%, versus 52.11% for released LDA.UniFlowWAM is the strongest policy on every real-world task.Source: §4.2.2 · p. 7, Table 2 · p. 8, §4.3 · p. 10
4.4 Ablation Studies
- The full objective grid is still improving after 120k steps, reaches 87.0% on LIBERO-Plus at 170k steps, and outperforms policy-only training on the real robot at equal steps.The reduced objective sets level off within the initial 120k-step budget.Source: §4.4 · p. 11
- On the real robot, the full ten-objective grid gains 12–18 points over policy-only training on chili transfer, pen manipulation, and towel folding.Relative to the no-flow variant, 3D-flow objectives add 10 points on pen manipulation and 8 on towel folding, versus 2 on chili transfer.Source: §4.4 · p. 11, Table 4 · p. 11
4.5 Discarded Designs
- The image VAE flow design collapsed to near-identical average flows, while gradual flow gating did not help RoboCasa-GR1 and was dropped from later experiments.The authors replaced the image VAE with a 4D VAE trained on geometry and flow.Source: §4.5 · p. 12
5 Limitations
The limitations concern noise in offline flow pseudolabels and the comparatively limited scale of flow pretraining.
5 Limitations
- Offline Track4World flow pseudolabels can contain temporal inconsistencies; SAM-based segmentation retains foreground flow, which the pretrained 4D VAE compresses effectively.Source: §5 · p. 12
- Flow mid-training uses approximately 1,200 hours of demonstrations, substantially less than the EI-30k corpus used for LDA pretraining [24]; larger-scale joint pretraining is left to future work.Source: §5 · p. 12