pi05_piper_stacking / README.md
rvsagar9567's picture
Document step_5000 + step_7500 checkpoints
c00a2f1 verified
|
Raw
History Blame Contribute Delete
3.65 kB
---
license: other
library_name: openpi
pipeline_tag: robotics
tags:
- robotics
- vla
- manipulation
- bimanual
- lerobot
- openpi
datasets:
- axiboai/piper_stacking
---
# π₀.₅ · piper_stacking (bimanual cube stacking)
π₀.₅ policy fine-tuned with [openpi](https://github.com/Physical-Intelligence/openpi) on
**[axiboai/piper_stacking](https://huggingface.co/datasets/axiboai/piper_stacking)** — a bimanual PiperX
robot picking up a red cube and stacking it on a blue cube.
**Prompt:** `stack red cube on blue cube`
## Task & embodiment
- Robot `piperx_bimanual` — 14-D joint state/action (per arm: 6 joints + gripper × 2)
- 3 RGB cameras: `cam_front`, `cam_left_wrist`, `cam_right_wrist` (resized to 224×224)
- Data: 60 demos / 23,087 frames @ 30 Hz, LeRobot v2.1
## Training
- Base: Physical Intelligence `pi05_base` · full fine-tune (no LoRA) · 3.35B params
- 10,000 steps · batch 32 · action_horizon 50 · AdamW · cosine LR (warmup 600 → peak 5e-5 → 5e-6) · EMA 0.99
- State input: discrete tokens · normalization: quantile (q01–q99)
- **Final flow-matching train loss: 0.0012**
- Hardware: 1 × NVIDIA RTX PRO 6000 Blackwell (96 GB)
## Contents
- `params/` — EMA inference weights (orbax)
- `assets/piperx_bimanual/norm_stats.json` — normalization stats (baked from this dataset)
- `_CHECKPOINT_METADATA`
- The 29 GB optimizer `train_state/` is intentionally **excluded** (not needed for inference).
## Load & run (openpi)
Requires the piperx-openpi fork (provides `piperx_policy` + the `pi05_piper_stacking` config).
```python
import pathlib
from huggingface_hub import snapshot_download
from openpi.training import config as _config
from openpi.policies import policy_config
ckpt = pathlib.Path(snapshot_download("axiboai/pi05_piper_stacking"))
cfg = _config.get_config("pi05_piper_stacking")
policy = policy_config.create_trained_policy(cfg, ckpt)
# observation = {"observation.images.cam_front": img, ...cam_left_wrist, ...cam_right_wrist,
# "observation.state": state_14d, "prompt": "stack red cube on blue cube"}
action_chunk = policy.infer(observation)["actions"] # (50, 14)
```
If `pi05_piper_stacking` is not yet in your `src/openpi/training/config.py`, add:
```python
TrainConfig(
name="pi05_piper_stacking",
model=pi0_config.Pi0Config(pi05=True),
weight_loader=weight_loaders.CheckpointWeightLoader("gs://openpi-assets/checkpoints/pi05_base/params"),
data=LeRobotPiperXBimanualDataConfig(
repo_id="axiboai/piper_stacking",
base_config=DataConfig(prompt_from_task=True),
assets=AssetsConfig(asset_id="piperx_bimanual"),
default_prompt="stack red cube on blue cube",
),
lr_schedule=_optimizer.CosineDecaySchedule(warmup_steps=600, peak_lr=5e-5, decay_steps=9400, decay_lr=5e-6),
num_train_steps=10000, batch_size=32, save_interval=2500, keep_period=2500,
),
```
## Provenance / license
Fine-tuned from Physical Intelligence's `pi05_base` checkpoint via openpi; the weights inherit the
upstream openpi / base-model license. Dataset: axiboai/piper_stacking.
## Checkpoints in this repo
Checkpoints from the **same run** (for overfitting comparison):
| Path | Step | Train loss |
|---|---|---|
| `/` (root) | 10,000 (final) | 0.0012 |
| `step_7500/` | 7,500 | 0.0021 |
| `step_5000/` | 5,000 | 0.0031 |
| `step_2500/` | 2,500 | 0.0048 |
Load any with `create_trained_policy(cfg, ckpt_dir / "step_XXXX")` (root = final). Training loss fell monotonically (2,500 → final); compare **real-rollout success** to tell whether the longer run overfit the 60 demos.