flex-pi/flexpi-pretrain-rgb-preview

Flex-π · Pretrain RGB (preview) — RGB + DINO

Preview. An early release of the pretrained Flex-π model; a stronger pretrained model will follow.

Flex-π is a 6B-parameter world-action model that supervises three visual futures — RGB, 3D pointmaps and DINOv3 semantics — alongside actions. A single checkpoint runs on any subset of its streams, so the compute/accuracy operating point is a runtime flag.

This checkpoint is a pretrain for bimanual 6-DoF YAM-family robots with one head camera and two wrist cameras. It trains on RGB + DINO modalities, no depth. It serves the action-only path and joint video and DINO generation. It is a starting point for fine-tuning; zero-shot use is possible on a matching rig. The sibling with optional depth, trained on a larger mixture, is flex-pi/flexpi-pretrain-preview.

Model details

Architecture	Mixture-of-Transformers: Wan2.2-TI2V-5B video expert (5.00 B) + ActionDiT action expert (1.02 B), 30 layers, HBridge 7 / 16 / 7
Cameras	cam_high 256×320, cam_left_wrist and cam_right_wrist 224×224, composited at 384×320 (robotwin_uniform14), 33 frames
Modalities	RGB + DINO (no depth)
Visual streams	RGB through the frozen Wan2.2 VAE; DINOv3 ViT-B/16 (frozen), temporal stride 2 (keep-far); pointmap stream off (p_present_pointmap = p_jp = 0)
Flex regime	p_present_video = p_present_dino = p_jv = p_jd = 0.5; cross-modal forcing on for video and DINO
Action head	34-D, shared by 6- and 7-joint arms. Channels 26 and 33 are each arm's 7th joint, which a YAM arm does not have: they are left out of the loss on YAM data and held at 0 when serving a YAM robot, which sends and receives 32-D
Wire format	32-D YAM end-effector state in, 32 × 32 absolute action chunk out
Normalization	per-episode q01/q99 over 32 channels (dataset_stats.json), normalized targets clamped to ±1
Training data	~760 hours of robot data: MolmoAct2 and in-house YAM teleoperation data
Checkpoint	step_286545, bf16, 12,082,213,301 bytes, sha256 273745e56bdbd559aa1852788c3ff6e5117109312fb63fa68a63494715410f9d

Usage

The serving and fine-tuning code is in the Flex-π repository; the zero-shot guide is docs/PRETRAINED.md.

huggingface-cli download flex-pi/flexpi-pretrain-rgb-preview --local-dir ./checkpoints/flexpi-pretrain-rgb-preview

config.yaml and dataset_stats.json stay beside the weights: the server reads the architecture back from the saved config, and --dataset-stats-path names the statistics the checkpoint was trained with.

Zero-shot server

D=./checkpoints/flexpi-pretrain-rgb-preview
python -u scripts/serve_yam_flexpi.py \
  --ckpt-path $D/checkpoints/weights/step_286545.pt \
  --dataset-stats-path $D/dataset_stats.json \
  --default-prompt "Pick up the plates from the table one by one and place each into a slot on the dish rack." \
  --host 0.0.0.0 --port 8000 --no-torch-compile --offload-text-encoder \
  --action-horizon 32 --mixed-precision bf16 --num-inference-steps 10 \
  <regime flags>

--num-inference-steps is 10 in the pre-release robot deploy of this architecture and 4 in the repository's configs/real_yam.yaml; neither has been measured on this checkpoint.

Regime	Regime flags
Action only	--infer-present-video true --infer-present-dino true --infer-present-pointmap false --infer-joint-video false --infer-joint-dino false --infer-joint-pointmap false
Joint video and/or DINO	--infer-present-pointmap false with --infer-joint-video true and/or --infer-joint-dino true, --infer-joint-pointmap false

The client sends RGB and state only. With no regime flags the model defaults to joint video and DINO generation with the pointmap stream off, and the connect metadata reports depth_required: false; --infer-present-pointmap true is not a trained regime and makes the server require depth. Video and DINO both absent is not a trained regime either. The server wraps the bare instruction in the training prompt template. A correct start logs action_dim_mask: channels [26, 33] zero-filled, Pointmap stream GLOBALLY OFF and a FlexPi inference regime: line. The client contract (state layout, frames, gripper sense) is in docs/YAM.md.

Fine-tuning

docs/FINETUNE.md covers fine-tuning from this checkpoint with the task preset configs/task/yam_finetune_pretrain_rgb_preview_3cam_34d_rel_1e-4.yaml: RGB + DINO, no depth (pointmap stream off), DINO temporal stride 2 (keep-far), and this checkpoint's own dataset_stats.json.

Scope

• Zero-shot use needs a bimanual 6-DoF YAM-family rig with one head and two wrist cameras speaking the 32-D wire format; LIBERO and RoboTwin need fine-tuning.

License

The weights are released under the MIT license, as is the code. The training data carries its own terms: MolmoAct2 BimanualYAM is Apache-2.0 and the released flex-pi YAM datasets are CC BY 4.0. At inference the model downloads DINOv3 timm/vit_base_patch16_dinov3.lvd1689m, which is under the DINOv3 License, and Wan2.2 components, which are Apache-2.0.

Citation

@article{yan2026flexpi,
  title   = {Flex-$\pi$: A Multi-Stream World-Action Model with Compute Flexibility},
  author  = {Yan, Ge and Liu, Jinghao and Fan, Yuzhi and Cai, Lei and Liao, Minwen
             and Zhang, Jesse and Fox, Dieter},
  journal = {arXiv preprint arXiv:2608.10860},
  year    = {2026},
  url     = {https://arxiv.org/abs/2608.10860}
}