File size: 5,142 Bytes
dffdd06 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 | ---
license: apache-2.0
tags:
- robotics
- vla
- behavior-1k
- openpi
- lora
library_name: openpi
---
# pi0.5 LoRA SFT β BEHAVIOR-1K `turning_on_radio`
LoRA fine-tune of the BEHAVIOR-1K challenge Ο0.5 baseline on the single task
`turning_on_radio`, trained with **openpi's PyTorch path + peft** (not the JAX path).
Checkpoints are uploaded **during** training, one folder per step, each with its
optimizer state β so a run can be resumed exactly, LR schedule included, on a
different host.
## Checkpoints
| folder | contents |
|---|---|
| `step-<N>/model.safetensors` | 7.73 GB β full model (base + merged-in LoRA params) |
| `step-<N>/optimizer.pt` | 2.63 GB β AdamW state, needed to resume |
| `step-<N>/metadata.pt` | step counter, config |
| `step-<N>/assets/` | dataset norm stats used at train time |
```bash
hf download madokalif/pi05-b1k-turning-on-radio-lora --include "step-6000/*" --local-dir ./ckpt
```
## Training setup
| | |
|---|---|
| base | Ο0.5 BEHAVIOR-1K challenge baseline, converted to PyTorch **in the patched venv** |
| LoRA | r=32, alpha=32, dropout=0, gaussian init, applied to the VLM only |
| trainable | 758.3M / 3.682B = **20.60%** (the action expert stays fully trainable, per RLinf's recipe) |
| batch | 16 per step, 4Γ A6000, DDP |
| schedule | 50,000 steps, cosine, peak LR 2.5e-5 |
| throughput | 2.85 s/it |
| dataset | `turning_on_radio`, LeRobot v3.0, RGB only (Ο0.5 does not consume depth) |
Loss by 1,000-step bucket: 0.417 β 0.176 β 0.156 β 0.144 β 0.142 β 0.141 β 0.130.
Flow-matching loss resamples the noise level every step, so its floor is high and the
value alone does not indicate policy quality β evaluate with rollouts.
## THE PART THAT BREAKS IF YOU SKIP IT
These weights **cannot be loaded by a stock transformers install.** openpi ships a
`transformers_replace/` tree that injects **AdaRMS** (adaptive RMSNorm) into gemma β
it changes the architecture, not just the API. Without it you get:
```
Missing key(s): ...gemma_expert.model.layers.0.input_layernorm.dense.weight
Unexpected key(s): ...gemma_expert.model.layers.0.input_layernorm.weight
```
So: copy `src/openpi/models_pytorch/transformers_replace/*` into
`site-packages/transformers/`, and fix the install gate β
`transformers/models/siglip/check.py` hardcodes `transformers.__version__ == "5.5.4"`,
which must become `"4.53.2"`. The gate only compares a version string, so passing it
proves nothing; a missing patch surfaces later as
`'PaliGemmaForConditionalGeneration' object has no attribute 'language_model'`.
Corollary: convert any base checkpoint **inside the venv that will train it**. A
conversion done in an unpatched venv is 6.52 GiB and unusable; the patched one is
6.74 GiB. The 0.22 GiB difference is the AdaRMS `dense` layers.
## Environment
```
Python 3.11.15
torch==2.7.1 transformers==4.53.2 lerobot==0.4.4
torchcodec==0.4.0 numpy==2.4.6 peft==0.20.0
safetensors==0.8.0 accelerate==1.14.0
openpi @ 0cc8e355f7bac0976db1cc3139b1ff0379feea60 (branch behavior)
```
Every pin is load-bearing and none is openpi's default:
- **transformers 4.53.2** β `transformers_replace/` targets this API. openpi's own pin
(5.5.4, since commit fd7e261) dropped `utils.LossKwargs` and `cache_utils.HybridCache`
without porting the replace files.
- **lerobot 0.4.4** β declares `CODEBASE_VERSION = v3.0` (so it reads this dataset) *and*
`huggingface-hub<0.36`. 0.5.x requires `huggingface-hub>=1.0`, which collides with
transformers 4.53.2. Stock PyPI lerobot is fine: Ο0.5 is RGB-only, so the b1k fork's
depth handling is irrelevant to SFT.
- **numpy>=1.26** β lerobot 0.4.4 requires rerun-sdk, which wants numpy>=2, while openpi
pins <2. Resolved with a `[tool.uv] override-dependencies` entry.
- **torchcodec 0.4.0** β 0.10.0 raised `Could not load libtorchcodec` against the
training host's FFmpeg. Host-specific; re-test rather than assuming.
### GPU architecture
The training environment is **torch 2.7.1+cu126**, whose arch list is
`sm_50 β¦ sm_90`. On **Blackwell (B200, sm_100)** there is no kernel image and it fails
at the first CUDA op. Swap the build, not the version:
```bash
uv pip install torch==2.7.1 --index-url https://download.pytorch.org/whl/cu128
python -c "import torch; print(torch.cuda.get_arch_list())" # must contain sm_100
```
torchvision and torchcodec are compiled against torch and may need matching rebuilds.
## Setup files
`setup/` in this repo carries what the Hub weights alone don't:
`openpi-sft.patch` (diff vs the base commit), `pyproject.toml` + `uv.lock`,
`pip_freeze.txt` (all 211 resolved versions), `check.py.fixed`,
`norm_stats.json`, and the training/upload scripts.
**Do not recompute norm stats** β `compute_norm_stats.py` has no exists-check and takes
6.5 hours on this dataset. Place the file at
`outputs/assets/pi05_b1k_pytorch_local/turning_on_radio/norm_stats.json`.
## Verifying a fresh setup
Startup must log:
```
LoRA enabled (r=32): trainable 758.3M / 3.682B (20.60%)
```
A different fraction means the LoRA wrap landed somewhere other than the VLM.
|