pi05 / README.md
renbin1's picture
Publish RoboTwin Pi05 final inference checkpoint and loading instructions
dba8cd8 verified
|
Raw History Blame Contribute Delete
4.43 kB
metadata
language:
  - en
  - zh
tags:
  - robotics
  - pi05
  - openpi
  - robotwin
  - aloha-agilex

Pi05 fine-tuned on RoboTwin (900 demonstrations)

This repository contains the inference checkpoint from a Pi05 fine-tuning run on 900 RoboTwin Aloha AgileX demonstrations (159,568 frames), covering place_object_basket, place_object_scale, and place_object_stand. Training completed 15,000 steps on October 7, 2026 (UTC+08:00). The final checkpoint is 14999, because steps are indexed from zero. The last logged training loss was 0.0013 at step 14900. No evaluation success rate is reported here; training loss does not establish task success.

Files

  • 14999/params/: JAX/Orbax model parameters (approximately 12 GiB).
  • 14999/assets/: normalization statistics required for inference.
  • 14999/_CHECKPOINT_METADATA: original checkpoint metadata.
  • load_policy.py: portable inference configuration and loading helper.
  • openpi_inference_source.tar.gz: snapshot of the local OpenPI source, client package, dependency declarations, lockfile, and upstream code license.

Optimizer/training state and the demonstration dataset are not included. This is a JAX/Orbax checkpoint, not a Transformers from_pretrained model. The model repository hosts downloadable files, not a hosted inference endpoint.

Download and install

python -m pip install -U huggingface_hub
hf download renbin1/pi05 --local-dir pi05
cd pi05
mkdir -p openpi_runtime
tar -xzf openpi_inference_source.tar.gz -C openpi_runtime
cd openpi_runtime
uv sync --frozen --no-dev
uv pip install huggingface_hub
cd ..

Install uv first if it is not already available. Use Python 3.11 and a CUDA-compatible NVIDIA GPU environment. The training environment used JAX/JAXlib 0.5.3, Flax 0.10.2, Orbax Checkpoint 0.11.13, NumPy 1.26.4, and Torch 2.10.0. The bundled lockfile supplies the remaining dependencies. First use may download OpenPI tokenizer assets.

Run inference

Run the following Python code using openpi_runtime/.venv/bin/python, from the downloaded repository directory:

import numpy as np
from load_policy import load_policy

policy = load_policy("14999")
observation = {
    "state": state.astype(np.float32),  # real robot state, shape (14,)
    "images": {
        "cam_high": head_rgb,
        "cam_left_wrist": left_wrist_rgb,
        "cam_right_wrist": right_wrist_rgb,
    },
    "prompt": instruction,
}
result = policy.infer(observation)
actions = result["actions"]  # shape (50, 14), joint/gripper action sequence

Supply actual observations: state, images, and instruction are placeholders. Images must be RGB uint8 arrays in HWC or CHW layout. Training images were 480x832 and the model transform resizes them to 224x224. The 14 state/action dimensions are ordered as left arm joints (6), left gripper (1), right arm joints (6), right gripper (1). Match the RoboTwin Aloha AgileX joint conventions. The loader preserves joint delta transforms internally and returns absolute actions; grippers remain absolute. adapt_to_pi=False matches training. Inference does not require the training dataset.

Configuration and provenance

Training configuration: pi05_base_aloha_full_sim_robotwin_900_defaults. Normalization asset ID: RoboTwin-pi05_900-aloha_agilex-joint. Architecture: Pi0Config(pi05=True), action horizon 50, internal action width 32, with 14 robot action dimensions returned by the Aloha output transform. The source snapshot includes the custom training configuration. Its git base was XPolicyLab commit 9c98a3aaf02d05c6f9999a5a0a7a42090555ddf3; local changes are captured in the snapshot, so the commit alone is not the complete source. The portable helper loads normalization from the downloaded checkpoint and does not depend on paths from the training machine.

Upstream projects: OpenPI, XPolicyLab, and RoboTwin. The code license is included in the source archive. This model card does not assign a new license to the weights; consult the applicable base-model and dataset terms before reuse or redistribution.

Validation for this release checks file presence, normalization dimensions, and configuration construction. GPU inference and task evaluation have not been validated as part of the upload preparation.