| --- |
| 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 |
| |