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

2026-08-11

AI Summary

Existing world models typically predict only RGB latents, lacking explicit modeling of 3D geometry and object semantics, which limits their capability in fine-grained manipulation tasks. This work proposes a novel approach based on a frozen video VAE, demonstrating for the first time that a general-purpose VAE can losslessly encode 3D point clouds without additional training or specialized sensors. By leveraging a shared latent space to jointly denoise visual signals and actions, and introducing a multi-stream hybrid Transformer architecture, the model enables single-model multimodal inference from arbitrary subsets of inputs. Integrating object-centric DINO semantics, cross-modal constraints, and stream-level Dropout, the method achieves substantially improved sample efficiency and outperforms the strongest baseline by 2–7× in both in-distribution and out-of-distribution settings on real-world bimanual dexterous manipulation tasks, while also surpassing π₀.₅ in inference speed.

Abstract

World-action models (WAMs) predict the future to act better, but nearly all of them predict only RGB latents, trained purely for pixel reconstruction, with no explicit signal for the 3D geometry or object semantics manipulation needs. We find a surprising free lunch: the same frozen video-generation VAE that encodes RGB also encodes 3D pointmaps almost losslessly, with no pointmap-specific training at all. This lets us supervise Flex-$π$, a 6B-parameter WAM, on 3D geometry and object-centric DINO semantics alongside RGB, at no cost in new sensors, new pre-training, or inference latency. Every visual signal is projected into this shared latent space and denoised jointly with actions inside a Mixture-of-Transformers backbone; per-stream dropout with cross-modality forcing then lets a single trained checkpoint run on any subset of these streams, from a fast action-only mode to full joint generation. The result is a policy that is exceptionally demonstration-efficient and generalizes well, beating the strongest baselines by up to 2-7$\times$ on dexterous, precise, real-world bimanual manipulation tasks both in and out of distribution, all while running faster than $π_{0.5}$. Our project website: https://flex-pi.github.io/

Problem

• world-action models
• 3D geometry
• object semantics
• visual representation
• manipulation tasks

Innovation

• world-action model
• multi-stream learning
• 3D geometry from VAE
• Mixture-of-Transformers
• compute flexibility