超越惩罚机制：基于扩散模型的离线强化学习分布外检测与选择性正则化

# 超越惩罚：离线强化学习中的基于扩散的分布外检测与选择性正则化

来源: https://arxiv.org/html/2605.08202  
Qingjun Wang$^1$, Hongtu Zhou$^1$, Hang Yu$^1$, Junqiao Zhao$^{1,2}$, Yanping Zhao$^1$, Chen Ye$^1$, Ziqiao Wang$^1$, Guang Chen$^{1,3}$  
$^1$ 同济大学计算机科学与技术学院，中国上海  
$^2$ 教育部嵌入式系统与服务计算重点实验室，同济大学，中国上海  
$^3$ 上海人工智能实验室  
$\{2432069, \text{zhouhongtu}, 2432034, \text{zhaojunqiao}\}@\text{tongji.edu.cn}$  
$\{2534018, \text{yechen}, \text{ziqiaowang}, \text{guangchen}\}@\text{tongji.edu.cn}$

###### 摘要

离线强化学习（RL）面临的一个关键挑战是对分布外（Out-of-Distribution, OOD）动作价值的过高估计。现有方法通过对未见样本施加惩罚来缓解这一问题，但它们往往无法准确识别OOD动作，并且可能会抑制行为支持范围之外有益的探索。尽管已有若干方法被提出以区分具有不同特性的OOD样本，但它们通常依赖于对数据分布的严格假设，且在判别能力上仍然有限。为了解决这个问题，我们提出了 **DOSER**（基于**D**iffusion的**O**OD**D**etection and **S**elective **R**egularization），这是一种新颖的框架，超越了统一的惩罚机制。DOSER训练两个扩散模型以捕捉行为策略和状态分布，使用单步去噪重建误差作为可靠的OOD指标。在策略优化过程中，它通过评估预测的转移来区分有益和有害的OOD动作，选择性地抑制风险动作，同时鼓励对高潜力动作的探索。理论上，我们证明了DOSER是一个 $\gamma$-收缩映射，因此存在唯一的不动点，且具有有界的值估计。我们进一步提供了在模型近似和OOD检测误差下相对于最优策略的渐近性能保证。在广泛的离线RL基准测试中，DOSER consistently 取得了优于先前方法的性能，特别是在次优数据集上。

## 1 引言

离线强化学习（RL）已成为一种强大的范式，它仅从静态数据集中学习策略，无需进行可能昂贵或危险的在线交互（Levine et al., 2020）。这一特性使其在探索受限的现实世界领域（如机器人、医疗保健和自主系统）中特别具有吸引力。然而，直接将标准的离策（off-policy）RL算法应用于离线数据集会引发*分布偏移*的根本性挑战。当学习到的策略生成明显偏离训练数据分布的动作时，值函数往往会错误地外推，导致严重的价值高估，最终造成性能急剧下降（Fujimoto et al., 2019）。

现有方法主要分为两类：

1.  **策略约束方法**：强制学习到的策略保持接近行为策略，以避免OOD动作（Kumar et al., 2019; Wu et al., 2019; Fujimoto and Gu, 2021; Kostrikov et al., 2021），通常依赖变分自编码器（VAEs）（Kingma and Welling, 2013）进行行为建模。虽然在原理上有效，但这些方法难以捕捉现实世界行为的多模态特性，往往在低密度区域内将多样化的动作分布坍缩为次优的平均输出（Wang et al., 2022）。
2.  **值正则化方法**：通过学习保守的Q函数来惩罚OOD动作，提供了一种替代方案（Kumar et al., 2020; Wu et al., 2021; Bai et al., 2022; Mao et al., 2023）。它们的有效性取决于底层的OOD识别机制，由于用于表征数据分布的模型表征能力有限，这是一项具有挑战性的任务。此外，它们通常在整个非支持区域施加统一的惩罚，而不考虑能够提升策略性能的有价值探索（图1，左）。

近期的工作试图通过细粒度地控制保守水平来缓解过度的悲观情绪。CCVL（Hong et al., 2022）将Q函数条件化于置信水平，以学习一系列保守值估计，使得策略能够在线评估期间动态调整悲观程度。ACL-QL（Wu et al., 2024）将行为策略建模为高斯分布，并引入可学习的权重函数，在状态-动作层面上自适应地调节保守性。DoRL-VC（Huang et al., 2024）采用基于VAE的检测器将OOD与分布内（ID）动作分离，并进一步区分具有不同特性的OOD动作。然而，此类方法要么依赖Q集成学习来实现不同程度的保守性，从而增加了额外的训练开销，要么继承了关于行为策略的强烈高斯假设，这从根本上限制了它们可靠识别OOD样本的能力。

