robometer-4b-full-threading-d0

Robometer-4B Full Fine-tune — MimicGen threading_d0

Full-parameter fine-tune of robometer/Robometer-4B (Qwen3-VL-4B-Instruct) on MimicGen threading_d0 rollouts. Emits, per frame, a progress value in [0,1] and a success probability.

This is the no-PEFT counterpart of chomeed/robometer-4b-lora-threading-d0: identical data, schedule, and eval config — the only change is model.use_peft=false.

• 4.03B trainable params (90.7%) — language model + all three prediction heads; the vision tower stays frozen (train_vision_encoder: false), same as the LoRA run
• 1000 steps, batch 8, lr 2e-5, cosine, warmup 0.1, weight decay 0.01, grad-norm clip 10
• bf16 full fine-tuning via Unsloth; 34 GB peak on one B200 (the LoRA run peaked at 14.7 GB)
• Uploaded checkpoint is step 500, the best of the run by mean(pearson, kendall_last)
• Data: chomeed/mimicgen_threading_d0_train_rfm (160 traj: 80 successful + 80 failure), evaluated on ..._test_rfm (40 traj, disjoint seeds)

Results (held-out, 40 trajectories with unseen seeds)

Best checkpoint of each run, scored by the trainer's own eval:

Matched-step comparison at step 1000, so the two runs are compared at equal training budget:

Full fine-tuning wins consistently but by a small margin — Pearson is higher at every matched eval step, and eval loss is lower throughout. Kendall is saturated at 1.000 for both, so this test set does not separate them on trajectory ranking.

Success-head calibration: both models emit very low absolute success probabilities (peak ≈0.002 on successful held-out trajectories). Ranking is excellent (AUROC 0.99) but the probabilities are not calibrated — threshold well below 0.01, or just use the progress signal.

Usage

Important: this is a full-weights checkpoint, like the upstream robometer/Robometer-4B. It keeps fp32 pixel_values against bf16 vision blocks, which is fine under the Trainer's autocast but raises RuntimeError: expected scalar type BFloat16 but found Float at plain inference. Wrap the forward pass in torch.autocast. (LoRA checkpoints work either way.)

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")
exp_config, tokenizer, processor, model = load_model_from_hf(model_path="chomeed/robometer-4b-full-threading-d0", device=device)
model.eval()
collator = setup_batch_collator(processor, tokenizer, exp_config, 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()}

with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):   # <-- required
    out = compute_batch_outputs(model, tokenizer, inputs, sample_type="progress",
                                is_discrete_mode=True, num_bins=10)   # progress head is discrete, 10 bins

progress = np.array(out["progress_pred"][0])                     # (T,) in [0,1]
success  = np.array(out["outputs_success"]["success_probs"][0])  # (T,) in [0,1]

Serving: python robometer/evals/eval_server.py model_path=chomeed/robometer-4b-full-threading-d0

Reproducing

TASK=mimicgen_threading_d0 GPU=0 STEPS=1000 bash run_full_task.sh

Requires robometer/data/dataset_success_cutoff.txt to contain mimicgen_threading_d0_train_rfm,1.0 and mimicgen_threading_d0_test_rfm,1.0 (simulation trajectories have exact endpoints), and the RFM datasets preprocessed via robometer/configs/preprocess_threading_d0.yaml.