[Paper Notes] Flex-π: A Multi-Stream World-Action Model with Compute Flexibility
Published:
Flex-π turns one robot checkpoint into a family of deployment-time policies. It jointly learns four future streams—actions, RGB latents, 3D pointmap latents, and object-centric DINO features—inside a shared 6B-parameter world-action model. At inference, an operator can generate actions alone for low latency or activate visual futures for higher task completion. The key training device is a pair of independently sampled stream masks: one chooses which current modalities the model observes, while the other chooses which predicted futures the action stream may read. Every future remains supervised, including modalities absent from the input, forcing appearance, geometry, and semantics to become mutually predictable.
The paper’s strongest contribution is this coupling of richer world supervision with compute-flexible deployment. On the authors’ real bimanual platform, action-only inference runs at about 60 ms per policy call and reaches 76% average task completion; full joint generation takes about 193 ms and reaches 83%. The gains are largest on gripper self-repair and soft-bag zipping, where long horizons, contact, small clearances, and deformable objects make RGB-only prediction especially brittle.
Paper Info
“Flex-π: A Multi-Stream World-Action Model with Compute Flexibility” is by Ge Yan, Jinghao Liu, Yuzhi Fan, Lei Cai, Minwen Liao, Jesse Zhang, and Dieter Fox from the University of Washington and the Allen Institute for AI. It is an August 2026 preprint: arXiv:2608.10860. The project page contains interactive stream configurations and real-robot videos, while the code repository provides training and deployment material.
Why RGB Futures Are an Incomplete World Model
Recent world-action models generate a future visual representation together with an action chunk. This can improve policy learning in two ways: future prediction shapes a representation around scene dynamics, and a video-generation backbone contributes priors learned from large video corpora. Most WAMs, however, predict latents from an RGB reconstruction model. Such latents preserve appearance and motion but have no explicit objective for metric geometry or object-level semantics.
Manipulation needs all three views of the same scene:
| Stream | Information supplied to control | Typical question |
|---|---|---|
| RGB | appearance and spatiotemporal video priors | What will the scene look like? |
| DINO | object- and part-level semantic structure | Which entities and regions matter? |
| Pointmap | explicit 3D scene geometry | Where are surfaces and clearances? |
Flex-π derives DINO features and pointmaps from RGB using frozen DINOv3 and Depth Anything 3 models. Its surprising engineering observation is that the frozen Wan-2.2 video VAE, although trained on RGB, can also encode and reconstruct image-shaped 3D pointmaps accurately. RGB and geometry can therefore share one pretrained latent interface, avoiding a separately pretrained 3D VAE.
One Backbone, Four Generated Streams
At time (t), the policy receives language (l), proprioception (s_t), and any non-empty subset of RGB (o_t), pointmap (p_t), and DINO features (d_t). Its maximal conditional model is
[ \pi_\theta!\left( a_t,o_{t+1},p_{t+1},d_{t+1} \mid o_t,p_t,d_t,s_t,l \right), ]
where (a_t) represents an action chunk and each future visual output is also a temporal chunk.
flowchart TD
O["Current RGB"] --> VAE["Frozen Wan-2.2 VAE"]
P["Current pointmap"] --> VAE
O --> DINO["Frozen DINOv3"]
VAE --> R["RGB latent stream"]
VAE --> G["Pointmap latent stream"]
DINO --> S["Semantic token stream"]
R --> M["5B shared visual transformer"]
G --> M
S --> M
L["Language + proprioception"] --> M
M --> FR["Future RGB latent"]
M --> FG["Future pointmap latent"]
M --> FS["Future DINO features"]
M --> A["1B action expert"]
FR --> A
FG --> A
FS --> A
A --> AC["Action chunk"]
The backbone follows a Mixture-of-Transformers design. A 5B visual transformer initialized from Wan-2.2 processes all three visual streams with shared transformer weights and stream-specific adapters. A narrower, roughly 1B action expert has separate attention and feed-forward parameters. Cross-stream attention occurs in the middle 16 of 30 blocks, leaving early encoding and late decoding stream-specific.
The attention direction is important. Action tokens may attend to current observations and active future visual tokens, so the evolving imagined future can influence control. Visual tokens never attend to action tokens. This one-way dependency prevents the future predictor from simply encoding the target action and gives the action stream a clean removable fast path.
Joint Flow Matching Objective
For a target latent (z_1), Flex-π samples Gaussian noise (\epsilon) and interpolates along
[ z_\tau=\tau z_1+(1-\tau)\epsilon, \qquad \tau\sim\mathcal U[0,1]. ]
The flow network learns the constant velocity from noise to data:
[ \mathcal L_{\mathrm{FM}}(z_1) = \mathbb E_{z_1,\epsilon,\tau} \left|v_\theta(z_\tau\mid \tau,c)-(z_1-\epsilon)\right|_2^2. ]
Actions, RGB latents, pointmap latents, and DINO features are optimized together:
[ \mathcal L(\theta) = \lambda_a\mathcal L^a_{\mathrm{FM}}(a_t) +\sum_{i\in{o,d,p}} \lambda_i\mathcal L^i_{\mathrm{FM}}(i_{t+1}), \qquad \lambda_a=\lambda_o=\lambda_d=\lambda_p=1. ]
The DINO head uses clean-feature prediction instead of velocity prediction because folding each (2\times2) patch neighborhood creates high-dimensional 3,072-D tokens. At inference, active output streams are generated with four Euler integration steps in the main experiments.
Two Masks Create Compute Flexibility
Flex-π samples two independent binary masks over RGB, DINO, and pointmap streams for every training example.
The input presence mask (m^{\mathrm{in}}\in{0,1}^3) selects the current modalities available to the model. Each visual input is dropped with probability 0.5, while at least one stream is retained. The output attention mask (m^{\mathrm{out}}\in{0,1}^3) selects the future visual streams visible to the action tokens and controls attention among future streams.
Crucially, (m^{\mathrm{out}}) is an attention mask, not a loss mask. All three visual futures are denoised and supervised on every training example. When the pointmap input is absent, for example, the model may still need to predict future geometry from RGB and DINO. The authors call this cross-modality forcing.
This scheme produces (7\times8=56) deployment configurations: seven non-empty subsets of current visual inputs and eight subsets of future visual outputs, with actions always generated. A single checkpoint can therefore run as:
- an action-only policy;
- an RGB-future WAM;
- a geometry-and-action model;
- a full RGB + DINO + pointmap + action generator;
- or another intermediate configuration.
Cross-modality forcing contributes more than missing-sensor robustness. Removing it in the five-task RoboTwin ablation reduces average success by 21%. Reconstructing an absent modality from the remaining ones encourages the shared trunk to represent correspondences among appearance, object identity, and geometry.
Pretraining and Deployment Recipe
Flex-π is initialized from Wan-2.2-5B and pretrained on roughly 500 hours from 100 tasks in AGIBOT World-Beta. The dataset contains a head camera and two wrist cameras recorded at 30 Hz; Depth Anything 3 supplies offline pointmap annotations. Domain-specific finetuning follows for RoboTwin, LIBERO, and each real-robot task.
On the real bimanual YAM platform, the three camera views are tiled into one canvas. The model conditions on a single observation with no visual history, predicts 32 robot actions, and optionally generates a nine-frame latent window containing the current frame plus eight futures. The controller executes the complete 32-step chunk and replans every 1.07 seconds.
The distinction between input and output streams matters when reading the latency claims. Action-only means that no future visual stream is generated; it does not inherently mean RGB-only observation. Flex-π can also omit pointmap input, and the paper’s real-world ablation reports no measurable loss on plate placement when depth is withheld. Thus geometry can act as training supervision without becoming a mandatory deployment sensor.
Real-World Results
The real-robot suite contains five bimanual tasks. Every method receives the same per-task dataset, and each task uses its own finetuned policy. Evaluation reports both normalized partial-credit task completion and binary full success over 10–20 trials per task.
| Task | Flex-π action-only | Flex-π full joint | Strongest baseline |
|---|---|---|---|
| Put Plate on Rack | 84.2 | 95.0 | ManiFlow: 75.8 |
| Sort Utensils | 70.0 | 75.0 | ManiFlow: 55.0 |
| Kitchen Organization | 96.3 | 98.8 | ManiFlow: 93.8 |
| Self-Repair Gripper | 66.9 | 76.0 | ManiFlow: 33.3 |
| Soft-Bag Zipping | 64.9 | 70.0 | π₀.₅: 42.8 |
| Five-task average | 76.0 | 83.0 | ManiFlow: 58.0 |
The two hardest tasks reveal what the extra streams buy. Self-Repair Gripper is an eight-stage sequence in which the robot picks up and inserts its own replacement gripper, places a screw, uses an electric screwdriver, and clears the workspace. Critical insertions have only ±0.25–0.5 mm clearance. Full-joint Flex-π completes the entire sequence in 11 of 20 rollouts; the strongest baseline completes it once.
Soft-Bag Zipping requires opening a deformable pencil case, placing a pen inside, reacquiring the zipper pull, and closing the loaded pouch. The object’s geometry changes after every contact. Full-joint task completion reaches 70.0%, compared with 42.8% for π₀.₅ and 31.9% for ManiFlow. On an unseen bag, Flex-π falls from 70.0% to 63.3%, while π₀.₅ falls to 17.2% and ManiFlow to 6.9%.
These results use substantial task-specific data. The datasets range from 152 demonstrations / 1.2 hours for utensil sorting to 802 demonstrations / 11.8 hours for gripper repair, plus 570 DAgger correction episodes / 5.6 hours for the repair task. “Demonstration-efficient” is therefore a comparative claim against the evaluated baselines, not a claim of few-shot real-world learning.
Simulation Results and Ablations
On the 50-task RoboTwin benchmark, both action-only and full-joint Flex-π average 94.6% success under the clean and randomized settings, slightly above the reported VLA and WAM baselines. The low-data regime is more informative: with 50 randomized demonstrations per task, full-joint Flex-π reaches 78.8%, compared with 41.9% for Fast-WAM, 31.4% for π₀.₅, and 17.2% for LingBot-VA. At 100 demonstrations, the corresponding scores are 87.0%, 68.1%, 44.7%, and 32.2%.
LIBERO is close to saturation. The flexible checkpoint scores 98.4% action-only and 98.5% full-joint. A task-fitting variant trained without stream dropout reaches 98.7% and 99.2%, illustrating the tradeoff: dropout improves one-checkpoint flexibility and robustness, while a fixed input/output regime can fit a closed benchmark slightly better.
The five-task RoboTwin ablations isolate the mechanism:
- adding DINO to video input raises success by 6.8%;
- adding pointmaps on top of video and DINO adds another 20%;
- with the same checkpoint and RGB-only input, action-only inference obtains 40.2% at about 60 ms;
- generating video raises success to 60.4%;
- generating RGB, DINO, and pointmap futures reaches 63.8% at about 193 ms;
- removing cross-modality forcing lowers success by 21%.
The same weights span more than a threefold latency range and a 23.6-point success range. Compute flexibility is therefore an empirical speed–accuracy frontier, not just an architectural option.
Strengths
The work connects representation supervision directly to deployment behavior. Its mask design creates meaningful operating points from one checkpoint, and the ablations separately measure input modalities, generated outputs, cross-modality forcing, and denoising steps. The real-robot comparison also controls the training data: Flex-π, π₀.₅, ManiFlow, and Fast-WAM see the same demonstrations for a given task.
The method uses existing visual priors efficiently. DINOv3 supplies semantics, Depth Anything 3 supplies pointmap annotations, and Wan-2.2 supplies both the latent visual interface and the transformer initialization. The action policy can retain benefits from geometric and semantic supervision even when expensive future streams are disabled.
Finally, the task suite tests more than coarse pick-and-place. Sequential self-repair, screwdriver use, tight insertion, zipper manipulation, clutter, and unseen objects provide evidence that future representations help where scene state must be tracked through contact and long horizons.
Limitations and Open Questions
Flex-π remains a large and data-hungry system. It has roughly 6B parameters, uses 500 hours of robot pretraining, requires at least 10 finetuning epochs on the reported real tasks, and is benchmarked on an RTX 5090. Full joint generation costs about three times the latency of the action-only path.
The real-world policies are finetuned separately per task. The experiments establish a reusable pretrained backbone and flexible inference mechanism, while open-vocabulary multi-task deployment on the YAM platform remains untested. Trial counts are 10–20 per task, and the headline “task completion” metric awards partial credit; binary end-to-end success is lower, especially for long sequences.
The geometry story also deserves careful interpretation. AGIBOT pointmaps are generated offline by Depth Anything 3, and real hardware can provide depth. The masking results show that pointmap input can be removed, yet the complete training pipeline still pays for geometric annotation and multi-stream optimization. The paper reports longer convergence from the added modalities and forcing objective.
On LIBERO-Plus, π₀.₅ and Qwen-RobotManip slightly outperform Flex-π in aggregate. The authors attribute this gap to stronger VLM semantics and much larger robot pretraining corpora. This suggests that multi-stream world prediction complements language-semantic scale; it does not replace it.
The paper is also a recent preprint whose hardware results come from the authors’ platform. Independent replication, tests on additional robot embodiments, and comparisons under equal pretraining scale would strengthen the generality claim.
Takeaways
Flex-π makes three ideas concrete:
- A useful robot future has multiple representations. RGB captures appearance, DINO supplies object semantics, and pointmaps expose geometry.
- Missing-modality prediction can shape the policy trunk. Cross-modality forcing turns each stream into supervision for the others and improves action prediction even when some streams are absent at deployment.
- World-model compute can be a runtime choice. One checkpoint covers a fast action policy and several richer WAM configurations, producing a measurable speed–accuracy frontier.
The broader lesson is that a world-action model need not commit to one fixed definition of “the future.” Appearance, semantics, and geometry can share a latent dynamics model, while the deployment system chooses how much of that future to instantiate for the task and hardware at hand.