Learning Multimodal One-step Flow Policy via Value-weighted Optimal Transport
cs.LG, cs.AI
Submitted: 2026-09-14
Updated: 2026-09-14
Comments: Preprint, 38 pages
Code: https://github.com/Yonsei-DILLab/OptiFlow
License: http://creativecommons.org/licenses/by/4.0/
The gist: Offline reinforcement learning aims to learn a policy solely from fixed datasets, which often contain multimodal action distributions.
Terminology
Abstract
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.
Sources
- Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems
- Behavior Regularized Offline Reinforcement Learning
- OptiDICE: Offline Policy Optimization via Stationary Distribution Correction Estimation
- AlgaeDICE: Policy Gradient from Arbitrary Experience
- Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow
- Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning
- Progressive Distillation for Fast Sampling of Diffusion Models
- One Step Diffusion via Shortcut Models
- Advantage-Weighted Regression: Simple and Scalable Off-Policy Reinforcement Learning
- Optimal Transport for Offline Imitation Learning
- Imitation Learning from Observation through Optimal Transport
- Building Normalizing Flows with Stochastic Interpolants
- Playing Atari with Deep Reinforcement Learning
- D4RL: Datasets for Deep Data-Driven Reinforcement Learning
- Gaussian Error Linear Units (GELUs)
- Layer Normalization
- Adam: A Method for Stochastic Optimization
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: Local-Consistency Optimisation for Bayesian Neural Networks
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks