# Diffusion policy: Visuomotor policy learning via action diffusion 詳讀筆記
###### tags: `Imitation learning`,`visuomotor policy`, `manipulation`

>最初於 2023 年發佈在 arXiv，2024發表在 The International Journal of Robotics Research (IJRR)。引用次數1000+

https://arxiv.org/abs/2303.04137v5
https://github.com/real-stanford/diffusion_policy

## 廢話、注意事項
* 這篇論文似乎不適合入門(也可能我不適合研究，我好廢)
* 建議具備機器學習、深度學習、生成式AI基礎
* 建議先了解[DDPM、Diffusion Model](https://hackmd.io/@1I0zSdF_QSi-tPfJXnBEXA/r1fBrnrgZe)，本文艱澀詞彙太多
* 此筆記只是供我自己學習使用，有空會持續更新，還有很大的改進空間，可以去找其他簡中資源來讀

## Introduction
* 從示範中學習策略在最簡單的形式可表述為一個監督式迴歸任務(通過學習將觀測結果映射到動作)
* But 實踐中預測機器人動作的特性(多模態分佈(多峰性)、動作間的順序相關性(序列相關性)，以及對高精度的要求)使得這項任務與其他監督學習問題相比變得不同且更具挑戰

所以作者先用下面這張圖(其實是三張圖)介紹古人是怎麼解決的(a)(b)，以及自己提出來的方法

總之在策略生成或模仿學習有如下分類：
* Explicit Policy：直接觀測映射到動作，精度差，會維度爆炸，解決不了多峰性(拉基)
* Implicit Policy：不直接輸出具體動作，透過優化能量函數找出最佳動作，但解決不了負採樣問題，訓練不穩定(還是拉基)
* Diffusion Policy：在機器人動作空間上進行條件去噪擴散，在給定視覺觀測的條件下，推斷動作的得分（score）函數梯度，上面古人的問題都解決了，實在是太棒了！

## Diffusion Policy Formulation
作者將視覺機器人策略制定為去噪擴散概率模型（DDPM）。
Diffusion Policy能夠表達複雜的多模態動作分佈，同時在訓練中表現穩定
### Denoising Diffusion Probabilistic Models
DDPM是一種上古時期的經典擴散模型，下圖為原論文找的圖(建議先找DDPM來看)
https://arxiv.org/abs/2006.11239

* 由右至左為正(前)向過程，為馬爾可夫鏈，逐漸向數據中添加高斯噪聲，數據呈高斯分布。(Training)
* 由左至右為反向過程，為參數化的瑪爾可夫鏈，逐步從噪聲中恢復原始數據。(Sampling)

>論文本文簡要說明了DDPM的公式和作用，下章節會來比較此論文是怎麼應用的

DDPM基本思想是從高斯噪聲${x}^{k}$開始，不斷迭代去噪來逐步生成多階段的中間動作，噪聲水平逐漸降低，經過$K$次迭代後得到期望的無噪聲輸出$x^0$

$$\mathbf{x}^{k-1} = \alpha(\mathbf{x}^k - \gamma \varepsilon_\theta(\mathbf{x}^k, k)) + \mathcal{N}(0, \sigma^{2} I))$$

其中$\varepsilon_\theta$是噪聲預測網路（具有參數$\theta$，需要通過學習進行優化），${N}(0, \sigma^{2} I)$是每次迭代加入的高斯隨機噪聲，

$\alpha$、$\gamma$、$\sigma$則是與迭代步驟$k$相關的係數，論文中稱其為noise schedule(不確定中文怎麼命名)。分別控制去噪過程中的學習率（$\alpha$）、噪聲預測的影響（$\gamma$），所添加的高斯噪聲量（$\sigma$）。研究發現$\alpha$略小於1可以提升穩定性(去看論文，我也不是很清楚怎麼發現的)

