Quantile-Coupled Flow Matching for Distributional Reinforcement Learning

• URL: http://arxiv.org/abs/2605.08515v1
• Date: Fri, 08 May 2026 21:49:37 GMT
• Title: Quantile-Coupled Flow Matching for Distributional Reinforcement Learning
• Abstract summary: Conditional Flow Matching (CFM) critics have attracted attention for modelling continuous, multi-modal return distributions.<n> FlowIQN sorts source and Bellman target samples within each mini-batch to approximate the monotone optimal transport coupling.<n>We prove that FlowIQN is the first flow-matching distributional critic with an explicit Wasserstein-aligned projection guarantee.
• Score: 14.19561152578512
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Unlike standard expected-return Reinforcement Learning (RL), Distributional RL (DRL) models the full return distribution, making it better-suited for uncertainty-aware and risk-sensitive decision-making. Conditional Flow Matching (CFM) critics have recently attracted attention for modelling continuous, multi-modal return distributions. Despite this interest, there remains a substantial metric mismatch: DRL theory relies on the distributional Bellman operator being contractive in the $p$-Wasserstein distance, yet existing CFM critics are trained with arbitrary source-target couplings, so their flow-matching losses are not Wasserstein-aligned surrogates for matching Bellman target return distributions. In this work, we address this mismatch by proposing FlowIQN, a CFM critic that sorts source and Bellman target samples within each mini-batch to approximate the monotone optimal transport coupling, replacing arbitrary pairings with quantile-aligned flow paths. We prove that the loss of our quantile-coupled CFM critic yields a Wasserstein-aligned approximate projection compatible with the foundations of DRL. To our knowledge, FlowIQN is the first flow-matching distributional critic with an explicit Wasserstein-aligned projection guarantee. We further extend FlowIQN with shortcut models for efficient inference. Empirical results show that FlowIQN improves Wasserstein return-distribution accuracy over other CFM critics. It also yields competitive performance on offline RL benchmarks across multiple policy extraction methods, providing a theoretically grounded CFM critic that is readily compatible with DRL pipelines. Code: https://github.com/ori-goals/flowIQN.

Related papers

• DISA: Offline Importance Sampling for Distribution-Matching LLM-RL [56.9445657766829]
  We introduce DISA, which moves this calibration problem outside the RL loop.<n> DISA draws proposal trajectories offline, estimates the partition function via importance sampling, and freezes the resulting partition-function estimate.<n>On two open-weight backbones across six math and three code benchmarks, DISA matches or exceeds the online-coupled distribution-matching baseline Flow.
  arXiv  Detail & Related papers  (2026-05-17T07:14:44Z)
• Multi-Fidelity Flow Matching: Cascaded Refinement of PDE Solutions [6.144605787421899]
  Multi-Fidelity Flow Matching is a cascade refinement framework for PDE solutions.<n>We exploit this in Multi-Fidelity Flow Matching (MFFM), a cascade refinement framework for PDE solutions.<n>MFFM is a learned analog of multigrid refinement that reaches the finest grid in $L$ deterministic network evaluations per query.
  arXiv  Detail & Related papers  (2026-05-15T16:02:18Z)
• Path-Coupled Bellman Flows for Distributional Reinforcement Learning [6.000836623151053]
  Path-Coupled Bellman Flows (PCBF) learns return distributions with flow matching.<n>PCBF couples current and successor return flows through shared base noise.<n>Experiments on analytically tractable MRPs, OGBench, and D4RL show improved distributional fidelity and training stability.
  arXiv  Detail & Related papers  (2026-05-07T19:05:01Z)
• PARM: Pipeline-Adapted Reward Model [60.769414637325326]
  Reward models (RMs) are central to aligning large language models (LLMs) with human preferences, powering advanced decoding strategies.<n>While most prior work focuses on single-step generation, real-world applications increasingly adopt multi-stage pipelines.<n>We investigate this through code generation for optimization, constructing a pipeline that integrates reward models into both formulation and solution stages.
  arXiv  Detail & Related papers  (2026-04-20T14:29:08Z)
• Distributional Reinforcement Learning with Diffusion Bridge Critics [57.70134665595571]
  We propose a novel distributional reinforcement learning method with Diffusion Bridge Critics (DBC)<n>DBC directly models the inverse cumulative distribution function (CDF) of the Q value.<n>We derive an analytic integral formula to address discretization errors in DBC.
  arXiv  Detail & Related papers  (2026-02-05T15:40:14Z)
