Offline reinforcement learning policy improves multimodal action selection
Learning Multimodal One-step Flow Policy via Value-weighted Optimal Transport
Machine LearningArtificial Intelligence
Summary
Offline reinforcement learning tries to teach computers to make decisions by learning only from already collected data, which often has multiple good choices for each situation. The authors found that existing methods either miss these multiple options or get confused by unreliable predictions. They introduced a new approach called OptiFlow, which matches high-value actions with learned policies in a careful way to avoid these problems. This method helps computers learn diverse and effective decision-making strategies from fixed data sets.
What this means in practice
- •For robotics developers: Develop robot control policies from fixed datasets to handle multiple valid action choices reliably.
- •For autonomous vehicle engineers: Create safer driving policies by modeling complex multimodal decisions from prior driving data without online learning.
Authors
Jaehun Shon, Jinha Choi, Jongwook Jeon, Jongmin Lee
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.