IMPALA の V-trace ってなんだ

はじめに

この記事は DeepMind の IMPALA（Importance Weighted Actor-Learner Architectures） について、V-trace がオフポリシー学習の安定化にどのように寄与しているのかについて、論文を読んだときに釈然としなかった部分について勉強を兼ねて補足したような記事です。

そもそも、オンポリシーとオフポリシーって何？

強化学習では行動ポリシー $\mu$ と学習ポリシー $\pi$ を使い分けて学習を行います。

• 行動ポリシー（μ）：実際に環境で行動してデータを集める方策
• 学習ポリシー（π）：最終的に学習したい目標の方策

この2つの関係性によって、学習方法が分類されます。

オンポリシー学習（π = μ）
同じ方策でデータ収集と学習を行う方法です。シンプルで安定していますが、以下の問題があります：

• データ収集→学習→データ収集... という逐次的な処理が必要
• 並列化が難しく、学習速度に限界がある

オフポリシー学習（π ≠ μ）
異なる方策でデータ収集と学習を行う方法です。例えば、複数のアクターが並列にデータを集めつつ非同期に学習に用いるといった場合に該当します。効率的ですがこれにより新たな問題が発生します。

オフポリシー学習の問題

$\mu$ と $\pi$ が異なる、つまり $\mu(a|s) \neq \pi(a|s)$ であることにより、$\mu$ が収集した経験が「的外れ」となってしまい学習に利用しにくくなります。

1. 分布のミスマッチ
状態 $s$ の下で行動 $a$ の分布が異なることから、$\mu$ に従った場合と $\pi$ に従った場合では異なる状態に到達する可能性が高くなります。その結果、行動ポリシーと学習ポリシーが遭遇する状態の分布が異なってしまいます。

2. TD 学習のバイアス
時間差分（TD）学習において価値推定に偏りが生じます。以下ではこの問題について触れていきます。

オンポリシー時の TD 誤差

ポリシー $\pi$ についての TD 誤差 $\delta_t^{\pi}$ は次のようになります。

$$\delta_t^{\pi} = r_t + \gamma V^{\pi}(x_{t+1}) - V^{\pi}(x_t)$$

このとき、ポリシー $\pi$ についての条件付き期待値を取ると 0 になります。つまりオンポリシーでは TD 誤差にバイアスがありません。

$$\begin{align*} \mathbb{E}_{\pi}[\delta_t^{\pi} | x_t] &= \mathbb{E}_{\pi}[r_t + \gamma V^{\pi}(x_{t+1}) | x_t] - V^{\pi}(x_t) \\ &= \sum_a \pi(a|x_t) Q^{\pi}(x_t, a) - V^{\pi}(x_t) \\ &= V^{\pi}(x_t) - V^{\pi}(x_t) \\ &= 0 \end{align*}$$

オフポリシー時の TD 誤差

異なるポリシー $\mu$ でデータを収集しながら、ポリシー $\pi$ の価値関数を学習しようとする場合その期待値は 0 にならず、バイアスが生じます。

$$\begin{align*} \mathbb{E}_{\mu}[\delta_t^{\pi} | x_t] &= \mathbb{E}_{\mu}[r_t + \gamma V^{\pi}(x_{t+1}) | x_t] - V^{\pi}(x_t) \\ &= \sum_a \mu(a|x_t) Q^{\pi}(x_t, a) - V^{\pi}(x_t) \\ &\neq 0 \end{align*}$$

Importance Sampling とその限界

この問題を解決する古典的な手法が Importance Sampling（重要度サンプリング） です。
これは TD 誤差 $\delta_t$ をそのまま使うのではなく、重み付け係数 $π(a|s) / μ(a|s)$ でポリシー間の差を補正した $(π(a|s) / μ(a|s)) × \delta_t$ を使う方法です。
ところが、式から予想される通りこの手法には欠点があります。$μ(a|s)$ が小さいとき、重み付け係数 $π(a|s) / μ(a|s)$ が爆発的に大きくなってしまいます。その結果バイアスは 0 になるものの分散が大きくなってしまい学習が不安定になります。

IMPALA においては Importance Sampling の拡張として V-trace を導入しています。

V-trace の定義

V-trace（n-step ターゲット）

$$\begin{align*} v_s^{(n)}(V) &\stackrel{def}{=} V(x_s) + \sum_{t=s}^{s+n-1} \gamma^{t-s} \Big(\prod_{i=s}^{t-1} c_i\Big)\, \delta_t V,\\ \delta_t V &= \rho_t \big(r_t + \gamma V(x_{t+1}) - V(x_t)\big),\\ \rho_t &= \min\!\left(\bar{\rho}, \frac{\pi(a_t|x_t)}{\mu(a_t|x_t)}\right),\quad c_t = \min\!\left(\bar{c}, \frac{\pi(a_t|x_t)}{\mu(a_t|x_t)}\right). \end{align*}$$

