Efficient, stable and scalable reinforcement of diffusion models

Text-to-image  ·  Text-to-video  ·  Image-to-video

Method  /  Quick Start  /  Documentation

DiffusionART is an on-policy diffusion RL algorithm via Aggregation-First Reward Tilting with Adaptive Restoration. It combines the advantages of forward-process and reverse-process reinforcement learning while avoiding their limitations. DiffusionART is highly efficient, stable and simple, requires minimal hyperparameter tuning and no differentiable reward model, naturally supports pairwise feedback, and is robust to small rollout groups for scalable diffusion reinforcement learning.

This repository provides a comprehensive framework for fine-tuning visual generation models, supporting widely used open-source base models, reward models and state-of-the-art algorithms across text-to-image, text-to-video and image-to-video generation.

Fine-tuning HunyuanVideo (13B) (left) and Wan2.1-I2V-14B-480P, CFG/step-distilled (14B) (right) using VideoReward, with Visual Quality (VQ):Motion Quality (MQ) = 1:1. Mean train reward is the mean of raw VQ and MQ. T2V ratios compare GPU hours to reach each baseline's final reward level.

1. NFT-style forward-process RL contains a hidden forward-reverse consistency force in its objective, slowing and destabilizing training.
2. Off-policy strategies stabilize reinforcement learning of continuous-space generative models through a strong restoration force. Even a small output-space restoration MSE can produce a large aggregated parameter gradient.

DiffusionART is simple: roll out K samples per prompt with an ODE solver, compute the reward-endpoint covariance as guidance, apply it to the visited states and actions, then perform two disjoint minibatch updates. We call this reward-improving direction guidance, bringing a classifier-free guidance (CFG)-like idea into diffusion RL: the model learns to internalize reward guidance, much like distilling CFG. For a reward-tilted endpoint distribution with tilt strength eta,

Here c is a prompt, x₀ is a clean latent endpoint, and R is its reward. DiffusionART systematically removes the self-consistency force: it aggregates reward evidence before applying the normalized guidance to reverse-process rollout states, without relying on an ideal forward-process velocity field. The second update retains guidance while adaptively restoring velocity drift toward the rollout policy, without an EMA off-policy strategy.

Guidance geometry and state assignment	Pseudocode
	Implementation: covariance guidance, ART loss, and disjoint update objectives.

DiffusionART-Pairwise is a variant in which the reward model gives preference scores for two visual outputs: 1 for the winner, 0 for the loser, and 0.5 for a tie. It accommodates pairwise feedback by simply replacing scalar rewards with these labels when computing the shared guidance. See the preference-label and pairwise-guidance implementation.

Pairwise T2V with K = 2 or 8 videos per prompt. Left: VQ train reward versus GPU hours. Right: final evaluation VQ and total GPU hours for pairwise K = 2, 4, 8 and scalar K = 8, each trained for 400 updates.

• ✅ ODE solver: deterministic rollouts, with no reverse-transition likelihood.
• ✅ Rollout efficiency: reuse rollout states and cached predictions, without adding noise again or recomputing predictions on newly noised states.
• ✅ Sample efficiency: fewer samples per prompt with minimal performance loss and lower GPU-hour cost.
• ✅ Training efficiency: remove forward-process RL's unstable self-consistency force and address reverse-process RL's inefficient credit assignment.
• ✅ Simple and scalable: minimal hyperparameter tuning, no differentiable rewards required, and natural support for pairwise reward models.

Pre-trained versus DiffusionART. Video strips show four frames.

• 🎉 [2026.10.12] Our paper, code, models and project page are now public! DiffusionART and six baselines share one framework for T2I, T2V and I2V. Explore the training recipes!

Each model card includes its inference setup, base-model requirements and examples.

Use Python 3.10+ and uv. Pretrained training requires Linux and CUDA.

git clone https://github.com/vanzll/DiffusionART.git
cd DiffusionART
uv sync --locked --python 3.10 --extra diffusers

uv run --locked visual-rl --config configs/examples/cpu_smoke.yaml --validate
uv run --locked visual-rl --config configs/examples/cpu_smoke.yaml

Accept the model license on Hugging Face, then configure HF_TOKEN and WANDB_API_KEY, or use existing Hub/W&B logins. Paper recipes enable W&B by default. Set VISUAL_RL_ASSET_ROOT to choose a writable asset cache. Videos use the pinned video environment.

Run the complete SD3 PickScore experiment on eight GPUs:

