はじめに
CoRL2025より以下の論文のまとめ
[1] A. L. Chandra, et. al. "DiWA: Diffusion Policy Adaptation with World Models". CoRL2025.
概要
- ロボットアームなどを精度よく学習させる手法
- DPPOをworld model で学習させたもの
手法
以下では
- 方策勾配による強化学習
- log-derivative-trickによる勾配の変形
- REINFORCE型に変形
- PPO
- DPPO
- WiWA（本手法）
という順でまとめる。
１）方策勾配による強化学習
上図左のようにロボットアームが青い缶をピッキングする状況を考える。強化学習によってロボットアームを賢くし、ちゃんと缶をピッキングできるようにしたい。
側においたカメラから取得した画像を状態 $S_t$ とし、それに対する行動を $a_t$ とする。 行動 $a_t$ は上図右のように確率的に与えられるような状況を考える。
上図左のようには状態 $S_t$ を何ニューラルネットのようなものに入れて、行動 $a_t$ を出力させる場合を考える。このニューラルネットをどんどん賢くすると、その状況に応じた行動を出力するだろう。それに沿って動けば、上図右のように最終的にはすばらしい軌跡 $\tau$ を辿って、ピッキングしてくれるだろう。
この状態 $S_t$ でこの行動 $a_t$ を出す方策を $\pi_{\theta}$ とする。
ある方策 $\pi_{\theta}$ に従って軌道 $\tau$ を描き、上図右のように t 時に缶をピッキングできたら報酬 $r_t$ が与えられる。
そうするとこの報酬の総和の割引現在価値（割引報酬和）を最大化させたい。
\sum_{t=0}^T \gamma^t \cdot r_t \rightarrow max
ただ、行動は確率的に決まるので期待値で考えると
J(\theta) = \mathbb{E}_{\tau \sim p_{\theta}(\tau)} \left[\sum_{t=0}^T \gamma^t \cdot r_t \right] \rightarrow max
と最大化させたい。
上図左のような $J(\theta)$ がmaxとなる $\theta$ の組み合わせを取得したい。よって上図右のように勾配を求めてちょっとずつ頂上に近づくことを考える。
\theta \leftarrow \theta + \alpha \cdot \nabla_\theta J(\theta)
ただ目的関数の勾配は
\nabla_\theta J(\theta) = \nabla_\theta \left( \mathbb{E}_{\tau \sim p_{\theta}(\tau)} \left[\sum_{t=0}^T \gamma^t \cdot r_t \right] \right)
だが、これは何？
方策 $\pi_\theta$ に従って軌道 $\tau$ を進み、そこで得られる割引報酬和の期待値を微分している。これを求めるのは簡単でないので、普段我々がやっているニューラルネット部分（$\pi_{\theta}$）だけを微分する形にしたい。
２）log-derivative-trickによる勾配の変形
\begin{eqnarray}
\nabla_\theta J(\theta) &=& \nabla_\theta \left( \mathbb{E}_{\tau \sim p_{\theta}(\tau)} \left[\sum_{t=0}^T \gamma^t \cdot r_t \right] \right) \\
&=& \nabla_\theta \int p_{\theta}(\tau) R (\tau) d\tau, \ \because R(\tau) = \left[\sum_{t=0}^T \gamma^t \cdot r_t \right] \\
&=& \int \nabla_\theta p_{\theta}(\tau) R (\tau) d\tau
\end{eqnarray}
ここで以下の log-derivative-trick を使う。
\nabla_\theta p_{\theta}(\tau) = p_{\theta} \nabla_\theta \log p_{\theta} (\tau)
これは合成関数の対数関数を微分するところからきてる。
\begin{eqnarray}
&& \nabla_\theta (\tau) \log p_{\theta}(\tau) = \frac{1}{p_{\theta}(\tau)} \nabla_\theta p_{\theta}(\tau) \\
&\iff& p_{\theta}(\tau) \nabla_\theta (\tau) \log p_{\theta}(\tau) = \nabla_\theta p_{\theta}(\tau)
\end{eqnarray}
よって勾配は
\begin{eqnarray}
\nabla_\theta J(\theta) &=& \int \nabla_\theta p_{\theta}(\tau) R (\tau) d\tau \\
&=& \int p_{\theta}(\tau) \nabla_\theta (\tau) \log p_{\theta}(\tau) R(\tau) d\tau \\
&=& \mathbb{E}_{\tau \sim p_{\theta}(\tau)} \left[ \nabla_\theta (\tau) \log p_{\theta}(\tau) R(\tau) \right]
\end{eqnarray}
と変形できる。
３）REINFORCE型に変形
上図左のように行動 $a_t$ は方策 $\pi(a_t | s_t)$ によって確率的に決まるが、環境によって次の状態 $s_{t+1}$ も確率的に変わると考える。（上図右）
そうすると軌道確率 $p_\theta (\tau)$ は
p_\theta (\tau) = p(s_0) \Pi^{T-1}_{t=0} \pi(a_t | s_t) \cdot p(s_{t+1} | s_t, a_t)
となる。対数をとると
\log p_\theta (\tau) = \log p(s_0) + \sum^{T-1}_{t=0} \log \pi_0 (a_t | s_t) + \sum^{T-1}_{t=0} \log p (s_{t+1} | s_t, a_t)
なので $\theta$ で微分すると、$\theta$ と関係の無い１項目と３項目が消えて
\nabla_{\theta} \log p_\theta (\tau) = \sum^{T-1}_{t=0} \nabla_{\theta} \log \pi_0 (a_t | s_t)
となる。よって勾配は
\begin{eqnarray}
\nabla_\theta J(\theta)
&=& \mathbb{E}_{\tau \sim p_{\theta}(\tau)} \left[ \nabla_\theta (\tau) \log p_{\theta}(\tau) R(\tau) \right] \\
&=& \mathbb{E}_{\tau \sim p_{\theta}(\tau)} \left[ \left( \sum^{T-1}_{t=0} \nabla_{\theta} \log \pi_0 (a_t | s_t) \right) R(\tau) \right] \\
&=& \mathbb{E}_{\tau \sim p_{\theta}(\tau)} \left[ \sum^{T-1}_{t=0} \nabla_{\theta} \log \pi_0 (a_t | s_t) G(t) \right] \tag{3.1} \\
\end{eqnarray}
こうすると、この
方策 $\pi_\theta$ がニューラルネットなど微分可能であれば、(3.1)式の $\log \pi_\theta$ も微分でき、勾配が求まって学習できる。
\theta \leftarrow \theta + \alpha \cdot \sum^{T-1}_{t=0} \nabla_{\theta} \log \pi_0 (a_t | s_t) G(t)
PPO
(3.1)式は分散が大きく学習が不安定。なので価値関数をベースラインとして利用する。
\begin{eqnarray}
\nabla_\theta J(\theta)
&=& \mathbb{E}_{\tau \sim p_{\theta}(\tau)} \left[ \sum^{T-1}_{t=0} \nabla_{\theta} \log \pi_0 (a_t | s_t) \left( G(t) - V(s_t) \right) \right] \tag{4.1} \\
\end{eqnarray}
価値関数はactor-critic的に同時に学習させる。
方策勾配は「良かった行動の確率を増やす」ので、学習率やデータのばらつき次第で 方策が一気に変わりすぎて性能が崩壊しがち、だそう。
そこで PPO は「方策の変化量を抑えながら」更新する。
確率比 $r_t(\theta)$ を導入する。
r_t(\theta) = \frac{\pi_\theta(a_t | s_t)}{\pi_{\rm old}(a_t | s_t)}
として
\max_\theta \mathbb{E} \left[ r_t(\theta) \left( G(t) - V(s_t) \right) \right]
を最大化させる。
ただ $r_t$ が大きくなりすぎると更新が危険（？）だそうで、PPOは clipping を入れて
J^{\rm CLIP} (\theta) = \mathbb{E} \left[ \min(r_t(\theta)\left( G(t) - V(s_t) \right), clip(r_t(\theta), 1-\epsilon , 1+\epsilon) \left( G(t) - V(s_t) \right) ) \right]
を最大化させる。 $\epsilon$は0.1など。
DPPO (Diffusion Policy Policy Optimization)
DPPOは [2] . Ren, et. al. "Diffusion Policy Policy Optimization"（https://arxiv.org/pdf/2409.00588）で提示されている仕組み。
$\pi_\theta$ の部分をDiffusionモデルに置き換えたい。しかし簡単では無い。
Diffusionモデルは一般的に
\min_\theta \mathbb{E} \left[ \| \epsilon - \epsilon_\theta \|^2 \right]
のようにして付与したノイズ $\epsilon$ を推定するように学習する。一方でDPPOは
\max_\theta \mathbb{E} \left[ r_t(\theta) \left( G(t) - V(s_t) \right) \right]
をCLIP付きで最大化させる。つまりdenoiseが報酬最大化に一致するようなモデリングが必要。
そこで、まず事前学習で「デモ分布に沿った行動を生成する」学習を行う。
次にDPPOな学習。具体的にはdenoiseの１過程をMDPの１課程とみなし、
a_t^K \rightarrow a_t^{K-1} \rightarrow a_t^{K-2} \rightarrow \cdots \rightarrow a_t^0
と考える。
各denoise stepは報酬最大化に向けた１過程と考えられる。
DiWA（本手法）
DPPOからDiWAへの最大の変更は、DPPOの実環境がWorld Modelになった点。
上図中央のDPPOでは下側の実環境から状態 $S_t$ 等を受け取り Diffusion policy でdenoiseし、行動 $a_t$ を生成して実環境へ反映させていた。
一方で上図右のDiWAは実環境がWorld modelに置き換わっている。
これを実現させるために下図のような仕組みを導入している。
a) World Model自体を学習させる
実環境から取得した動画からWorld Modelを用いて
(s_t, a_t) \rightarrow s_{t+1}
を学習させる。
b) Diffusion Policyをexpert dataで事前学習させる
これはDPPOでも実施していたprocess。
c) 報酬モデルを学習させる
World Modelからそれらしい報酬を生み出すよう、報酬モデルを学習させる。
d) DiWAの学習
World Model、報酬モデルを用いて、仮想空間の中で Diffusion Policyを学習させる。
定量的評価
DPPOがDiWAと同程度に達するには10万回を超えるstepが必要。実機でこれは時間がかかるだろう。
一方で、回数さえ回せば実機で学習しているDPPOがDiWAを上回るか？
感想
- 結局、World Modelを学習の高速化に用いる系の話か？