上式也可理解為單步帶噪的梯度下降過程：
$$\mathbf{x}' = \mathbf{x} - \gamma \nabla E(\mathbf{x})$$
其中噪聲預測網路 $\varepsilon_{\theta}$ 可以有效的預測梯度場 $\nabla E(x)$；$\gamma$ 為學習率

### DDPM Training

訓練步驟如下：
1. 首先從資料集中隨機取出未加擾動的樣本$x^0$
2. 對於每個樣本，隨機選擇一個去噪迭代步數$k$
3. 根據該步數對應的noise schedule參數生成一個隨機噪聲$\varepsilon_k$
4. 接著讓噪聲預測網路$\varepsilon_{\theta}$對添加了噪聲$\varepsilon^k$的資料樣本$x^0 + \varepsilon^k$進行預測

訓練目標是讓網路輸出還原出原來加入的噪聲$\varepsilon^k$
其損失函數為：
$$\mathcal{Loss} = MSE(\varepsilon^{k}, \varepsilon_{\theta}( \mathbf{x}^{0} + \varepsilon^{k}, k))$$
其中$\mathsf{MSE}$表示均方誤差損失
最小化上述損失函數，等價於最小化資料分佈 $p(x^{0})$ 與由 DDPM $q(x^{0})$
所產生之樣本分佈之間的 KL 散度（KL divergence）的變分下界。

### Diffusion for Visuomotor Policy Learning
以前的DDPM多用於圖像生成，這篇論文開創性的使用DDPM去做機器人視覺運動策略生成。這需要在框架上做兩點主要修改：
1. 將擴散模型的輸出圖像$\mathbf{x}$，改為機器人動作$A_t$
2. 去噪過程加入視覺觀測$O_t$作為條件
>基本上公式與DDPM一模一樣，所以搞懂DDPM就能很輕鬆去看Diffusion Policy


講完公式回到圖片和本文做說明

作者用下面這張圖(論文是三張圖放一起擠在最上方)進行概述
可以看到input有圖像與機器人的robot pose，output在$A_t^k$時刻(也就是噪聲最大的時刻)生成了一堆動作點(全是可能動作點，類似於高斯分布)，在不斷迭代後生成比較密集的軌跡($A_t^3$)，最後生成了一條軌跡($A_t^0$)


#### Closed-loop action-sequence prediction
>簡單來說就是先輸入一段觀測，他會預測一大段動作，然後執行其中一個動作，以此類推、周而復始


閉迴路動作序列預測：一個有效的動作生成形式應當在長視野規劃中鼓勵動作的時間連貫和平滑，同時又允許面對非預期觀測時的即時反應。為此，作者將擴散模型產生動作序列的能力與滾動視界控制相結合，以實現穩健的動作執行。
具體而言，在時間步長$t$，策略將最近$T_o$步的觀測資料$O_t$作為輸入，預測未來$T_p$步的動作序列$A_t$，其中僅取前$T_a$步動作不重新規劃直接在機器人上執行。

$T_o$定義為觀測視界長度，$T_p$為動作預測視界長度，$T_a$為動作執行視界長度
透過這種設計，在保持動作時間一致性的同時保留了對新觀測的響應能力，並在長視野規劃與實時響應之間取得了平衡
關於參數$T_a$對性能的影響後面章節會討論

#### Visual observation conditioning
>觀測數據是條件分布而不是聯合分布

視覺觀測條件：他們使用DDPM來近似條件分佈$p(A_t \mid O_t)$，而非其他論文使用的聯合分佈$p(A_t, O_t)$，這種改動允許模型在給定觀測的條件下預測動作，而無需推斷未來狀態，從而加快擴散過程並提高生成動作的準確性。(接著論文寫了剛才提到公式修改的地方，講過了所以不再說明)

