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.