adipotnis/m2-train-time-cf-robowarp

M2 — Train-Time Counterfactual Flow (RoboWarp)

Training-time internalization of RoboWarp's test-time CF-Steer: π₀.₅ finetuned so that the flow's source distribution is tilted toward the policy's own counterfactual (attention- erased) behavior, making the regressed velocity the prior→evidence correction. 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 (M2)

1. Evidence map — RoboWarp CF-Steer masking on the frozen released pi05_libero: layer-12 suffix→image attention, register-mask + spike cleaning (ratio 2.0, 15 iters, σ 0.8), smoothed, thresholded at the 0.85 quantile (top-15% pixels; measured coverage 15.2%).
2. Counterfactual prior — a_prior = pi05_libero(o₋) with the attended pixels greyed (fill 128, camera still present), 10-step flow sample, normalized action space. prior_norm_mean = 1.453.
3. Tilted source — with probability p, x₁ = α·a_prior + √(1−α²)·ε (else ε ~ N(0,I)), so the CFM target on tilted samples is α(a_prior − a) + √(1−α²)ε: prior behavior earns zero loss reduction; only evidence-dependent corrections remain learnable. Inference is unchanged single-pass from N(0, I).

Hyperparameters

Init & drafter	gs://openpi-assets/checkpoints/pi05_libero (released; drafter = init)
Model	π₀.₅ (pi05=True, action_horizon 10, discrete_state_input=False — state-blind, matching the released checkpoint and the RoboWarp paper 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)
Coupling	α = 0.6, p = 0.5
Batch / steps	160 (2× GH200, FSDP) / 2,300 (≈6.9 epochs)
LR	3.95e-5 → 3.95e-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
Prior sampling	10 Euler steps, seed 7000, batch 256
Precision	bfloat16

Caveats

• Training loss is not comparable to a standard-CFM baseline — tilted sources shrink targets on prior-correct states. Evaluate by simulator success rate (LIBERO-PRO spatial/swap) only.
• Not evaluated in sim at upload time. ~6.9 epochs (vs ~48 in the reference recipe).
• State-blind by design (matches the RoboWarp test-time baseline for comparability).
• Comparison runs: spatial_run1 (plain finetune of pi05_base on the same data, HF: adipotnis/pi05-libero-spatial-run1); CF-Steer test-time numbers in the RoboWarp report.