RFC: image/video diffusion (flow-matching) RL — flow-GRPO training
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
**Is your feature request related to a problem? Please describe.**
NeMo RL currently covers LLMs and VLMs (understanding), but has no support for RL post-training of image or video generation models. Flow-matching base models such as Qwen-Image benefit substantially from online RL against task rewards (OCR accuracy, human preference) — see [Flow-GRPO](https://arxiv.org/abs/2505.05470). Today users must leave the NeMo RL ecosystem to do this.
**Describe the solution you'd like**
Add a diffusion-GRPO training path: convert the flow-matching ODE into an equivalent SDE so per-step transition log-probs are computable in closed form, then apply the standard GRPO recipe (group-relative advantages, clipped policy gradients, optional Gaussian KL vs. the LoRA-disabled reference).
Naming note: to avoid confusion with diffusion language model work (e.g. #3167), everything here is visual generation (continuous flow-matching); the two efforts are independent and do not overlap in code. The initial work is scoped to text-to-image; video flow-matching models are a natural follow-up — they reuse the same SDE/GRPO machinery through the pipeline-adapter protocol.
I have a working implementation (Qwen-Image via Diffusers, LoRA + Ray DP workers, pluggable image reward environment, OCR task converging) and will upstream it as a series of focused, self-contained PRs. Each PR introduces its own scope, notes how it relates to the previously merged ones, and links back to this issue.
**Describe alternatives you've considered**
A single monolithic PR (~6.6k lines) — rejected in favor of reviewable, independently-testable increments.
Building on an inference engine for rollouts — not applicable: the training framework itself samples the SDE trajectory (log-probs must match the training graph), so no vLLM/refit machinery is involved.
**Additional context**
Reference implementation parity: hyperparameters and SDE semantics mirror the Flow-GRPO authors' repo and verl-omni's flowgrpo trainer; the exemplar config reproduces their Qwen-Image OCR LoRA recipe.
Scope/limitations of the initial series: Qwen-Image (Diffusers) only, LoRA + single-GPU-per-worker data parallelism (multi-node via Ray), no FSDP/Megatron sharding. The pipeline adapter is a Protocol, so other flow-matching pipelines (image or video) can be added incrementally.
Contributor guide
Assessment
This issue has not been assessed yet.