つまり、Importance Sampling の中で出てきた重み付けに上限を設けた形になっています。これにより $\mu(a_i|x_i)$ が小さいケースにおいても重み付け係数が大きな値を取ることを防げます。

以下では一貫して $\bar{\rho} \ge \bar{c}$ を仮定します。これにより、各時刻で $\rho_t \ge c_t \ge 0$ が保たれます。

V-trace による収束先

V-trace では重み付け係数が大きな値を取らないようにできるのは確かですが、このターゲットは妥当なのでしょうか？　論文中では次の点について述べられていますが、ここでは主に 1 と 2 について説明します。

1. V-trace 作用素が縮小写像であること
2. V-trace 作用素の不動点の同定
3. オンライン学習の収束

以下の証明では

$$\begin{align*} v_s^{(\infty)}(V) &\equiv V(x_s) + \sum_{t=s}^{\infty} \gamma^{t-s} \Big(\prod_{i=s}^{t-1} c_i\Big)\, \rho_t\big(r_t+\gamma V(x_{t+1})-V(x_t)\big), \\ (\mathcal{R}V)(x) &\equiv \mathbb{E}_\mu\!\left[v_s^{(\infty)}(V)\mid x_s=x\right] \end{align*}$$

を用いる。以降は $\gamma\in(0,1)$ とし、$\bar{\rho}\ge \bar{c}$（ゆえに $0\le c_t\le \rho_t$）、
および coverage 仮定 （$\exists\,\beta\in(0,1] \, s.t. \, \forall s \,\ \mathbb{E}_\mu[\rho_s]\ge\beta$） を前提とします。

V-trace 演算子は縮小写像

V-trace 演算子 $\mathcal{R}$ を以下のように定義します：

$$\begin{align*} (\mathcal{R}V)(x) &\equiv \mathbb{E}_\mu\!\left[v_s^{(\infty)}(V)\mid x_s=x\right] \\ &= V(x) +\mathbb{E}_{\mu} \left[\sum_{t=s}^{\infty} \gamma^{t-s} \left(\prod_{i=s}^{t-1} c_i\right) \rho_t (r_t + \gamma V(x_{t+1}) - V(x_t)) \Big| x_s = x \right] \tag{1} \end{align*}$$

V-trace 演算子が縮小写像であることを示します。

仮定

$$\gamma \in (0, 1),\, \bar{\rho} \ge \bar{c},\, \exists{\beta} \, s.t. \, \forall{s} \, \mathbb{E}_{\mu}[\rho_s] \ge \beta \,(\text{coverage仮定})$$

定理（V-trace 演算子の縮小性）：
任意の価値関数 $V_1, V_2$ に対して次の不等式を示せばよい。（ノルムの等価性から状態が有限次元ベクトルであるなら $L^{\infty}$ ノルムの場合を示せば十分）

$$\|\mathcal{R} V_1 - \mathcal{R} V_2\|_\infty \lt \eta \|V_1 - V_2\|_\infty, \eta \lt 1$$

証明

$(1)$ の $\mathbb{E}_{\mu}$ の中身を 3 つの項に分けます。

$$\begin{align*} \sum_{t=s}^{\infty} \gamma^{t-s} \left(\prod_{i=s}^{t-1} c_i\right) \rho_t (r_t + \gamma V(x_{t+1}) - V(x_t)) &= \sum_{t=s}^{\infty} \gamma^{t-s} \left(\prod_{i=s}^{t-1} c_i\right) \rho_t r_t \\ &+ \sum_{t=s}^{\infty} \gamma^{t-s+1} \left(\prod_{i=s}^{t-1} c_i\right) \rho_t V(x_{t+1}) \\ &- \sum_{t=s}^{\infty} \gamma^{t-s} \left(\prod_{i=s}^{t-1} c_i\right) \rho_t V(x_{t}) \end{align*}$$

• 第1項：$\sum_{t=s}^{\infty} \gamma^{t-s} \left(\prod_{i=s}^{t-1} c_i\right) \rho_t r_t$
• 第2項：$\sum_{t=s}^{\infty} \gamma^{t-s+1} \left(\prod_{i=s}^{t-1} c_i\right) \rho_t V(x_{t+1})$
• 第3項：$- \sum_{t=s}^{\infty} \gamma^{t-s} \left(\prod_{i=s}^{t-1} c_i\right) \rho_t V(x_{t})$

