robometer-4b-lora-threading-d0-seg

Robometer-4B LoRA — MimicGen threading_d0, segmented clips

LoRA fine-tune of robometer/Robometer-4B on grasp → terminal clips of MimicGen threading_d0. Emits, per frame, a progress value in [0,1] and a success probability.

Distinct from chomeed/robometer-4b-lora-threading-d0, which trains on full episodes; that adapter is untouched. Two things differ here: the data is cut to the segment between grasp and outcome, and the success cutoff is 0.95 instead of the 1.0 default.

Comparison against the original adapter

Both adapters evaluated on the same held-out set, chomeed/mimicgen_threading_d0_test_seg_rfm (n=40: 20 successful, 20 failure). Metrics are deliberately cutoff-independent — the two models were trained with different success cutoffs, so anything defined against the training target would bake that difference into the score.

The two models rank equally well; this one is far better calibrated. Ranking accuracy, monotonicity and AUPRC are within noise at n=40 — both order successes above failures about equally. What changed is the scale: the original assigns final progress above 0.5 to 85% of failure trajectories, declaring nearly-complete on runs that failed. This adapter halves that and separates success from failure by 2.3× the margin.

That distinction matters for a reward model specifically, because an RL agent consumes the value, not the ranking. A signal that says "0.9 done" on a failure is misleading even when its ordering is perfect.

Caveats, stated plainly

• n=40. neg_acc 3/20 vs 10/20 is significant but marginal (Fisher p ≈ 0.04). The tie-level metrics are genuinely ties, not small wins — do not read ranking_acc 0.922 vs 0.917 as a difference.
• _seg is this model's native domain. The original trained on full episodes including approach frames, so part of the margin may be domain match rather than better reward modelling. The reverse-direction evaluation (both models on _test_rfm) would separate those and has not been run.
• save_best could not rank checkpoints: it is configured with the reference run's metric names, which do not match the ones logged for this dataset. The published weights are the final step-1000 checkpoint, not a selected best.

Training

• LoRA r=32, α=64, dropout=0.05 on all attention + MLP projections (75.9M trainable, 1.68%)
• 1000 steps, batch 8, lr 2e-5, cosine, warmup 0.1, weight decay 0.01 — matching the reference recipe
• 54m34s on one B200; loss 0.55 → 0.085, grad-norm 10 → 0.5
• Success cutoff 0.95 for both data sources

Data

chomeed/mimicgen_threading_d0_train_seg_rfm (156 trajectories: 78 successful / 78 failure) and ..._test_seg_rfm (40: 20/20). Balanced by construction — 98 failure clips is the hard ceiling, so successes were held to match, drawn half from policy rollouts and half from expert demos so both success distributions are represented.

On the 0.95 cutoff

Measured on the full episodes, success fires at ~97.6% of episode length (p10 = 0.953), so 1.0 is too strict there. In these _seg clips the cut is made at the outcome frame, so success sits at exactly 1.0 and a cutoff of 1.0 would be exact. 0.95 was chosen deliberately anyway: it widens the success plateau from 1 frame to 5 of 64, raising the positive rate ~5×, which is what the reference run's degenerate success head needed. It trades a little label precision for class balance.

Usage

import numpy as np, torch
from robometer.data.dataset_types import ProgressSample, Trajectory
from robometer.evals.eval_server import compute_batch_outputs
from robometer.utils.save import load_model_from_hf
from robometer.utils.setup_utils import setup_batch_collator

device = torch.device("cuda")
cfg, tokenizer, processor, model = load_model_from_hf(
    model_path="chomeed/robometer-4b-lora-threading-d0-seg", device=device)
model.eval()
collator = setup_batch_collator(processor, tokenizer, cfg, is_eval=True)

frames = ...  # uint8 (T, H, W, C), RGB, third-person view
traj = Trajectory(frames=frames, frames_shape=tuple(frames.shape),
                  task="Pick up the needle and thread it through the hole in the tripod.",
                  id="0", metadata={"subsequence_length": int(frames.shape[0])},
                  video_embeddings=None)
batch = collator([ProgressSample(trajectory=traj, sample_type="progress")])
inputs = {k: (v.to(device) if hasattr(v, "to") else v)
          for k, v in batch["progress_inputs"].items()}

# is_discrete_mode MUST be True: 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. That
# does not raise -- it silently yields plausible but meaningless numbers.
out = compute_batch_outputs(model, tokenizer, inputs, sample_type="progress",
                            is_discrete_mode=True, num_bins=10)
progress = np.array(out["progress_pred"][0])   # (T,) in [0,1], monotone