hqfang/pi05-so100_101

pi05 SO100/101 — step 150,000

Final checkpoint from training pi05 on the accessible SO100/101 mixture: 1,209 datasets, 36,877 episodes, and 19,227,195 frames. All included datasets were in LeRobot v3 format. Training used global batch size 256 for 150,000 updates (38.4 million sampled frames), four GB200 GPUs, FP32 master weights and Adam moments, and BF16 autocast. W&B was disabled. The final loss was 0.00433; mean training loss over the last 100 updates was 0.00482. No task-success evaluation or robot rollout evaluation is included.

Files and training

The original checkpoint files are preserved exactly. model.safetensors contains FP32 weights. The saved processors contain shared pooled q01/q99 state/action normalization; tokenizer/tokenizer.model is the native PaliGemma SentencePiece tokenizer. Optimizer and RNG state are excluded. The inference checkpoint is approximately 16.6 GB.

training/run.json records the actual training schedule: 1,000 warmup steps, cosine decay over 150,000 steps, peak LR 2.5e-5 and final LR 2.5e-6. The original config.json retains the policy's default scheduler length of 30,000; the custom trainer used the 150,000-step schedule in training/run.json. Training used max-autotune compilation, fused AdamW, layer checkpointing, and eager attention.

Compatible LeRobot code

This checkpoint uses a custom SentencePiece processor and camera padding support. Recreate the compatible source using the included patch:

git clone https://github.com/huggingface/lerobot.git
cd lerobot
git checkout b6ec0060779550c0a157ae34feb89e0cf86012a8
# Set CHECKPOINT to the downloaded model repository directory.
git apply "$CHECKPOINT/code/lerobot.patch"
# Install the patched package and its pi extra in your Python environment.
pip install -e '.[pi]'

The tested training environment used Python 3.12, PyTorch 2.11.0, Transformers 5.5.4, and SentencePiece 0.2.2. The included lockfile patch records dependency resolution.

Load the saved processors rather than constructing new normalization statistics:

from pathlib import Path
from huggingface_hub import snapshot_download
from lerobot.policies.pi05.configuration_pi05 import PI05Config
from lerobot.policies.pi05.modeling_pi05 import PI05Policy
import lerobot.policies.pi05.processor_pi05  # register the saved custom processor
from lerobot.policies.factory import make_pre_post_processors

root = Path(snapshot_download("hqfang/pi05-so100_101"))
cfg = PI05Config.from_pretrained(str(root))
cfg.device = "cuda"
cfg.compile_model = False
cfg.text_tokenizer_name = str(root / "tokenizer/tokenizer.model")
policy = PI05Policy.from_pretrained(str(root), config=cfg).eval()
pre, post = make_pre_post_processors(
    cfg, pretrained_path=str(root),
    preprocessor_overrides={"device_processor": {"device": "cuda"}},
)

The tokenizer path in the original policy config refers to the training workspace; the example overrides it. The saved processor resolves its bundled tokenizer relative to the checkpoint automatically.

Observation and action contract

Provide task text and observation.state in this six-joint order: shoulder_pan, shoulder_lift, elbow_flex, wrist_flex, wrist_roll, gripper. Actions use the same order and the original dataset units, with absolute joint targets. The action chunk length is 50. Do not normalize state manually before the saved processor.

Use one to four RGB views in observation.images.camera_0 through camera_3, with float values in [0,1]. The model letterboxes to 224 by 224. Camera assignments were randomized per episode during training; retain a consistent assignment during a rollout. For explicit four-slot batches, use zero images for missing views and set the corresponding observation.images.camera_N_is_pad boolean to true (false for real views). Camera keys identify input slots, not fixed physical camera identities.

Apply pre(observation) before policy.select_action, then apply post to the predicted action to restore joint units. Reset the policy between episodes. This input flexibility does not establish task success for every camera setup.