---
license: apache-2.0
language:
- en
tags:
- robotics
- vla
- pi05
- subtask
- openpi
- lerobot
- orbax
datasets:
- physical-intelligence/libero
pipeline_tag: robotics
---

# pi0.5 subtask fine-tune

A 100-step fine-tune of `pi05_base` for subtask generation from the original [pi05 paper](https://www.pi.website/download/pi05.pdf).
We reproduced steps from a community issue thread on openpi that studies this [#701](https://github.com/Physical-Intelligence/openpi/issues/701).

## TL;DR

- **Start weights**: `gs://openpi-assets/checkpoints/pi05_base/params`
- **Config**: `pi05_subtask_libero` (adds `Pi05Subtask` head: joint flow-matching + CE-on-subtask-tokens loss)
- **Training**: 100 steps × batch 8 on 30 LIBERO episodes, 1× H100 on Modal
- **Final loss**: 3.04 → 0.23

## Loading

```python
from pathlib import Path
import jax
import jax.numpy as jnp
import flax.nnx as nnx
from huggingface_hub import hf_hub_download
import tarfile

from openpi.models import model as _model
from openpi.models.pi0 import Pi0
from openpi.models.pi0_config import Pi0Config

# 1. Download + extract
tar = hf_hub_download("swatery/pi05-subtask",
                     "jax/pi05_subtask.tar")
tarfile.open(tar).extractall(".")
ckpt = Path("99")

# 2. Build model and restore weights
config = Pi0Config(pi05=True)
model = config.create(jax.random.key(0))
params = _model.restore_params(ckpt / "params", dtype=jnp.bfloat16)
nnx.update(model, nnx.State(params))
model.eval()
```

For end-to-end subtask generation (JIT-compiled AR decode with ASCII vocab mask over PaliGemma's LM head), see the `SubtaskGenerator` implementation in [openpi/hosting](https://github.com/Hebbian-Robotics/openpi) `src/hosting/subtask_generator.py`.
That module loads a checkpoint like this one and calls `.generate(prompt, images)`.

## Training details

| | |
|---|---|
| Architecture | pi0.5 — PaliGemma + Gemma action expert, with `Pi05Subtask` head |
| Loss | Flow-matching (action) + cross-entropy (subtask tokens) |
| Knowledge insulation | Yes — LM backbone receives only CE gradients |
| Steps | 100 |
| Batch size | 8 (global, single device) |
| Optimizer | AdamW, cosine schedule, peak LR 5e-5, warmup 10k (only 100 steps used, so effectively constant warmup) |
| EMA decay | 0.999 |
| Precision | bfloat16 |
| Hardware | 1× NVIDIA H100 80GB (Modal) |
| Wall-clock | ~10 min training + ~5 min data/weight fetch |

### Data

- **Dataset**: first 30 episodes of `physical-intelligence/libero` chunk-000 (~8,294 frames)
- **Norm stats**: reused `pi05_libero`'s precomputed full-dataset stats from `gs://openpi-assets/checkpoints/pi05_libero/assets/`
- **Subtask annotation**: **identity** — `high_prompt = low_prompt = task_prompt`
  (real hierarchical subtask annotations for LIBERO are not publicly available)

## References

- https://www.pi.website/blog/pi05
- https://github.com/Physical-Intelligence/openpi (upstream pi0.5 implementation)
- https://github.com/Physical-Intelligence/openpi/issues/701 (community issue thread reproducing subtask generation)
- https://github.com/LisavilaLee/openpi_with_subtask (fork with training example)

## License

- Code & fine-tuned weights: Apache 2.0 (inherited from openpi)
- Gemma dependency: this checkpoint is derived from Google's Gemma via PaliGemma. Usage is subject to the Gemma Terms of Use in addition to Apache 2.0.