# Quantile-Coupled Flow Matching for Distributional Reinforcement Learning

arXiv: https://arxiv.org/abs/2605.08515 · HTML: https://mlsanity.com/paper/2605.08515/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion and flow policies for robotics, one of 14 directions in Diffusion models.
- Builds on 27 of its 32 references that are in the map.
- First posted on arXiv 8 May 2026.

## Placed in

- Diffusion models / [Diffusion and flow policies for robotics](https://mlsanity.com/diffusion/diffusion-and-flow-policies-for-robotics/)
- Deep reinforcement learning methods / [RL with diffusion and flow models](https://mlsanity.com/deeprl/rl-with-diffusion-and-flow-models/)

## Builds on (most cited first, 27 of 27)

- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [Conservative Q-Learning for Offline Reinforcement Learning](https://mlsanity.com/paper/2006.04779/) (2020)
- [Off-Policy Deep Reinforcement Learning without Exploration](https://mlsanity.com/paper/1812.02900/) (2018)
- [Stochastic Interpolants: A Unifying Framework for Flows and Diffusions](https://mlsanity.com/paper/2303.08797/) (2023)
- [Improving and generalizing flow-based generative models with minibatch optimal transport](https://mlsanity.com/paper/2302.00482/) (2023)
- [A Distributional Perspective on Reinforcement Learning](https://mlsanity.com/paper/1707.06887/) (2017)
- [Distributional Reinforcement Learning with Quantile Regression](https://mlsanity.com/paper/1710.10044/) (2017)
- [One Step Diffusion via Shortcut Models](https://mlsanity.com/paper/2410.12557/) (2024)
- [Implicit Quantile Networks for Distributional Reinforcement Learning](https://mlsanity.com/paper/1806.06923/) (2018)
- [Flow Matching Guide and Code](https://mlsanity.com/paper/2412.06264/) (2024)
- [OGBench: Benchmarking Offline Goal-Conditioned RL](https://mlsanity.com/paper/2410.20092/) (2024)
- [Flow Q-Learning](https://mlsanity.com/paper/2502.02538/) (2025)
- [Conservative Offline Distributional Reinforcement Learning](https://mlsanity.com/paper/2107.06106/) (2021)
- [Distributional Reinforcement Learning for Efficient Exploration](https://mlsanity.com/paper/1905.06125/) (2019)
- [1000 Layer Networks for Self-Supervised RL: Scaling Depth Can Enable New Goal-Reaching Capabilities](https://mlsanity.com/paper/2503.14858/) (2025)
- [The Benefits of Being Distributional: Small-Loss Bounds for Reinforcement Learning](https://mlsanity.com/paper/2305.15703/) (2023)
- [More Benefits of Being Distributional: Second-Order Bounds for Reinforcement Learning](https://mlsanity.com/paper/2402.07198/) (2024)
- [floq: Training Critics via Flow-Matching for Scaling Compute in Value-Based RL](https://mlsanity.com/paper/2509.06863/) (2025)
- [Value Flows](https://mlsanity.com/paper/2510.07650/) (2025)
- [Theoretical guarantees in KL for Diffusion Flow Matching](https://mlsanity.com/paper/2409.08311/) (2024)
- [Unleashing Flow Policies with Distributional Critics](https://mlsanity.com/paper/2509.23087/) (2025)
- [Expressive Value Learning for Scalable Offline Reinforcement Learning](https://mlsanity.com/paper/2510.08218/) (2025)
- [FlowCritic: Bridging Value Estimation with Flow Matching in Reinforcement Learning](https://mlsanity.com/paper/2510.22686/) (2025)
- [Formalizing the Sampling Design Space of Diffusion-Based Generative Models via Adaptive Solvers and Wasserstein-Bounded Timesteps](https://mlsanity.com/paper/2602.12624/) (2026)
- [What Does Flow Matching Bring To TD Learning?](https://mlsanity.com/paper/2603.04333/) (2026)

## Built on by (most cited first, 1 of 1)

- [Retimed Bellman Flows: Escaping the Impossible Triangle of Velocity Bootstrapping](https://mlsanity.com/paper/2609.32828/) (2026)

## Abstract (arXiv)

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.