pi0-pretrain-vlabench-primitive-aligned

pi0-pretrain-vlabench-primitive-aligned

Pi0 official implementation trained on VLABench datasets.

This repository provides a Pi0 checkpoint trained with the VLABench pretrain primitive tasks dataset. The dataset contains 2,000 trajectories for each task. This config uses the aligned delta chunk setting.

The uploaded checkpoint includes both inference parameters and the full training state:

• params/: Orbax parameters for inference/evaluation.
• train_state/: full Orbax training state for resuming training.
• assets/: normalization statistics and checkpoint assets.
• _CHECKPOINT_METADATA: Orbax checkpoint metadata.

Evaluation

To run this checkpoint, please clone this repo: https://github.com/Shiduo-zh/openpi, and checkout to the branch main. Assume that you download this checkpoint and put it in the directory checkpoints, to run the policy as server, please run:

bash serve_policy.sh pi0_pretrain_vlabench_primitive_aligned checkpoints/pi0-pretrain-vlabench-primitive-aligned/

After serving the policy, open another terminal and run:

bash multi_run_vlabench.sh <Your path to store the evaluate results>

Train

To reproduce the training result, please run the training script with the config pi0_pretrain_vlabench_primitive_aligned.

XLA_PYTHON_CLIENT_MEM_FRACTION=0.95 uv run scripts/train.py pi0_pretrain_vlabench_primitive_aligned \
  --exp-name=pi0_pretrain_vlabench_primitive_aligned \
  --batch-size=32 \
  --save_interval=10000 \
  --overwrite

This checkpoint is trained for 200k iterations on the VLABench pretrain primitive tasks dataset with the aligned delta chunk setting.

Reference Results

Reference success rates are not included in this checkpoint card. Please run the evaluation script above to reproduce results for your environment.

Citation

If you use this checkpoint, please consider citing VLABench:

@article{zhang2024vlabench,
  title={Vlabench: A large-scale benchmark for language-conditioned robotics manipulation with long-horizon reasoning tasks},
  author={Zhang, Shiduo and Xu, Zhe and Liu, Peiju and Yu, Xiaopeng and Li, Yuan and Gao, Qinghui and Fei, Zhaoye and Yin, Zhangyue and Wu, Zuxuan and Jiang, Yu-Gang and others},
  journal={arXiv preprint arXiv:2412.18194},
  year={2024}
}