Quantile Geometry Regularization for Distributional Reinforcement Learning¶

会议: NeurIPS 2026
arXiv: 2605.08182
OpenReview: https://openreview.net/forum?id=nZTmdGMiNl
领域: 强化学习
关键词: 分布式强化学习、分位数回归、Wasserstein 分布鲁棒优化、分位数几何正则化、隐式分位数网络

一句话总结¶

针对基于分位数回归的分布式强化学习因有限统计量非 Bellman 闭包而导致的分布坍缩问题，本文提出基于 Wasserstein 分布鲁棒优化的隐式分位数网络（RQIQN），通过闭式、分位数依赖的局部位置校正正则化回报分布几何，在不改变期望值目标的前提下有效恢复分布扩展度。

研究背景与动机¶

分布式强化学习（DistRL）通过对贴现累积回报的完整概率分布进行建模，为随机环境与风险敏感决策提供了比传统标量期望强化学习更丰富的表征基础。在各类分布参数化方案（分类原子、分位数、期望分位数、显式样本集）中，以 QR-DQN 和隐式分位数网络（IQN）为代表的基于分位数的方法尤为突出，因为分位数函数天然与 Wasserstein 几何对齐、无需预设固定的价值支撑集，且能通过扭曲风险测度（如 CVaR）灵活实现风险敏感控制。然而，有限维分位数统计量在数学上并不是 Bellman 闭合的——这意味着对后继状态分位数进行 Bellman 备份后生成的分布，其真实分位数无法仅通过后继分位数和即时奖励做有限代数运算精确得到。在实际时序差分（TD）更新中，模型将当前分位数估计直接对齐到自举的分位数目标上，导致逼近偏差随训练迭代不断累积，引发严重的分布退化与几何畸变。

这种退化在深度强化学习中被进一步放大：为了优化过程的数值稳定性，实际算法普遍采用分位数 Huber 损失替代理论上的精确分位数检验损失（Check loss），使得回归过程产生局部平滑效应，导致学到的回报分布方差急剧缩减并向均值坍缩，极值尾部信息大面积丢失。先前解决非闭包问题的工作通常依赖复杂的结构改造，例如通过求解高代价的非线性反演系统从期望分位数重构样本（如 ER-DQN）、引入多头预测与辅助解耦损失联合学习（如 IEQN），或是借助 EM 投影与高斯混合进行样本集运算（如 MMD、MWG），这些方法带来了极高的计算复杂度与训练不稳定性。

本文的切入视角是回归分位数估计的本质：将 IQN 损失的单次更新解构为针对采样分位数分量的若干独立局部经验分位数估计槽位。既然自举目标是由有限采样构成的经验分布，那么分布退化本质上就是经验目标扰动下的估计偏差。核心 idea：将 IQN 的每个局部经验分位数拟合槽位建模为 Wasserstein 模糊集下的分布鲁棒分位数估计问题，推导出具有中位数反对称性与单调性的闭式加性校正项，以均值中性的方式主动拉伸上下分位数间距，从几何层面消除分布坍缩。

方法详解¶

整体框架¶

RQIQN 的核心思想是在时序差分更新的局部插槽层面引入分布鲁棒性，而非直接对无限维的回报分布空间强加全局先验约束。对于智能体与环境交互收集的转移样本 $(s, a, r, s')$，算法首先通过均值贪婪准则确定最优后继动作 $a^* = \arg\max_{a'} \frac{1}{K}\sum_{k=1}^K Z_{\tilde{\tau}_k}(s', a')$。随后，算法分别从均匀分布 $\mathcal{U}([0,1])$ 中独立采样 $N$ 个当前分位数分量 $\{\tau_i\}_{i=1}^N$ 与 $N'$ 个后继目标分位数分量 $\{\tau'_j\}_{j=1}^{N'}$。对于任意固定的当前分量 $\tau_i$，后继自举样本集合构成了一个经验目标测度 $\hat{\mu}_{s,a}$。RQIQN 考虑以 $\hat{\mu}_{s,a}$ 为中心、半径为 $\epsilon_t$ 的 Wasserstein 概率球，通过求解最坏情况扰动下的分位数回归问题，解析得到一个无需数值求解的闭式位置偏移量 $\Delta_p(\tau_i, \epsilon_t)$，并将其直接注入到时序差分残差中。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入转移样本 (s, a, r, s')<br/>采样分位数分量 τ_i, τ'_j ~ U([0,1])"] --> B["局部经验分位数估计重构<br/>构建经验 Bellman 目标律 μ̂_{s,a}"]
    B --> C["闭式 Wasserstein 分布鲁棒分位数校正<br/>计算中位数反对称偏移 Δ_p(τ_i, ϵ_t)"]
    C --> D["动态退火半径与鲁棒 TD 误差构建<br/>结合逆向 Logistic 调度与 Huber 代理"]
    D --> E["输出优化损失并更新网络参数<br/>重塑回报分布几何宽度，抑制分布坍缩"]

