VLABench/pi0-primitive-10task

Pi0 official implementation trained on VLABench datasets.

This repository provides the official release of the Pi0 model trained with the whole VLABench's official primitive tasks dataset.

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 checkpoints and put it in the directory checkpoints, to run the policy as server, please run:

bash vla_bench_scipts/serve_policy.sh pi0_ft_vlabench_primitive checkpoints/VLABench/pi0-primitive-10task/29999/

After serving the policy, open another terminal and run:

bash vla_bench_scipts/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 pi05_ft_vlabench_primitive.

XLA_PYTHON_CLIENT_MEM_FRACTION=0.95 uv run scripts/train.py pi0_ft_vlabench_primitive --exp-name=pi0_ft_vlabench_primitive --overwrite

Our checkpoint is trained on 8 H100 for 30k iterations, with 5000 episodes data acrossing 10 tasks.

Reference Results

The reference success rate of this model is:

Track	add_condiment	insert_flower	select_book	select_chemistry_tube	select_drink	select_fruit	select_mahjong	select_painting	select_poker	select_toy	Avg_SR
track_1_in_distribution	0.66	0.18	0.694	0.52	0.52	0.38	0.25	0.46	0.54	0.5	0.47
track_2_cross_category	0.14	0.04	0.064	0.12	0.224	0.46	0.02	0.26	0.26	0.36	0.212
track_3_common_sense	0.34	0.22	0.417	0.7	0.08	0.083	0.125	0.5	0.06	0.38	0.291
track_4_semantic_instruction	0.26	0.02	0.311	0.06	0.1	0.06	0.12	0.56	0.12	0.12	0.173
track_6_unseen_texture	0.56	0.1	0.714	0.28	0.44	0.30	0.02	0.3	0.28	0.18	0.322