chomeed/robometer-4b-lora-square-d0-seg

Robometer-4B LoRA — MimicGen square_d0

LoRA fine-tune of robometer/Robometer-4B on MimicGen square_d0 (nut assembly). Emits, per frame, a progress value in [0,1] and a success probability.

Task: Pick up the square nut and place it on the peg.

Results — against the zero-shot base

Both models on the same held-out set, chomeed/mimicgen_square_d0_test_seg_rfm (n=40: 20 successful, 20 failure). Metrics are cutoff-independent, so the training-target definition does not leak into the score.

metric	base Robometer-4B	this LoRA
ranking_acc — P(success final > failure final)	0.877	0.988	+13%
kendall — progress monotonicity vs frame index	0.567	0.774	+36%
success_auprc	0.895	0.990	+11%
succ_fail_gap — mean final progress, success − failure	+0.302	+0.415	+37%
pos_acc — successes scored > 0.5	0.90	0.95
neg_acc — failures scored ≤ 0.5	0.55	0.85	+55%

Every metric improves. Ranking, monotonicity and calibration all move together — unlike the threading counterpart (..-lora-threading-d0-seg), where only calibration improved and ranking stayed flat.

Read this as a zero-shot-base comparison, not a finetune-vs-finetune one. A large gain over an un-adapted base is expected; it does not say this recipe beats an alternative finetune. n=40 also means small differences would be noise — these gaps are large enough to be real, but the test set is not big enough to rank close variants.

Both models were cast to bfloat16 for evaluation: the base ships mixed dtype (483 bf16 / 249 fp32 params) and its LayerNorms reject the activations otherwise. Both get identical treatment.

Training

• LoRA r=32, α=64, dropout=0.05 on all attention + MLP projections
• 1000 steps, batch 8, lr 2e-5, cosine, warmup 0.1, weight decay 0.01; 50m36s on one B200
• loss 0.42 → 0.13; success cutoff 0.95

Data

chomeed/mimicgen_square_d0_train_seg_rfm — 160 trajectories, 80 successful / 80 failure, balanced by construction.

Failures come from ..._40k_failure_seg: tail cut at the outcome, approach kept, with the cut point depending on how far the attempt got (lifted_then_failed → closest approach to the seated pose; failed_grasp → the moment the grasp failed; never_reached_nut → uncut). Successes are drawn half from ..._40k_success and half from the expert demos ..._224x224, unsegmented.

Successes are unsegmented because they already end at their outcome, so both classes start at frame 0 and end at the outcome — the framing does not leak the label.

Usage

out = compute_batch_outputs(model, tokenizer, inputs, sample_type="progress",
                            is_discrete_mode=True, num_bins=10)   # MUST be True
progress = np.array(out["progress_pred"][0])   # (T,) in [0,1], monotone

is_discrete_mode=True is required: this model trains with progress_loss_type=discrete, and decoding with False returns raw per-bin logits (T×bins values, range ~[-3,5]) instead of progress. It does not raise — it silently yields plausible but meaningless numbers.

Caveat

save_best could not rank checkpoints (its configured metric names do not match the ones logged for this dataset), so the published weights are the final step-1000 checkpoint, not a selected best.