Shanghai Jiao Tong University and Partners Introduce UniFlowWAM, Elevating 3D Scene Flow to a Core Modality

By Ethan

Shanghai Jiao Tong University, Shanghai Chuangzhi Academy, and Qianzhi Intelligent introduce UniFlowWAM. Over the past six months, research on world-action models has increasingly trended toward extracting physical signals into abstract modalities, feeding them into encoders to train embodied AI models.

Shanghai Jiao Tong University, Shanghai Zhi Chuang Academy, and Agibot Jointly Release UniFlowWAM: When 'Flow' Becomes a Modality

Over the past six months, research on World-Action Models (WAMs) has increasingly trended toward extracting physical world signals into abstract modalities. These signals are processed as data and fed into encoders to train embodied models. As more research teams investigate which encoding formats maximize data absorption efficiency, modalities have expanded beyond world knowledge, action information, and world-action coupling. They now increasingly encompass time, causality, and future states. Ultimately, the concept of a modality itself is becoming more abstract.

As early as September, Bilawal Sidhu, former product manager and technical lead for Google Maps Immersive View, demonstrated a 3D "flow" modality. By combining time and space based on the trajectories of objects, changes in the world and object actions are reconstructed in 4D. The multifaceted nature of the "flow" concept reveals numerous new possibilities.

Prior to UniFlowWAM, the MVIG Lab at Shanghai Jiao Tong University, Shanghai Zhi Chuang Academy, and Agibot published ChronoFlow-Policy and LaMP at ECCV 2026. These works established 3D interactive flow as a mediating signal between observation understanding and action prediction. UniFlowWAM is their latest achievement.

Paper link: https://www.alphaxiv.org/abs/2610.30991

Project homepage: https://kksham111.github.io/UniFlowWAM_ProjectPage/

The Vision and Intelligence (MVIG) Lab at Shanghai Jiao Tong University, together with Shanghai Zhi Chuang Academy and Agibot, has just released UniFlowWAM. Reportedly, this unified World-Action Model integrates visual semantics and 4D geometry, elevating 3D scene flow to a third modality on par with vision and action. The co-first authors of this research are Jialin Tian and Zhuofan Li, both members of the MVIG Lab at Shanghai Jiao Tong University. Key scholars on the team include Professor Cewu Lu, Vice Dean of the Institute of AI at Shanghai Jiao Tong University and Co-founder and CEO of Agibot, with Assistant Professor Lixin Yang from the Institute of AI at Shanghai Jiao Tong University and Agibot serving as the corresponding author.

How should we understand this 3D scene flow? As a completely new modality, how does the model allocate attention across different modalities?

From "Light" to 3D Scene Flow

In classical methods, flow represents the trajectory of moving objects in the physical world, essentially acting as a line with physical significance. Where does a specific pixel in the current frame move in the next frame, and how does it move? Consequently, flow is used to describe two-dimensional displacement on the image plane.

However, using two-dimensional flow to depict object displacement in three-dimensional space has clear limitations. The true motion of objects in 3D space, along with projection changes caused by camera perspectives, cannot always be fully reconstructed into action trajectories using flow alone. In cases such as object rotation or localized trembling, two-dimensional flow is largely inapplicable.

Flow has been utilized in some VLA or WAM studies. In approaches like Im2Flow2Act and PointAction, for instance, flow serves directly as an explicit interface between perception and control. It acts as an intermediate representation that is predicted first and subsequently converted into actions in a separate stage. Flow can also be integrated within a single policy. Models such as TraceVLA, MolmoAct, 3D Flow Diffusion Policy, and LaMP use tracked or predicted flow as a conditioning signal for actions during inference to enhance policy performance. Alternatively, flow may simply serve as a training signal without functioning as an interface or a policy component. Additionally, frameworks like UMA jointly learn 3D keypoint trajectories alongside actions.

Flow holds significant value for unified future prediction and action learning. Models such as UWM and UVA learn both to predict the visual consequences of actions and to infer actions from future visual features. While these features capture the semantics of future target images, they fail to explicitly describe how the three-dimensional scene evolves into that state.

What, then, distinguishes the 3D scene flow proposed by UniFlowWAM from the aforementioned concepts of flow?

For UniFlowWAM, learning the flow trajectories of three-dimensional points improves action prediction. By explicitly supervising the 3D displacement of scene points, UniFlowWAM transcends the geometric states supervised by depth-based objectives. It combines future visual semantics with the 4D geometry derived from scene-level 3D flow, elevating flow to a modality on par with action, vision, and semantics. To integrate these elements effectively, 3D scene flow cannot merely serve as an input.

3D Scene Flow-Based Transformer

Capturing this insight, the UniFlowWAM team proposed a unified world-action model based on a diffusion Transformer architecture. The model assigns a dedicated expert to each modality, with the three experts exchanging information via joint attention.

