shrg7's picture
step 18000 (best val 0.0041); full training README
c037506 verified
|
Raw
History Blame Contribute Delete
4.79 kB
---
license: apache-2.0
library_name: openpi
pipeline_tag: robotics
tags:
- robotics
- vla
- pi0
- paligemma
- droid
- manipulation
---
# MolmoBot-Pi0 β€” DROID absolute-joint, reasoning-init full fine-tune
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.
**Checkpoint: step 18000** β€” best held-out validation loss (`val_action_loss = 0.0041`).
---
## 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