ここで第 2 項において $u=t+1$ とおくと、

$$\sum_{u=s+1}^{\infty} \gamma^{u-s} \left(\prod_{i=s}^{u-2} c_i\right) \rho_{u-1} V(x_u)$$

また、第 3 項で $t=s$ の項と $t \ge s+1$ の項を分けると

$$\begin{align*} - \sum_{t=s}^{\infty} \gamma^{t-s} \left(\prod_{i=s}^{t-1} c_i\right) \rho_t V(x_{t}) &= - \rho_{s}V(x_s) - \sum_{t=s+1}^{\infty} \gamma^{t-s} \left(\prod_{i=s}^{t-1} c_i\right) \rho_t V(x_{t}) \\ &= - \rho_{s}V(x_s) - \sum_{t=s+1}^{\infty} \gamma^{t-s} \left(\prod_{i=s}^{t-2} c_i\right) c_{t-1} \rho_t V(x_{t}) \end{align*}$$

これらを $(1)$ に当てはめると

$$\begin{align*} (\mathcal{R} V)(x) &= (1 - \mathbb{E}_{\mu}[\rho_s]) V(x) + \mathbb{E}_{\mu} \left[ \sum_{t=s}^{\infty} \gamma^{t-s} \left(\prod_{i=s}^{t-1} c_i\right) \rho_t r_t + \sum_{t=s+1}^{\infty} \gamma^{t-s} \left(\prod_{i=s}^{t-2} c_i\right) (\rho_{t-1} - c_{t-1} \rho_t) V(x_{t}) \right] \\ &= (1 - \mathbb{E}_{\mu}\rho_s) V(x) + \mathbb{E}_{\mu} \left[ \sum_{t=s}^{\infty} \gamma^{t-s} \left(\prod_{i=s}^{t-1} c_i\right) (\rho_t r_t + \gamma(\rho_t - c_t \rho_{t+1})V(x_{t+1})) \right] \end{align*}$$

ここで、$\rho_{s-1}=c_{s-1}=1, 空積=1$ としています。

2 つの価値関数 $V_1$ と $V_2$ に対して $\mathcal{R}$ を作用させたときの差を考えると、報酬は打ち消し合います。

$$(\mathcal{R}V_1- \mathcal{R}V_2)(x) = (1 - \mathbb{E}_{\mu}[\rho_s]) [V_1(x) - V_2(x)] + \mathbb{E}_{\mu} \left[ \sum_{t=s+1}^{\infty} \gamma^{t-s} \left(\prod_{i=s}^{t-2} c_i\right) (\rho_{t-1} - c_{t-1} \rho_t) [V_1(x) - V_2(x)]) \right]$$

ここで
$t=s$ のとき

$$\begin{align*} \gamma^{t-s} \left(\prod_{i=s}^{t-2} c_i\right) (\rho_{t-1} - c_{t-1} \rho_t) &= \gamma^0 \left(\prod_{i=s}^{s-2} c_i\right) (\rho_{s-1} - c_{s-1} \rho_s) \\ &= 1 - \rho_s \end{align*}$$

$$\begin{equation*} (1 - \mathbb{E}_{\mu}[\rho_s]) [V_1(x) - V_2(x)] + \mathbb{E}_{\mu} \left[ \sum_{t=s+1}^{\infty} \gamma^{t-s} \left(\prod_{i=s}^{t-2} c_i\right) (\rho_{t-1} - c_{t-1} \rho_t) [V_1(x) - V_2(x)]) \right] \\ = \mathbb{E}_{\mu} \left[ \sum_{t=s}^{\infty} \gamma^{t-s} \left(\prod_{i=s}^{t-2} c_i\right) (\rho_{t-1} - c_{t-1} \rho_t) [V_1(x) - V_2(x)]) \right] \end{equation*}$$

この係数部分を $\eta$ として、

