🎉🎉🎉 MaskGRPO is accepted by ICLR 2026. We will be updating this repository with the newest results and checkpoints, stay tuned!

This is the official implementation of MaskGRPO: Consolidating Reinforcement Learning for Multimodal Discrete Diffusion Models，an open-sourced framework for SFT and GRPO training of discrete diffusion models (DDMs) including LLaDA and MMaDA.

In this repo, we release:

1. Improved importance estimation for reinforcing DDMs with controlled randomness across devices.
2. AR-like reversing for RL training on math reasoning and coding tasks.
3. Emerge sampler for image generation and RL training (the principled DDM sampler modified from MDLM).
4. Detailed SFT, RL and Evaluation scripts.

🚀 Quick Start

First, set up the enviroment with python==3.10 (supposing you are using cu128. For other versions, please read the comments in requirements):

conda create -n maskgrpo python==3.10
conda activate maskgrpo
pip install torch==2.7.0 torchvision==0.22.0 torchaudio==2.7.0 --index-url https://download.pytorch.org/whl/cu128
pip install -r requirements.txt
pip install git+https://github.com/huggingface/trl.git@0f88c179e30b3439467942a08c3190f624d5c423

Then, download the MMaDA and LLaDA checkpoints, and update related paths in configs/*.yaml and scripts/train/*.sh accordingly.

You may also need to pre-download or prepare the datasets you want. Please refer to data or data_utils for data management details.

📖 Math reasoning and coding (T2T tasks):

For math reasoning, run

bash scripts/train/t2t_grpo_train.sh

For coding, you have to first launch a sandbox using

bash scripts/run_e2b.sh

then you can run the training script.

🎨 Image generation (T2I tasks):

For SFT, use

bash scripts/train/t2i_sft.sh

For GRPO training with CLIP or HPSv2, use

bash scripts/train/t2i_grpo_train.sh

If you want to run with Unified-Reward, HPSv3 or other MLLM as servers, please launch the server (see reward_func for communication protocols, reward_server for detailed configuration) and run

bash scripts/train/t2i_grpo_train_w_server.sh

🧑‍🏫 Evaluation

Just run the scripts in scripts/eval and get the results.

As additional enviroments may need to be installed for image generation tasks, please refer to the original benchmark repository for instructions.

Here are more comparisons between the default MaskGIT-style results (before RL) and ours (after RL).

🧷 PS

Theoretically, Lumina-Dimoo can also be supported, as it shares a similar structure to MMaDA. We welcome any PR or issue post on efficiency improvement and compatibility!

📜 Acknowledgement

This repository is built off of MMaDA, diffu-GRPO, DiffuCoder, MDLM and ReDDiT. Thanks to all the authors for their great work.

📄 BibTex

@article{ma2025MaskGRPO,
      title={Consolidating Reinforcement Learning for Multimodal Discrete Diffusion Models.}, 
      author={Tianren Ma and Mu Zhang and Yibing Wang and Qixiang Ye},
      year={2025},
      eprint={2510.02880},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}