## Key Design Decisions
本節描述Diffusion Policy中幾個關鍵的設計決策，以及作者在實作時為$\varepsilon_{\theta}$（噪聲預測網路）所選用的具體神經網路架構。
### Network Architecture Options
作者選擇了兩種常見的架構類型：卷積神經網路（CNN）和Transformer，並比較了兩者之間的性能和特性。
須注意一點，噪聲預測網路$\varepsilon_{\theta}$的選擇與視覺編碼器(Visual Encoder)相互獨立
>總結一下兩者架構優缺點
>CNN-based:超參數較為穩定，但高頻變化表現差
>Transformer-based:複雜任務和高頻動作表現很好，但很難訓練
>作者建議用他們的東西先用CNN開始實現，再來才考慮Transformer

#### CNN-based Diffusion Policy
CNN部分作者使用了一維時間卷積網路(1D temporal CNN)，並做了以下修改：
1. 只對條件分佈$p(A_t \mid O_t)$建模，透過FiLM（特徵線性調製）機制將觀測特徵和擴散迭代步數$k$作為條件融入動作生成過程
2. 網路僅預測未來的動作序列$A_t$
3. 移除了原方法中基於圖像修補（inpainting）的目標狀態調節，因為本論文採用的滾動視界預測框架不相容該方法

在實驗中，他們發現基於CNN的主架構在大多數任務上開箱即用即可獲得良好效果，幾乎不需要大量的超參數調整。然而，當期望的動作序列隨時間快速且劇烈變化時（例如速度指令類的動作空間），CNN主架構表現不佳。
推測可能是由於時間卷積的歸納偏差傾向於低頻訊號，也就是CNN更擅長建模平滑變化，對需要高頻變化的情況適應較差。

#### Time-series diffusion transformer
為減輕CNN模型的過度平滑效應，作者提出了一種新穎的Transformer式DDPM。該模型採用minGPT的Transformer架構來預測動作序列。
加入噪聲的動作序列$A_t^k$作為Transformer解碼器的輸入token，各動作的嵌入會加上擴散迭代步數$k$對應的正弦位置編碼。
觀測$O_t$經共享的多層感知機(MLP)轉換為觀測嵌入序列，作為解碼器中交叉注意力層的輸入特徵。
在Transformer的解碼堆疊中，每個輸出token（對應每個動作）都會預測“梯度”值$\varepsilon_{\theta}(O_t, A^k_t, k)$
在純狀態輸入實驗中，大多數性能最好的策略都來自Transformer架構，特別是當任務的複雜度和動作變化速率較高時。然而，他們發現Transformer對超參數更加敏感，訓練難度較高，這並非Diffusion Policy所獨有的問題，而是Transformer訓練本身的常見挑戰。
未來透過改進Transformer訓練技術或增大數據規模，可能緩解這一問題。

### Visual Encoder
視覺編碼器將原始圖像序列映射為潛在嵌入$O_t$，並與擴散策略一起端到端訓練。當有多個攝影機視角時，我們為每個視角使用單獨的編碼器，對每一步的圖像獨立編碼，然後將每步的圖像特徵拼接形成$O_t$。
作者採用經典的ResNet-18（無預訓練）作為編碼器骨幹，並進行了以下修改：
1. 用空間softmax池化(spatial softmax pooling)替代了全局平均池化(global average pooling)，以保留圖像的空間資訊
2. 用群組歸一化（GroupNorm）取代批次歸一化（BatchNorm），以獲得更穩定的訓練。

這些改進在擴散模型中尤為重要：例如在DDPM中常使用指數移動平均（EMA）來提升穩定性，此時BatchNorm可能帶來不利影響，因此我們改用不依賴批量統計的GroupNorm。

### Noise Schedule
>作者試了DDPM方法不好用，所以用iDDPM的方法實現，效果絕佳

由$\sigma$、$\alpha$、$\gamma$及每步附加的高斯噪聲$\varepsilon_k$隨迭代步$k$的函數所構成。
底層的噪聲時程決定了Diffusion Policy在多大程度上捕捉動作訊號的高頻和低頻特性。在控制類任務中，作者發現平方cosine時程（square cosine schedule）在DDPM框架下效果最佳(iDDPM)，在所有任務上均取得了最佳效果。

