pi0-sponge-marker-10k-nullstate-3k (PyTorch)
pi0 policy for the SO-101 sponge/marker task, with a classifier-free-guidance-style null-state capability so it can be run with or without proprioceptive state.
Format: PyTorch (model.safetensors), converted from the JAX/orbax training checkpoint.
Provenance
- Base:
mkpongm/pi0-sponge-marker-10k(stage-1 pi0, 10k steps) - Stage-2: action-head-only fine-tune, 3,000 steps,
state_dropout_prob=0.2 - Null token scope: dropout/null token restricted to the real state dims
[0:6](state_dropout_dim=6); the padded dims[6:32]are exactly zero, matching how real state is zero-padded from 6 →action_dim=32. - Precision: bf16 (frozen backbone was bf16 in the stage-2 checkpoint; inference computes in bf16).
The null-state embedding is baked into the model
Unlike a stock JAX→PyTorch conversion (which silently drops it), this checkpoint carries the
learned null-state embedding as a real parameter, null_state_token — shape (32,), with
dims [6:32] zero. Verified identical to the JAX source (max abs diff ~1.7e-6).
Dual inference (with / without proprioception)
The state token is fed through state_proj. Both branches use the same converted weights:
| branch | state fed to the model |
|---|---|
| with proprioception | the real robot state, zero-padded 6 → 32 |
| without proprioception | model.null_state_token (32-dim), broadcast to batch |
# null (no-proprioception) branch:
state = model.null_state_token.unsqueeze(0).expand(batch_size, -1)
Normalization
Z-score (mean/std) — this is a plain pi0 model (model_type == PI0), so openpi uses the
mean/std branch of Normalize, not quantile or min-max. Stats come from
assets/mkpongm/sponge_marker_merged/norm_stats.json (shipped here).
- normalize (state in):
(x - mean) / (std + 1e-6) - un-normalize (action out):
x * (std + 1e-6) + mean
Un-normalize the action outputs with the inverse before sending to the robot.