> **图1**：基于VAE的行为建模方法（左）误判OOD动作，而统一惩罚抑制了高潜力的OOD动作。DOSER（右）通过扩散模型建模多模态行为策略，并使用重建误差作为OOD指标，进一步区分有害与有益动作以进行选择性正则化。

为了解决这些挑战，我们提出了 **DOSER**（基于**D**iffusion的**O**OD**D**etection and **S**elective **R**egularization），通过两项关键创新推进OOD处理（图1，右）。首先，我们利用扩散模型实现OOD检测。通过部署两个扩散模型用于行为策略近似和状态分布建模，我们建立了重建误差作为理论严谨的度量指标，在避免强烈高斯假设的同时保持校准良好的检测性能。其次，我们引入了一种自适应判别机制，超越了分布内（ID）和分布外（OOD）的二元分类。通过集成学习到的动力学模型，我们区分有益的OOD动作（那些有潜力在保持状态分布不变的情况下提高性能的动作）和有害的OOD动作（那些可能引起状态分布偏移或价值退化的动作）。这种细粒度的判别使得选择性正则化成为可能，既劝阻危险动作，又鼓励有希望的探索，从而形成一个在保持必要保守性的同时促进策略改进的鲁棒框架。

本文的主要贡献如下：

1.  我们提出了一种基于扩散的离线RL中OOD检测方法，使用重建误差作为有理论依据的度量。
2.  我们引入了一种双重正则化策略，根据预测结果自适应地调整对OOD动作的处理，抑制有害动作同时鼓励有益动作。
3.  在D4RL基准上的大量实验表明，与先前方法相比，我们的方法取得了 superior 或具有竞争力的性能，详细的消融实验验证了每个组件的有效性。

## 2 预备知识

