pi0-base-two-wu

pi0-base-two-wu

This repository contains the complete resumable checkpoints for a sequential two-Wu PI0 model trained jointly on the four LIBERO suites (libero_spatial, libero_object, libero_goal, and libero_10).

Checkpoint stages

• VLM-Wu: 60,000 optimizer steps. The VLM-Wu checkpoint is frozen when training the second stage.
• Action-Wu: 85,000 optimizer steps. This is the current resumable Action-Wu checkpoint.

The included optimizer.pt, rng_state.pt, metadata.pt, split files, and metric histories are intentional: this upload is meant to support continued training, not inference only.

Architecture and training configuration

The PI0 base checkpoint is not bundled. Convert the same PI0 base checkpoint to PyTorch and verify that model.safetensors has SHA-256:

9e8013b0d353ed4b45c42aa0999731ef2667547af8518e397d5d38844ce443ef

Repository layout

checkpoints/pi0_two_wu_libero/
  pi0_two_wu_4suite_vlm_r16_d1152_n12_lr1e5_60k_6gpu_b30_no_wu_ac_final/
    episode_split.json
    metrics.jsonl
    best_checkpoint.json
    60000/
      vlm_wu_nnt.pt
      optimizer.pt
      rng_state.pt
      metadata.pt
      vlm_wu_config.json
      action_wu_config.json
  pi0_two_wu_4suite_action_r16_d768_n12_lr1p5e5_200k_6gpu_b30_no_wu_ac_final/
    episode_split.json
    metrics.jsonl
    best_checkpoint.json
    85000/
      action_wu_nnt.pt
      optimizer.pt
      rng_state.pt
      metadata.pt
      vlm_wu_config.json
      action_wu_config.json

Download

from huggingface_hub import snapshot_download

snapshot_download(
    repo_id="WendaYu/pi0-base-two-wu",
    local_dir="/path/to/pi0-base-two-wu",
)

Prepare the PI0 base checkpoint

Using the matching openpi-wu source tree:

uv run examples/convert_jax_model_to_pytorch.py \
  --checkpoint_dir /path/to/pi0_base_jax \
  --config_name pi0_libero \
  --output_path /path/to/pi0_base_pytorch

sha256sum /path/to/pi0_base_pytorch/model.safetensors

The hash must be 9e8013b0d353ed4b45c42aa0999731ef2667547af8518e397d5d38844ce443ef. The resume code intentionally rejects a different base fingerprint.

Resume Action-Wu training from 85k

Use the matching source code from SunnyYWD/openpi-wu. The original run used six GPUs (0,1,2,3,4,7 on the source machine); any six local GPU IDs can be exposed in the same logical order.

cd /path/to/openpi-wu

HF_ROOT=/path/to/pi0-base-two-wu
VLM_EXP=pi0_two_wu_4suite_vlm_r16_d1152_n12_lr1e5_60k_6gpu_b30_no_wu_ac_final
ACTION_EXP=pi0_two_wu_4suite_action_r16_d768_n12_lr1p5e5_200k_6gpu_b30_no_wu_ac_final

export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5
export JAX_PLATFORMS=cpu
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export PYTHONUNBUFFERED=1

.venv/bin/torchrun \
  --standalone --nnodes=1 --nproc-per-node=6 \
  scripts/train_pi0_two_wu_4suite_pytorch.py \
  --stage action \
  --exp-name "$ACTION_EXP" \
  --checkpoint-base-dir "$HF_ROOT/checkpoints" \
  --base-pytorch-weight-path /path/to/pi0_base_pytorch \
  --vlm-checkpoint "$HF_ROOT/checkpoints/pi0_two_wu_libero/$VLM_EXP/60000" \
  --resume \
  --num-train-steps 200000 \
  --lr-decay-steps 200000 \
  --batch-size 6 \
  --gradient-accumulation-steps 5 \
  --learning-rate 1.5e-5 \
  --decay-learning-rate 1.5e-6 \
  --warmup-steps 4000 \
  --adam-b1 0.9 --adam-b2 0.98 --adam-eps 1e-8 \
  --weight-decay 0.01 --clip-gradient-norm 5 \
  --validation-interval 5000 --validation-batches-per-suite 40 \
  --save-interval 5000 --log-interval 100 \
  --require-cuda

For the most faithful continuation, keep the original six-rank setup because rng_state.pt stores one RNG state per original rank. The action checkpoint metadata contains the source machine's absolute VLM path; the explicit --vlm-checkpoint argument above intentionally overrides that machine-specific path.

Inference

Pass both checkpoint directories explicitly:

.venv/bin/python3 scripts/serve_pi0_two_wu_policy.py \
  --config-name pi0_wu_libero \
  --base-pytorch-weight-path /path/to/pi0_base_pytorch \
  --vlm-wu-checkpoint-dir /path/to/pi0-base-two-wu/checkpoints/pi0_two_wu_libero/pi0_two_wu_4suite_vlm_r16_d1152_n12_lr1e5_60k_6gpu_b30_no_wu_ac_final/60000 \
  --action-wu-checkpoint-dir /path/to/pi0-base-two-wu/checkpoints/pi0_two_wu_libero/pi0_two_wu_4suite_action_r16_d768_n12_lr1p5e5_200k_6gpu_b30_no_wu_ac_final/85000 \
  --pytorch-device cuda \
  --port 8000

License

The openpi code is Apache-2.0. The separately obtained Gemma-derived PI0 base checkpoint is also subject to its applicable Gemma terms; users are responsible for complying with those terms.