flex-pi/flexpi-pretrain-preview

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

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 large-scale pretrain for bimanual 6-DoF YAM-family robots with one head camera and two wrist cameras. It trains on RGB + DINO + depth modalities; depth is optional at inference. Training depth is DA3 for ABC-130k, MolmoAct2 and AgiBot World, and sensor depth for the YAM data. It is a starting point for fine-tuning; zero-shot use is possible on a matching rig. The sibling without depth is flex-pi/flexpi-pretrain-rgb-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 + depth (depth optional at inference)
Visual streams	RGB and 3D pointmap through the frozen Wan2.2 VAE; DINOv3 ViT-B/16 (frozen), temporal stride 1
Flex regime	p_present_* = p_j* = 0.5 for video, DINO and pointmap; cross-modal forcing on
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	frame-pooled q01/q99 over 34 channels (dataset_stats.json), normalized targets unclamped (action_norm_clamp: inf)
Training data	~4,600 hours of robot data: ABC-130k, MolmoAct2, AgiBot World and in-house YAM teleoperation data
Training depth	DA3 for ABC-130k, MolmoAct2 and AgiBot World; sensor depth for YAM
Checkpoint	step_190000, bf16, 12,082,212,409 bytes, sha256 b168eb68f255b98bf4805ad762dd870e06dd5f0f6f044f3e4280622338799eb2

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-preview --local-dir ./checkpoints/flexpi-pretrain-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-preview
python -u scripts/serve_yam_flexpi.py \
  --ckpt-path $D/checkpoints/weights/step_190000.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	Client sends depth + K
no depth, 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	no
with depth, action only	same, with --infer-present-pointmap true	yes
no depth, joint video + DINO	--infer-present-pointmap false --infer-joint-video true --infer-joint-dino true --infer-joint-pointmap false	no
with depth, full joint	none (the trained default)	yes

With no regime flags the model runs full joint generation and requires depth; the connect metadata reports this as depth_required. Depth is uint16 millimetres for all three cameras or none, on the per_cam_hw grid from the metadata, with intrinsics [fx, fy, cx, cy] on the same grid. The server wraps the bare instruction in the training prompt template. A correct start logs action_dim_mask: channels [26, 33] zero-filled 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_preview_3cam_34d_rel_1e-4.yaml, which reads this checkpoint's statistics from dataset_stats_rel32.json (the same statistics without channels 26 and 33, the 7th-joint slots a YAM arm does not use).

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, RoboTwin and AgiBot need fine-tuning.

License

The weights are released under the MIT license, as is the code. The training data carries its own terms: ABC-130k and MolmoAct2 BimanualYAM are Apache-2.0, AgiBot World is gated and AgiBot's terms are CC BY-NC-SA 4.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}
}