**离线强化学习**。我们考虑由马尔可夫决策过程（MDP） formulate 的RL问题，其定义为元组 $(\mathcal{S}, \mathcal{A}, \mathcal{P}, R, \gamma, d_0)$，其中 $\mathcal{S}$ 为状态空间，$\mathcal{A}$ 为动作空间，$\mathcal{P}: \mathcal{S} \times \mathcal{A} \times \mathcal{S} \rightarrow [0,1]$ 为转移动力学，$R: \mathcal{S} \times \mathcal{A} \rightarrow [R_{\mathrm{min}}, R_{\mathrm{max}}]$ 为奖励函数，$\gamma \in [0,1)$ 为折扣因子，$d_0: \mathcal{S} \rightarrow [0,1]$ 为初始状态分布（Sutton et al., 1998）。RL的目标是学习一个策略 $\pi: \mathcal{S} \rightarrow \Delta(\mathcal{A})$，以最大化期望折扣回报 $J(\pi) = \mathbb{E}[\sum_{t=0}^{\infty} \gamma^t R(\bm{s}_t, \bm{a}_t)]$。对于任意策略 $\pi$，我们定义值函数为 $V^\pi(\bm{s}) = \mathbb{E}[\sum_{t=0}^{\infty} \gamma^t R(\bm{s}_t, \bm{a}_t) | \bm{s}_0=\bm{s}]$，Q函数为 $Q^\pi(\bm{s}, \bm{a}) = \mathbb{E}[\sum_{t=0}^{\infty} \gamma^t R(\bm{s}_t, \bm{a}_t) | \bm{s}_0=\bm{s}, \bm{a}_0=\bm{a}]$。鉴于奖励是有界的，Q函数必须位于 $Q_{\mathrm{min}} = R_{\mathrm{min}}/(1-\gamma)$ 和 $Q_{\mathrm{max}} = R_{\mathrm{max}}/(1-\gamma)$ 之间。在离线RL中，智能体仅限于从由行为策略 $\pi_\beta$ 收集的静态数据集 $\mathcal{D}=\{(\bm{s}, \bm{a}, r, \bm{s}')\}$ 中学习，不与环境进行任何交互（Lange et al., 2012）。我们将经验行为策略表示为 $\hat{\pi}_\beta$，它描述了在 $\mathcal{D}$ 中观察到的条件动作分布。

**扩散模型**。扩散模型（Sohl-Dickstein et al., 2015; Ho et al., 2020; Song et al., 2020）已成为一类强大的生成模型，擅长捕捉复杂的数据分布。其核心思想围绕一个前向扩散过程，该过程逐渐将数据扰动为噪声，以及一个学习重构原始数据的反向过程。给定一个标准差为 $\sigma_{\mathrm{data}}$ 的干净样本 $\bm{x}_0 \sim p_{\mathrm{data}}(\bm{x}_0)$，前向过程通过添加独立同分布（i.i.d.）的高斯噪声（标准差 $\sigma_t$ 沿调度 $\sigma_{\mathrm{min}}=\sigma_0 < \sigma_1 < \dots < \sigma_N=\sigma_{\mathrm{max}}$ 增加）构建一系列越来越嘈杂的样本 $\bm{x}_t \sim p(\bm{x}_t; \sigma_t)$。通常，$\sigma_{\mathrm{min}}$ 选择得足够小，使得 $p_{\mathrm{min}}(\bm{x}) \approx p_{\mathrm{data}}(\bm{x})$，而 $\sigma_{\mathrm{max}}$ 足够大，使得最终分布近似各向同性高斯噪声，即 $p_{\mathrm{max}}(\bm{x}) \approx \mathcal{N}(\bm{x}; 0, \sigma_{\mathrm{max}}^2 \bm{I})$。在原始的DDPM（Ho et al., 2020）公式中，此过程被建模为离散马尔可夫链。后续工作通过随机微分方程（SDEs）的视角重新诠释了这一过程（Song et al., 2020），描述 $\bm{x}_t$ 在连续时间 $t \in [0, T]$ 上的演化为：

$$
d\bm{x}_t = f(\bm{x}_t, t) \, dt + g(t) \, d\bm{w}_t \quad (1)
$$

其中 $f(\cdot, t)$ 和 $g(t)$ 是漂移和扩散系数，$\bm{w}_t$ 是标准维纳过程。EDM框架（Karras et al., 2022）通过使用可微噪声调度 $\sigma(t)$ 重新参数化扩散路径来完善这一范式。反向过程由从前向SDE导出的相应概率流ODE控制，其公式为：

$$
d\bm{x}_t = -\dot{\sigma}(t) \sigma(t) \nabla_{\bm{x}_t} \log p_t(\bm{x}_t) dt \quad (2)
$$

其中 $\dot{\sigma}(t) = \frac{d\sigma}{dt}$ 是控制噪声变化率的噪声调度的时间导数，$\nabla_{\bm{x}_t} \log p_t(\bm{x}_t)$ 是边际分布 $p_t(\bm{x}_t)$ 的分数函数。该分数由通过去噪分数匹配（Vincent, 2011）训练的神经网络 $\bm{\epsilon}_\theta(\bm{x}_t; \sigma_t)$ 近似。去噪模型 $\bm{\epsilon}_\theta$ 被训练以从其噪声版本 $\bm{x}_t = \bm{x}_0 + \sigma_t \bm{\epsilon}$ 中预测真实的干净样本 $\bm{x}_0$，通过最小化加权 $L_2$ 损失：

$$
\mathcal{L}(\theta) = \mathbb{E}_{\sigma_t, \bm{x}_0 \sim p(\bm{x}_0), \bm{\epsilon} \sim \mathcal{N}(0, \bm{I})} \left[ \lambda(\sigma_t) \| \bm{x}_0 - \bm{\epsilon}_\theta(\bm{x}_t, \sigma_t) \|_2^2 \right], \quad (3)
$$

其中 $\lambda(\sigma_t)$ 是损失权重。与需要数千步去噪步骤的原始DDPM相比，EDM通过引入优化的噪声调度和高阶ODE求解器加速采样，仅在几十步内即可实现高质量生成。

## 3 基于扩散的OOD检测与选择性正则化

在本节中，我们介绍DOSER的技术框架。我们首先介绍三个使OOD动作的精确检测和分类成为可能的主要组件，然后展示这些组件如何集成到一个统一的算法框架中，并详细说明实际实现。图2提供了所提方法的概览。关于全面的理论分析，请参阅附录A。

> **图2**：所提方法概览：(a) 基于扩散的OOD动作检测，(b) 集成检测器以实现OOD动作分类。

### 3.1 基于扩散的行为与状态建模

我们方法的基础是建立两个扩散模型，以共同捕捉离线数据集的底层分布。我们首先构建一个条件扩散模型，通过训练去噪网络 $\bm{\epsilon}_{\theta_a}(\bm{a}_t, \sigma_t, \bm{s})$ 来重构原始动作 $\bm{a}_0$，从而学习经验行为策略分布 $\hat{\pi}_\beta(\bm{a}|\bm{s})$，优化目标如下：

$$
\mathcal{L}(\theta_a) = \mathbb{E}_{\sigma_t, (\bm{s}, \bm{a}_0) \sim \mathcal{D}, \bm{\epsilon} \sim \mathcal{N}(0, \bm{I})} \left[ \lambda(\sigma_t) \| \bm{a}_0 - \bm{\epsilon}_{\theta_a}(\bm{a}_t, \sigma_t, \bm{s}) \|_2^2 \right].
$$