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

TL;DR

Flex-π is a 6B-parameter World-Action Model (UW / Allen AI, Yan et al.) that jointly denoises RGB, 3D pointmap geometry, and object-centric DINOv3 semantics with actions in a single shared latent space, using per-stream dropout with cross-modality forcing so one checkpoint runs on any subset of streams at inference. This produces a runtime speed-vs-accuracy dial (fast action-only vs full joint generation) without new sensors or visual priors. Reports beating strongest baselines by 2–7× on dexterous, precise, real-world bimanual YAM tasks in and out of distribution while running faster than π0.5 — including finishing all eight stages of a self-repair gripper task 11/20 vs the best baseline’s 1/20.

Key claims

A shared latent space carrying RGB, 3D pointmap, and DINO semantic streams through one 6B backbone can be jointly denoised with actions under a single flow-matching objective, using a frozen Wan-2.2 VAE for both RGB and 3D pointmaps and a frozen DINOv3 encoder for semantics [§How it works].
Per-stream dropout with cross-modality forcing — training to predict each stream even when other streams are missing — yields a single checkpoint deployable on all 56 input/output stream combinations without per-configuration finetuning [§How it works, architecture diagram].
Cross-modality forcing alone raises RoboTwin success by 47% relative on their controlled ablation [§How it works].
On five real-world bimanual YAM tasks in distribution, Flex-π (full joint) averages 83.0% task completion vs 58.0 (ManiFlow) / 52.1 (π0.5) / 31.7 (Fast-WAM) [Table 1 in-distribution].
Out-of-distribution average drops to 76.1% (Flex-π) vs 31.5 (ManiFlow) / 43.2 (π0.5) / 16.9 (Fast-WAM) — a widened gap on OOD, with Flex-π losing only 6.9 points from the ID average vs 26.5 for ManiFlow [Table 1 OOD].
On the eight-stage self-repair gripper task with a sub-millimeter insertion, Flex-π (full joint) finishes all eight stages in 11/20 rollouts vs 1/20 for the best baseline [§01 Precision and long horizon, ordered-completion bars].
Flex-π with half the training data (action only) performs similarly to π0.5 with all the data on Put Plate on the Rack; with full joint prediction on 50% data, it outperforms every baseline on full data [§03 Robustness, 50%-data panel].
Joint generation trades latency for accuracy: full joint mode is 193 ms vs 60 ms for action-only on an RTX 5090 (~3×); the action-only path is still faster than every baseline compared [§Conclusion and limitations].
On simulation, Flex-π leads RoboTwin at 94.6% in either mode; LIBERO is saturated with everything but π0 within 2.5 points [§Simulation benchmarks].

Method

Flex-π encodes three latent streams from a single RGB observation through frozen encoders: RGB latents from the Wan-2.2 video VAE, DINO object-semantic tokens from a frozen DINOv3 encoder (linearly projected into the backbone’s dimension), and 3D geometry latents by pushing a per-pixel 3D pointmap through the same Wan-2.2 VAE (rather than a geometry-specific encoder). Proprioception and language instructions condition all streams. The shared backbone denoises action chunks jointly with future latents in each stream under one flow-matching objective; actions are read out under shared self-attention with the predicted futures, so the policy conditions on the imagined future without ever decoding it to pixels.

Training uses per-stream latent dropout combined with cross-modality forcing: each modality must be predictable from any subset of the others, which prevents the shared backbone from splitting into three weakly-coupled channels and pushes it toward a mutually-predictive appearance/geometry/semantics representation. At deployment, a runtime argument picks the input/output stream combination (56 total from the same weights). The fast action-only path skips every future-stream computation and matches VLA-level latency (~60 ms); the full joint path runs 3× slower (~193 ms on RTX 5090) but adds accuracy.

Results

Real-world YAM bimanual, 5 in-distribution tasks: Flex-π full joint 83.0% avg vs ManiFlow 58.0 / π0.5 52.1 / Fast-WAM 31.7. Action-only path is 76.4% — beats every baseline while running faster than π0.5. On the self-repair-gripper eight-stage sequence (tightest stage: 4 mm bit into a 4.5 mm socket, ±0.25 mm), Flex-π completes all stages in-order in 11/20 rollouts vs 1/20 for the strongest baseline. On soft-bag zipping, unseen bags drop π0.5 from 42.8 → 17.2 and ManiFlow 31.9 → 6.9; Flex-π drops only 70.0 → 63.3.

On cluttered / unseen-object generalization: 95.0% on Put Plate and 70.0% on Sort Utensils under unseen conditions — 2.5 and 5.0 points below in-distribution; ManiFlow (with 3D inputs) drops 32.5 and 22.5 points. On 50% training data: Flex-π half-data action-only matches π0.5 full-data; Flex-π half-data full-joint beats every baseline full-data. Simulation: RoboTwin 94.6% (both modes leading); LIBERO saturated across baselines. Ablations show removing cross-modality forcing costs 47% relative RoboTwin success.

Why it’s interesting

Flex-π sits at the intersection of two active levers on the wiki’s board and stakes out a distinct point on each. On the World Foundation Models side, it extends the “WFM-as-policy-internal-substrate” pattern — established by Causal World Modeling for Robot Control (LingBot-VA) LingBot-VA and the “frozen trace-WM + small action expert” bet of μ₀: A Scalable 3D Interaction-Trace World Model µ₀ — with the sharpest filed instance of runtime compute flexibility: one checkpoint, 56 deployable combinations, a live speed-vs-accuracy dial rather than a per-configuration checkpoint. On the VLA Models side, it adds a new axis to the recipe-lever board: rather than the RGB-only latent futures of RynnWorld-4D: 4D Embodied World Models for Robotic Manipulation RynnWorld-4D (RGB+depth+flow with tri-branch DiT) or the wrist-only latent futures of World-to-Wrist: Task-Conditioned Future Wrist Modeling for Fine-Grained Robot Manipulation W²-VLA, Flex-π uses the same VAE weights for RGB and 3D pointmaps and pushes semantics through a linear DINOv3 projection — a compute-parsimonious answer to “which future modalities to predict” that keeps the encoder budget flat. The load-bearing empirical finding — cross-modality forcing alone lifts RoboTwin 47% relative — is a concrete recipe piece that transfers regardless of the specific encoder choices, and complements See like a Robot: Robot-Centric Pointmaps for Vision-Language-Action Models and WAM-TTT: Steering World-Action Models by Watching Human Play at Test Time WAM-TTT as a training-side answer to the “how do we make the future prediction actually help control” open question this concept keeps re-raising.