capvector-pi05

1. Environment Setup

We use uv to manage Python dependencies. See the uv installation instructions to set it up. Once uv is installed, run the following to set up the environment:

GIT_LFS_SKIP_SMUDGE=1 uv sync
GIT_LFS_SKIP_SMUDGE=1 uv pip install -e .
cp -r ./src/openpi/models_pytorch/transformers_replace/* .venv/lib/python3.11/site-packages/transformers/
source .venv/bin/activate

NOTE: GIT_LFS_SKIP_SMUDGE=1 is needed to pull LeRobot as a dependency.

2. Data Preparation

Here we take the real-world Aloha data as example, more detail simulation data could be refered in the official openpi repo.

First, you need to collect the task-specific raw data with your own robot, and save it in the .hdf5 format.

Then, convert the data to LeRobot dataset format.

uv run examples/aloha_real/convert_aloha_data_to_lerobot.py --raw-dir /path/to/raw/data --repo-id <org>/<dataset-name>
# By default, The converted data is stored in ~/.cache/huggingface/lerobot/<org>/<dataset-name>/

3. Obtain the capability vectors and merge it to obtain $\theta_{meta}$

First, define your task-specific config in config.py. And we provide an example of our real-world task here.

Then, convert a JAX model checkpoint to PyTorch format:

uv run examples/convert_jax_model_to_pytorch.py \
    --checkpoint_dir gs://openpi-assets/checkpoints/pi05_base \
    --config_name <config_name> \
    --output_path checkpoints/pytorch_pi05_base
# This command will automatically download pi05_base checkpoint to ~/.cache/openpi/openpi-assets/checkpoints/pi05_base/
# Otherwise you can download it manually and modify the --checkpoint_dir

⭐ If you don't use the regularization strategy, you could download the capability-merged meta model we provided, place it at ./checkpoints/vector_init/pi05SF-LIBEROspatial_minus_pi05-LIBEROspatial/, and directly jump to the next Training step.

Then, the capability vectors are obtained by simply conducting parameter arithmetic between two models finetuned with different strategies. Therefore, we need to prepare these two trained models, e.g., Pi0.5 on LIBERO-Spatial) and Pi0.5-SF on LIBERO-Spatial). The directory structure is as below:

capvector-pi05
    ├── checkpoints
    ·   ├── pi05-LIBEROspatial
        │   ├── model.safetensors
        │   └── ...
        ├── pi05SF-LIBEROspatial
        │   ├── model.safetensors
        │   └── ...
        ├── diff
        ├── vector_init
        ·

Next, conduct parameter arithmetic between these two models:

CONFIG=pi05_capvector_aloha_place_block && \
EXT=pi05SF-LIBEROspatial && \
DOWN=pi05-LIBEROspatial && \
uv run capvector/compute_param_diff.py \
  --config $CONFIG \
  --a.dir checkpoints/$EXT \
  --b.dir checkpoints/$DOWN \
  --out checkpoints/diff/${EXT}_minus_${DOWN}.pth \
  --strict-keys \
  --dtype fp32

Finally, merge these diff parameters to obtain $\theta_{meta}:

DIFF=pi05SF-LIBEROspatial_minus_pi05-LIBEROspatial && \
uv run capvector/apply_param_diff.py \
  --base-safetensors checkpoints/pytorch_pi05_base/model.safetensors \
  --diff-pth checkpoints/diff/${DIFF}.pth \
  --out-safetensors checkpoints/vector_init/${DIFF}/model.safetensors \
  --scale 1.0 \
  --no-strict-keys \
  --dtype fp32 \
  --device cpu

4. Training

First, you need to compute the normalization statistics for the training data.

uv run scripts/compute_norm_stats.py --config-name <config_name>

Finally, launch training using one of these modes:

# Single GPU training:
uv run scripts/train_regular_loss_pytorch.py <config_name> --exp_name <run_name> --save_interval <interval>
# Example:
uv run scripts/train_regular_loss_pytorch.py pi05_capvector_aloha_place_block --exp_name pytorch_test
uv run scripts/train_regular_loss_pytorch.py pi05_capvector_aloha_place_block --exp_name pytorch_test --overwrite  # Overwrite existing checkpoints

# Multi-GPU training (single node):
uv run torchrun --standalone --nnodes=1 --nproc_per_node=<num_gpus> scripts/train_regular_loss_pytorch.py <config_name> --exp_name <run_name>

# Multi-Node Training:
uv run torchrun \
    --nnodes=<num_nodes> \
    --nproc_per_node=<gpus_per_node> \
    --node_rank=<rank_of_node> \
    --master_addr=<master_ip> \
    --master_port=<port> \
    scripts/train_regular_loss_pytorch.py <config_name> --exp_name=<run_name> --save_interval <interval>

5. Inference

Real-world inference is executed in the server-client form.

First, launch a model server (we use the checkpoint for iteration 20,000 for this example, modify as needed):

uv run scripts/serve_policy.py policy:checkpoint --policy.config=<config_name> --policy.dir=checkpoints/<config_name>/<run_name>/20000

This will spin up a server that listens on port 8000 and waits for observations to be sent to it.

Then, We can then run an client robot script that queries the server.

You need to write your client script according to your robot. A simple client exmaple is as below:

uv run examples/simple_client/main.py --env ALOHA