UniGRPO: Unified Policy Optimization for Reasoning-Driven Visual Generation

TL;DR

This paper introduces UniGRPO, a unified reinforcement learning framework that enhances reasoning-driven visual generation by jointly optimizing text and image policies, enabling scalable, multi-round interleaved generation.

Contribution

It proposes a novel unified RL approach for interleaved text and image generation, with specific modifications to improve scalability and robustness in multi-turn scenarios.

Findings

Improved image quality through reasoning-based generation.

Scalable approach for multi-round interleaved generation.

Effective mitigation of reward hacking with new regularization.

Abstract

Unified models capable of interleaved generation have emerged as a promising paradigm, with the community increasingly converging on autoregressive modeling for text and flow matching for image generation. To advance this direction, we propose a unified reinforcement learning framework tailored for interleaved generation. We validate our approach on its fundamental unit: a single round of reasoning-driven image generation, where the model first expands the user prompt through reasoning, followed by image synthesis. Formulating this multimodal generation process as a Markov Decision Process with sparse terminal rewards, we introduce UniGRPO to jointly optimize text and image generation policies using GRPO. Adopting a minimalist methodology to avoid over-design, we leverage established training recipes for both modalities by seamlessly integrating standard GRPO for reasoning and FlowGRPO for visual synthesis. To ensure scalability to multi-round interleaved generation, we introduce two critical modifications to the original FlowGRPO: (1) eliminating classifier-free guidance to maintain linear, unbranched rollouts, which is essential for scaling to complex scenarios involving multi-turn interactions and multi-condition generation (e.g., editing); and (2) replacing the standard latent KL penalty with an MSE penalty directly on the velocity fields, providing a more robust and direct regularization signal to mitigate reward hacking effectively. Our experiments demonstrate that this unified training recipe significantly enhances image generation quality through reasoning, providing a robust and scalable baseline for the future post-training of fully interleaved models.