PI0.5#
Environment Setup#
We use uv to manage Python dependencies.
conda activate robosyn
# Install uv
pip install uv
Once uv is installed, run the following commands to set up the environment:
cd policy/pi05
apt install -y git-lfs build-essential pkg-config
GIT_LFS_SKIP_SMUDGE=1 uv sync
cp -r ./src/openpi/models_pytorch/transformers_replace/* .venv/lib/python3.11/site-packages/transformers/
Generate RoboSynChallenge Data#
See Collect Data Section for more details.
Prepare openpi0 Data for Training#
RoboSynChallenge depends on the EmbodiChain emulator, which by default only supports acquiring Lerobot 3.0 data. We provide a script to convert Lerobot 3.0 data to Lerobot 2.1 data; this script is located in scripts/convert_lerobot3.0_to_2.1.py. Furthermore, our data collection script already supports one-click conversion; please refer to Collect Data Section for details.
convert_lerobot3.0_to_2.1.py script will give you a processed dataset with the same name as the original {repo_id}, but in version 3.0, the dataset will be named {repo_id}_v3.0. Usage examples:
python scripts/convert_lerobot3.0_to_2.1.py --repo-id {repo_id} --root /path/to/datasets
If you want to train on multiple datasets together (e.g., multi-task, mixed training with simulated and real data), merge them with the lerobot-edit-dataset tool or launch/collect_combined_dataset.sh before placing the result in this policy’s training data folder.
After preparing the data in lerobot2.1 format, create the training_data folders in the policy/pi05 directory:
mkdir training_data
Then copy all the data you wish to use for training into training_data/.
Train with Downloaded Data#
Download instructions are in Download Data. policy/pi05/finetune.sh sets HF_LEROBOT_HOME to policy/pi05/training_data, so downloaded datasets used by this policy should be placed under that folder.
For the click_bell real dataset:
cd /root/workspace/RoboSynChallenge/policy/pi05
huggingface-cli download RoboSynChallenge/cobotmagic_Real_click_bell \
--repo-type dataset \
--local-dir training_data/RoboSynChallenge/cobotmagic_Real_click_bell
Set repo_id="RoboSynChallenge/cobotmagic_Real_click_bell" in src/openpi/training/config.py for pi05_base_robosynchallenge_full, then run:
uv run scripts/compute_norm_stats.py --config-name pi05_base_robosynchallenge_full
bash finetune.sh pi05_base_robosynchallenge_full click_bell_real 0
bash eval.sh click_bell clear pi05_base_robosynchallenge_full click_bell_real 0 --max_episodes 50
Write the Corresponding train_config#
In src/openpi/training/config.py, there is a dictionary called _CONFIGS. You can modify a pre-configured PI0.5 configurations I’ve written:
pi05_base_robosynchallenge_full.
Set repo_id to the dataset you want to train on, e.g., RoboSynChallenge/cobotmagic_Sim_click_bell or RoboSynChallenge/cobotmagic_Real_click_bell.
Finetune model#
# compute norm_stat for dataset
uv run scripts/compute_norm_stats.py --config-name ${train_config_name}
# uv run scripts/compute_norm_stats.py --config-name pi05_base_robosynchallenge_full
# train_config_name: The name corresponding to the config in _CONFIGS, such as pi05_base_robosynchallenge_full
# model_name: You can choose any name for your model
# gpu_use: if not using multi gpu,set to gpu_id like 0;else set like 0,1,2,3,4,5,6,7
bash finetune.sh ${train_config_name} ${model_name} ${gpu_use}
#bash finetune.sh pi05_base_robosynchallenge_full click_bell 0,1,2,3,4,5,6,7
| Training mode | Memory Required | Example GPU |
|---|---|---|
| Fine-Tuning (LoRA) | > 46 GB | A6000(48G) |
| Fine-Tuning (Full) | > 100 GB | 2*A100 (80GB) / 2*H100 |
If your GPU memory is insufficient, please set the fsdp_devices parameter according to the following GPU memory reference, or reduce the batch_size parameter.
Or you can try setting XLA_PYTHON_CLIENT_PREALLOCATE=false in finetune.sh, it will cost lower gpu memory, but make training speed slower.
The default batch_size is 32 in the table below.
| GPU memory | Model type | GPU num | fsdp_devices | Example GPU |
|---|---|---|---|---|
| 40G | full | 4 | 4 | A100(40G) |
| 80G | full | 2 | 2 | A100(80G) |
Eval on RoboSynChallenge#
Checkpoints will be saved in policy/pi05/checkpoints/${train_config_name}/${model_name}/${checkpoint_id}
You can modify the deploy_policy.yml file to change the checkpoint_id you want to evaluate.
# ckpt_path like: policy/pi05/checkpoints/pi05_base_robosynchallenge_full/click_bell/30000
bash eval.sh ${task_name} [random | clear] ${train_config_name} ${model_name} ${gpu_id}
# bash eval.sh click_bell random pi05_base_robosynchallenge_full click_bell 0
# This command evaluate the policy($model_name) trained by the `click_bell_random` setting using the same `click_bell_random` setting.
The evaluation results, including videos, will be saved in the eval_result/{task_name}/pi05/{setting}/{train_config_name}/{model_name}/{checkpoint_id}/ directory under the project root.