UniFlowWAM extracts point maps and scene flow from the same demonstration clips used for policy learning via a monocular 3D tracker, then encodes them with a frozen 4D VAE to output flow latents. Using a DINOv3 encoder and an action encoder, the semantic and visual modalities each have their own dedicated expert, and all three experts exchange information exclusively through shared joint attention. During the mid-training and post-training phases, the flow modality is integrated into a pre-trained vision-action model. The vision and action experts initialize from pre-trained weights, while the flow expert is trained from scratch. The three experts are then jointly trained to adapt the pre-trained representations to this new modality.

Each function of the model corresponds to a specific training objective, with each column representing one objective. Ten objectives—including policy, visual prediction, and flow prediction—form a role grid that assigns one of three roles to each modality: target, condition, or missing. Solid squares denote clean condition inputs, noisy squares indicate noisy targets to be predicted, and blank squares represent learned placeholders. Because each objective has only one target modality, the grid covers all four combinations of the presence or absence of flow and future visual features for action prediction, while reserving three objectives each for flow prediction and visual prediction.

This intentional asymmetric design prioritizes condition selection for control variables. Simply put, when the target is the action, the presence or absence of flow and visual features forms exactly four combinations. The task objectives cover these combinations while reserving three objectives each for flow and visual prediction, leaving more condition combinations for the action modality. Meanwhile, the four objectives that exclude flow exactly recover the original task set of the pre-trained model, allowing the new modality to be integrated without degrading existing capabilities. Under the direct policy objective, because the flow branch already carries learned weights, deployment requires neither scene flow extraction nor flow encoding. Consequently, in the policy column, only the action is a noisy target, while both flow and vision are blank.

Empirical Results

The team evaluated UniFlowWAM across simulation environments, out-of-distribution settings, and real-robot dual-arm tasks to assess its learning capabilities, generalization ability, and real-world performance.

In the RoboCasa-GR1 simulation tests, UniFlowWAM achieved a peak success rate of 58.9%, performing slightly below OpenWAM-alpha. The version of UniFlowWAM augmented only with flow, without mid-training, achieved a 54.84% success rate, which improved to 56.38% with the addition of mid-training. This means that compared to the publicly released LDA, integrating flow alone yielded a 2.7 percentage point improvement, and incorporating mid-training further expanded this gain to 4.27 percentage points.

Researchers familiar with OpenWAM-α will recognize that the openness of its architecture enables theoretically optimal data absorption efficiency, with UniFlowWAM following closely behind. An analysis of the performance gains reveals that although UniFlowWAM's base model may be slightly weaker than OpenWAM's, the 3D scene flow modality compensates for this to a certain extent, yielding corresponding improvements in handling spatial variations dictated by geometric relationships.

However, UniFlowWAM underwent only 1,200 hours of mid-training. Given that LDA pre-training utilized the EI-30k corpus, which comprises approximately 30,000 hours of data, extending the training duration further could potentially yield even better performance for UniFlowWAM.

UniFlowWAM demonstrated highly impressive performance on both the LIBERO and LIBERO Plus benchmarks. It achieved an average success rate of 98.1% across eight tasks, significantly outperforming π0.5 and FastWAM, though it trailed OpenWAM-alpha and 4D WAM. On the LIBERO Plus benchmark, during a zero-shot evaluation involving seven types of perturbations across 10,030 episodes, UniFlowWAM attained an overall success rate of 87.7%. This result surpassed π0.5 and established a clear lead over Track4Action, FLEX-π, and LaMP. A breakdown of the results revealed that UniFlowWAM holds a remarkable advantage against camera viewpoint perturbations, lighting changes, and sensor noise. However, its average success rate dropped significantly when subjected to perturbations in the robot's initial state and object layouts.

This also highlights a unique characteristic inherent in extending 2D flow to 3D scene flow. As previously discussed, while the "flow" modality facilitates strong object understanding, the continuity of this representation within a given space grants it exceptional robustness against camera, lighting, and noise perturbations. However, if neither the agent undergoes significant changes nor the scene changes, the flow modality becomes ineffective or rapidly collapses on a specific task, despite being 3D.

The team deployed UniFlowWAM on an AgileX Cobot Magic for real-robot testing, executing four tasks: picking up peppers, transferring cabbage, placing pens into a drawer, and folding towels. Each task was evaluated independently based on specific milestones. For the cabbage and pepper task, the milestones included picking up the items, placing them into a pot, and picking up and placing each individual pepper. For the pen task, the milestones were picking up and placing each pen, stabilizing the drawer with the right gripper, and pushing the drawer closed. For the towel task, the milestones were picking up, lifting, folding in half, folding again, and flattening. Placing pens into the drawer and folding towels posed the greatest challenges for bimanual dexterous manipulation, with towel folding additionally testing the robot's visual understanding of deformable objects. If the model failed at a specific milestone during a trial but succeeded at a subsequent one, it still received credit for the latter milestone, with the final score calculated as an average.

