提出量化对齐的流匹配方法,让分布强化学习更准确地建模回报分布。
Quantile-Coupled Flow Matching for Distributional Reinforcement Learning
- 通过排序源与目标样本实现分位数对齐的流路径。
- 在多个离线强化学习基准上提升分布精度并保持竞争力性能。
- 首个具有明确沃尔什距离对齐保证的流匹配分布批判器,适合理论严谨的研究者。
与标准期望回报强化学习不同,分布强化学习(DRL)建模完整的回报分布,更适合不确定性感知和风险敏感决策。条件流匹配(CFM)批判器因其能建模连续、多模态回报分布而受到关注。然而,现有方法存在显著度量不匹配:DRL理论依赖分布贝尔曼算子在p-沃尔什距离下为压缩映射,但现有CFM批判器使用任意源-目标配对训练,其流匹配损失并非沃尔什对齐的代理。本文提出FlowIQN,通过在每个小批量内对齐源与贝尔曼目标样本的分位数,近似单调最优传输耦合,以量化对齐的流路径替代任意配对。我们证明该方法的损失函数可提供沃尔什对齐的近似投影,符合DRL理论基础。据我们所知,FlowIQN是首个具有显式沃尔什对齐投影保证的流匹配分布批判器。此外,我们还引入快捷模型以提升推理效率。实验表明,FlowIQN在回报分布的沃尔什精度上优于其他CFM批判器,并在多种策略提取方法下的离线强化学习基准中表现优异,提供一个理论坚实且易于集成到DRL流程中的批判器。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。