# Learning Multimodal One-step Flow Policy via Value-weighted Optimal Transport

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Diffusion and flow policies for offline RL, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 32 of its 37 references that are in the map.
- First posted on arXiv 14 Sep 2026.

## Placed in

- Offline RL and imitation learning from data / [Diffusion and flow policies for offline RL](https://mlsanity.com/offlinerl/diffusion-and-flow-policies-for-offline-rl/)

## Builds on (most cited first, 30 of 32)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2403.03206/) (2024)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Gaussian Error Linear Units (GELUs)](https://mlsanity.com/paper/1606.08415/) (2016)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [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)
- [D4RL: Datasets for Deep Data-Driven Reinforcement Learning](https://mlsanity.com/paper/2004.07219/) (2020)
- [Offline Reinforcement Learning with Implicit Q-Learning](https://mlsanity.com/paper/2110.06169/) (2021)
- [Building Normalizing Flows with Stochastic Interpolants](https://mlsanity.com/paper/2209.15571/) (2022)
- [A Minimalist Approach to Offline Reinforcement Learning](https://mlsanity.com/paper/2106.06860/) (2021)
- [IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures](https://mlsanity.com/paper/1802.01561/) (2018)
- [Advantage-Weighted Regression: Simple and Scalable Off-Policy Reinforcement Learning](https://mlsanity.com/paper/1910.00177/) (2019)
- [Reinforcement Learning with Deep Energy-Based Policies](https://mlsanity.com/paper/1702.08165/) (2017)
- [Behavior Regularized Offline Reinforcement Learning](https://mlsanity.com/paper/1911.11361/) (2019)
- [Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning](https://mlsanity.com/paper/2208.06193/) (2022)
- [One Step Diffusion via Shortcut Models](https://mlsanity.com/paper/2410.12557/) (2024)
- [Critic Regularized Regression](https://mlsanity.com/paper/2006.15134/) (2020)
- [OGBench: Benchmarking Offline Goal-Conditioned RL](https://mlsanity.com/paper/2410.20092/) (2024)
- [AlgaeDICE: Policy Gradient from Arbitrary Experience](https://mlsanity.com/paper/1912.02074/) (2019)
- [Flow Q-Learning](https://mlsanity.com/paper/2502.02538/) (2025)
- [Revisiting the Minimalist Approach to Offline Reinforcement Learning](https://mlsanity.com/paper/2305.09836/) (2023)
- [OptiDICE: Offline Policy Optimization via Stationary Distribution Correction Estimation](https://mlsanity.com/paper/2106.10783/) (2021)
- [Is Value Learning Really the Main Bottleneck in Offline RL?](https://mlsanity.com/paper/2406.09329/) (2024)
- [Optimal Transport for Offline Imitation Learning](https://mlsanity.com/paper/2303.13971/) (2023)
- [Policy Regularization with Dataset Constraint for Offline Reinforcement Learning](https://mlsanity.com/paper/2306.06569/) (2023)
- [Rethinking Optimal Transport in Offline Reinforcement Learning](https://mlsanity.com/paper/2410.14069/) (2024)

## Abstract (arXiv)

Offline reinforcement learning aims to learn a policy solely from fixed datasets, which often contain multimodal action distributions. Flow policies can naturally represent such multimodal behaviors, but learning an efficient one-step flow policy remains challenging: standard value guidance often leads to mode collapse or exploits overestimation bias in out-of-distribution regions. To address this, we introduce One-step Flow policy via Optimal Transport (OptiFlow), a framework for one-step flow policy learning as a structured sample-allocation problem. OptiFlow jointly trains a value-aware reference flow policy and an efficient one-step policy, coupling their action samples through state-wise entropic optimal transport. For each state, critic-estimated values define the priority of distillation target actions, while the action-distance cost ensures geometrically compatible pairings. By avoiding direct critic maximization, our transport-guided approach enables in-distribution exploitation by anchoring the one-step policy to high-value, dataset-supported modes without the risk of out-of-distribution divergence. Experimental results demonstrate that OptiFlow effectively captures optimal multimodal behaviors and achieves strong performance across diverse offline RL benchmarks. Our code is available at https://github.com/Yonsei-DILLab/OptiFlow.