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.