Online Reinforcement Learning (RL) holds immense potential for advancing instruction-guided image editing, but its progress has been severely hindered by a critical perception gap we term "Attention Collapse". Existing reward models frequently neglect cross-image comparisons and fail to capture fine-grained editing details, leading to inaccurate evaluations and unstable RL training.
1git clone https://github.com/Kwai-Keye/SpatialReward.git
2cd SpatialReward
3
4conda create -n spatialreward python=3.11 -y
5conda activate spatialreward
6
7pip install torch==2.8.0 torchvision --extra-index-url https://download.pytorch.org/whl/cu126
8pip install -r requirements.txt
1# Start reward server
2cd example/reward/server
3bash start_servers.sh
4bash start_proxy.sh
5
6# Query from client
7from example.reward.client.reward_client_edit import RewardClient
8
9client = RewardClient(proxy_host="127.0.0.1", proxy_port=23456)
10scores, rewards, reasoning, meta_data = client.evaluate(
11 input_images=[input_img],
12 output_image=[output_img],
13 meta_datas=[{"instruction": "Remove the dog"}]
14)
Model and data are loaded directly from HuggingFace by default.
1# MER-Bench
2bash eval/MERBench/run.sh
3
4# MMRB2
5bash eval/MMRB2/run.sh
6
7# EditReward-Bench
8bash eval/EditReward-Bench/run.sh
1# Replace ORM first
2cp example/SpatialReward-train/rl/orm.py <ms-swift>/swift/plugin/orm.py
3bash example/SpatialReward-train/rl/run_mater.sh
1@article{long2026spatialreward,
2 title={SpatialReward: Bridging the Perception Gap in Online RL for Image Editing via Explicit Spatial Reasoning},
3 author={Long, Yancheng and Yang, Yankai and Wei, Hongyang and Chen, Wei and Zhang, Tianke and Fan, Haonan and Liu, Changyi and Jiang, Kaiyu and Chen, Jiankang and Tang, Kaiyu and Wen, Bin and Yang, Fan and Gao, Tingting and Li, Han and Yang, Shuo},
4 journal={arXiv preprint arXiv:2602.07458},
5 year={2026}
6}
