HWC
← Note

No date

Diffusion-QL | Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning

PaperRLOffline RLDiffusion

快速解釋

這篇論文提出 Diffusion-QL，核心主張是：offline RL 不只需要「保守」，還需要一個足夠有表達能力的 policy class。作者認為很多 policy-regularized offline RL 方法效果受限，不是因為 regularization 這個方向錯了，而是因為它們常用的 Gaussian 或簡化 generative policy 太弱，無法準確表示真實的 behavior policy，特別是多峰、多策略混合的人類或 exploratory dataset。

Diffusion-QL 的做法是把 policy 直接建模成一個 conditional diffusion model，讓 policy 本身就能表示複雜的 action distribution；接著再把 Q-learning 的目標直接加進 diffusion policy 的訓練 loss，讓 policy 一方面維持在資料分佈附近，一方面往高價值動作的方向偏移。這不是先做 behavior cloning、再額外做 policy constraint 的兩階段流程，而是在同一個生成模型裡同時完成 regularization 與 policy improvement。

問題設定

Offline reinforcement learning 的設定是：agent 只能從固定資料集學習，而不能再與環境互動。令資料集記為：

D={(s,a,r,s′)}.\mathcal{D} = \{(s,a,r,s')\}.

目標仍然是學得一個 policy，使其在原始 MDP 中的期望折扣回報最大。但 offline RL 的核心困難在於：只要 policy improvement 把 actor 推到資料集外的 action 區域，critic 就得對 out-of-distribution action 做外插估計，這些估計往往不可靠，進而讓 policy 偏好被高估的動作。

過去不少 offline RL 方法會透過 policy regularization 來限制 actor 不要離 behavior policy 太遠，例如 BCQ、BEAR、BRAC、TD3+BC、IQL 等。但這篇論文指出，這類方法常有兩個更深層的問題：

  • policy class 不夠 expressive：若真實 behavior policy 是多峰的、偏態的，或不同 action 維度之間有相依結構，對角 Gaussian 很難正確擬合。
  • regularization 是兩階段且有誤差傳遞：先學一個 cloned behavior policy，再把 current policy regularize 到它附近；如果第一步沒學好，第二步的 regularization 反而會誤導 policy improvement。

作者用一個 2D bandit 例子說明這件事：若資料分佈有多個 mode，Gaussian policy 往往會在 mode 中間的低密度區域放上較高機率，這在 offline RL 裡非常危險，因為 regularization 會把 policy 拉向其實不屬於 dataset support 的位置。

核心想法

1. 用 conditional diffusion model 當 policy

Diffusion-QL 不再用 Gaussian actor 或 CVAE actor，而是把 policy 表示成一條條件式 reverse diffusion chain。給定 state 之後，模型從高斯雜訊開始，逐步去噪，最後得到 action。這種 implicit policy distribution 的優點是：

  • 可以表示多峰 action distribution
  • 不需要顯式 density 才能運作
  • 可以自然對應到 behavior cloning
  • 在 human demonstrations 或多策略混合資料中，能比 Gaussian 類 policy 更忠實地保留資料結構

2. 把 behavior cloning 與 policy improvement 耦合在同一個 loss

Diffusion policy 的基本 loss 是標準 denoising loss，本質上就是一個 behavior cloning objective：它讓 policy 學會從資料集中的 state-action distribution 生成合理動作。

但單靠 cloning 不可能超越資料來源 policy，所以作者再加入 Q-learning 項，直接鼓勵 diffusion policy 生成高價值動作。最重要的是，這個 Q 項不是事後拿來重排序或額外約束，而是直接加在 diffusion policy 的訓練目標裡，讓 policy 在學習「像資料」與「高價值」時同步進行。

這使得 Diffusion-QL 和 BCQ / BEAR 這類方法的差別很明顯：

  • BCQ / BEAR 先學 behavior model，再對 policy improvement 做限制
  • Diffusion-QL 則讓 policy 自己就是 expressive behavior model，並在同一個訓練過程中接受 Q-guidance

3. 用 diffusion policy 的表達能力，改善 regularization 的品質

這篇論文真正想證明的不是 diffusion 一定比別的生成模型神奇，而是 policy regularization 的效果高度依賴 behavior policy 的建模品質。如果 policy class 本身不夠 expressive，那 regularization 就會把 policy 拉到錯的地方；反之，若 policy 能準確覆蓋 dataset support，Q-learning 才有機會在一個可信的局部區域內進行 improvement。

因此 Diffusion-QL 的成功，來自兩個元件的配合：

  • diffusion model 提供高表達能力的 behavior regularization
  • Q-learning guidance 把這個分佈往高價值區域推

作者的 ablation 也支持這件事：BC-Diffusion 比 BC-CVAE 更能擬合行為資料；而把 diffusion policy 放進同樣的 policy improvement 框架後，也比 CVAE-based 版本穩定得多。

4. 用小步數 diffusion 換取可行的訓練成本

Diffusion model 的缺點是 sampling 需要多步反覆去噪，計算成本會隨 diffusion steps 增加。作者因此把步數 N 視為 expressiveness 與效率的折衷：

  • 大的 N 更能精細擬合資料分佈
  • 小的 N 計算更便宜、訓練更快

實驗顯示在 D4RL 上，N = 5 已足以提供很好的效果，讓 Diffusion-QL 在效能與成本之間取得合理平衡。

關鍵公式

1. Diffusion policy 的表示

給定 state 之後，policy 被表示成一個條件式 reverse diffusion chain：

πθ(a∣s)=pθ(a0:N∣s)=p(aN)∏i=1Npθ(ai−1∣ai,s),\pi_\theta(a \mid s) = p_\theta(a_{0:N} \mid s) = p(a_N) \prod_{i=1}^{N} p_\theta(a_{i-1} \mid a_i, s),

其中最終得到的 denoised action 就是實際輸出的動作，而最初的 latent sample 來自標準高斯噪聲。

每一步 reverse kernel 通常建模為：

pθ(ai−1∣ai,s)=N(ai−1;μθ(ai,s,i),βiI).p_\theta(a_{i-1} \mid a_i, s) = \mathcal{N}\bigl(a_{i-1}; \mu_\theta(a_i,s,i), \beta_i I\bigr).

這裡會用到 diffusion schedule 的標準記號，例如 step-wise 的 alpha、beta 與 cumulative alpha-bar；請注意，它們和後面 Q-guidance 裡的 loss 權重不是同一個量。

2. Reverse denoising update

作者使用 noise prediction parameterization，令：

μθ(ai,s,i)=1αi(ai−βi1−αˉi ϵθ(ai,s,i)).\mu_\theta(a_i,s,i) = \frac{1}{\sqrt{\alpha_i}} \left( a_i - \frac{\beta_i}{\sqrt{1-\bar\alpha_i}}\, \epsilon_\theta(a_i,s,i) \right).

因此在 sampling 時，每一步的更新可寫成：

ai−1=1αi(ai−βi1−αˉi ϵθ(ai,s,i))+βi ϵ,ϵ∼N(0,I).a_{i-1} = \frac{1}{\sqrt{\alpha_i}} \left( a_i - \frac{\beta_i}{\sqrt{1-\bar\alpha_i}}\, \epsilon_\theta(a_i,s,i) \right) + \sqrt{\beta_i}\,\epsilon, \quad \epsilon \sim \mathcal{N}(0,I).

也就是說，policy 是從噪聲逐步去噪，最後產生動作，而不是像一般 actor 一次前向傳播就直接輸出 action。

3. Diffusion behavior cloning loss

Diffusion policy 的基本訓練目標是 denoising loss：

Ld(θ)=Ei∼U, ϵ∼N(0,I), (s,a)∼D[∥ϵ−ϵθ(αˉia+1−αˉi ϵ,s,i)∥22].L_d(\theta) = \mathbb{E}_{i \sim U,\, \epsilon \sim \mathcal{N}(0,I),\, (s,a) \sim \mathcal{D}} \left[ \left\| \epsilon - \epsilon_\theta\bigl(\sqrt{\bar\alpha_i}a + \sqrt{1-\bar\alpha_i}\,\epsilon, s, i\bigr) \right\|_2^2 \right].

這個 loss 在 offline RL 裡扮演的角色其實就是 behavior cloning：它要求 diffusion policy 去重建資料集中的 action distribution，因此會自然把 policy 壓在 dataset support 附近。

4. Q-learning guidance 與最終 policy objective

只做 behavior cloning 不足以超越行為資料，所以作者加入 Q-learning guidance，得到最終的 policy objective：

L(θ)=Ld(θ)+Lq(θ)=Ld(θ)−α Es∼D, a0∼πθ(⋅∣s)[Qϕ(s,a0)].L(\theta) = L_d(\theta) + L_q(\theta) = L_d(\theta) - \alpha\, \mathbb{E}_{s \sim \mathcal{D},\, a_0 \sim \pi_\theta(\cdot\mid s)} \bigl[Q_\phi(s,a_0)\bigr].

這裡的第二項會鼓勵 diffusion policy 生成更高 Q 的 action。由於最終 action 是透過整條 reverse diffusion chain 取樣得到，所以 Q 的梯度可以沿著整個 denoising chain 反向傳播回 diffusion policy。

為了避免不同資料集上的 Q-value 尺度差異太大，作者另外用下式對這個權重做 normalization：

α=ηE(s,a)∼D[∣Qϕ(s,a)∣].\alpha = \frac{\eta}{\mathbb{E}_{(s,a) \sim \mathcal{D}}\bigl[|Q_\phi(s,a)|\bigr]}.

其中 eta 是真正需要調的超參數，而分母只用來做尺度正規化，不參與梯度更新。

5. Critic 的 Bellman loss

Q-function 本身仍然是標準 actor-critic 架構中的 twin critics，使用 double Q trick 與 target networks。對每個 critic，可寫成：

LQ(ϕi)=E(st,at,rt,st+1)∼D, a′∼πθ′(⋅∣st+1)[(rt+γmin⁡j∈{1,2}Qϕj′(st+1,a′)−Qϕi(st,at))2].L_Q(\phi_i) = \mathbb{E}_{(s_t,a_t,r_t,s_{t+1}) \sim \mathcal{D},\, a' \sim \pi_{\theta'}(\cdot\mid s_{t+1})} \left[ \Bigl( r_t + \gamma \min_{j \in \{1,2\}} Q_{\phi'_j}(s_{t+1}, a') - Q_{\phi_i}(s_t,a_t) \Bigr)^2 \right].

也就是說，Diffusion-QL 並不是用 diffusion model 取代 critic，而是把 diffusion model 放在 actor 端，critic 仍然維持一般離策略 Q-learning 的形式。

模型結構

Diffusion-QL 的整體架構可以拆成三個部分：

1. Conditional diffusion policy

核心 policy network 是一個條件式 noise predictor。輸入包含：

  • 當前 state
  • 某一步的 noisy action
  • diffusion timestep

輸出則是該步需要移除的 noise。透過這個網路反覆執行 reverse denoising，就能從高斯噪聲逐步生成最終 action。

2. Twin critics

和 TD3、SAC 類的 actor-critic 方法一樣，Diffusion-QL 使用兩個 Q-networks 來降低 overestimation bias，並搭配 target networks 提供穩定的 bootstrap target。

3. 聯合訓練流程

整體訓練流程可以概括成：

  1. 從固定資料集取出 mini-batch
  2. 用 target diffusion policy 在 next state 上取樣 action
  3. 用 Bellman loss 更新 twin critics
  4. 用當前 diffusion policy 取樣 action
  5. 用 diffusion BC loss + Q-guidance loss 聯合更新 policy
  6. 用 Polyak averaging 更新 target networks

從方法定位上看，Diffusion-QL 介於純 behavior cloning 與傳統 policy-constraint offline RL 之間：

  • 它不是只模仿資料，因為有明確的 Q-learning guidance
  • 它也不是單純靠顯式 divergence constraint 限制 policy，而是讓 expressive diffusion policy 自然提供高品質的 in-distribution regularization

這也是它能在 D4RL 的 Gym、AntMaze、Adroit、Kitchen 四類任務上都維持強表現的主要原因。