armanakbari4/imagewam-ur3-ethernet

ImageWAM-FLUX.2-4B — real-world UR3 ethernet-cable insertion

Real-world fine-tune of ImageWAM (FLUX.2 [klein] base 4B editing DiT + ActionDiT action expert) on a bimanual dual-arm UR3, for the task "insert the ethernet cable on the slot".

Initialized from yuyangalin/ImageWAM-FLUX.2-4B-InternData-A1-EE (the InternData-A1 pretrain, step 60k) — loaded with 0 missing / 0 unexpected keys.

Files

File	What
ur3_ethernet_ee16_step4000.pt	step 4,000
ur3_ethernet_ee16_step5000.pt	step 5,000 — action_l1 0.0122
ur3_ethernet_ee16_step7000.pt	step 7,000 (final) — best held-out action_l1, 0.0106
ur3_ethernet_ee16_dataset_stats.json	z-score normalization stats — required for inference
train_config.yaml	resolved training config, for provenance

Each .pt holds only the trained parts (mot ≈ 8.2 B params + proprio_encoder), ~9.0 GB. The FLUX.2 klein-base-4B base weights and autoencoder must be prepared separately, as for any ImageWAM checkpoint.

Data

EmbodyX/UR3 → ethernet_2.0_lerobot — LeRobot v2.1, 103 episodes / 59,796 frames at 15 fps.

• 3 cameras (camera_top, camera_wrist_left, camera_wrist_right) at 240×320, AV1-encoded
• 14-dim absolute joint action/state, [L_arm(6) | L_grip | R_arm(6) | R_grip]

Recipe

Upstream's fine-tune-from-pretrain recipe: lr 2.5e-5 cosine w/ 5% warmup, AdamW(0.9, 0.95), wd 1e-2, grad-clip 1.0, bf16, DeepSpeed ZeRO-1, global batch 192 (12/GPU × 2 GPUs × 8 grad-accum), 7,000 steps (~22 epochs), ~14 h on 2×H100. num_frames=17, action_video_freq_ratio=1 → 16-step action horizon, endpoint_frames_only=true.

Inference — three things that will bite you

1. Mosaic is 288×256 (compact_288x256): camera_top 192×256 on top, the two wrists 96×128 side-by-side below, order fixed [top, left, right], pixels normalized to (−1, 1).
2. Actions come back 16D. This checkpoint uses the ee16 layout: the robot's 14D vector is zero-padded at dims 7 and 15 (nothing is converted to end-effector poses). Un-pad back to 14D — concat(x[..., 0:7], x[..., 8:15]) — before sending to the controller.
3. Use ur3_ethernet_ee16_dataset_stats.json for denormalization. These stats are computed on this dataset, not the pretrain's.

Checkpoint selection

Unlike our other UR3 fine-tune, held-out action error kept improving through the whole run rather than plateauing. By thirds of training, action_l1 went 0.0303 → 0.0209 → 0.0169; regressing action_l1 on step over 14 evals gives a slope of −0.0027 per 1k steps (t = −4.26), so the downtrend is real rather than eval noise. The best evals are the latest: step 7000 (0.0106), 5000 (0.0122), 5500 (0.0132).

Two caveats. Per-eval noise (sd ≈ 0.0071 over 32 clips) is larger than the gap between adjacent checkpoints, so this ranks trends, not neighbours. And offline action_l1 has not been validated against real-robot success rate on this setup. The curve had not flattened at 7,000 steps, so training longer may help.

Citation

Built on ImageWAM:

@misc{zhang2026imagewam,
      title={ImageWAM: Do World Action Models Really Need Video Generation, or Just Image Editing?},
      author={Yuyang Zhang and Wenyao Zhang and Zekun Qi and He Zhang and Haitao Lin and Jingbo Zhang and Yao Mu and Xiaokang Yang and Wenjun Zeng and Xin Jin},
      year={2026},
      eprint={2606.19531},
      archivePrefix={arXiv},
      url={https://arxiv.org/abs/2606.19531},
}