$$\begin{align*} \eta &\stackrel{\mathrm{def}}{=} \mathbb{E}_{\mu} \left[ \sum_{t=s}^{\infty} \gamma^{t-s} \left(\prod_{i=s}^{t-2} c_i\right) (\rho_{t-1} - c_{t-1} \rho_t) \right] \\ &= \mathbb{E}_{\mu} \left[ \sum_{t=s}^{\infty} \gamma^{t-s} \left(\prod_{i=s}^{t-2} c_i\right) \rho_{t-1} \right] - \mathbb{E}_{\mu} \left[ \sum_{t=s}^{\infty} \gamma^{t-s} \left(\prod_{i=s}^{t-1} c_i\right) \rho_t \right] \\ &= \mathbb{E}_{\mu} \left[ \sum_{t=s}^{\infty} \gamma^{t-s} \left(\prod_{i=s}^{t-2} c_i\right) \rho_{t-1} \right] - \gamma^{-1} \left( \mathbb{E}_{\mu} \left[ \sum_{t=s}^{\infty} \gamma^{t-s} \left(\prod_{i=s}^{t-2} c_i\right) \rho_{t-1} \right] - 1 \right) \\ &= \gamma^{-1} - (\gamma^{-1} - 1) \mathbb{E}_{\mu} \left[ \sum_{t=s}^{\infty} \gamma^{t-s} \left(\prod_{i=s}^{t-2} c_i\right) \rho_{t-1} \right] \end{align*}$$

ここで

$$\sum_{t=s}^{\infty} \gamma^{t-s} \left(\prod_{i=s}^{t-2} c_i\right) \rho_{t-1} = 1 + \gamma \rho_s + \gamma^2 c_s \rho_{s+1} + \cdots$$

$\gamma \in (0, 1), \rho_t \ge 0, c_t \ge 0$ であることから、この各項は非負。したがって

$$\sum_{t=s}^{\infty} \gamma^{t-s} \left(\prod_{i=s}^{t-2} c_i\right) \rho_{t-1} \ge 1 + \gamma \rho_s$$

期待値を取ると

$$\mathbb{E}_{\mu} \left[ \sum_{t=s}^{\infty} \gamma^{t-s} \left(\prod_{i=s}^{t-2} c_i\right) \rho_{t-1} \right] \ge 1 + \gamma \mathbb{E}_{\mu} \left[\rho_s \right] \gt 0$$

$$\begin{align*} \gamma^{-1} - (\gamma^{-1} - 1) \mathbb{E}_{\mu} \left[ \sum_{t=s}^{\infty} \gamma^{t-s} \left(\prod_{i=s}^{t-2} c_i\right) \rho_{t-1} \right] &\le \gamma^{-1} - (\gamma^{-1} - 1)(1 + \gamma \mathbb{E}_{\mu}[\rho_s]) \\ &= 1 - (1 - \gamma) \mathbb{E}_{\mu}[\rho_s] \end{align*}$$

以下では、$\rho_s=\min(\bar{\rho},\,\pi/\mu)$ の期待値に一様の下限 $\beta \in (0, 1]$ があると仮定します（単にサポートが重なるだけでは十分でない）。この前提のもとで

$$\eta \le 1 - (1 - \gamma) \mathbb{E}_{\mu}[\rho_s] \le 1 - (1 - \gamma) \beta \lt 1$$

となるため、$\eta \lt 1$ 。

$$\|\mathcal{R} V_1 - \mathcal{R} V_2\|_\infty \lt \eta \|V_1 - V_2\|_\infty$$

$\mathcal{R}$ は縮小写像であり不動点を持つ。

V-trace 演算子による収束先の価値関数

仮定：$\bar{\rho} \ge \bar{c}$
このとき V-trace の不動点は、以下のポリシー $\pi_{\bar{\rho}}$ に対する価値関数 $V^{\pi_{\bar{\rho}}}$ となります。

