Flex-$π$: A Multi-Stream World-Action Model with Compute Flexibility

TL;DR: Flex-π is a 6-billion-parameter robot policy that predicts not just future RGB frames but also future 3D geometry and object semantics — all squeezed into the same latent space, using a frozen video-generation VAE that turns out to encode 3D pointmaps almost losslessly despite never having seen one. Random per-stream dropout during training yields a single checkpoint that runs anywhere from a 60 ms action-only mode (faster than π0.5) to full "imagine everything" mode, and it beats strong VLA and world-action baselines by 2–6× on genuinely hard real-world bimanual tasks — sub-millimeter screw insertion, zipping a floppy fabric pouch — with striking data efficiency.

The setup: world-action models and what they're missing

A quick orientation for anyone who hasn't tracked this corner of robot learning. The dominant recipe for generalist manipulation is the vision-language-action model (VLA): take a pretrained vision-language model, fine-tune it to map images and instructions to action chunks. π0.5 is the canonical example. VLAs inherit web-scale semantic priors but, during action fine-tuning, they only predict actions — they never have to model what happens next.

World-action models (WAMs) are the emerging alternative: initialize from a pretrained video generation model and train it to jointly denoise future observations and actions. The intuition is that predicting the future forces a richer internal representation than imitation alone, and the video backbone brings strong spatiotemporal priors. Empirically this buys demonstration efficiency and generalization.

But there's a mismatch. Existing WAMs predict future RGB latents from a video VAE — a representation trained purely for pixel reconstruction. Pixel-faithful latents care a lot about texture and lighting and not at all, explicitly, about the two things manipulation actually needs: 3D geometry (where is the surface I'm about to contact, in metric space?) and object identity (which of these latents is "the plate"?). You could supervise those directly, but every obvious route has a tax: depth sensors mean new hardware, a 3D-specific encoder means training a new visual prior from scratch (forfeiting the video pretraining), and generating extra output streams means slower inference on the robot.

The paper's question: can a WAM get geometric and semantic supervision without paying any of those three taxes?

The free lunch: a video VAE that speaks pointmap

Here is the observation the whole paper pivots on. A pointmap is an H \times W \times 3 tensor where each pixel stores its 3D coordinate (x, y, z) — the standard output format of modern monocular geometry models like Depth Anything 3. Structurally, it's just a three-channel image. The authors tried the dumbest possible thing: feed pointmaps through the frozen VAE of Wan-2.2, a video generation model trained exclusively on RGB pixels, and see what comes out the other side.

It reconstructs them almost perfectly (their Figure 3). No fine-tuning, no adapter on the VAE itself.

Why should this work? A good VAE for natural video learns a latent code for smooth, spatially-structured 3-channel signals — and pointmaps are smoother than natural images in most respects (depth varies gradually except at object boundaries, which are exactly where RGB has edges too). The VAE's inductive machinery for compressing "images" generalizes to this out-of-distribution-but-image-shaped signal. It's a small empirical finding, but it unlocks the architecture: RGB and 3D geometry now live in the same latent space, produced by the same frozen encoder, and both inherit the video model's pretrained priors. No new sensor is needed either, since the pointmaps are estimated from RGB offline (Depth Anything 3 for annotation; real robot depth cameras optionally at deployment — more on "optionally" below).

Object semantics come from a third frozen encoder, DINOv3, whose patch tokens are famously object-centric. To keep the token count manageable, each 2×2 patch neighborhood is losslessly folded into one 3072-dim token (pixel-unshuffle), cutting DINO tokens 4× with zero information loss.

The architecture: three visual streams, one action expert, one trunk

Denoised futures RGB t+1 3D t+1 DINO t+1 32 actions

language + proprioception (cross-attention, always visible)

RGB and pointmap share one frozen VAE, so both streams inherit the video pretraining prior.

