强化学习入门（五）

强化学习+diffusion模型

算法 Diffusion-QL （伪代码）

输入：

• 经验池$\mathcal{D}$
• 迭代步数 $N_{iter}$
• 批量大小$B$
• 折扣因子$\gamma$
• 软更新系数$\tau$
• 行为克隆/Q-learning 权衡系数 $\eta$
• 其他超参数（${\beta }_{schedule}$, $n_{timesteps}$, grad-norm等）

初始化：

• Actor ${\pi }_{\theta }$ ← 条件扩散模型 Diffusion
• EMA-Actor ${\pi }_{\bar{\theta }}$ ← $\theta$ 的深拷贝
• Critic Qϕ={Q1,Q2}Q_\phi=\{Q_1,Q_2\}Qϕ​={Q1​,Q2​} ← 双 Q 网络
• Target-Critic $Q_{\bar{\varphi }}$ ← $\varphi$ 的深拷贝
• 优化器 & 学习率调度器

──────────────────────────────────────────────────────────────
FOR $step=1\dots N_{\text{iter}}$ DO

1. 采样小批量
     $(s,a,r,s^{\prime },d)\leftarrow \mathcal{D}.\text{sample}(B)$

2. 更新 Critic（值函数）
      2.1 目标动作
     ${\hat{a}}^{\prime }\leftarrow {\pi }_{\bar{\theta }}(s^{\prime })$  # 用 EMA-Actor 生成
     IF $max\_q\_backup$ THEN  # 可选：10 次重复取最大 Q
      $\bar{y}\leftarrow \underset{j}{\min }\underset{k}{\max }{Q}_{{\bar{\varphi }}_j}(s^{\prime },{\hat{a}}_k^{\prime })$
     ELSE
      $\bar{y}\leftarrow \underset{j}{\min }{Q}_{{\bar{\varphi }}_j}(s^{\prime },{\hat{a}}^{\prime })$
     ENDIF
     $y=r+\gamma (1-d)\bar{y}$
     2.2 损失与梯度
     $L_{\text{critic}}=\sum _j(Q_{{\varphi }_j}(s,a)-y{)}^2$
     $g_{\varphi }\leftarrow {\nabla }_{\varphi }{L}_{\text{critic}}$
     $\varphi \leftarrow \text{Adam}(\varphi ,g_{\varphi },\text{clip}=gra{d}_{\text{norm}})$

3. 更新 Actor（策略）
      3.1 行为克隆损失（扩散噪声重建）
     $L_{\text{BC}}=\Vert ϵ-{ϵ}_{\theta }(a_t,t,s){\Vert }^2$
     3.2 价值提升项
     $\tilde{a}\leftarrow {\pi }_{\theta }(s)$
     随机选择 $(i,j)\in \{(1,2),(2,1)\}$
     $L_{\text{QL}}=-\frac{Q_{{\varphi }_i}(s,\tilde{a})}{|Q_{{\varphi }_j}(s,\tilde{a})|}$
     $L_{\text{actor}}=L_{\text{BC}}+\eta L_{\text{QL}}$
     $g_{\theta }\leftarrow {\nabla }_{\theta }{L}_{\text{actor}}$
     $\theta \leftarrow \text{Adam}(\theta ,g_{\theta },\text{clip}=gra{d}_{\text{norm}})$

4. 软更新目标 Critic
     $\bar{\varphi }\leftarrow \tau \varphi +(1-\tau )\bar{\varphi }$

5. EMA 更新 Actor (每 $update\_\text{ema\_every}$ 步且 $step\ge step\_\text{start\_ema}$)
     $\bar{\theta }\leftarrow \alpha \bar{\theta }+(1-\alpha )\theta$  # $\alpha =em{a}_{\text{decay}}$

6. 学习率调度 （可选）
     scheduler.step()

END FOR
──────────────────────────────────────────────────────────────

推断（采样动作）：
  给定状态 $s$
  重复 50 次生成候选动作 {ak∼πθ(s)}\{a_k\sim\pi_\theta(s)\}{ak​∼πθ​(s)}
  计算 $q_k=\min (Q_{{\bar{\varphi }}_1},Q_{{\bar{\varphi }}_2})(s,a_k)$
  按 $\text{softmax}(q_k)$ 采样 1 个动作作为输出

