molmobot-pi0-toys-absjoint-visl2sp-10000

MolmoBot-Pi0 — droid-toys, absolute joint, vision tower L2-SP (step 10000)

Full fine-tune of allenai/MolmoBot-Pi0-DROID on a 31-episode real-robot toy-manipulation set, with the SigLIP vision tower trainable but held near its pretrained weights by an L2-SP penalty instead of being hard-frozen.

This checkpoint has not been evaluated. No rollouts, no success rates, no offline action-error numbers on held-out data. Everything below is training-time telemetry.

What the run was testing

The usual choice when fine-tuning a VLA on a small dataset is to freeze the vision tower outright. This run replaces that binary with a penalty: the tower gets gradients, but a term lambda * (p - p_init) is added to them, so it can move only as far as the task loss can justify against the pull back to init.

Result: at lambda = 0.1 the tower reached equilibrium by ~step 250 and stayed there for the remaining 9,750 steps.

(relative L2 over the 412.4M vision parameters; min 0.0000%, max 0.0429%, median 0.0291% across 1,000 logged points)

For scale, a comparable unregularized full fine-tune of the same base moved its vision tower 1.86% by step 5000 — about 65x further, and further than its own LLM moved (0.85%). So lambda = 0.1 behaves as a soft freeze rather than a soft constraint. If the goal is to let the tower genuinely adapt, a smaller lambda (~1e-2) is the setting to reach for; the drift scales roughly as 1/lambda at equilibrium.

Training

Loss:

step=0      0.5065
step=1000   0.0300
step=2500   0.0168
step=5000   0.0121
step=7500   0.0126
step=9990   0.0078

Caveats worth reading before you trust anything here

• The val_action_loss in the logs is not held out. The trainer reported [val] no VAL_DATA_ROOT -- batches drawn from the TRAIN pool. Those numbers (0.0058 at step 8500, 0.0060 at step 10000) are comparable across steps but are not a generalization estimate.
• Roughly 700 passes over the data. 10,000 steps x 512 effective batch is ~5.1M sample draws against 6,952 frames. The training loss is in memorization range and should be read that way. (The epoch counter in wandb reads 99.9 because it divides by the micro-batch and ignores gradient accumulation.)
• No EMA. ema_decay=0.99 is in the config but the PyTorch trainer ignores it — these are raw weights, not an EMA average. This matters if you compare against a JAX-arm checkpoint, where params/ is the EMA.
• Optimizer state is not included. optimizer.pt (12.5 GB) was left out, so this checkpoint can be loaded for inference or as an init, but not resumed exactly.
• Two sibling checkpoints exist locally at steps 8500 and 8750; their train-pool loss is within noise of this one, so step 10000 is not a validated "best".

Contents

model.safetensors                      7.0 GB, PI0Pytorch state dict
config.json                            model shape (action_dim, horizon, variants)
metadata.pt                            full TrainConfig + global_step
assets/droid_equad/norm_stats.json     quantile norm stats used in training

Loading

These are not transformers-loadable. The state dict is openpi's PI0Pytorch layout (paligemma_with_expert.paligemma.* / paligemma_with_expert.gemma_expert.*) and building the model requires openpi's transformers_replace patch (transformers 4.53.2 with a patched siglip). Use the MolmoBot-Pi0 / openpi PyTorch loader, pointing pytorch_weight_path at a directory containing model.safetensors and assets/.