M3 — Train-Time Counterfactual Velocity Steering (RoboWarp)
Training-time distillation of RoboWarp's test-time CF-Steer guidance geometry: π₀.₅
finetuned so that the flow-matching regression target is steered away from the frozen
base checkpoint's counterfactual (attention-erased) velocity field. Companion to M2
(adipotnis/m2-train-time-cf-robowarp), which instead tilts the flow source; init,
teacher, data and normalization are identical across the two, so they are directly
comparable. Trained with openpi (JAX)
on 2× NVIDIA GH200.
Contents: params/ (Orbax weights) + assets/ (the released pi05_libero quantile
normalization stats — required at inference). No train_state/, so this cannot be resumed.
Method (M3)
Per training sample (openpi convention, t=1 noise):
eps ~ N(0, I), t ~ Beta(1.5, 1)·0.999 + 0.001
x_t = t·eps + (1−t)·a u = eps − a (standard CFM pair)
v_neg = v_teacher(x_t, t | o_cf) frozen released pi05_libero on the
attention-ERASED base image
s* = ⟨u, v_neg⟩ / ‖v_neg‖² CFG-Zero* scale
v_tgt = u + ω·(u − s*·v_neg) w.p. p, else u
loss = ‖v_θ(x_t, t | o_factual) − v_tgt‖²
This is CF-Steer's guidance formula v₊ + ω(v₊ − s*v₋) with the ground-truth conditional
velocity in the factual slot: the student is pushed away from what the evidence-blind view
dictates while staying anchored to the demonstrations (not capped at the teacher). The
teacher runs frozen, forward-only, under stop-gradient. Erasure masks are RoboWarp CF-Steer
masking on the frozen released checkpoint: layer-12 suffix→image attention, register-mask +
spike cleaning (ratio 2.0, 15 iters, σ 0.8), 0.85-quantile threshold (top-15 % pixels,
measured coverage 15.2 %), fill 128. Inference is unchanged: single pass from N(0, I),
no masking, no second forward.
Hyperparameters
| Init & teacher | gs://openpi-assets/checkpoints/pi05_libero (released; teacher = init, frozen) |
| Model | π₀.₅ (pi05=True, action_horizon 10, discrete_state_input=False — state-blind, matching the released checkpoint and the RoboWarp baseline) |
| Data | physical-intelligence/libero v2.0, LIBERO-Spatial suite only: 432/1693 episodes, 52,970 frames (task indices 30–39) |
| Normalization | released pi05_libero quantile stats (from the checkpoint's own assets) |
| Steering | ω = 0.5, p = 0.5, CFG-Zero* scale on |
| Batch / steps | 112 (2× GH200, FSDP; reduced from M2's 160 to fit the resident teacher) / 2,300 (≈4.9 epochs) |
| LR | 3.3e-5 → 3.3e-6 cosine, warmup 200 (sqrt-rescaled from 256@5e-5) |
| Optimizer | AdamW β=(0.9, 0.95), eps 1e-8, wd 1e-10, grad-clip 1.0, EMA 0.999 |
| Masking params | layer 12, erase_pct 0.85, fill 128, ratio 2.0, iters 15, σ 0.8, attn from 2 denoise steps |
| Precision | bfloat16 (teacher also bfloat16) |
Caveats
- Training loss is not comparable to a standard-CFM baseline or to M2 — the steered target has larger norm on disagreement states. Evaluate by simulator success rate (LIBERO-PRO spatial/swap) only.
- Not evaluated in sim at upload time. ~4.9 epochs (vs ~48 in the reference recipe).
- ω = 0.5 is an initial estimate, not swept.
- State-blind by design (matches the RoboWarp test-time baseline for comparability).
- Comparison runs:
adipotnis/m2-train-time-cf-robowarp(source-tilted twin),adipotnis/pi05-libero-spatial-run1(plain finetune of pi05_base on the same data); CF-Steer test-time numbers in the RoboWarp report.