1. 代码结构概览

• Critic

  • 两个 Q 网络 $Q_{{\varphi }_1},Q_{{\varphi }_2}$，采用相同的全连接结构（激活函数为 Mish）。
  • 前向输入拼接状态 $s$ 与动作 $a$，输出对应的 Q 值。
  • 提供辅助函数 q1、q_min 方便单独取 $Q_1$ 或 $\min (Q_1,Q_2)$。

• Actor

  • 调用 agents.diffusion.Diffusion，本质是条件扩散模型：给定状态逐步去噪得到动作。
  • 模型骨干是 MLP（见 agents.model），输出噪声或直观动作表征。

• EMA-Actor

  • 用 agents.helpers.EMA 维护 $\bar{\theta }$：
    $$\bar{\theta }\leftarrow \alpha \bar{\theta }+(1-\alpha )\theta$$
  • 仅在 step ≥ step_start_ema 且满足步频 update_ema_every 时更新。

• 优化与调度

  • Actor、Critic 均用 Adam；可选余弦退火学习率 (CosineAnnealingLR)。
  • 梯度裁剪 $\Vert g{\Vert }_2\le \text{grad\_norm}$。

2. 数学推导

2.1 扩散策略

• 正向加噪
  $$q({\mathbf{a}}_t\mid {\mathbf{a}}_{t-1})=\mathcal{N}(\sqrt{1-{\beta }_t}{\mathbf{a}}_{t-1},{\beta }_t\mathbf{I})$$
• 反向去噪（策略）
  $$p_{\theta }({\mathbf{a}}_{t-1}\mid {\mathbf{a}}_t,s)=\mathcal{N}({\mu }_{\theta },{\sigma }_t^2\mathbf{I}),{\mu }_{\theta }=\frac{1}{\sqrt{1-{\beta }_t}}({\mathbf{a}}_t-{\beta }_t{ϵ}_{\theta }({\mathbf{a}}_t,t,s))$$
• 行为克隆损失
  $${\mathcal{L}}_{\text{BC}}={\mathbb{E}}_{(s,a_0)\sim \mathcal{D},t\sim \mathcal{U}[1,T],\epsilon \sim \mathcal{N}(0,\mathbf{I})}\Vert \epsilon -{ϵ}_{\theta }({\mathbf{a}}_t,t,s){\Vert }_2^2$$

2.2 双 Q 学习

• 目标动作
  ${\hat{a}}^{\prime }={\pi }_{\bar{\theta }}(s^{\prime })$（或 10×重复取最大 Q，取决于 max_q_backup）。
• TD 目标
  $$y=r+\gamma (1-d)\underset{i=1,2}{\min }{Q}_{{\bar{\varphi }}_i}(s^{\prime },{\hat{a}}^{\prime })$$
• Critic 损失
  $${\mathcal{L}}_{\text{critic}}=\sum _{i=1}^2{\mathbb{E}}_{(s,a)\sim \mathcal{D}}(Q_{{\varphi }_i}(s,a)-y{)}^2$$

2.3 策略改进项

• 生成新动作 $\tilde{a}={\pi }_{\theta }(s)$
• 随机交换索引 $(i,j)\in \{(1,2),(2,1)\}$，计算
  $${\mathcal{L}}_{\text{QL}}=-\frac{Q_{{\varphi }_i}(s,\tilde{a})}{|Q_{{\varphi }_j}(s,\tilde{a})|}$$
• 总 Actor 损失
  $${\mathcal{L}}_{\text{actor}}={\mathcal{L}}_{\text{BC}}+\eta {\mathcal{L}}_{\text{QL}}$$

2.4 软更新

$$\bar{\varphi }\leftarrow \tau \varphi +(1-\tau )\bar{\varphi }$$

3. 训练循环摘要