• Boosting Maximum Entropy Reinforcement Learning via One-Step Flow Matching [8.665369041430969]
  Flow Matching (FM) enables one-step generation, but integrating it into Entropy Reinforcement Learning (MaxEnt RL) is challenging.<n>We propose textbfFlow-based textbfLog-likelihood-textbfAware textbfMaximum textbfEntropy RL (textbfFLAME), a principled framework that addresses these challenges.
  arXiv  Detail & Related papers  (2026-02-02T03:54:11Z)
• Reverse Flow Matching: A Unified Framework for Online Reinforcement Learning with Diffusion and Flow Policies [4.249024052507976]
  We propose a unified framework, reverse flow matching (RFM), which rigorously addresses the problem of training diffusion and flow models without direct target samples.<n>By adopting a reverse inferential perspective, we formulate the training target as a posterior mean estimation problem given an intermediate noisy sample.<n>We show that existing noise-expectation and gradient-expectation methods are two specific instances within this broader class.
  arXiv  Detail & Related papers  (2026-01-13T01:58:24Z)
• Whatever Remains Must Be True: Filtering Drives Reasoning in LLMs, Shaping Diversity [13.211627219720796]
  Reinforcement Learning (RL) has become the de facto standard for tuning LLMs to solve tasks involving reasoning.<n>We argue that RL implicitly optimize the "mode-seeking" or "zero-forcing" Reverse KL to a target distribution causing the model to concentrate mass on certain high-probability regions of the target while others.<n>In this work, we instead begin from an explicit target distribution, obtained by filtering out incorrect answers while neglecting the relative probabilities of correct ones.
  arXiv  Detail & Related papers  (2025-12-05T18:56:40Z)
• Reweighted Flow Matching via Unbalanced OT for Label-free Long-tailed Generation [9.488297561764211]
  We propose Unbalanced Optimal Transport Reweighted Flow Matching (UOT-RFM) for long-tailed distributions.<n>Our method mitigates majority bias through a principled inverse reweighting strategy.<n>Our model outperforms existing flow matching baselines on long-tailed benchmarks.
  arXiv  Detail & Related papers  (2025-09-30T03:19:52Z)
• Unleashing Flow Policies with Distributional Critics [15.149475517073258]
  We introduce the Distributional Flow Critic (DFC), a novel critic architecture that learns the complete state-action return distribution.<n>DFC provides the expressive flow-based policy with a rich, distributional Bellman target, which offers a more stable and informative learning signal.
  arXiv  Detail & Related papers  (2025-09-27T03:51:06Z)
• FlowRL: Matching Reward Distributions for LLM Reasoning [69.88820066093798]
  We propose FlowRL: matching the full reward distribution via flow balancing instead of maximizing rewards in large language model (LLM) reinforcement learning (RL)<n>We transform scalar rewards into a normalized target distribution using a learnable partition function, and then minimize the reverse KL divergence between the policy and the target distribution.
  arXiv  Detail & Related papers  (2025-09-18T17:56:36Z)
• On the Wasserstein Convergence and Straightness of Rectified Flow [54.580605276017096]
  Rectified Flow (RF) is a generative model that aims to learn straight flow trajectories from noise to data.<n>We provide a theoretical analysis of the Wasserstein distance between the sampling distribution of RF and the target distribution.<n>We present general conditions guaranteeing uniqueness and straightness of 1-RF, which is in line with previous empirical findings.
  arXiv  Detail & Related papers  (2024-10-19T02:36:11Z)
• Policy Evaluation in Distributional LQR [70.63903506291383]
  We provide a closed-form expression of the distribution of the random return. We show that this distribution can be approximated by a finite number of random variables. Using the approximate return distribution, we propose a zeroth-order policy gradient algorithm for risk-averse LQR.
  arXiv  Detail & Related papers  (2023-03-23T20:27:40Z)
• Distributional Reinforcement Learning for Multi-Dimensional Reward Functions [91.88969237680669]
  We introduce Multi-Dimensional Distributional DQN (MD3QN) to model the joint return distribution from multiple reward sources. As a by-product of joint distribution modeling, MD3QN can capture the randomness in returns for each source of reward. In experiments, our method accurately models the joint return distribution in environments with richly correlated reward functions.
  arXiv  Detail & Related papers  (2021-10-26T11:24:23Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.