ImageWAM-FLUX.2-4B-RoboTwin-C2R-FT

This repository contains the ImageWAM FLUX.2 4B checkpoint for OOD RoboTwin Clean2Random finetune experiment from ImageWAM: Do World Action Models Really Need Video Generation, or Just Image Editing?

ImageWAM is a family of world action models built on image-editing foundation models. This checkpoint is intended for evaluation and research use with the accompanying ImageWAM codebase.

Model Details

• Model family: ImageWAM
• Image-editing backbone: FLUX.2 [klein] base
• Variant: FLUX.2 klein-base-4B
• Benchmark: RoboTwin
• Training code: yuyangalin/ImageWAM
• Base model weights: Users must separately prepare the FLUX.2 klein-base-4B weights and FLUX.2 autoencoder as described in the ImageWAM README.

Files

Expected file layout:

.
├── model.pt
├── dataset_stats.json
└── config.yaml

• model.pt: ImageWAM checkpoint used by the evaluation scripts.
• dataset_stats.json: normalization statistics required for policy evaluation.
• config.yaml: original training configuration for provenance and reproducibility.

Usage

Install and prepare the ImageWAM repository following the project README. Then download this model repository:

mkdir -p checkpoints/imagewam_release/robotwin/flux2_klein_4b_c2r

huggingface-cli download yuyangalin/ImageWAM-FLUX.2-4B-RoboTwin-C2R-FT \
  --repo-type model \
  --local-dir checkpoints/imagewam_release/robotwin/flux2_klein_4b_c2r

Prepare FLUX.2 4B weights and set:

export FLUX2_VARIANT=4b
export FLUX2_MODEL_PATH=/path/to/flux-2-klein-base-4b.safetensors
export FLUX2_AE_MODEL_PATH=/path/to/ae.safetensors
export FLUX2_QWEN3_MODEL_SPEC=Qwen/Qwen3-4B

Prepare the RoboTwin benchmark environment and assets following the ImageWAM README. Then evaluate on RoboTwin:

export CKPT_PATH="$(pwd)/checkpoints/imagewam_release/robotwin/flux2_klein_4b_c2r/model.pt"
export DATASET_STATS_PATH="$(pwd)/checkpoints/imagewam_release/robotwin/flux2_klein_4b_c2r/dataset_stats.json"

NUM_GPUS=8 FLUX2_VARIANT=4b bash scripts/flux2/run_eval_flux2_robotwin.sh

RoboTwin evaluation uses EVALUATION.skip_get_obs_within_replan=true by default to speed up evaluation. Set SKIP_GET_OBS_WITHIN_REPLAN=false if you need fully rendered videos.

Intended Use

This checkpoint is intended for:

• Reproducing ImageWAM RoboTwin evaluations.
• Research on robot policy learning, world action models, and image-editing-based action generation.
• Comparison against other RoboTwin policy models under the same evaluation setup.

This checkpoint is not intended for safety-critical or real-world robot deployment without additional validation.

Limitations

• Evaluation requires the ImageWAM codebase, RoboTwin benchmark environment, and RoboTwin assets.
• The checkpoint assumes the same model variant and configuration used during training. See train_config.yaml.
• Users must separately prepare the matching FLUX.2 4B base model and autoencoder weights.
• Performance may differ if the simulator version, RoboTwin assets, dataset preprocessing, camera layout, action normalization statistics, or evaluation settings differ from the release setup.

Citation

If you use this checkpoint, please cite the ImageWAM paper:

@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},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2606.19531}, 
}

Acknowledgements

ImageWAM builds on several open-source projects and model families, including FLUX.2, FastWAM, LIBERO, LIBERO-plus, and RoboTwin. Please also follow the licenses and citation requirements of the corresponding upstream projects.