The results show that UniFlowWAM significantly outperformed π0.5, DreamZero, and LDA, achieving an average success rate of 87.0% and a process score of 92.3. UniFlowWAM delivered the strongest performance across all four tasks. Notably, there was a stark discrepancy between the success rate and the process score for UniFlowWAM in the cabbage transfer task. The team noted that UniFlowWAM picked up and placed the cabbage just as the pot lid was being placed, causing the cabbage to roll out. In this specific trial, although UniFlowWAM experienced a minor failure at the very end, it still successfully completed 92.5% of the milestones.

The primary reason for this likely lies in the representational capacity of "flow" combined with multiple modalities reaching its limit. Because 3D flow records how points move rather than capturing forces, friction, mass, and deformation, it is also possible that the 4D geometric representations learned during training were not subjected to short-horizon reasoning, or that the mid-training duration for flow was insufficient. In contact-intensive phases, bimanual dexterous manipulation remains highly challenging. As tracking models continue to improve, flow extraction is expected to become faster and more accurate, making flow as easily scalable for robot training as video. Looking ahead, the team plans to conduct larger-scale joint pre-training.

What modality is the most effective

When selecting the "flow" modality, 4D VAE was not the initial choice.

During the exploration of joint training with 3D flow, the team experimented with several ineffective designs. One approach involved encoding flow using an image VAE, extending the Im2Flow2Act method—which writes 2D point trajectories into image channels before encoding—to 3D. Although the VAE faithfully reconstructed these images and the flow denoising loss decreased as expected, the generated flows were nearly identical across all inputs. The model collapsed to the dataset's average flow, failing to produce any usable flow.

Another approach involved gating the flow modality. To preserve the knowledge the LDA backbone acquired during pre-training, the team attempted to gradually unblock the connections from existing tokens to flow tokens by applying layer-wise gating to the attention logits. However, this yielded minimal improvement.

As an emerging internal trend in embodied AI, 'flow' enables models to encode transitions between states. Yet, what are its limitations? Across modalities, policies, and explicit interfaces, 'flow' appears to have exhausted the architectural forms it can take. The 4D VAE encoder it employs operates by having a tracker output a 3D point and its displacement to the next frame for every pixel in each frame. It then encodes the resulting position and displacement sequences separately before concatenating them into a single latent variable. In other words, it functions as a kinematic language—specifying where points are, where they move, and by how much—to reconstruct the analytical form of a state without incorporating actual physical quantities. Similarly, the recently released Natura-1 introduced another form akin to 'flow': a multimodal unified sequence based on flow matching.

Furthermore, because this flow is tracked from demonstration videos using Track4World, it can sometimes introduce errors, such as temporal inconsistencies, sudden direction reversals within continuous trajectories, or the loss of background and context under SAM segmentation. More importantly, flow-based predictions struggle in regions where actual physical contact occurs, as occlusion frequently arises at the boundary between the gripper and the object. Additionally, articulated constraints like those of a drawer, the precise fit between a pen and a drawer, and the folding topology of a towel lack explicit variables. Not only is this representation insufficient on its own, but supporting the input of these variables would also require massive amounts of data.

Overall, this representation still feels somewhat incomplete, leaving room for further advancement. A key question remains: since the 'flow' used in UniFlowWAM focuses primarily on policy, is it too simplistic for the physical world? Should improvements to 'flow' occur within the flow representation itself, or outside the model?

This brings us back to the fundamental question: for embodied models, what constitutes the most effective modality? Or, delving deeper, what is the most critical component within a modality? Clearly, more information is not inherently better. Building an effective modality depends on its explanatory power, its utility, and empirical data support. It must satisfy at least four conditions: it must explicitly describe state transitions rather than just the states themselves; it must possess its own supervision signal rather than serving merely as a byproduct of other tasks; it must be queryable as a condition and predictable as a target; and it must share the same spatiotemporal coordinate system with other modalities, otherwise joint attention becomes impossible.

UniFlowWAM demonstrates that flow-based modeling holds significant practical applicability and explanatory power for physical reality, potentially offering long-term implications for embodied AI. Forcing a model to 'predict flow'—requiring it to reconstruct 3D displacements within its own latent space—essentially enables the machine to internalize and understand the geometric dynamics of objects. Advancing the 'world model' from visual imagination to geometric understanding is far more efficient than merely 'reading flow.' However, every modality has its boundaries and cannot be flawless across all dimensions. A more pressing consideration may be how to combine physics-grounded modalities with others. Looking ahead, an interesting direction would be upgrading this 'flow' into a unified state encompassing contact, force, and deformation, or a representation incorporating foldability. This would allow explicit contact states, force and tactile feedback, deformation gradients, and object and joint structures to interact seamlessly during real-world training and inference. Achieving millisecond-level computation and highly accurate machine understanding could ultimately give 'flow' an entirely new form.