for each gradient step:
    1. 从经验池采样 (s,a,r,s',d)
    2. 更新 Critic ：最小化 L_critic
    3. 更新 Actor  ：最小化 L_actor
    4. 软更新 Critic_target
    5. EMA 更新 Actor_bar   (按频率与起始步数触发)
    6. 可选：学习率调度器步进

加噪过程(Diffusion模型损失函数)

加噪过程由 q_sample 函数实现，数学公式为： $$x_t=\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ$$ 其中，$x_0$ 是初始动作，$ϵ$ 是高斯噪声，${\bar{\alpha }}_t={\prod }_{i=1}^t{\alpha }_i$，${\alpha }_i=1-{\beta }_i$，${\beta }_i$ 是扩散系数。

损失计算
若 self.predict_epsilon 为 True，模型预测噪声，损失函数为： $$L=\text{Loss}(f(x_t,t,s),ϵ)\cdot w$$ 若 self.predict_epsilon 为 False，模型直接预测 $x_0$，损失函数为： $$L=\text{Loss}(f(x_t,t,s),x_0)\cdot w$$ 其中，$f(x_t,t,s)$ 是模型的输出，$\text{Loss}$ 是损失函数（如 L2 损失），$w$ 是权重。

loss 函数最终公式
综合以上步骤，loss 函数的数学公式为：
$$L=\frac{1}{B}\sum _{i=1}^B\{\begin{array}{ll}& \text{Loss}(f(x_{t_i},t_i,s_i),{ϵ}_i)\cdot w_i,\text{if}\text{self.predict\_epsilon}=\text{True}\\ & \text{Loss}(f(x_{t_i},t_i,s_i),x_{0_i})\cdot w_i,\text{if}\text{self.predict\_epsilon}=\text{False}\end{array}$$
其中，$B$ 是批量大小，$x_{t_i}$ 是第 $i$ 个样本在时间步 $t_i$ 的加噪动作，$s_i$ 是第 $i$ 个样本的状态，${ϵ}_i$ 是第 $i$ 个样本的高斯噪声，$x_{0_i}$ 是第 $i$ 个样本的初始动作，$w_i$ 是第 $i$ 个样本的权重。

参考文献翻译

相关工作：策略正则化（Policy Regularization）
离线强化学习中，大多数基于策略正则化的方法依赖行为克隆 (behavior cloning, BC) 来实现策略正则化：

• BCQ (Fujimoto 等人，2019) 将策略建模为在单独学习得到的条件变分自编码器 (Conditional-VAE，CVAE；Sohn 等人，2015) 行为克隆模型附近、且受到最大 Q 值约束的可训练偏移；
• BEAR (Kumar 等人，2019) 在策略改进步骤中加入加权的行为克隆损失，通过最小化 MMD 来实现；
• TD3+BC (Fujimoto & Gu，2021) 采用与 BEAR 相同的技巧，但使用最大似然估计 (MLE)；
• BRAC (Wu 等人，2019) 评估了多种行为克隆正则化方法，如 KL 散度、MMD 以及 Wasserstein 对偶形式；
• IQL (Kostrikov 等人，2021b) 是一种基于优势 (advantage) 加权的行为克隆方法，使用“样本内”学习得到的 Q 值；
• Goo & Niekum (2022) 强调在离线 RL 中显式执行行为克隆的必要性，而 Ajay 等人 (2022) 则肯定了条件生成模型在决策中的威力。

相关工作：扩散模型在强化学习中的应用

• Pearce 等人 (2023) 提出了利用具有强表达能力且稳定的扩散模型更好地模仿人类行为；
• Diffuser (Janner 等人，2022) 将扩散模型用作轨迹生成器，将完整的状态-动作序列视为单个样本，并单独训练一个回报模型来预测每条轨迹的累积奖励，然后在反向采样阶段注入该回报模型的引导。该方法与 Decision Transformer (Chen 等人，2021) 类似，后者也在 GPT-2 (Radford 等人，2019) 上训练轨迹生成器，并借助真实轨迹回报；
  • 在在线使用时，序列模型无法再自回归地由状态预测动作（因为状态是环境输出）。因此在评估阶段，需要为每个状态预测整条轨迹，但只执行第一步动作，计算开销较大。
• 我们的方法以完全不同的方式在 RL 中使用扩散模型：
  • 我们将扩散模型作用于动作空间，并将其构造成条件扩散模型，以状态作为条件；
  • 该方法是无模型 (model-free) 的，扩散模型一次仅采样一个动作；
  • Q 值函数的引导在训练期间注入，在我们的实验中带来了良好的性能。

虽然 Diffuser (Janner 等人，2022) 和我们的工作都在离线 RL 中应用扩散模型，但 Diffuser 立足于基于模型的轨迹规划视角，而我们的方法则属于离线、无模型的策略优化视角。

DIFFUSION Q-LEARNING

下面说明我们如何将条件扩散模型用作行为克隆的一种高表达力策略；随后介绍如何在训练阶段把 Q-learning 指导融入扩散模型的学习过程，并把行为克隆项视作策略正则化的一种形式。

3.1 扩散策略（Diffusion Policy）

符号约定： 本文中存在两类时间步——扩散过程的时间步与强化学习 (RL) 的轨迹时间步。我们使用上标 i∈{1,…,N}i\in\{1,\dots,N\}i∈{1,…,N} 表示扩散时间步，下标 t∈{1,…,T}t\in\{1,\dots,T\}t∈{1,…,T} 表示轨迹时间步。

策略表示

我们将 RL 策略表示为条件扩散模型反向过程的终点分布：

$${\pi }_{\theta }(a\mid s)=p_{\theta }(a_{0:N}\mid s)=\mathcal{N}(a_N;0,I)\prod _{i=1}^N{p}_{\theta }(a_{i-1}\mid a_i,s),$$

其中反向链的最终样本 $a_0$ 就是 RL 执行时采用的动作。一般而言，

$$p_{\theta }(a_{i-1}\mid a_i,s)=\mathcal{N}(a_{i-1};{\mu }_{\theta }(a_i,s,i),{\Sigma }_{\theta }(a_i,s,i)).$$

噪声预测参数化

按照 Ho 等人 (2020) 的做法，固定协方差为 ${\Sigma }_{\theta }(a_i,s,i)={\beta }_{i}I$，并将均值写成

$${\mu }_{\theta }(a_i,s,i)=\sqrt{\frac{1}{{\alpha }_i}}\left(a_i-\sqrt{\frac{{\beta }_i}{1-{\bar{\alpha }}_i}}{ϵ}_{\theta }(a_i,s,i)\right).$$

反向采样

1. 先从高斯先验采样 $a_N\sim \mathcal{N}(0,I)$；
2. 然后按下式逐步去噪，$i=N,\dots ,1$：

$$a_{i-1}|a_i=\sqrt{{\alpha }_i}{a}_i-\sqrt{\frac{{\beta }_i}{{\alpha }_i(1-{\bar{\alpha }}_i)}}{ϵ}_{\theta }(a_i,s,i)+\sqrt{{\beta }_i}ϵ,ϵ\sim \mathcal{N}(0,I).\text{\hspace{1em}(1)}$$

在 $i=1$ 时，遵循 DDPM (Ho 等人，2020) 的做法将 $ϵ$ 置零以提升采样质量。

训练目标

借鉴 Ho 等人 (2020) 的简化目标，训练条件 $ϵ$-模型的行为克隆损失为

$$L_d(\theta )={\mathbb{E}}_{\begin{array}{c}i\sim U,ϵ\sim \mathcal{N}(0,I),(s,a)\sim \mathcal{D}\end{array}}\Vert ϵ-{ϵ}_{\theta }(\sqrt{{\bar{\alpha }}_i}a+\sqrt{1-{\bar{\alpha }}_i}ϵ,s,i){\Vert }_2^2.\text{\hspace{1em}(2)}$$

该损失驱动 ${\pi }_{\theta }$ 在行为克隆意义上逼近离线数据分布，同时为后续加入 Q-learning 指导奠定基础。

其中 $U$ 表示离散集合 {1,…,N}\{1,\dots,N\}{1,…,N} 上的均匀分布，$D$ 表示离线数据集，由行为策略 ${\pi }_b$ 收集。该扩散模型损失 $L_d(\theta )$ 属于行为克隆损失，其目标是学习行为策略 ${\pi }_b(a\mid s)$（即从与训练数据同分布的动作中进行采样）。需要注意的是，反向扩散链的边缘分布提供了一种隐式而富有表现力的分布，能够捕捉离线数据集中表现出的偏态、 多峰等复杂特性。此外，该正则化是基于采样的，只需从 $D$ 和当前策略中随机取样即可（即不必显式获得行为策略；当数据集来源于人类演示时，行为策略往往无法确定）。与常见的“两步式”方法不同，我们的策略为在灵活策略上实施正则化提供了一种简洁且有效的途径。

$L_d(\theta )$ 可通过对每个数据点仅采样一次扩散步 $i$ 来高效优化，但式 (1) 中的反向采样需要迭代计算 ${ϵ}_{\theta }$ 网络 $N$ 次，这可能成为运行时间瓶颈。因此，我们倾向于将 $N$ 控制为较小值。为了在较小 $N$ 下工作，在 ${\beta }_{\min }=0.1$、${\beta }_{\max }=10.0$ 的设置下，我们遵循 Xiao 等人 (2021) 定义

$${\beta }_i=1-{\alpha }_i=1-e^{-{\beta }_{\min }\left(\frac{1}{N}\right)}-0.5({\beta }_{\max }-{\beta }_{\min })\frac{2i-1}{N^2},$$

该噪声调度取自 Song 等人 (2021) 的方差保持随机微分方程 (SDE) 框架。

3.2 Q-学习（Q-LEARNING）

行为正则化损失 $L_d(\theta )$ 本质上是一个行为克隆项，只能让策略复现生成训练数据的行为策略，而无法超越该行为策略。为了提升策略性能，我们在训练阶段向反向扩散链中注入 Q 值函数 的引导，使模型倾向于采样高价值动作。最终的策略学习目标是策略正则化与策略改进的线性组合：

$$\pi =\underset{{\pi }_{\theta }}{\mathrm{argmin}}L(\theta )=L_d(\theta )+L_q(\theta )=L_d(\theta )-\alpha \cdot {\mathbb{E}}_{s\sim \mathcal{D},a_0\sim {\pi }_{\theta }}[Q_{\varphi }(s,a_0)].\text{\hspace{1em}(3)}$$

其中 $a_0$ 通过公式 (1) 重新参数化，因此 Q 值对动作的梯度能够反向传播穿过整条扩散链。

由于不同离线数据集中的 Q 值量级可能差异较大，为了归一化，我们遵循 Fujimoto & Gu (2021) 的做法，将

$$\alpha =\frac{\eta }{{\mathbb{E}}_{(s,a)\sim \mathcal{D}}[|Q_{\varphi }(s,a)|]},$$

其中 $\eta$ 为平衡两项损失的超参数；分母中的 Q 仅用于归一化，不参与求导。

Q 值函数本身按照常规方法学习——在双 Q-learning 技巧 (Hasselt, 2010) 下最小化 Bellman 误差 (Lillicrap et al., 2015; Fujimoto et al., 2019)。我们构建了两组 Q 网络 $Q_{{\varphi }_1},Q_{{\varphi }_2}$ 以及对应的目标网络 $Q_{{\varphi }_1^{\prime }},Q_{{\varphi }_2^{\prime }}$，同时建立目标策略 ${\pi }_{{\theta }^{\prime }}$。优化 ϕi  (i∈{1,2})\phi_i\;(i\in\{1,2\})ϕi​(i∈{1,2}) 的目标为

$${\mathbb{E}}_{(s_t,a_t,s_{t+1})\sim \mathcal{D},a_{t+1}^0\sim {\pi }_{{\theta }^{\prime }}}[||(r(s_t,a_t)+\gamma \underset{i=1,2}{\min }{Q}_{{\varphi }_i^{\prime }}(s_{t+1},a_{t+1}^0))-Q_{{\varphi }_i}(s_t,a_t)|{|}^2].\text{\hspace{1em}(4)}$$

在第 4 节和第 5 节的大量实验中，我们验证了 $L_d$ 与 $L_q$ 的协同作用可获得最佳性能。实现细节见算法 1。