The backbone is a Mixture-of-Transformers (MoT): one shared 5B-parameter transformer trunk, initialized directly from Wan-2.2-5B, processes all visual tokens with per-modality projection adapters, while a narrower ~1B-parameter "action expert" carries its own attention and feedforward parameters (initialized by resampling Wan-2.2's weights via 1D interpolation, since the widths differ). Cross-stream attention happens only in the middle 16 of 30 blocks — early encoding and late decoding stay stream-private. Everything is trained with flow matching: each future stream's latent z_{t+1} is noised along the linear path z^\tau = \tau z_{t+1} + (1-\tau)\epsilon, and the model regresses the velocity z_{t+1} - \epsilon. (One wrinkle: the 3072-dim folded DINO tokens use x-prediction — predicting the clean features and converting to velocity — because velocity prediction degrades at high feature dimension.)

The attention is deliberately asymmetric: action tokens read the current observations and cross-attend to whichever generated visual futures are active, but no visual token ever attends to the action stream. This one-way street is what makes action-only inference possible at all — the visual generation can be dropped from the graph entirely and nothing upstream of the actions changes.

The training trick: stream dropout plus cross-modality forcing

The second idea is where the "Flex" comes from, and it has a subtlety worth understanding precisely. Two independent binary masks are sampled per training example, each entry Bernoulli(0.5):

• \mathbf{m}^{\text{in}}: which of the three visual streams are observed at time t (at least one, by rejection sampling).
• \mathbf{m}^{\text{out}}: which generated futures the action tokens are allowed to read, and which futures are mutually visible to each other.

Crucially, \mathbf{m}^{\text{out}} is an attention mask, not a loss mask. Every future stream is denoised and supervised on every training sample, no matter what. The masks only shape who can look at whom.

Because the masks are independent, you routinely get training samples where, say, the pointmap is absent from the input but its future must still be generated — from RGB and DINO alone. The authors call this cross-modality forcing, and it's the conceptual cousin of masked autoencoding: the model must learn that appearance, geometry, and semantics are mutually predictive, and internalize each representation well enough to synthesize it from the others. Their Appendix Figure 13 shows a checkpoint generating accurate scene geometry with pointmaps entirely withheld from the input.

The payoff is threefold:

1. One checkpoint, any inference regime. At deployment you choose which streams to feed in and which to generate — action-only, RGB+action, everything — as runtime flags, not retraining decisions.
2. The depth sensor becomes optional. On the real robot, withholding depth input costs 95.0% → 91.7% task completion on Put Plate on Rack. The geometry supervision lives in the weights, not the sensor.
3. It improves the policy outright. Ablating cross-modality forcing — same inputs, same architecture, just no forced synthesis — drops RoboTwin success by 21%. This is the strongest evidence that the mechanism isn't merely robustness plumbing; the mutual-prediction constraint builds a better representation.

The evidence

Real robot (the headline). Five bimanual tasks on a YAM platform, pre-trained on ~500 hours of AGIBOT World, per-task fine-tuned, scored with partial-credit rubrics over 10–20 rollouts. Two tasks are legitimately hard: Self-Repair Gripper (the robot re-installs its own gripper and drives a screw with a powered screwdriver, eight sequential stages, ±0.25–0.5 mm insertion clearances) and Soft-Bag Zipping (unzip a deformable fabric pouch, insert pens, zip it shut). On Self-Repair Gripper, full-joint Flex-π completes the entire eight-stage sequence in over half its rollouts; ManiFlow finishes 1 in 20 and π0.5 finishes zero. Averaged across tasks, Flex-π's success rate is 3.5× π0.5's and 2.3× ManiFlow's, with margins of +42.7 and +27.2 task-completion points over the strongest baseline on the two hardest tasks.

Speed. Action-only inference runs at ~60 ms per 32-step action chunk on an RTX 5090 (4 Euler steps), faster than π0.5, ManiFlow, and even a torch-compiled Fast-WAM (90 ms) — while beating all of them on every task. Full joint generation costs ~193 ms and buys a further +13% average success. One checkpoint spans a 3× latency range and ~24 points of success in the RoboTwin sweep.

The data-scaling result above is, to me, the most informative chart in the paper. At 50 demos per task, Flex-π nearly doubles the best baseline (78.8% vs 41.9%) — this is where the richer supervision signal should matter most, and it does. At 500 demos the gap shrinks to a few points, exactly as you'd expect if the auxiliary prediction targets are substituting for demonstrations.

Ablations (Section 4.5 — read this one). On five RoboTwin tasks with 50 demos each: adding DINO to video-only training gives +6.8% success; adding pointmaps on top gives a further +20%. So geometry is the bigger lever, semantics the smaller — but both matter, and both matter even when not generated at inference, which is the crux of the "supervision is free" claim.

Simulation leaderboards are less dramatic but respectable: 94.6% on RoboTwin's 50 tasks (edging out Qwen-RobotManip, whose pre-training corpus is 76× larger), 99.2% LIBERO average (tied for best), 80.9% on LIBERO-Plus perturbations — behind Qwen-RobotManip's 91.4% and π0.5's 84.7%, which the authors attribute to those baselines' stronger VLM backbones and vastly larger robot-data pre-training.

What to make of it, and where to be skeptical

The genuinely new bits. The frozen-VAE-encodes-pointmaps observation is small but generative: it means any latent-space world model built on a video VAE can be extended with geometric supervision at essentially zero architectural cost, and I'd expect it to be adopted quickly. The independent input/output masking with cross-modality forcing is the deeper contribution — it converts multi-stream prediction from a fixed architectural commitment into a deployment-time dial, and the 21% ablation delta suggests the forcing objective is doing real representational work, not just handling missing sensors. Together they support a clean thesis: what a world model is trained to predict matters more than what it consumes at test time. The OOD experiments back this up pointedly — action-only Flex-π (no depth input, no visual generation) drops only ~4 points under heavy visual shift while ManiFlow, which consumes depth at inference, drops 26.7. Supervision baked into weights beats sensors bolted onto inputs.

Reasons for caution. All real-robot baselines were trained by the authors (no published YAM numbers exist), which is unavoidable but means baseline tuning effort is uncontrolled — they do document identical per-task datasets and native configs for each baseline, and their retrained π0.5 numbers on RoboTwin look sane. The abstract's "2–7×" (2–6× in the paper body) refers to success-rate ratios on hard tasks where baselines are near zero, so the multiplier is partly a small-denominator effect; the task-completion margins (+27 to +43 points) are the more honest statistic. Real-world evaluation is 10–20 rollouts per task per method — enough to see the big gaps, not the small ones. The latency comparison explicitly does not hold engineering effort constant across methods (Appendix F.3), though they did optimize Fast-WAM too. The "almost losslessly" VAE claim is supported qualitatively (Figure 3) rather than with reconstruction-error numbers, which is a gap worth noting given how central it is. And LIBERO-Plus shows the ceiling: without a strong VLM backbone, Flex-π trails on language-heavy and camera-shift perturbations. Convergence is also slower — at least 10 fine-tuning epochs on real tasks, which the authors flag themselves.

Where to spend your time. Section 4.5 (ablations) plus Section 3.2 (the masking scheme) contain the ideas; Appendix G.3 (Self-Repair Gripper) is the best read on what the real-world evaluation actually demanded; and Appendix I is an unusually candid, quantitative inference-optimization writeup (TensorRT ladders, prefill/decode splits, what didn't work) that's useful well beyond this paper.

If the result holds up, the practical takeaway for anyone building WAM-style policies is simple: annotate your RGB data with off-the-shelf geometry and semantics, push all of it through the video VAE you already have, denoise everything jointly with independent input/output dropout — and then throw the extra streams away at deployment if you need the speed. The supervision was the point, not the prediction.