step 18000 (best val 0.0041); full training README

README.md

@@ -1,26 +1,133 @@
 ---
 license: apache-2.0
+library_name: openpi
+pipeline_tag: robotics
 tags:
   - robotics
   - vla
   - pi0
+  - paligemma
   - droid
+  - manipulation
 ---
 
-# MolmoBot-Pi0-DROID-absjoint-reasoninit
+# MolmoBot-Pi0 — DROID absolute-joint, reasoning-init full fine-tune
 
-MolmoBot-Pi0 (PaliGemma-decoder VLA) **full fine-tune** on the equad DROID data
-(`shrg7/droid-equad-stage-1`), with **absolute joint-pose** action targets
-(`arm[t] = q[t+1]`, 7 joints + gripper).
+MolmoBot-Pi0 (a PaliGemma-decoder VLA with a flow-matching action expert),
+**fully fine-tuned** on equad DROID with **absolute joint-pose** action targets,
+starting from a **reasoning-initialized** backbone.
 
-**Reasoning-init variant:** before full fine-tuning, a frozen reasoning LoRA
-(trained reasoning-only on the same DROID data) is folded into the PaliGemma
-decoder weights in-place, so the full-FT starts from a reasoning-initialized
-backbone (a weight-space prior), rather than plain base weights.
+**Checkpoint: step 18000** — best held-out validation loss (`val_action_loss = 0.0041`).
 
-- Config: `molmobot_pi0_lerobot_droid_absjoint`
-- **Checkpoint: step 2000 (intermediate — training in progress)**, held-out
-  `val_action_loss = 0.0199`.
-- Contents: `model.safetensors` (flat, full model), `metadata.pt`,
-  `assets/droid_equad/norm_stats.json`, `assets/train_config.pkl`.
-  Optimizer state omitted.
+---
+
+## What "reasoning-init" means
+
+Standard full fine-tuning starts from the base VLA weights. Here we first bake a
+**frozen reasoning LoRA** into the PaliGemma decoder, then full-fine-tune from that
+point. The reasoning adapter is *not* kept as a separate module — its delta is folded
+directly into the base weights:
+
+```
+W  <-  W + (alpha / r) * (B @ A)      # applied in-place, per target module
+```
+
+with `r = 64`, `alpha = 128` (scaling `2.0`), folded across **207 modules**. The
+adapter came from a reasoning-only training run on the *same* DROID data
+(Gemini-annotated chain-of-thought, PaliGemma prompt format).
+
+So reasoning enters as a **weight-space prior on the initialization**, not as an
+auxiliary loss and not as a runtime module. After the fold the model is a plain flat
+checkpoint — inference is identical to any other pi0 checkpoint, with no reasoning
+code path and no extra parameters.
+
+On resume the fold is deliberately skipped (the delta is already inside the saved
+weights), so it is applied exactly once, at step 0.
+
+---
+
+## Action space
+
+Absolute joint pose, 8-dim per step:
+
+| dims | meaning |
+|---|---|
+| `0:7` | absolute future joint positions, `arm[t] = q[t+1]` (7 arm joints) |
+| `7` | gripper command (0–1) |
+
+Action horizon 16. At execution the prediction is used **directly** as the target
+joint pose — no current-position offset (contrast with the `joint_delta` variant,
+where `target = current_q + pred`).
+
+Normalization statistics are bundled in `assets/droid_equad/norm_stats.json`; arm
+stats follow the `observation.state` distribution, gripper from the action gripper dim.
+
+---
+
+## Training setup
+
+| | |
+|---|---|
+| Base checkpoint | `MolmoBot-Pi0-DROID` |
+| Reasoning adapter | reasoning-only LoRA, r=64 / α=128, same DROID data |
+| Dataset | `droid-equad-stage-1` (LeRobot v3.0) |
+| Cameras | `exterior_1_left`, `wrist_left` |
+| Config | `molmobot_pi0_lerobot_droid_absjoint` |
+| Trainable | **entire model** (full FT, no LoRA, no freezing) |
+| Hardware | 4 × NVIDIA L40 (46 GB), DDP |
+| Batch | 64 × grad-accum 8 → **effective 512** |
+| Optimizer | AdamW, weight_decay 1e-10, grad clip 1.0 |
+| LR | peak 5e-5, warmup 1000 steps, constant thereafter |
+| Precision | bf16 |
+| Throughput | ~30 s / step |
+| Steps | **20000** (early-stopped; 30000 was the cap) |
+
+### Checkpoint selection
+
+Validation runs every 1000 steps on **8 fixed held-out batches** (flow noise reseeded
+so the metric is deterministic and comparable across steps). Retention keeps the
+**latest** checkpoint (for resume) plus the **2 best** by `val_action_loss`.
+
+### Early stopping
+
+Training was supervised by a self-chaining job that, each time a training job ended,
+read the cumulative validation series and either resumed training (if val was still
+improving) or stopped. It resumed 3 times and stopped at step 20000 once no new best
+appeared for 2 consecutive validation points — leaving 10000 of the budgeted steps
+unspent.
+
+---
+
+## Validation curve
+
+`val_action_loss`, every 1000 steps:
+
+| step | val | step | val | step | val | step | val |
+|---|---|---|---|---|---|---|---|
+| 1000 | 0.0346 | 6000 | 0.0074 | 11000 | 0.0057 | 16000 | 0.0060 |
+| 2000 | 0.0199 | 7000 | 0.0067 | 12000 | 0.0063 | 17000 | 0.0046 |
+| 3000 | 0.0128 | 8000 | 0.0071 | 13000 | 0.0089 | **18000** | **0.0041** |
+| 4000 | 0.0112 | 9000 | 0.0064 | 14000 | 0.0057 | 19000 | 0.0059 |
+| 5000 | 0.0096 | 10000 | 0.0082 | 15000 | 0.0054 | 20000 | 0.0045 |
+
+Roughly an 8× reduction from the first validation point, with most of the gain by
+step ~6000 and a slow, noisy tail after ~11000.
+
+---
+
+## Contents
+
+```
+model.safetensors                    # flat full model (7.0 GB), reasoning delta folded in
+metadata.pt                          # global_step and training bookkeeping
+assets/droid_equad/norm_stats.json   # state + action normalization statistics
+assets/train_config.pkl              # serialized openpi TrainConfig
+```
+
+Optimizer state is not included (10.9 GB) — this checkpoint is for inference and
+evaluation, not for resuming training.
+
+## Related
+
+- `shrg7/MolmoBot-Pi0-DROID-absjoint` — same data and action space, **without** reasoning-init (the control)
+- `shrg7/MolmoAct2-DROID-absjoint` — MolmoAct2 on the same absolute-joint DROID targets

assets/train_config.pkl

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:025f507e7c8f0d7cb753f9bf1c6eca541b90e6ec7b549fbcac683328f75739a8
+oid sha256:1aa91058f7970777d4acfd8235eacdbf9cde38ac5e5e0fd0131d41ea5ed79df6
 size 1691

metadata.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c56cf90a614fdd2e32ea6e352c05b545f484360f11de7e943c3b5e49d8a9b19d
+oid sha256:21ad27a7a6daed6d6c9b6c5777d2f92d077c310972878e543419304856d18317
 size 2931

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:60c394e7d6a0218a31460fb7caf68510f46bbe66f9be692cfcbd555d4d17a2ab
+oid sha256:cca36677d1d8ca6f6772b9edfbd71e931a4053e254e267cbf87424be63ade8c8
 size 7011248944