zhuyuze/zero-wam-robot-hm-phase

Zero-WAM robot-branch NEXT / GRASP / PLACE heatmap readouts

These are small readout heads (4.5M parameters each, fp32) trained on frozen Zero-WAM activations. Each head predicts where the robot's gripper will next make contact in the head camera image. The input is the robot's most recent deployment chunk.

head	target
NEXT	The next gripper event strictly after the chunk end, with both arms merged. During approach this is the grasp (or press) point; while holding it is the release point. This is the default readout.
GRASP	The upcoming grasp point. Rows are approach-phase only.
PLACE	The arm's next gripper-release point, including before the grasp (it skips over the upcoming grasp).

Grasp and release points follow one definition: the gripper TCP plus a per-task finger-contact offset measured in simulation. The release point is taken at the onset of gripper opening, so it is not the object's landing point.

Input features

• Model: Zero-WAM, the RoboTwin post-train checkpoint, run in deployment mode. That means the robot stream uses the empty text embedding (TARGET_TEXT_CFG=-1), sees the paired human demo (ICL) and past actions, and is teacher-forced.
• Layer: the residual stream after 18 transformer blocks (forward-hook output of blocks[17]).
• Tokens: the clean cam_high tokens of latent frame kf = min(2c+1, F-1) for deployment chunk c. That is a 7 × 9 grid (token index r*9+c) of 3072-d vectors, shape [63, 3072].
• Image frame: the 640 × 480 head camera is bilinearly resized to 288 × 224 "model px", with u_m = 0.45 u and v_m = v * 224/480.

Output

• A softmax probability map over 28 × 36 cells of 8 px each (convex ×4 upsampling of the token grid).
• Up to two decoded peaks in model px. The second peak counts as a real second target (the other arm's event within 16 frames) when count_score >= tau (tau is in config.json).

from modeling import load_head, predict_points
head = load_head("NEXT/seed0")                       # model.safetensors + config.json
prob, peaks, count_score = predict_points(head, feats)   # feats [B, 63, 3072]

model.safetensors stores the head (head.*) and the feature standardizer (std.*: per-channel mean and sd fitted on the training rows). modeling.py has no dependencies beyond torch and safetensors.

Training

• Data: RoboTwin 2.0, 2,499 episodes. Readout trained on 43 seen tasks, tested on 7 unseen tasks: open_microwave, place_empty_cup, stamp_seal, stack_blocks_three, place_bread_basket, move_stapler_pad, place_object_scale.
• Loss: cross-entropy against a σ = 8 px Gaussian target. AdamW, lr 1e-3, wd 5e-2, batch 64, cosine schedule with warmup.
• Epochs: chosen by early stopping on an episode-grouped inner validation split (patience 10, max 60). The head was then refit on all training rows for best_epoch + 1 epochs.
• Seeds: 0, 1 and 2 (<HEAD>/seed<k>/). Full hyperparameters are in each config.json.

Results (unseen-7 test rows, PCK@8 / PCK@16 in model px)

head	test rows	PCK@8 (these weights, 3 seeds)	PCK@16	PCK@8 in the original study
NEXT	1847	0.714 ± 0.006	0.894 ± 0.004	0.716 ± 0.003
GRASP	1031	0.848 ± 0.003	0.933 ± 0.001	0.849 ± 0.015
PLACE	1815	0.473 ± 0.009	0.839 ± 0.002	0.478 ± 0.017

Two notes on these numbers:

• PLACE before the grasp: the original study's PCK@8 on pre-grasp rows was 0.43 with the demo visible, 0.31 with demo attention masked, and 0.28 for raw VAE latents.
• Reproducibility: the original study kept only predictions, not weights. These heads were re-trained with the identical code, configuration and seeds. GPU training is not bit-deterministic, so the chosen epochs and weights differ slightly from the original runs, but the metrics agree within seed spread. Each config.json records the comparison with the original run's predictions under reproduces_original_run. The peak differences there include the arbitrary second peak on single-target rows.