$$\pi_{\bar{\rho}}(a|x) = \frac{\min(\bar{\rho}\mu(a|x), \pi(a|x))}{\sum_{a'} \min(\bar{\rho}\mu(a'|x), \pi(a'|x))} \tag{2}$$

これは正規化因子 $\sum_{a'} \min(\bar{\rho}\mu(a'|x), \pi(a'|x))$を両辺にかけることで次のように変形できます。

$$\min(\bar{\rho}\mu(a|x), \pi(a|x)) = \pi_{\bar{\rho}}(a|x) \sum_{a'} \min(\bar{\rho}\mu(a'|x), \pi(a'|x)) = \pi_{\bar{\rho}}(a|x) Z(x) \tag{2'}$$

$$\begin{align*} (\mathcal{R} {V^{\pi_{\bar{\rho}}}})(x) &= {V^{\pi_{\bar{\rho}}}}(x) +\mathbb{E}_{\mu} \left[\sum_{t=s}^{\infty} \gamma^{t-s} \left(\prod_{i=s}^{t-1} c_i\right) \rho_t \delta_t^{V^{\pi_{\bar{\rho}}}} \Big| x_s = x \right] \tag{3} \\ \delta_t^{V^{\pi_{\bar{\rho}}}} &= r_t + \gamma V^{\pi_{\bar{\rho}}}(x_{t+1}) - V^{\pi_{\bar{\rho}}}(x_t) \end{align*}$$

$(3)$ の第 2 項について、$\left(\prod_{i=s}^{t-1} c_i\right)$ は時刻 $t$ より前の情報だけで決まるので次のように変形できます。

$$\begin{align*} \mathbb{E}_{\mu} \left[\sum_{t=s}^{\infty} \gamma^{t-s} \left(\prod_{i=s}^{t-1} c_i\right) \rho_t \delta_t^{V^{\pi_{\bar{\rho}}}} \Big| x_s = x \right] = \mathbb{E}_{\mu} \left[ \sum_{t=s}^{\infty} \gamma^{t-s} \left(\prod_{i=s}^{t-1} c_i\right) \mathbb{E}_{\mu} \left[ \rho_t \delta_t^{V^{\pi_{\bar{\rho}}}} \right] \Big| x_s = x \right] \tag{4} \end{align*}$$

$$\begin{align*} \mathbb{E}_{\mu} \left[ \rho_t \delta_t^{V^{\pi_{\bar{\rho}}}} \Big| x_t \right] &= \mathbb{E}_{\mu} \left[ \rho_t (r_t + \gamma V^{\pi_{\bar{\rho}}}(x_{t+1}) - V^{\pi_{\bar{\rho}}}(x_t)) \right] \\ &= \sum_a \mu(a | x_t) \min \left(\bar{\rho}, \frac{\pi(a | x_t)}{\mu(a | x_t)} \right) ( r_t + \gamma V^{\pi_{\bar{\rho}}}(x_{t+1}) - V^{\pi_{\bar{\rho}}}(x_t) ) \\ &= \sum_a \pi_{\bar{\rho}}(a|x_t) Z(x_t) ( r_t + \gamma V^{\pi_{\bar{\rho}}}(x_{t+1}) - V^{\pi_{\bar{\rho}}}(x_t) ) \\ &= Z(x_t) \sum_a \pi_{\bar{\rho}}(a|x_t) ( r_t + \gamma V^{\pi_{\bar{\rho}}}(x_{t+1}) - V^{\pi_{\bar{\rho}}}(x_t) ) \\ &= \underbrace{Z(x_t)}_{正規化因子} \underbrace{\mathbb{E}_{\pi_{\bar{\rho}}} \left[ r_t + \gamma V^{\pi_{\bar{\rho}}}(x_{t+1}) - V^{\pi_{\bar{\rho}}}(x_t) \right ]}_{\pi_{\bar{\rho}} についてのベルマン誤差の期待値=0} \\ &= 0 \end{align*}$$

これを $(4)$ に代入すると

$$\mathbb{E}_{\mu} \left[ \sum_{t=s}^{\infty} \gamma^{t-s} \left(\prod_{i=s}^{t-1} c_i\right) \underbrace{\mathbb{E}_{\mu} \left[ \rho_t \delta_t^{V^{\pi_{\bar{\rho}}}} \right]}_{=0} \Big| x_s = x \right]=0$$

さらにこの結果から $(3)$ は

$$(\mathcal{R} {V^{\pi_{\bar{\rho}}}})(x) = V^{\pi_{\bar{\rho}}}(x) + \underbrace{\mathbb{E}_{\mu} \left[\sum_{t=s}^{\infty} \gamma^{t-s} \left(\prod_{i=s}^{t-1} c_i\right) \rho_t \delta_t^{V^{\pi_{\bar{\rho}}}} \Big| x_s = x \right]}_{=0} = V^{\pi_{\bar{\rho}}}(x)$$

となるので、$V^{\pi_{\bar{\rho}}}$ は $\mathcal{R}$ の不動点。またこの結果から $\mathcal{R}$ の収束先は $\bar{\rho}$ には依存しますが $\bar{c}$ には依存しないということもわかります。

オンライン学習

論文中の式 $(7)$：

$$V_{k+1}(x_s) = V_k(x_s) + \alpha_k(x_s) \underbrace{ \sum_{t \ge s}\gamma^{t-s}(c_s \dots c_{t-1})\rho_t(r_t + \gamma V_k(x_{t+1}) - V_k(x_t)) }_{V-trace ターゲットとの差}$$

有限 MDP・Robbins–Monro の学習率条件・全状態への無限回訪問・$\bar{\rho} \ge \bar{c}$ および coverage 条件と適切な $\alpha$ のもとで $V_k$ の収束先が $V^{\pi_{\bar{\rho}}}$ になるらしいです。