uv run --locked --extra diffusers visual-rl \
  --config configs/paper/image/sd3_pickscore_art.yaml --nproc-per-node 8

The command prepares pinned models and prompt data before starting 1000 training updates. Missing models are downloaded from Hugging Face automatically; existing assets are reused. No manual model download is required for this recipe. GenEval and OCR also download benchmark prompts from our HF dataset. First complete the GenEval scorer setup or OCR scorer setup, then use the full training command in that guide. See all image recipes for other algorithms.

Resume the latest complete checkpoint, optionally extending the total budget:

uv run --locked --extra diffusers visual-rl \
  --config configs/paper/image/sd3_pickscore_art.yaml --nproc-per-node 8 \
  --resume latest --max-updates 1500

Video recipes use the same entrypoint inside the video environment. For shared multi-node preparation, generate a fresh launch ID once:

python -c 'import uuid; print(uuid.uuid4().hex)'

Export that same value as PREPARATION_RUN_ID on every node. This lets peers reuse one full asset validation within this launch; use a new ID on restart. T2V (400 updates):

visual-rl --config configs/paper/video/hunyuan_art_721.yaml \
  --nproc-per-node 8 --nnodes 4 --node-rank "$NODE_RANK" \
  --master-addr "$MASTER_ADDR" --preparation-run-id "$PREPARATION_RUN_ID" \
  --rendezvous-timeout-seconds 7200

I2V (400 updates, automatically generated evaluation references):

visual-rl --config configs/examples/wan_art_online_eval.yaml \
  --nproc-per-node 8 --nnodes 4 --node-rank "$NODE_RANK" \
  --master-addr "$MASTER_ADDR" --preparation-run-id "$PREPARATION_RUN_ID" \
  --rendezvous-timeout-seconds 7200

Run on each node with its rank and a shared asset/output root. Models and prompt data are prepared from HF automatically. I2V generates FLUX evaluation images on first use and caches them for later evaluations and resume; no pre-existing image manifest is needed. This is a new evaluation protocol, separate from historical fixed-image results. See the I2V guide for baseline recipes and shared reference caches. Add --prepare-only to prepare without training, or --validate for network-free checks. See distributed execution for multi-node launches and checkpoint/resume for recovery and evaluation.

Extend an example or paper recipe. Save this as configs/my-pickscore.yaml:

extends: examples/sd3_pickscore.yaml
name: my-pickscore
training:
  max_updates: 1000
runtime:
  output_dir: outputs/my-pickscore
  checkpoint_keep: 2

uv run --locked --extra diffusers visual-rl --config configs/my-pickscore.yaml

Select components through model.name, algorithm.name and reward.name; compose objectives with reward.weights. Sampling, microbatches, evaluation and checkpoint cadence live in the same file. Required asset paths use ${ENV_VAR}.

Metrics are saved to metrics.jsonl; evaluation images, videos and scores stay on disk. Set runtime.log: wandb for scalar logging.

Reward quality matters. As with ReFL and other reward-based fine-tuning methods, DiffusionART depends on how well the reward model captures the desired behavior. Our experiments use open-source reward models. For practical applications, we recommend higher-quality reward models better aligned with the intended use case: DiffusionART improves rewards quickly, which can also expose weaknesses in the reward model and lead to reward hacking sooner.

In our experiments, when different algorithms reach a comparable score on an optimized dimension, they tend to exhibit similar levels of degradation in other dimensions. This suggests a shared limitation of the reward model rather than a failure unique to DiffusionART.

Documentation  /  Development  /  Model interfaces  /  Reward interfaces

Thanks to the authors of Flow-GRPO, DiffusionNFT, DanceGRPO and UniRL for open-sourcing their work.

The code is released under the Apache License 2.0. Third-party code, model weights and datasets retain their respective licenses; see Third-Party Notices.

If you find DiffusionART inspiring or useful for your projects, please consider starring this repository and citing our paper:

@misc{wan2026diffusionartefficientstablescalable,
  title={DiffusionART: Efficient, Stable and Scalable Reinforcement of Diffusion Models},
  author={Zhenglin Wan and Jiayu Chen and Zhengfeiyang Zhang and Pengfei Zhou and Wangbo Zhao and Fangyuan Kong and Pengfei Wan and Kun Gai and Yang You},
  year={2026},
  eprint={2610.12572},
  archivePrefix={arXiv},
  primaryClass={cs.LG},
  url={https://arxiv.org/abs/2610.12572},
}