[論文解説] Diffusion Policy

Diffusion Policyについて解説します。

元論文
Diffusion Policy: Visuomotor Policy Learning via Action Diffusion

公式のサイト
https://diffusion-policy.cs.columbia.edu/

公式のgithub
https://github.com/real-stanford/diffusion_policy

概要

ロボット行動方策を「条件付きノイズ除去拡散過程 (Conditional Denoising Diffusion Process)」として定式化したDiffusion Policyを提案している。これにより、行動分布を拡散モデルで表現し、マルチモーダル（多峰性）な行動分布を安定的に学習できる。

先行研究との違いは

模倣学習のアルゴリズムとはして、
BC(Behavior Cloning): 目標行動への直接回帰として扱う。マルチモーダルな行動分布を扱うことは困難。下図(a)
IBC(Implicit Behavior Cloning): エネルギーベースモデルとして方策を学習する。マルチモーダルな行動分布を表現可能であるが、正規化定数推定の難しさや負例サンプリングによる学習の不安定さがある。下図(b)

提案手法のDiffusion Policyは、Diffusionを方策学習に適用することで、IBCにおける不安定要素を回避しながらマルチモーダルで高次元な行動分布を安定して学習できる点が特徴となる。下図(c)

＊ ここでいうマルチモーダルとはVLMのような異なる種類の入力データという意味ではなく、異なる行動（左から掴む、右から掴む）というモードの違いだと理解している。

技術のキモは？

拡散過程を用いた行動生成

• ノイズ予測ネットワーク$\epsilon_{\theta}$は、データ上の行動分布に付随するノイズ（環境変動や観測誤差などによる外乱成分）を逐次推定・除去する役割を担う。
• 方策学習を拡散過程として定式化し、行動に付与されたガウスノイズを段階的に除去する手順をK回繰り返すことで、モデルは外乱や不確実性を考慮しつつ最終的に安定かつロバストな行動を出力可能となる。
• この過程で得られる行動分布は、初期状態を確率的にサンプリングし、ノイズ除去を反復的に行うことで複数の安定点（モード）へと収束しうるため、多峰性を有するマルチモーダルな特性を同時に学習できる。

時刻$t$におけるロボットの行動は以下の式を$A_{t}^{K}, A_{t}^{K-1}, ... , A_{t}^{0}$まで繰り返すことでノイズが少ない$A_{t}^{0}$が得られる。

$$A_{t}^{k-1} = \alpha\bigl(A_{t}^{k} - \gamma\epsilon_{\theta}(O_t, A_{t}^{k}, k) + \mathcal{N}(0, \sigma^{2}I)\bigr)$$

$\mathcal{N}(0, \sigma^{2}I))$はガウシアンノイズ
$α, γ, \sigma$はパラメータである
$\epsilon_{\theta}$はノイズ予測ネットワークで、外乱$\mathcal{N}(0, \sigma^{2}I))$を予測していることになる。
$\epsilon_{\theta}$は以下のように求める。

$$\mathcal{L} = \mathrm{MSE}\bigl(\epsilon^k, \epsilon_{\theta}(O_t, A_{t}^{0} + \epsilon^k, k)\bigr)$$

後退ホライズン制御

• マルチステップ先まで行動をまとめて生成（後退ホライズン制御の形で出力）することで、行動が時間的に一貫性を持ち、ロボットが安定した操作を継続できる。
• 一時的に停止するようなアイドルアクションがある。これらの行動の予測は容易であるため、BCなどではオーバーフィットを起こす。これに対して、行動シーケンス全体を考慮する後退ホライズン制御では、長期的整合性が求められるため、そのような不適切なオーバーフィットが緩和される。

• 後退ホライズン制御とは、

時刻kにおいて予測ホライズン先までの時刻区間(k+H_p)で最適化を行い， 次の時刻k+1では(k+H_p+1)までの時刻区間で最適化を行う． このように， 現時刻から一定時刻先までの最適化を各時刻で繰り返していくことを 後退ホライズン（receding horizon）方策と呼ぶ．

結果

複数のロボット制御タスク（2DoF～6DoFアクション）において、ベースライン手法(LSTM-GMM, BET, IBCなど)を上回る性能が示された。

LSTM-GMMは、行動系列をモデル化する際に、LSTM（Long Short-Term Memory）とGMM（Gaussian Mixture Model）を組み合わせた手法です。行動生成をガウス混合分布で表すことで、単純な平均回帰よりも多様な行動分布を表現することができる。
IBC(Implicit Behavior Cloning): エネルギーベースモデルとして方策を学習する。マルチモーダルな行動分布を表現可能であるが、正規化定数推定の難しさや負例サンプリングによる学習の不安定さがある。
BET(Behavior transformers: Cloning $k$ modes with one stone): Transformerを用いたBehavior Cloning。

ph: 熟練者のデモデータ
mh: 非熟練者のデモデータ
記号入力による結果

画像入力による結果

まとめ

Diffusion Policyは拡散過程を用いた方策学習によって不安定性を回避し、マルチモーダルで高次元な行動分布を安定的に表現・学習する手法である。後退ホライズン制御と組み合わせることで長期的な時間的一貫性を保ち、単純な回帰方策では困難なタスクにも柔軟かつロバストに対処できることを示した。

論文に書かれた知見

• 事前学習済みモデルをfreezeした状態では精度が上がらなかった。ResnetとViTを試した。
• vision encoderを学習率を低くして学習すると精度が上がる。
• CNNとTransformerアーキテクチャの両方を試したが、CNNのほうが良い？表4だけ違う気がする
• ノイズスケジューラーはQuare Cosine Schedulerを使う。
• カメラの前で手を振る、ブロックを動かしている最中に邪魔をするなどノイズを入れてもロバストだった。