Psi0/baselines/pi05/README.md at main · physical-superintelligence-lab/Psi0

OpenPI $\pi_{0.5}$

Set-up Envrionment

ℹ️ The following commands assume the variables eg., PSI_HOME are loaded by running cd /path/to/psi0 && source .env.

Set up seperate environment for $\pi_{0.5}$:

ℹ️ We manage the $\Psi_0$ environment and all the baselines through uv and they all share the same src/ code. See Environment Management for more details.

uv venv .venv-openpi --python 3.10
source .venv-openpi/bin/activate
VIRTUAL_ENV=.venv-openpi uv pip install -e .
VIRTUAL_ENV=.venv-openpi uv pip install -e src/openpi/openpi-client
VIRTUAL_ENV=.venv-openpi GIT_LFS_SKIP_SMUDGE=1 uv pip install -r baselines/pi05/requirements-openpi.txt

Apply the transformers library patches:

See also the official OpenPI README.md

cp -r src/openpi/models_pytorch/transformers_replace/* .venv-openpi/lib/python3.10/site-packages/transformers/

Download Pre-Trained Weights

Download pretrained pi05_droid model of torch variant:

hf download USC-PSI-Lab/psi-model \
	--local-dir=$PSI_HOME/cache/checkpoints \
	--include="openpi/pi05_droid/*" \
	--repo-type=model

[Optional] Expand to see how we produced the `pi05_droid` pytorch checkpoint.

Download `pi05-droid` checkpoints

```
python src/openpi/shared/download.py
```

Convet from `jax` to `pytorch` checkpoint
```
uv run examples/convert_jax_model_to_pytorch.py \
	--checkpoint_dir=/home/user/.cache/openpi/openpi-assets/checkpoints/pi05_droid \
	--config_name=pi05_droid \
	--output_path $PSI_HOME/cache/checkpoints/openpi/pi05_droid
```

Download Psi0 Task Data

Download the task data, for example

export task=G1WholebodyXMovePick-v0

hf download USC-PSI-Lab/psi-data simple/$task.zip --local-dir=$PSI_HOME/data --repo-type=dataset
unzip "$PSI_HOME/data/simple/$task.zip" -d "$PSI_HOME/data/simple"

Create a new TrainConfig for the task in src/openpi/training/config.py:

Skip this step if you are finetuning the same SIMPLE/real tasks provided by $Psi_0.

vim src/openpi/training/config.py

for example

	# ...
	TrainConfig(
        name="simple_bend_pick_v1",
        project_name="psi",
        num_workers=8,
        model=pi0_config.Pi0Config(
            pi05=True,
            action_dim=36,
            action_horizon=30,
            max_token_len=250,
        ),
        data=LeRobotHFMDataConfig(
            repo_id=f"{os.environ['PSI_HOME']}/data/simple/<-- replace with $task -->",
            base_config=DataConfig(prompt_from_task=True),
        ),
        weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_droid/params"),
        num_train_steps=40_000,
        batch_size=128,
        lr_schedule=_optimizer.CosineDecaySchedule(
            warmup_steps=1_000,
            peak_lr=1e-4,
            decay_steps=40_000,
            decay_lr=1e-8,
        ),
        pytorch_weight_path=f"{os.environ['PSI_HOME']}/cache/checkpoints/pi05_droid",
        policy_metadata={"dataset": "<-- replace with $task -->"},
        checkpoint_base_dir=f".runs/openpi-05"
    ),
	# ...

Compute stats for the task using official script (which is slow, see below for another option)

python src/openpi/compute_norm_stats.py --config-name $task

Or you can rewrite our precomputed stats to openpi format: python src/openpi/rewrite_norm_stats.py --task_path=$PSI_HOME/data/simple/$task. The calculations are slightly different but it's ok.

Fine-Tune $\pi_{0.5}$

Launch the training script

bash baselines/pi05/train_pi05.sh $task

Eval $\pi_{0.5}$

export port=9000
export step=40000
bash baselines/pi05/serve_pi05.sh $task $step $port

Open-loop evaluation (on train data)

python baselines/pi05/eval_openloop.py --port=$port --task=$task

[Optional] Expand to see more implementation details.

Below are the gist of how to adpat OpenPI on humaonoid loco-manipulation tasks:

change action dim from 32 to 36

# vim src/openpi/models_pytorch/pi0_pytorch.py
self.action_in_proj = nn.Linear(36, action_expert_config.width)
self.action_out_proj = nn.Linear(action_expert_config.width, 36)

adapt the loading of pretrained pi05_torch

model_path = os.path.join(config.pytorch_weight_path, "model.safetensors")
# Psi-0: adapt action dim to 36
from safetensors.torch import load_file
state_dict = load_file(model_path)
pad_dim = config.model.action_dim - state_dict["action_in_proj.weight"].shape[1]
if pad_dim > 0:
	# eg., torch.Size([1024, 32]) -> torch.Size([1024, 36])
	# Replicate the last 4 columns instead of padding with zeros
	w = state_dict["action_in_proj.weight"]
	to_pad = w[:, -pad_dim:]
	state_dict["action_in_proj.weight"] = torch.cat([w, to_pad], dim=1)

	b = state_dict["action_out_proj.bias"]
	state_dict["action_out_proj.bias"] = torch.cat([b, b[-pad_dim:]], dim=0)

	w = state_dict["action_out_proj.weight"]
	to_pad = w[-pad_dim:, :]
	state_dict["action_out_proj.weight"] = torch.cat([w, to_pad], dim=0)

# https://github.com/Physical-Intelligence/openpi/issues/669
state_dict["paligemma_with_expert.paligemma.model.language_model.embed_tokens.weight"] = \
	state_dict["paligemma_with_expert.paligemma.lm_head.weight"]

_model = model.module if isinstance(model, torch.nn.parallel.DistributedDataParallel) else model
missing_keys, unexpected_keys = _model.load_state_dict(state_dict, strict=False)

create TrainConfig in config.py

TrainConfig(
    name="Pick_toys_into_box_and_lift_and_turn_and_put_on_the_chair_new_target_yaw",
    project_name="hfm",
    num_workers=8,
    model=pi0_config.Pi0Config(
        pi05=True,
        action_dim=36,
        action_horizon=16,
        max_token_len=250,
    ),
    data=LeRobotHFMDataConfig( # FIXME
        repo_id= f"{os.environ['DATA_HOME']}/Pick_toys_into_box_and_lift_and_turn_and_put_on_the_chair_new_target_yaw",
        base_config=DataConfig(prompt_from_task=True),
    ),
    weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_droid/params"),
    num_train_steps=40_000,
    batch_size=128,
    lr_schedule=_optimizer.CosineDecaySchedule(
        warmup_steps=1_000,
        peak_lr=1e-4,
        decay_steps=40_000,
        decay_lr=1e-8,
    ),
    pytorch_weight_path=os.environ["PYTORCH_WEIGHT_PATH"],
    policy_metadata={"dataset": "Pick_toys_into_box_and_lift_and_turn_and_put_on_the_chair_new_target_yaw"},
),

launch training

bash scripts/train/openpi/benchmark_pi05_nv_slurm.sh \
	Pick_toys_into_box_and_lift_and_turn_and_put_on_the_chair_new_target_yaw

[Optional] upload model weights

#export task=Hold_lunch_bag_with_both_hands_and_squat_to_put_on_the_coffee_table
#export task=Pick_toys_into_box_and_lift_and_turn_and_put_on_the_chair_new_target_yaw
export task=Pull_the_tray_out_of_chips_can_and_throw_the_can_into_trash_bin
export step=40000
hf upload USC-PSI-Lab/psi-model \
	.runs/openpi-05/$task/$task/$step/model.safetensors \
	benchmarks/openpi-05/$task/$step/model.safetensors \
	--repo-type=model

Serve Download:

export step=40000
python scripts/data/download.py \
	--repo-id=USC-PSI-Lab/psi-models \
	--remote-dir=benchmarks/openpi-05/$task/$step \
	--repo-type=model \
	--local-dir=.runs/openpi-05/$task/$task/$step

and Serve:

export port=9000
bash baselines/pi05/serve_pi05.sh $task $step $port

Open-loop evaluation

python baselines/pi05/eval_openloop.py --port=$port --task=$task

Eval in SIMPLE

TODO: migrate following instructions using SIMPLE third_party

cd <project root of SIMPLE>
source .venv/bin/activate

export task=G1WholebodyXMovePick-v0

Download eval data and extract it:

hf download USC-PSI-Lab/psi-data \
	simple-eval/$task.zip \
	--local-dir=data/evals \
	--repo-type=dataset

unzip data/evals/simple-eval/$task.zip -d data/evals/simple-eval

Now start SIMPLE eval in the SIMPLE environment:

We provide three domain randomization levels: level-0, level-1, level-2 for each task

export dr=level-0

We use two different entrypoints for evaluating different tasks:

set entrypoint and agent to eval_decoupled_wbc.py and pi05_decoupled_wbc if the evaluating task ends with Teleop, which means the task data is collected using teleoperation:

export entry=eval_decoupled_wbc.py
export agent=pi05_decoupled_wbc

and set entrypoint and agent to eval.py and pi05 if the evaluating task ends with MP, which means the task data is generated using CuRobo Motion planning:

export entry=eval.py
export entry=pi05

python src/simple/cli/$entry \
	simple/$task \
	$agent \
	$dr \
	--host=localhost \
	--port=9000 \
	--sim-mode=mujoco_isaac \
	--no-headless \
	--data-format=lerobot \
	--data-dir=data/evals/simple-eval/$task/$dr

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

OpenPI $\pi_{0.5}$

Set-up Envrionment

Download Pre-Trained Weights

Download Psi0 Task Data

Fine-Tune $\pi_{0.5}$

Eval $\pi_{0.5}$

Eval in SIMPLE

Uh oh!

FilesExpand file tree

README.md

Latest commit

History

README.md

File metadata and controls

OpenPI $\pi_{0.5}$

Set-up Envrionment

Download Pre-Trained Weights

Download Psi0 Task Data

Fine-Tune $\pi_{0.5}$

Eval $\pi_{0.5}$

Eval in SIMPLE