### Accelerating Inference for Real-time Control
>用DDIM的方法來實現正向過程和反向過程解耦，3080上迭代100次只需要0.1s

由於作者將擴散過程直接用作機器人控制策略，因此確保閉迴路控制下的推理速度足夠快至關重要。為加速推理，我們採用了去噪擴散隱式模型（DDIM）。DDIM的方法在於將訓練和推理所需的擴散迭代步數解耦，許在推理時使用更少的迭代步數來加快生成。在我們的真實實驗中，使用DDIM將訓練時的100步迭代減少為推理時的10步，實現了0.1秒級的推理延遲（基於Nvidia 3080 GPU）。透過DDIM大幅降低了推理的擴散步數，使得Diffusion Policy能滿足實時控制的需求。

## Intriguing Properties of Diffusion Policy
主要是在講Diffusion Policy的一些性質，以及與其他類型策略相比的優勢

### Model Multi-Modal Action Distributions
>動作多峰性(或者說動作多模態行為)：Diffusion model的底層隨機採樣過程和隨機初始化會導致動作出現多峰化。擴散過程引入隨機擾動對動作進行優化，讓生成結果可以聚合到不同動作模態上。(這裡的多模態是指一個任務可以有多種解決方案)


### Synergy with Position Control
>為什麼是動作控制而不是速度控制：位置控制有更好的多模態性質、速度控制更容易被複合誤差影響、位置控制能更好的預測動作序列


### Benefits of Action-Sequence Prediction
>說明了Diffusion Policy下動作序列預測的優點，一致性、平穩性、穩健性更強

### Training Stability
>理論上來說IBC應該達到一樣的效果，但因為隱式策略的動作分布依賴於基於能量模型EBM，要做負採樣，但負採樣準確性依賴於採樣品質，若樣本無法充分覆蓋動作空間，會導致歸一化係數的估計偏差，進而導致訓練不穩定
>Diffusion policy:老子直接建模求導，就穩定了

$$p_{\theta}(a \mid o) = \dfrac{e^{-E_{\theta}(o, a)}}{Z(o, \theta)}$$

$$Loss_{\text{infoNCE}} = -\log\!\left(\dfrac{e^{-E_{\theta}(o, a)}}{e^{-E_{\theta}(o, a)} + \sum_{j=1}^{N^{neg}} e^{-E_{\theta}(o, \tilde{a}^{j})}}\right)$$

$$\nabla_{a}\log p(a \mid o) = -\nabla_{a}E_{\theta}(a, o) - \underbrace{\nabla_{a}\log Z(o, \theta)}_{=0} \approx -\varepsilon_{\theta}(a, o)$$


## Evaluation、Realworld Evaluation
>第五節，作者在4個基準15個任務進行全面評估，全面超越、遙遙領先

>第六節，真實世界也是無機能敵

實驗很多很有趣，可以去看看

## 補充

### 平方餘弦噪聲排程（Square Cosine Schedule）

此方法來源於論文 《Improved Denoising Diffusion Probabilistic Models》。其核心目標是在取樣（sampling）過程的中段加入較多噪聲，而在起始與結尾階段加入較少噪聲。
傳統 DDPM 採用的是線性噪聲排程（linear noise schedule），即噪聲變異數會隨步數逐漸增加。

引入平方餘弦噪聲排程後，可在維持生成品質幾乎不變的情況下，大幅減少約一半的取樣步數，因而提升生成效率。

### Denoising Diffusion Implicit Models（DDIM）

此方法來自論文 《Denoising Diffusion Implicit Models》。DDIM 的主要貢獻在於提供了一種加速擴散模型生成過程的方法，且幾乎不會犧牲影像品質。

DDIM 的想法是將原本的擴散過程重新定義為非馬可夫過程（non-Markovian process），在不更動訓練過程的前提下，使得模型能以更少的推論步數生成高品質影像，達到明顯的加速效果。