Flow Models for Unbounded and Geometry-Aware Distributional Reinforcement Learning

📄 arXiv: 2505.04310v1 📥 PDF

作者: Simo Alami C., Rim Kaddah, Jesse Read, Marie-Paule Cani

分类: cs.AI, cs.LG, math.OC

发布日期: 2025-05-07

💡 一句话要点

提出基于Normalizing Flow的DistRL架构，解决回报分布建模中无界支持和几何感知问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 分布强化学习 Normalizing Flow 回报分布建模 无界支持 Cramèr距离

📋 核心要点

  1. 传统DistRL方法在回报分布建模上存在局限，如C51的固定边界和分位数方法的建模能力不足。
  2. 利用Normalizing Flow建模回报分布，实现灵活无界支持，并提升对多模态、偏度和尾部行为的建模能力。
  3. 提出一种新的Cramèr距离替代，直接从PDF计算，避免CDF计算，并在ATARI-5上验证了方法的有效性。

📝 摘要（中文）

本文提出了一种新的分布强化学习(DistRL)架构，该架构使用Normalizing Flow来建模回报分布。与依赖固定或有界表示的C51等分类方法相比，这种方法为回报分布提供了灵活的、无界的支持。与基于分位数的方法相比，它还提供了更丰富的建模能力，以捕获多模态、偏度和尾部行为。我们的方法比分类方法更具参数效率。用于训练现有模型的标准指标，如KL散度或Wasserstein距离，要么对尺度不敏感，要么具有有偏的样本梯度，尤其是在回报支持不重叠时。为了解决这个问题，我们提出了一种新的Cramèr距离的替代，它是几何感知的，可以直接从回报分布的PDF计算，避免了代价高昂的CDF计算。我们在ATARI-5子基准上测试了我们的模型，结果表明我们的方法优于基于PDF的模型，同时与基于分位数的方法保持竞争力。

🔬 方法详解

问题定义：现有的分布强化学习方法在建模回报分布时存在局限性。例如，C51等方法使用离散的分类表示，需要预先定义回报的上下界，这限制了其对无界回报分布的建模能力。基于分位数的方法虽然可以处理连续的回报分布，但在捕捉多模态、偏度和尾部行为方面能力有限。此外，训练这些模型常用的KL散度或Wasserstein距离等指标，在回报支持不重叠时，容易出现尺度不敏感或梯度偏差的问题。

核心思路：本文的核心思路是利用Normalizing Flow来建模回报分布。Normalizing Flow是一种生成模型，可以通过一系列可逆变换将一个简单的分布（如高斯分布）转换为复杂的目标分布。这种方法的优势在于，它可以提供灵活的、无界的回报分布支持，并且具有强大的建模能力，能够捕捉多模态、偏度和尾部行为。此外，Normalizing Flow可以直接生成回报分布的概率密度函数(PDF)，这为后续的优化提供了便利。

技术框架：该方法采用标准的Actor-Critic框架。Actor网络负责生成策略，Critic网络负责评估状态-动作对的价值分布。Critic网络使用Normalizing Flow来建模回报分布。具体来说，Critic网络接收状态作为输入，输出Normalizing Flow的参数。然后，通过Normalizing Flow将一个简单的基分布（如高斯分布）转换为回报分布。Actor网络的训练目标是最大化期望回报，Critic网络的训练目标是最小化预测回报分布与真实回报分布之间的距离。

关键创新：该论文的关键创新在于以下两点：一是使用Normalizing Flow来建模回报分布，从而实现了灵活的、无界的回报分布支持；二是提出了一种新的Cramèr距离的替代，该替代可以直接从回报分布的PDF计算，避免了代价高昂的CDF计算。传统的Cramèr距离需要计算累积分布函数(CDF)，而CDF的计算通常比较耗时。本文提出的替代方法利用PDF的几何信息，可以直接计算Cramèr距离，从而提高了训练效率。

关键设计：在Normalizing Flow的选择上，论文采用了RealNVP架构。RealNVP是一种常用的Normalizing Flow架构，具有计算效率高、易于训练等优点。在损失函数的设计上，论文使用了提出的Cramèr距离替代。此外，为了提高训练的稳定性，论文还采用了梯度裁剪等技巧。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

在ATARI-5子基准测试中，该方法优于基于PDF的模型，并与基于分位数的方法保持竞争力。具体来说，该方法在某些游戏中取得了显著的性能提升，表明其在建模复杂回报分布方面的优势。实验结果验证了Normalizing Flow在DistRL中的有效性，以及提出的Cramèr距离替代的优越性。

🎯 应用场景

该研究成果可应用于各种需要处理复杂回报分布的强化学习任务中，例如金融交易、机器人控制和游戏AI等。通过更准确地建模回报分布，可以提高智能体的决策能力和学习效率，从而在复杂环境中取得更好的表现。该方法在风险敏感型任务中具有潜在的应用价值，例如自动驾驶和医疗诊断。

📄 摘要（原文）

We introduce a new architecture for Distributional Reinforcement Learning (DistRL) that models return distributions using normalizing flows. This approach enables flexible, unbounded support for return distributions, in contrast to categorical approaches like C51 that rely on fixed or bounded representations. It also offers richer modeling capacity to capture multi-modality, skewness, and tail behavior than quantile based approaches. Our method is significantly more parameter-efficient than categorical approaches. Standard metrics used to train existing models like KL divergence or Wasserstein distance either are scale insensitive or have biased sample gradients, especially when return supports do not overlap. To address this, we propose a novel surrogate for the Cramèr distance, that is geometry-aware and computable directly from the return distribution's PDF, avoiding the costly CDF computation. We test our model on the ATARI-5 sub-benchmark and show that our approach outperforms PDF based models while remaining competitive with quantile based methods.