README.md

---
license: mit
library_name: pytorch
tags:
  - robotics
  - imitation-learning
  - world-model
  - image-editing
  - flux
  - imagewam
datasets:
  - EmbodyX/UR3
---

# 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`](https://huggingface.co/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`](https://huggingface.co/datasets/EmbodyX/UR3/tree/main/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](https://github.com/yuyangalin/ImageWAM):

```bibtex
@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},
}
```