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
Downloads last month

-

Downloads are not tracked for this model. How to track
Safetensors
Model size
4B params
Tensor type
F32
Β·
BF16
Β·
Video Preview
loading