面向分布强化学习的分位数几何正则化

# 分布强化学习中的分位数几何正则化

来源: https://arxiv.org/html/2605.08182
Zhaofan Zhang, Minghao Yang, Rufeng Chen, Sihong Xie, Hui Xiong
人工智能枢纽，香港科技大学（广州）
中国广东广州

###### 摘要

基于分位数的分布强化学习方法通过采样的分位数回归来学习回报分布，但其引导的目标分位数可能会诱导扭曲或退化的分布估计。我们提出了鲁棒基于分位数的隐式分位数网络（RQIQN），这是一种轻量级的 Wasserstein 分布鲁棒增强方法，从分位数估计的角度进行了改进。我们首先将 IQN 损失的快照重新解释为一组在采样当前分位比例上的局部经验分位数估计问题。然后，我们用 Wasserstein 分布鲁棒分位数估计公式对每个局部槽位进行鲁棒化，从而产生一个闭式的、依赖于分位比例的贝尔曼目标校正。该校正直接解决了分布退化问题：其中位数反对称性保持了风险中性分位数平均值，而其单调性扩大了上下分位数间隙并抵消了分布扩展的崩溃。因此，RQIQN 在不改变底层价值目标或需要额外样本集重建的情况下正则化了分位数几何。最后，我们通过实证表明，所提出的 RQIQN 在风险敏感导航和 Atari 游戏中优于其他现有的基于分位数的分布强化学习算法。

## 1 引言

分布强化学习（DistRL）对折扣回报的完整分布进行建模，而不仅仅是其期望值，从而为在随机性和风险下的学习提供了更丰富的基础。理解实用 DistRL 算法的一个有用方法是通过区分回报分布及其统计量。分布贝尔曼算子作用于完整的回报分布，而可扩展算法通常传播有限维摘要，如原子、分位数或期望值。现有的 DistRL 算法主要在如何表示和更新回报分布方面存在差异，包括分类原子、分位函数、期望值和样本集表示。在这些选择中，基于分位数的方法特别具有吸引力，因为分位函数自然地与 Wasserstein 几何对齐，不需要固定的值支持，并可以直接访问回报分布的不同区域。

