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

Problem Statement

World-action models (WAMs) have emerged as a powerful alternative to vision-language-action models (VLAs) by jointly predicting future observations and actions, thereby learning robust spatiotemporal representations and inheriting priors from large-scale video generation. However, current generalist WAMs almost exclusively predict future RGB image latents. These latents, trained purely for pixel reconstruction, capture appearance details but often fail to align with the explicit 3D geometry and object-level semantics required for precise robot manipulation.

Supervising WAMs directly on geometry and semantics typically incurs three significant costs:

1. Additional Sensors: Requiring new hardware modalities (e.g., depth sensors).
2. New Pre-training: Necessitating the training of new visual priors from scratch.
3. Inference Latency: Slowing down deployment due to complex multi-modal processing.

The paper asks: How can we amplify the strengths of WAMs (joint prediction and strong visual priors) to include 3D geometry and semantics without these sacrifices?

Methodology: FLEX-π

The authors introduce FLEX-π, a 6B-parameter WAM that predicts future RGB observations, 3D pointmaps, and object-centric DINO semantics simultaneously. The core innovation is achieving this "free lunch" of multi-modal grounding without new sensors or pre-training.

1. Shared Latent Space via Frozen VAE

FLEX-π utilizes a single, frozen Variational Autoencoder (VAE) from a pre-trained video generation model (Wan-2.2).

• RGB Encoding: The VAE encodes RGB images as usual.
• 3D Pointmap Encoding: Surprisingly, the authors find that this same frozen VAE, trained only on RGB pixels, can encode and decode 3D pointmaps almost losslessly. This allows 3D geometry to be projected into the exact same latent space as RGB without training a new encoder.
• Semantic Encoding: Object-centric semantics are derived using a frozen DINOv3 encoder.

2. Multi-Stream Mixture-of-Transformers (MoT) Backbone

The model embeds three visual streams (RGB, Pointmap, DINO) into a shared latent space of token streams.

• Architecture: A Mixture-of-Transformers backbone (5B parameters) processes these streams. It uses a shared visual transformer backbone but maintains independent feedforward projection networks for each modality.
• Action Expert: A separate, narrower Action Expert Transformer (1B parameters) cross-attends to the visual streams to generate action chunks.
• Causal Joint Generation: Action tokens attend to current observations and cross-attend to active future visual tokens. This allows action generation to benefit from the model's evolving representation of the future.

3. Flexible Training via Stream Dropout and Cross-Modality Forcing

To enable a single checkpoint to support any combination of inputs and outputs at inference, FLEX-π employs two key training strategies:

• Visual Stream Dropout: During training, input streams are randomly dropped (with probability 0.5), ensuring the model can operate with partial observations (e.g., RGB only).
• Cross-Modality Forcing: Crucially, even if a stream is dropped from the input, the model is still forced to denoise and predict that stream at the output. For example, if the 3D pointmap is missing from the input, the model must synthesize the future pointmap based on the RGB and DINO inputs. This forces the model to internalize a unified representation where appearance, geometry, and semantics are mutually predictive.

4. Inference Flexibility

At deployment, users can trade off latency and performance by selecting which streams to generate:

• Action-Only Mode: The model predicts only actions (fastest inference).
• Full Joint Generation: The model predicts actions plus future RGB, pointmaps, and DINO features (highest performance).
• Partial Generation: Any subset of streams can be generated.

Key Contributions

1. Free Geometric and Semantic Grounding: Demonstrating that a frozen RGB-trained video VAE can encode 3D pointmaps losslessly, allowing WAMs to be supervised on 3D geometry and DINO semantics without new sensors or pre-training.
2. Compute Flexibility: A single trained checkpoint supports any subset of input/output streams, allowing practitioners to choose their operating point on the speed–performance frontier at deployment time.
3. Cross-Modality Forcing: A training regime that forces the model to predict missing modalities from observed ones, significantly improving the internal representation and policy performance.

Experimental Results

The authors evaluate FLEX-π on simulation benchmarks (RoboTwin, LIBERO, LIBERO-Plus) and real-world bimanual manipulation tasks on a YAM robot.

• Real-World Performance: On five dexterous, precise, real-world tasks (including "Self-Repair Gripper" and "Soft-Bag Zipping"), FLEX-π outperforms strong baselines (π0.5, ManiFlow, Fast-WAM) by 2–7× in success rates.
  • Notably, the action-only variant of FLEX-π (which does not generate visual futures at inference) already outperforms all baselines and runs faster than π0.5.
  • Full joint generation further boosts performance, achieving up to a 42.7% absolute improvement in task completion on difficult tasks compared to the strongest baseline.
• Data Efficiency: FLEX-π is significantly more data-efficient. In RoboTwin, it outperforms baselines by 1.9–4.5× in low-data regimes (50 demonstrations). Even when trained on only 50% of the real-world data, FLEX-π outperforms baselines trained on the full dataset.
• Generalization (OOD): FLEX-π demonstrates robustness to out-of-distribution shifts (unseen objects, clutter, distractors). While baselines like ManiFlow and π0.5 suffer significant performance drops (25–26%) under distribution shift, FLEX-π drops by less than 5%.
• Ablations:
  • Removing DINO or Pointmap inputs during training reduces success rates significantly (by 6.8% and 20% respectively).
  • Removing Cross-Modality Forcing hurts success rates by 21%, confirming that forcing the model to synthesize missing modalities is critical for learning a strong shared representation.

Significance and Claims

The paper claims that FLEX-π demonstrates that a WAM can gain semantic and geometric grounding essentially for free. By leveraging the latent space of a pre-trained video model and employing cross-modality forcing, the model achieves:

• Better Policy Performance: Significantly higher success rates in complex, contact-rich, and long-horizon tasks.
• No Hardware Overhead: No need for additional sensors beyond RGB.
• No Pre-training Overhead: No need to train new visual priors.
• Deployment Flexibility: The ability to run in a fast, action-only mode or a high-performance, full-joint mode using the same checkpoint.

The authors conclude that while FLEX-π is highly effective, it still requires more fine-tuning epochs to converge than simpler models, and its full-joint mode is slower than parameter-comparable VLAs. They suggest that future improvements could come from combining FLEX-π with stronger VLM backbones and access to even larger robot datasets.