transfer

transfer — pi0.5 VACE object-swap training bundle

Everything needed to train the pi0.5 policy on the VACE object-swap augmented RoboCasa PickPlaceCounterToCabinet dataset (8 "hard mesh" neg-mesh objects swapped in), on 2× A100-80GB.

Contents

pi0.5 init weights are NOT bundled — the config's weight_loader auto-downloads pi05_base from gs://openpi-assets/checkpoints/pi05_base/params on first run (needs internet/GCS), or point it at a local copy. | code/openpi/assets/pi0_fast_robocasa_pretrain_human300/robocasa365_human300/norm_stats.json | normalization stats the config uses (state/action). Required. | tiny | | INSTALL.md | environment setup | — |

Train

To train on the action-augmentation dataset instead, use config pi05_robocasa_target_PickPlaceCounterToCabinet_actaug256 (same hyper-parameters; see actaug/README.md).

The config is already registered as pi05_robocasa_target_PickPlaceCounterToCabinet_vace_negmesh in code/openpi/src/openpi/training/config.py (full fine-tune, EMA on, nothing frozen, batch 64, 30k steps, cosine LR peak 2.5e-5 → 2.5e-6 over 30k, warmup 1k).

cd code/openpi
# init weights: the config's weight_loader auto-downloads pi05_base from gs://openpi-assets on first
# run (needs internet/GCS). To use a local copy instead, place it at
#   ~/.cache/openpi/openpi-assets/checkpoints/pi05_base/params
# or edit the weight_loader path in the config entry.

# point the config's data_dir at wherever you unpacked the dataset (edit the "path" in the
# pi05_...vace_negmesh entry of src/openpi/training/config.py if it differs from the bundled path)

XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 python scripts/train.py \
  pi05_robocasa_target_PickPlaceCounterToCabinet_vace_negmesh \
  --exp-name=vace_negmesh_bs64_30k \
  --batch-size=64 --num-train-steps=30000

• 2× A100-80GB, data-parallel (fsdp_devices=1), no gradient accumulation. Peak ~69 GiB/GPU.
• ETA ≈ 31 h for 30k @ bs64 (~3.6 s/step). bs32 ≈ 16 h.
• Checkpoints save every 2500 steps under code/openpi/checkpoints/<config>/<exp-name>/.

IMPORTANT: modality.json (state/action ordering)

dataset/.../meta/modality.json is required and is already correct — do not drop it. The ordering chain is subtle:

1. Raw parquet observation.state (16) / action (12) are base-first [base_pos, base_rot, eef_pos, eef_rot, gripper] (they come from the robocasa hdf5).
2. modality.json is a name→raw-slice lookup (base_position:[0:3], end_effector_position:[5:8], …). It does not by itself reorder; its key order is ignored by the openpi loader.
3. The exported loader groot_openpi_dataset.GrootOpenpiSingleDataset.__getitem__ hardcodes the assembly into arm-first:
   • state = [eef_pos, eef_rot, base_pos, base_rot, gripper]
   • action = [eef_pos, eef_rot, gripper, base_motion, control_mode]
4. norm_stats were computed through this same loader (see scripts/compute_norm_stats_robocasa.py), so training/inference are consistent arm-first.

Net: keep modality.json and norm_stats.json as shipped; the model expects arm-first, and this bundle already produces exactly that. If you serve/eval through a native GR00T loader (which assembles by modality.json key order), reorder the keys to arm-first — but for the openpi path here it is already correct.

Provenance

• Dataset built by vace_swap_negmesh/build_swap_dataset.py (stitches 81-frame/stride-80 VACE segments → full-length per-cam mp4 @20fps, h264/yuv420p, frame-exact to the parquet; copies actions/states verbatim; instruction = swapped object name).
• Validated: 256 contiguous episodes, 768 videos (decode-count == parquet length), action/state order == model-expected (arm-first, exact-equality checked), 2-GPU bs64 smoke trained + checkpointed clean.