Learning Multimodal One-step Flow Policy via Value-weighted Optimal Transport
OptiFlow learns one-step multimodal flow policies for offline RL via state-wise entropic optimal transport, avoiding critic overestimation and mode collapse.
The paper introduces OptiFlow, a framework that frames one-step flow policy learning as a structured sample-allocation problem in offline reinforcement learning. It jointly trains a value-aware reference flow policy and a one-step policy, coupling action samples through state-wise entropic optimal transport where critic values set distillation priority and action-distance cost ensures geometrically compatible pairings. By avoiding direct critic maximization, it anchors the policy to high-value dataset-supported modes without out-of-distribution divergence. Code is released on GitHub and the method performs strongly across diverse offline RL benchmarks.
- One-step flow policy learned as structured sample allocation via entropic optimal transport
- Avoids direct critic maximization, preventing mode collapse and OOD divergence
- Critic values prioritize distillation targets; action-distance cost ensures compatible pairings
- Code publicly available; strong results across offline RL benchmarks
Full article167 words · extracted from arxiv.org · click to collapse
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.
Text extracted automatically; images, tables and formatting may be missing. Original: https://arxiv.org/abs/2609.15883