在此范围内，隐式分位数网络（IQN）[Dabney et al. (2018a)](https://arxiv.org/html/2605.08182#bib.bib3) 进一步扩展了带有固定网格分位数的 QR-DQN [Dabney et al. (2018b)](https://arxiv.org/html/2605.08182#bib.bib4)，通过学习通过采样分位比例 $\tau \sim U([0,1])$ 的连续分位函数。通常，如果贝尔曼更新分布的统计量可以仅从下一状态分布和奖励的相同统计量中计算出来，而无需重建完整分布，则称一组统计量是贝尔曼封闭的。这一性质至关重要，因为它允许在所选统计空间中进行递归动态规划（公式 [1](https://arxiv.org/html/2605.08182#S2.E1)）。然而，有限分位统计量并非贝尔曼封闭 [Rowland et al. (2019)](https://arxiv.org/html/2605.08182#bib.bib7)，这可能导致分布退化和分布扩展的有偏估计（图 [1](https://arxiv.org/html/2605.08182#S1.F1)(a)）。

在基于分位数的 DistRL 中，深度 Q 学习 [Mnih et al. (2013)](https://arxiv.org/html/2605.08182#bib.bib1) 从标量值回归扩展到分布分位数拟合。对于每次转移，智能体在当前分位估计和引导的目标分位之间形成成对的 TD 误差，并最小化分位数回归损失（公式 [3](https://arxiv.org/html/2605.08182#S2.E3)）。因此，TD 更新具有明确的分位数估计解释，每个采样比例对应于一个局部分位数拟合问题。

在这项工作中，我们通过局部分位数估计的视角重新审视 IQN。具体而言，IQN 损失（公式 [3](https://arxiv.org/html/2605.08182#S2.E3)）在经验损失层面可以精确地解释为一组经验分位数估计槽位，每个采样的当前分位比例 $\tau_i$ 对应一个槽位。每个槽位拟合由目标分位比例 $\{\tau'_j\}_{j=1}^{N'}$ 诱导的引导贝尔曼目标的 $\tau_i$ 定位分位数。这一观点揭示了 IQN 在由有限数量的引导目标分位值形成的模型生成的经验目标分布上执行局部分位数估计。从这个角度来看，分布退化可以被理解为学习到的分位值中累积的偏差，这扭曲了所表示回报分布的几何形状。

我们提出 RQIQN（图 [1](https://arxiv.org/html/2605.08182#S1.F1)(b)）来解决分布退化问题，即有偏的分位数表示，方法是通过对经验贝尔曼目标分布内 Wasserstein 模糊集内的最坏情况扰动鲁棒化每个局部分位数拟合问题。所得的闭式分位校正通过稳定有限样本分位估计来正则化学习到的回报分布。与现有工作中通过修改分布贝尔曼算子或分布表示来引入计算密集型方案不同 [Rowland et al. (2019)](https://arxiv.org/html/2605.08182#bib.bib7); [Jullien et al. (2023)](https://arxiv.org/html/2605.08182#bib.bib5); [Nguyen et al. (2020)](https://arxiv.org/html/2605.08182#bib.bib17)，RQIQN 用围绕经验贝尔曼目标律的 Wasserstein 分布鲁棒分位数估计问题替换每个局部分位数回归槽位。对于检验损失公式[^1]，该鲁棒槽位接受一个闭式的、依赖于分位比例的校正 $\Delta_p(\tau; \epsilon)$，其中 $p$ 是 Wasserstein 阶数，$\epsilon$ 是随时间衰减的鲁棒性半径。在几何上，该校正关于中位数反对称且在 $\tau$ 上单调。因此，它在扩大上下分位数间隙的同时保持均匀分位平均值，提供了一种均值中立的分布扩展正则化。

[^1]: 即分位数损失。它与分位数 Huber 损失（Huberized 分位数回归损失）不同。

**图 1:** (a) 状态 0 处的分布退化以及 (b) 所提出的 RQIQN 校正如何调节分位数几何的图示。拟合回报分布的样本可视化来自一个具有确定性转移和单一动作的四状态链 MDP。状态转移是有方向且顺序的，从状态 0 进展到状态 3。除了终止状态外，奖励为零，在终止状态处 $r \sim \frac{1}{2}\mathcal{N}(-2,1) + \frac{1}{2}\mathcal{N}(2,1)$。在 (b) 中，每个智能体使用其默认训练损失。

### 1.1 相关工作

DistRL 已迅速发展为基于期望的价值学习的强大替代方案，代表性方法使用分类原子、分位数、期望值或基于样本的表示对回报分布进行建模 [Bellemare et al. (2017)](https://arxiv.org/html/2605.08182#bib.bib2); [Dabney et al. (2018b)](https://arxiv.org/html/2605.08182#bib.bib4), [a](https://arxiv.org/html/2605.08182#bib.bib3); [Rowland et al. (2019)](https://arxiv.org/html/2605.08182#bib.bib7); [Yang et al. (2019)](https://arxiv.org/html/2605.08182#bib.bib8)。其中，IQN [Dabney et al. (2018a)](https://arxiv.org/html/2605.08182#bib.bib3) 因其灵活的隐式分位函数表示而变得特别具有影响力，它支持任意分位数采样，并通过扭曲风险度量自然地实现风险敏感策略。IQN 已广泛应用于需要风险敏感决策的任务中，例如自主水面导航 [Lin et al. (2023)](https://arxiv.org/html/2605.08182#bib.bib13); [Zhang et al. (2025)](https://arxiv.org/html/2605.08182#bib.bib11) 和四足运动 [Shi et al. (2024)](https://arxiv.org/html/2605.08182#bib.bib14)。

然而，IQN 仍然继承了基于分位数的 DistRL 的一个局限性：有限分位统计量并非贝尔曼封闭。几项研究工作通过改变学习的统计量或引入显式的统计量到样本接口来解决由非贝尔曼封闭分布统计量引起的困难。期望值作为不对称最小二乘位置统计量被引入 [Newey and Powell (1987)](https://arxiv.org/html/2605.08182#bib.bib18)，它们很有吸引力，因为它们比检验损失分位数回归提供更平滑的 $L_2$ 基础拟合，且其对称情况 $\tau=1/2$ 与均值重合。这一特性使得期望值对于稳定的价值估计很有用，因为均值是用于风险中性动作选择的统计量。然而，有限的一组期望值通常也不是贝尔曼封闭的，且期望值不能直接解释为来自回报分布的样本。为了使期望值可用于贝尔曼备份，ER-DQN [Rowland et al. (2019)](https://arxiv.org/html/2605.08182#bib.bib7) 通过一个插补步骤恢复贝尔曼目标样本，这需要求解昂贵的非线性系统。IEQN [Jullien et al. (2023)](https://arxiv.org/html/2605.08182#bib.bib5) 则联合学习期望值和分位数，以保留高效的 $L_2$ 基础学习同时避免显式插补，但代价是额外的预测头和有辅助耦合损失。

样本集方法，如 MMD [Nguyen et al. (2020)](https://arxiv.org/html/2605.08182#bib.bib17) 和 MWG [Zhang et al. (2024)](https://arxiv.org/html/2605.08182#bib.bib16)，采取不同的路线，直接对显式回报样本进行操作，从而绕过将非封闭统计量作为样本传播的需要。然而，确定性样本集可能会限制随机目标的多样性，而高斯混合变体依赖于 EM 风格投影 [Dempster et al. (1977)](https://arxiv.org/html/2605.08182#bib.bib19) 和样本增强，引入了额外的计算和算法复杂性。

## 2 问题设置与主要分析

### 2.1 分布强化学习

我们考虑一个马尔可夫决策过程（MDP） $\mathcal{M}=(\mathcal{S},\mathcal{A},\mathcal{P},\mathbb{R},\gamma)$，其中 $\mathcal{S}$ 和 $\mathcal{A}$ 分别表示状态和动作空间，$\mathcal{P}(\cdot|s,a)$ 是转移核，$\mathbb{R}(\cdot|s,a)$ 表示奖励分布，$\gamma \in [0,1)$ 是折扣因子。对于策略 $\pi$，状态动作回报是随机变量 $Z^{\pi}(s,a) \overset{D}{=} \sum_{t=0}^{\infty}\gamma^{t}R_{t}$，其中 $S_{0}=s, A_{0}=a$，$\overset{D}{=}$ 表示分布相等，且转移过程可以解释为 $S_{t+1}\sim\mathcal{P}(\cdot|S_{t},A_{t})$，$R_{t}\sim\mathbb{R}(\cdot|S_{t},A_{t})$，$A_{t}\sim\pi(\cdot|S_{t})$ 对于 $t\geq 1$。

标准的强化学习通常只优化此回报分布的一阶矩。特别是，动作值函数定义为 $Q^{\pi}(s,a)=\mathbb{E}_{Z}\![Z^{\pi}(s,a)]$。最优动作值函数满足贝尔曼最优方程：
$$
Q(s,a)=\mathbb{E}_{R\sim\mathbb{R}(\cdot|s,a),\,S'\sim\mathcal{P}(\cdot|s,a)}\![R+\gamma\max_{a'\in\mathcal{A}}Q(S',a')\,|\,S=s,A=a]
$$
这种基于期望的公式将完整的回报分布压缩为一个标量，丢弃了由随机转移和奖励引起的分布信息。DistRL 通过对*回报分布* $Z$ 而不是仅对其期望进行建模来扩展标准 RL。正式地，$Z$ 可以通过动态规划进行更新，分布贝尔曼最优算子[^2]定义为：
$$
\mathcal{T}^{\pi}Z(s,a)\stackrel{\{\scriptstyle D\}}{\{=\}}R(s,a)+\gamma Z(S',A'). \tag{1}
$$
[^2]: 大写字母 $S$ 和 $A$ 表示随机变量，小写字母表示它们的实现值。

我们可以使用分布贝尔曼最优算子 $\mathcal{T}$ 计算最优回报分布，定义为：
$$
\mathcal{T}Z(s,a)\stackrel{\{\scriptstyle D\}}{\{=\}}R(s,a)+\gamma Z\left(S',a^*\right),\qquad a^*=\operatorname*{argmax}_{a'}\mathbb{E}_{Z}[Z(S',a')]. \tag{2}
$$
作为一种代表性的基于分位数的 DistRL 方法，隐式分位数网络（IQN）[Dabney et al. (2018a)](https://arxiv.org/html/2605.08182#bib.bib3) 通过回报分布的分位函数提供了一种灵活的分布表示。设 $F_{Z(s,a)}^{-1}(\tau)$ 表示随机回报 $Z(s,a)$ 在分位比例 $\tau$ 处的分位函数。为了符号简单，我们写 $Z_{\tau}(s,a):=F_{Z(s,a)}^{-1}(\tau)$，使得当 $\tau \sim U([0,1])$ 时，采样的分位值 $Z_{\tau}(s,a)$ 遵循回报分布 $Z(s,a)$。基于这种隐式分位表示，IQN 用神经网络参数化此分位函数 $Z_{\tau}(x,a;\theta) \approx Z_{\tau}(s,a)$。为简洁起见，我们直接使用 $Z_{\tau}(s,a)$。由于 IQN 显式采样分位比例，可以通过修改分位采样范围来诱导风险敏感行为。一个常见的选择是条件在险价值（CVaR）扭曲，它将均匀采样的比例 $\tau \sim U([0,1])$ 映射到 $\tilde{\tau}=\eta\tau$，从而将采样限制在下尾分位区域 $U([0,\eta])$ 以进行风险厌恶控制。在本工作中，除非明确声明，否则比例是无扭曲采样的。

在整个工作中，我们考虑与 DistRL 相关的两个分位数回归损失。第一个是检验损失：
$$
\rho_{\tau}(u)=u\left(\tau-\mathbbm{1}_{\{u<0\}}\right)=\left\|\tau-\mathbbm{1}_{\{u<0\}}\right\|\,\|u\|,
$$
它给出了精确的分位数回归目标。第二个是实际 QR-DQN/IQN 实现中使用的标准分位 Huber [Huber (1992)](https://arxiv.org/html/2605.08182#bib.bib15) 损失：
$$
\rho_{\tau}^{\kappa}(u)=\left\|\tau-\mathbbm{1}_{\{u<0\}}\right\|\frac{\mathcal{H}_{\kappa}(u)}{\kappa},\qquad\mathcal{H}_{\kappa}(u)=\begin{cases}\frac{1}{2}u^{2},&\|u\|\leq\kappa,\\ \\kappa\left(\|u\|-\frac{1}{2}\kappa\right),&\|u\|>\kappa.\end{cases}
$$
这里，$\kappa>0$ 是 Huber 阈值。检验损失保留了精确的分位数估计解释，而 Huberized 损失为数值稳定性提供了更平滑的替代方案。