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

Component	Configuration
Frozen backbone	PI0 base (gemma_2b VLM + gemma_300m action expert)
VLM-Wu	rank 16, width 1152, 12 NNT layers, 322.078M parameters
Action-Wu	rank 16, width 768, 12 NNT layers, 143.743M parameters
Training suites	Four LIBERO suites jointly
Original world size	6 GPUs
Effective batch size	30 (global micro-batch 6, accumulation 5)
Action-Wu LR	1.5e-5 to 1.5e-6, 4,000-step warmup, 200,000-step decay
AdamW	betas (0.9, 0.98), eps 1e-8, weight decay 0.01
Action gradient clipping	global norm 5.0

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.