Create README.md

README.md

@@ -0,0 +1,93 @@
+---
+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.