geyan21/flex-pi

TL;DR — Flex-π (UW + AI2, Dieter Fox's group) is a 6B world-action model for manipulation that co-denoises future RGB, 3D pointmaps, DINOv3 features, and actions as parallel token streams, then lets you pick at inference which streams to feed in and which to generate — 56 regimes from one checkpoint, spanning a 60 ms/call action-only VLA to a 193 ms full world model. The framing — compute flexibility as a first-class training objective rather than a distillation afterthought — is genuinely new, and the release is complete: code, checkpoints, datasets, and a real-robot serving stack, all MIT.

The idea

Most world-action models predict one thing: future RGB latents. That's a strong prior but a pixel-reconstruction objective carries no explicit signal for geometry or object identity — the things manipulation actually needs. Flex-π adds two more predicted futures: 3D pointmaps (annotated offline with Depth Anything 3, so no depth sensor is required at deployment) and DINOv3 semantic features.

The neat empirical finding enabling this: the Wan-2.2 video VAE, trained only on RGB, encodes 3D pointmaps almost losslessly. So RGB and geometry live in one shared latent space, DINOv3 tokens ride alongside, and a ~1B action DiT expert joins them inside a Mixture-of-Transformers backbone where everything is co-denoised jointly.

The training trick is cross-modality forcing: visual streams are dropped at random per-sample, and the model must generate the streams it never observed from the ones it did. This forces the backbone to internalize each modality's information redundantly — which is exactly what makes any input/output subset viable at inference. Two runtime flag families (infer_present_* for inputs, infer_joint_* for outputs) select the regime with no retraining.

Evidence

All numbers are the authors' own, but they're specific. Real-world: five bimanual YAM tasks, beating π0.5, ManiFlow, and Fast-WAM on all of them — 2.3× the success rate of the strongest baseline — including an eight-stage gripper self-repair with a ±0.25 mm-clearance insertion, plus the best robustness to unseen objects and clutter. Data efficiency is the more interesting claim: the world-model objectives act as a substitute for demonstrations — 1.9–4.5× baseline success at 50–100 RoboTwin demos/task, and on the real robot, half the demos beat all baselines at full data. Full-data results are near ceiling: 94.6% across RoboTwin's 50 tasks, up to 99.2% on LIBERO. The README gives no per-task breakdown; the paper (arXiv 2608.10860) presumably does.

What's released and what to expect

Unusually complete for a paper repo: training/eval/deployment code, RoboTwin and LIBERO checkpoints (including a no-dropout "full-joint-star" variant), all datasets on HuggingFace in LeRobot v2.1 format, a WebSocket policy server for real robots with a documented msgpack wire contract, and TensorRT engine builds. Docs cover install (conda/uv/Docker), per-benchmark training guides, and — refreshingly — "rules that have caused emergency stops on real hardware."

Gaps: the large-scale pretrained checkpoints (YAM, AgiBot World, DROID) are "planned," so what's out now is benchmark-specific. Training needs 4–8× 80 GB GPUs (LoRA fine-tuning also "planned"); inference fits one 4090/5090 at 16–26 GB. Repo is two weeks old with 3 forks, so expect rough edges. The main thing to verify independently is how much the visual streams actually buy over action-only mode — the README says joint generation "costs latency and wins more" but quantifies the win only via the data-efficiency claims.

Try it

git clone --recurse-submodules https://github.com/geyan21/flex-pi.git
cd flex-pi
# follow docs/INSTALL.md, then:
huggingface-cli download flex-pi/flexpi-libero-fulljoint-star \
  --local-dir runs/flexpi-libero-fulljoint-star
CKPT=$(ls runs/flexpi-libero-fulljoint-star/checkpoints/weights/*.pt) \
DATASET_STATS=runs/flexpi-libero-fulljoint-star/dataset_stats.json \
  bash scripts/eval_flexpi_libero_4suite.sh

Videos at flex-pi.github.io; datasets and checkpoints at huggingface.co/flex-pi.