关键设计¶

1. 局部经验分位数估计重构：将自举 Bellman 损失解构为独立分位数回归插槽

传统 IQN 的损失函数对所有当前分位数分量 $\tau_i$ 和目标分量 $\tau'_j$ 计算成对的时序差分误差： $$\mathcal{L}_{\mathrm{IQN}}(s,a,r,s') = \frac{1}{N'} \sum_{i=1}^N \sum_{j=1}^{N'} \rho_{\tau_i}^\kappa \left( r + \gamma Z_{\tau'_j}(s', a') - Z_{\tau_i}(s,a) \right)$$ 本文从统计学视角重新审视这一快照：若固定单个当前分位数分量 $\tau_i \in (0,1)$，记自举目标样本为 $y_j := r + \gamma Z_{\tau'_j}(s', a^*)$，待估分位数值为 $q_i := Z_{\tau_i}(s,a)$，则经验 Bellman 目标测度可表示为经验狄拉克测度的均匀混合 $\hat{\mu}_{s,a} = \frac{1}{N'}\sum_{j=1}^{N'} \delta_{y_j}$。在此测度下，针对 $\tau_i$ 的更新等价于求解经典的一维经验分位数回归问题： $$q_{\tau_i}^0 \in \arg\min_{q \in \mathbb{R}} \mathbb{E}_{Y \sim \hat{\mu}_{s,a}} \left[ \rho_{\tau_i}(Y - q) \right]$$ 其中 $\rho_\tau(u) = u(\tau - \mathbb{I}_{\{u<0\}})$ 为标准 Check 损失。该解严格满足经验覆盖条件 $\mathbb{P}_{Y \sim \hat{\mu}_{s,a}}(Y < q_{\tau_i}^0) \le \tau_i \le \mathbb{P}_{Y \sim \hat{\mu}_{s,a}}(Y \le q_{\tau_i}^0)$。这一重构明确了：分布退化的微观根源在于每个局部估计槽位所依赖的经验测度 $\hat{\mu}_{s,a}$ 包含有限采样方差与自举网络非平稳性引入的几何扭曲。

2. 闭式 Wasserstein 分布鲁棒分位数校正：均值中性与几何间距扩展的双重保证

为防止经验目标测度偏离真实目标律，RQIQN 引入基于 $p$ 阶 Wasserstein 距离（$p \in (1, \infty]$）的分布鲁棒优化框架，定义模糊集 $\mathcal{B}_p(\hat{\mu}_{s,a}, \epsilon) := \{ \nu \in \mathcal{P}(\mathbb{R}) : W_p(\nu, \hat{\mu}_{s,a}) \le \epsilon \}$，并考察最坏测度下的分位数极小极大问题： $$q_{\tau_i}^\epsilon \in \arg\min_{q \in \mathbb{R}} \sup_{\nu \in \mathcal{B}_p(\hat{\mu}_{s,a}, \epsilon)} \mathbb{E}_{Y \sim \nu} \left[ \rho_{\tau_i}(Y - q) \right]$$ 得益于 Check 损失对加性 Wasserstein 正则化的独特性质，最坏情况扰动下的稳健解具有解析形式的位置平移：$q_{\tau_i}^\epsilon = q_{\tau_i}^0 + \Delta_p(\tau_i, \epsilon)$。通用修正项定义为： $$\Delta_p(\tau, \epsilon) = \frac{\epsilon}{q} \left( \tau^q - (1-\tau)^q \right) c_{\tau,p}^{1-q}, \quad \frac{1}{p} + \frac{1}{q} = 1$$ 其中结构常数 $c_{\tau,p}$ 随 $p$ 变化。该校正项具备两大决定性几何性质： - 中位数反对称性（均值中性）：$\Delta_p(1-\tau, \epsilon) = -\Delta_p(\tau, \epsilon)$ 且 $\Delta_p(1/2, \epsilon) = 0$。在均匀采样 $\tau \sim \mathcal{U}(0,1)$ 下，积分均值为零 $\mathbb{E}_{\tau}[\Delta_p(\tau, \epsilon)] = 0$。因此，它严格保持了风险中性下的平均价值估计 $\mathbb{E}_\tau[Z_\tau(s,a) + \Delta_p(\tau, \epsilon)] = \mathbb{E}_\tau[Z_\tau(s,a)]$，不会引入类似传统 Q-learning 的系统性过估计或欠估计偏差。 - 单调递增性（扩展分布间距）：$\Delta_p(\tau, \epsilon)$ 关于 $\tau$ 严格非递减。对于任意下上分位数分量 $0 < \tau_l < \tau_h < 1$，校正后的区间宽度严格放大： $$[Z_{\tau_h} + \Delta_p(\tau_h, \epsilon)] - [Z_{\tau_l} + \Delta_p(\tau_l, \epsilon)] = (Z_{\tau_h} - Z_{\tau_l}) + \left(\Delta_p(\tau_h, \epsilon) - \Delta_p(\tau_l, \epsilon)\right)$$ 由于增量项 $\Delta_p(\tau_h, \epsilon) - \Delta_p(\tau_l, \epsilon) \ge 0$，该机制从数学上直接对抗了 Bellman 迭代导致的分布宽度萎缩。针对连续采样下端点发散的工程挑战，作者针对 $p=2$ 导出了保持单调性与反对称性的有界闭式变体： $$\Delta_2(\tau, \epsilon) = \frac{\epsilon}{2} \frac{1 - 2\tau}{\sqrt{\tau^2 + (1-\tau)^2}}$$

3. 动态退火半径与鲁棒 TD 误差构建：平衡探索期鲁棒性与收敛期平滑度

统计理论建议 Wasserstein 半径随样本容量按幂律收缩（$\epsilon_N = \mathcal{O}(N^{-1/2})$）。但在强化学习非平稳自举训练场景下，样本来自重放缓冲区且价值网络不断更新。若固定鲁棒半径，后期将持续引入外生方差。因此，RQIQN 将鲁棒半径解耦，采用时间依赖的逆向 Logistic 衰减调度（Reverse-Logistic Schedule）： $$\epsilon_t = \frac{\epsilon_0}{1 + \exp\left( k(t - t_0) \right)}$$ 其中 $\epsilon_0$ 为初始尺度，$k$ 控制衰减陡峭度，$t_0$ 为过渡中点。在训练初期，较大的 $\epsilon_t$ 赋予较强的几何正则化，支撑起分布宽度并促进风险感知探索；在训练后期，$\epsilon_t \to 0$ 使得算法渐近退火至标准一致估计，确保策略收敛。在损失函数实现上，算法将校正项平移并结合 Huber 损失构建平滑且稳健的训练目标： $$\mathcal{L}_{\mathrm{RQIQN\text{-}Huber}} = \frac{1}{N'} \sum_{i=1}^N \sum_{j=1}^{N'} \rho_{\tau_i}^\kappa \left( r + \gamma Z_{\tau'_j}(s', a^*) + \Delta_i - Z_{\tau_i}(s,a) \right)$$

损失函数 / 训练策略¶

RQIQN 的整体优化目标采用结合闭式 Wasserstein 校正的鲁棒 Huber 损失。训练过程保持与标准 IQN 相同的超参数体系，网络输入状态并经过全连接/卷积主干，分位数分量 $\tau$ 经过余弦嵌入层后与状态表征逐元素相乘并映射为分位数输出。训练时设置当前分位数采样数 $N=64$，自举目标分位数采样数 $N'=64$，贪婪动作选择评估数 $K=32$。初始鲁棒半径固定为 $\epsilon_0 = 1.0$。在复杂动力学与扰动环境中，智能体可进一步接入自适应 CVaR 采样机制：根据与最近障碍物的归一化距离动态调整下尾截断阈值 $\eta$，在危险区域专注下尾风险规避，在开阔区域保持全局分位数利用。

实验关键数据¶

主实验¶

论文在极具挑战性的安全关键型无人水面艇（USV）自主导航环境（包含 10 个静态障碍物和 8 个强涡流扰动源的困难模式）中进行了全面评估。对比基线包括经典标量 DQN、标准 IQN 以及结合自适应 CVaR 采样的变体（+Adaptive）。每个方法在 5 个随机种子上独立训练 3M 步，并在 500 个不同随机布局的测试轮次上评估平均表现。

方法 (Method)	成功率 (Success, %) ↑	碰撞率 (Collision, %) ↓	超时率 (Timeout, %) ↓	累积回报 (Return) ↑	成功时间 (Timesucc, s) ↓	成功能耗 (Energysucc) ↓
DQN	67.24 ± 4.79	32.72 ± 4.77	0.04 ± 0.09	12.97 ± 3.65	33.48 ± 0.33	98.95 ± 5.76
IQN	64.80 ± 22.66	21.92 ± 9.54	13.28 ± 29.47	-13.17 ± 51.78	39.08 ± 1.22	122.56 ± 17.29
IQN + Adaptive	71.04 ± 23.35	15.44 ± 6.13	13.52 ± 29.11	-12.19 ± 51.20	42.24 ± 1.55	133.31 ± 17.04
RQIQN (本文)	83.24 ± 1.19	16.48 ± 1.14	0.28 ± 0.27	19.20 ± 1.73	37.55 ± 1.25	111.24 ± 6.34
RQIQN + Adaptive (本文)	85.64 ± 2.27	14.20 ± 2.02	0.16 ± 0.26	18.57 ± 3.90	39.19 ± 1.87	115.83 ± 6.41

在 Atari 基准测试（包含 Atari-5 子集以及 Alien、BattleZone、Frostbite、Qbert 等 9 款经典游戏）中，RQIQN 基于 Dopamine 框架在 Sticky-actions 协议下训练 200M 帧，在绝大多数游戏上显著超越基线 IQN，并在多款游戏上超过了集成强化学习算法 Rainbow。

消融实验¶

为了分析 Wasserstein 模糊集阶数 $p$（几何形式）与回归损失类型（Check loss vs. Huberized loss）对性能的影响，论文在 Atari 游戏 Qbert 上开展了系统的消融实验。

配置 (Config)	损失函数 (Loss)	Wasserstein 几何阶数 ($p$)	最终评估得分 (Qbert Score)	说明 (Note)
RQIQN (默认配置)	Check loss	$p = 2$	~4,200	完整模型，兼顾理论分布鲁棒与平滑拓展
RQIQN-Huber	Huber loss	$p = 2$	~4,050	实用平滑替代，表现与 Check loss 相当
RQIQN ($p=\infty$)	Check loss	$p = \infty$	~2,100	$\ell_\infty$ 模糊集导致边缘过度惩罚，性能显著下降
RQIQN-Huber ($p=\infty$)	Huber loss	$p = \infty$	~2,000	几何形式不匹配引发较大退化
IQN 基准	Huber loss	无校正 ($\epsilon=0$)	~2,550	缺乏几何正则化，受分布坍缩制约

关键发现¶

分布鲁棒几何正则化极大地降低了训练方差：在强涡流 USV 导航中，标准 IQN 表现出极差的稳定性（成功率方差高达 ±22.66%，回报方差高达 ±51.78），存在高达 13.28% 的超时陷入概率；而 RQIQN 将方差收敛至 ±1.19%，超时率降至 0.28%，彻底解决了自举退化导致的策略震荡。
几何选择（$p=2$ vs $p=\infty$）比损失形式更为敏感：消融表明，采用 $p=2$ 阶的 Wasserstein 距离时，Check 损失与 Huber 损失均能带来优异表现；而一旦退化为 $p=\infty$，性能发生断崖式下跌。这证明 $p=2$ 能够平滑地对中高阶尾部几何进行比例修正，而 $p=\infty$ 的极值惩罚过于激进。
非有界支撑集与表达能力的权衡：在 NameThisGame 中，固定支撑集的 C51 表现更稳健，而无界分位数方法（IQN 与 RQIQN）因尾部估计误差放大而收敛于次优解。这揭示出无界分位数模型需要配合自适应风险扭曲（如 CVaR）来控制异常尾部方差。

亮点与洞察¶

从统计局部估计切入重构算法：巧妙地将全局非闭包的 Bellman 算子难题，降维转化为单个时序差分插槽内的经验分位数估计问题，避开了在无限维函数空间做投影运算的高难度数学困境。
闭式计算带来近乎零开销的性能飞跃：推导出的解析校正项仅由分位数分量 $\tau_i$ 和标量半径 $\epsilon_t$ 决定，只需一行公式即可直接融入 PyTorch/JAX 的张量批处理，无需辅助神经网络头，更无需任何反演或采样优化器。
均值中性与离散扩展的优美数学对齐：通过反对称函数设计，严格在期望意义下保持 $\mathbb{E}[\Delta_p]=0$，实现了“只调整分布形态与离散度，绝不干扰价值中枢”的理想正则化效果。

局限与展望¶

作者承认的局限：在某些极端环境（如 Atari 中的 NameThisGame），无界分位数函数对非平稳尾部波动的敏感性可能仍会导致次优策略，表现略逊于具有硬截断有界支撑集的 C51。
潜在不足：对于具有非连续跳跃或高度多峰特性的极端回报分布，一阶单调递增的加性偏移可能无法精细调制局部的峰谷结构；此外，逆向 Logistic 退火超参数（$t_0, k$）在不同动力学时长的环境中需要适当调谐。
改进方向：可进一步探索根据状态认知不确定性（Epistemic Uncertainty）自适应生成局部鲁棒半径 $\epsilon(s)$，实现“高置信区域精准拟合、高未知区域强力支撑”的自适应几何正则化。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次从 Wasserstein 分布鲁棒分位数估计角度重构 IQN 损失快照，推导出极具数学美感的均值中性闭式校正。
实验充分度: ⭐⭐⭐⭐ 涵盖从物理仿真的海洋无人艇安全导航到标准 Atari 游戏库，并附带针对几何度量与损失形式的深入消融。
写作质量: ⭐⭐⭐⭐⭐ 理论推导严丝合缝，几何直觉阐述清晰生动，公式表述规范克制。
价值: ⭐⭐⭐⭐⭐ 代码实现极为精简，几乎不增加任何计算壁垒即可无缝增强现有的分位数强化学习管线，实用价值巨大。

配置 (Config)	损失函数 (Loss)	Wasserstein 几何阶数 (\(p\))	最终评估得分 (Qbert Score)	说明 (Note)
RQIQN (默认配置)	Check loss	\(p = 2\)	~4,200	完整模型，兼顾理论分布鲁棒与平滑拓展
RQIQN-Huber	Huber loss	\(p = 2\)	~4,050	实用平滑替代，表现与 Check loss 相当
RQIQN (\(p=\infty\))	Check loss	\(p = \infty\)	~2,100	\(\ell_\infty\) 模糊集导致边缘过度惩罚，性能显著下降
RQIQN-Huber (\(p=\infty\))	Huber loss	\(p = \infty\)	~2,000	几何形式不匹配引发较大退化
IQN 基准	Huber loss	无校正 (\(\epsilon=0\))	~2,550	缺乏几何正则化，受分布坍缩制约