Authors: Maojiang Su, Po-Chung Hsieh, Weimin Wu, Mingcheng Lu, Jiunhau Chen, Jerry Yao-Chieh Hu, Han Liu
Organizations: Center for Foundation Models and Generative AI, Northwestern University, Evanston, IL 60208, USA · Department of Electrical Engineering, National Taiwan University, Taipei 10617, Taiwan · Department of Physics, National Taiwan University, Taipei 10617, Taiwan · Department of Statistics and Data Science, Northwestern University, Evanston, IL 60208, USA
We introduce Discrete flow Matching policy Optimization (DoMinO), a unified framework for Reinforcement Learning (RL) fine-tuning Discrete Flow Matching (DFM) models under a broad class of policy gradient methods. Our key idea is to view the DFM sampling procedure as a multi-step Markov Decision Process. This perspective provides a simple and transparent reformulation of fine-tuning reward maximization as a robust RL objective. Consequently, it not only preserves the original DFM samplers but also avoids biased auxiliary estimators and likelihood surrogates used by many prior RL fine-tuning methods. To prevent policy collapse, we also introduce new total-variation regularizers to keep the fine-tuned distribution close to the pretrained one. Theoretically, we establish an upper bound on the discretization error of DoMinO and tractable upper bounds for the regularizers. Experimentally, we evaluate DoMinO on regulatory DNA sequence design. DoMinO achieves stronger predicted enhancer activity and better sequence naturalness than the previous best reward-driven baselines. The regularization further improves alignment with the natural sequence distribution while preserving strong functional performance. These results establish DoMinO as an useful framework for controllable discrete sequence generation.
Table 1: Performance of DNA design methods on the HepG2 enhancer design task. We compare pre-trained generative models, prior reward-driven baselines (DRAKES and SEPO), and our policy-gradient fine-tuning methods for discrete flow matching. We do not apply regularization loss during fine-tuning in this setting. Pred-Activity and ATAC-Acc evaluate functional performance, while 3-mer Corr-All measures sequence naturalness relative to the HepG2 data distribution. Our methods outperform the state-of-the-art baseline SEPO on Pred-Activity and 3-mer Corr-All and achieve comparable performance on ATAC-Acc. Bold and underlined entries highlight the best and second-best reward-driven methods.
Method
Pred-Activity ↑
ATAC-Acc ↑ (%)
3-mer Corr - All ↑
DRAKES
6.37 (0.04)
96.1 (0.5)
-0.379 (0.009)
DRAKES with KL
5.61 (0.07)
92.5 (0.6)
-0.302 (0.011)
SEPO
7.55 (0.01)
99.5 (0.2)
-0.537 (0.002)
SEPO with GF
7.64 (0.01)
99.9 (0.09)
-0.496 (0.001)
DoMinO-REINFORCE
8.32 (0.01)
99.2 (0.2)
-0.285 (0.001)
DoMinO-REINFORCE with CE
8.22 (0.01)
94.1 (0.8)
-0.347 (0.004)
Table 2: Effect of regularization on DoMinO for the HepG2 enhancer design task. We compare our policy-gradient fine-tuning methods with and without regularization loss. Pred-Activity and ATAC-Acc measure functional performance, while 3-mer Corr-All measures sequence naturalness relative to the HepG2 data distribution. Regularization improves sequence naturalness and helps maintain strong functional performance. This yields a better balance between reward optimization and naturalness. Bold and underlined entries highlight the best and second-best regularized reward-driven methods.
Discrete flow models (DFMs) are a class of flexible generative models for generating discrete data, and diffusion large language models (dLLMs) can be viewed as a special case with a specific choice of a mixture path and a masked source distribution. While several recent works have explored reinforcement learning for dLLMs, its application to more general discrete flow models remains underexplored. In this work, we present discrete Flow-GRPO (dFlowGRPO), a unified reinforcement learning framework for discrete flow models that supports a broad family of probability paths and non-masked source distributions. We derive the full trajectory probability for DFMs and formulate the denoising process as a Markov decision process, enabling dFlowGRPO to incorporate information from both the associated conditional transition rates and the posterior model during reinforcement learning. We apply dFlowGRPO to FUDOKI, a recent multimodal discrete flow model, and evaluate it on both image generation and multimodal understanding tasks. Empirical results show that dFlowGRPO outperforms existing GRPO-type methods for dLLMs on text-to-image generation tasks and achieves performance competitive with continuous flow-based models trained using Flow-GRPO, while also demonstrating strong capabilities on understanding tasks.
Zhengyan Wan, Yidong Ouyang, Panwen Hu +1
Mohamed bin Zayed University of Artificial Intelligence · East China Normal University · University of California, Los Angeles +1
Flow Matching (FM) has recently emerged as a principled and efficient generative modeling framework for reinforcement learning (RL), enabling expressive, multimodal policy parameterizations via deterministic probability transport. Compared to diffusion-based policies that rely on stochastic denoising chains, FM uses sampling based on ordinary differential equations (ODEs), with learned velocity fields, which can substantially reduce inference latency and simplify the incorporation of RL objectives. As research in flow-based RL rapidly accelerates across offline continuous control, online fine-tuning, and foundation model alignment, the literature has become highly fragmented. In this survey, we provide a comprehensive taxonomy of flow-matching approaches in reinforcement learning. We organize the literature along two axes: the target distribution being modeled (e.g., action policies, value critics, transition dynamics) and the mechanism of RL signal integration (e.g., energy-weighted regression, flow-based policy gradients, and group relative policy optimization). Furthermore, we survey emerging frontiers such as discrete and non-Euclidean action spaces, provide a systematic comparative analysis against Gaussian and diffusion baselines, and outline critical open problems. Ultimately, this survey serves as a foundational roadmap for the next generation of generative reinforcement learning and alignment.
Nabuat Zaman Nahim, Fairoz Nower Khan, Peizhong Ju
Department of Computer Science, University of Kentucky, Lexington, Kentucky 40506, USA
Many reinforcement learning (RL) tasks have discrete action spaces, but most generative policy methods based on diffusion and flow matching are designed for continuous control. Meanwhile, generative policies usually rely heavily on offline datasets and offline-to online RL is itself challenging, as the policy must improve from new interaction without losing useful behavior learned from static data. To address those challenges, we introduce DRIFT, an online fine-tuning method that updates an offline pretrained continuous time Markov chain (CTMC) policy with an advantage-weighted discrete flow matching loss. To preserve useful pretrained knowledge, we add a path-space penalty that regularizes the full CTMC trajectory distribution, rather than only the final action distribution. For large discrete action spaces, we introduce a candidate-set approximation that updates the actor over a small subset of actions sampled from reference-policy rollouts and uniform exploration. Our theoretical analysis shows that focusing on a small set of high probability actions can reduce computation while keeping the approximation error small, making our method practical for large action spaces. We evaluate on macro-action MinAtar games where offline-pretrained DRIFT outperforms its online-only counterpart across all five games and achieves the highest mean return on three and on oracle free Jericho, DRIFT achieves a mean normalized score of 7.3% with a lightweight GRU encoder, outperforming the evaluated pretrained language-model baselines. Ablations demonstrate the importance of appropriately weighted path-space regularization and reference-guided candidate construction.
Fairoz Nower Khan, Nabuat Zaman Nahim, Peizhong Ju
Department of Computer Science, University of Kentucky