概要

　強化学習は、学習が不安定になりやすいことが知られています。学習の不安定性に関して１つの対処法として、報酬のクリッピング(Reward clipping)があります。

　Reward clipping は以下のように報酬を変えます。
・報酬が正なら +1
・報酬が負なら -1

　Reward clippingすることで、Q値の急激な増大が抑えられ、勾配爆発を抑えられるそうです。しかし、報酬の大小が区別できない欠点があります。

　簡単に、ボーリングを例にとり、倒したピンの数が本来の報酬とします。つまり、ピンを多く倒すように学習をさせたい。しかし、Reward clippingしてしまうと、ピンを一つ以上倒しても報酬は変わらないことを意味しています。

　今回の記事では、Reward clippingに替わる方法として、行動価値関数（Q 関数）に作用するベルマン演算子を変更する方法を紹介したいと思います。
　
　以下の論文の一部を引用しています。

　この方法は、MuZero(ミューゼロ)にも使われ、MuZeroを勉強して始めて知りました。

Transformed Bellman Operator

　通常のQ学習で使われるベルマン演算子は、行動空間を$\mathcal{A}$、状態空間を$\mathcal{S}$と表し、行動・状態・報酬を$a,s,R$で表せば

(\mathcal{T} Q)(s,a) =\mathbb{E}_{s^{\prime} \sim P(\cdot|s,a)}\left[ R(s,a) + \gamma \underset{a^{\prime} \in \mathcal{A}}{\max} Q(s^{\prime},a^{\prime}) \right]

と書ける。
　ベルマン演算子を関数$h$を使い、

(\mathcal{T}_h Q)(s,a) =\mathbb{E}_{s^{\prime} \sim P(\cdot|s,a)}\left[ h\left( R(s,a) + \gamma \underset{a^{\prime} \in \mathcal{A}}{\max} h^{-1}\left( Q(s^{\prime},a^{\prime})\right) \right) \right]

と変更する。

Contraction

　ベルマン演算子$\mathcal{T}_h$の収縮(Contraction)を見ていく。

命題 1

　関数$h$とその逆関数$h^{-1}$は、単調増加関数でLipschitz連続とする。$h$と$h^{-1}$に対応するLipschitz定数を$L_h$と$L_{h^{-1}}$とする。割引率$\gamma$が

\gamma < \frac{1}{L_h L_{h^{-1}}}

を満たすならば、ベルマン演算子$\mathcal{T}_h$は収縮する。

　収縮であれば、解の一意性が言える（そうだ）。$\gamma$ 収縮とは、$x,y\in\mathcal{S}$ および $0<\gamma<1$とし、演算子$\mathcal{T}$に対して、

\|\mathcal{T}x -\mathcal{T}y \| \leq \gamma \|x - y \|

が成立することである。演算子$\mathcal{T}$を作用し続ければ、$\mathcal{S}$上の固定点に収束する。

　ベルマン演算子$\mathcal{T}_h$の収縮ついて説明する（証明する？数学者じゃないので）。まず準備として、イェンセンの不等式

\mathbb{E}[h(x)] \leq h(\mathbb{E}[x])

そして、関数$h$はLipschitz連続であることから、距離関数$d$とし、Lipschitz定数$L_h$として

d(h(x_1),h(x_2)) \leq L_hd(x_1,x_2)

を使う。距離関数として一様ノルム

\|x\|_{\infty}=\max\{|x_1|,\ldots |x_n| \}

を使う。

任意の$Q,U$の関数を使い、計算すると

\begin{align}
\|\mathcal{T}_hQ &- \mathcal{T}_hU\|_{\infty}  \\
=&\underset{s,a \in \mathcal{A}\times\mathcal{S}}{\max} \left| \mathbb{E}_{s^{\prime} \sim P(\cdot|s,a)}\left[h\left( R(s,a) + \gamma \underset{a^{\prime} \in \mathcal{A}}{\max} h^{-1}\left( Q(s^{\prime},a^{\prime})\right) \right)\right]
- \mathbb{E}_{s^{\prime} \sim P(\cdot|s,a)}\left[ h\left( R(s,a) + \gamma \underset{a^{\prime} \in \mathcal{A}}{\max} h^{-1}\left( U(s^{\prime},a^{\prime})\right) \right)  \right] \right| \\

\leq & \underset{s,a \in \mathcal{A}\times\mathcal{S}}{\max} \left|h\left( \mathbb{E}_{s^{\prime} \sim P(\cdot|s,a)}\left[R(s,a) + \gamma \underset{a^{\prime} \in \mathcal{A}}{\max} h^{-1}\left( Q(s^{\prime},a^{\prime})\right) \right]\right)
- h\left( \mathbb{E}_{s^{\prime} \sim P(\cdot|s,a)}\left[R(s,a) + \gamma \underset{a^{\prime} \in \mathcal{A}}{\max} h^{-1}\left( U(s^{\prime},a^{\prime})\right) \right) \right]  \right| \\

\leq & \gamma L_h \underset{s,a \in \mathcal{A}\times\mathcal{S}}{\max} \mathbb{E}_{s^{\prime} \sim P(\cdot|s,a)}\left[\left| \underset{a^{\prime} \in \mathcal{A}}{\max} h^{-1}\left( Q(s^{\prime},a^{\prime})\right) - \underset{a^{\prime} \in \mathcal{A}}{\max} h^{-1}\left( U(s^{\prime},a^{\prime})\right) \right|\right]\\

\leq & \gamma L_h \underset{s,a \in \mathcal{A}\times\mathcal{S}}{\max} \mathbb{E}_{s^{\prime} \sim P(\cdot|s,a)}\left[\underset{a^{\prime} \in \mathcal{A}}{\max} \left| h^{-1}\left( Q(s^{\prime},a^{\prime})\right) -  h^{-1}\left( U(s^{\prime},a^{\prime})\right) \right|\right] \\ 

\leq & \gamma L_h L_{h^{-1}} \underset{s,a \in \mathcal{A}\times\mathcal{S}}{\max} \mathbb{E}_{s^{\prime} \sim P(\cdot|s,a)}\left[\underset{a^{\prime} \in \mathcal{A}}{\max} \left|  Q(s^{\prime},a^{\prime}) -  U(s^{\prime},a^{\prime})\right|\right] \\ 

\leq & \gamma L_h L_{h^{-1}} \|Q -U\|_{\infty} \\
\leq & \|Q -U\|_{\infty}

\end{align}

となり収縮する。
　上式の２行目～３行目は一様ノルムの定義、４行目～５行目はイェンセンの不等式、６行目は$h$がLipschitz連続であること、７行目は最大となるもの同士の引き算より、引き算してから最大となるものを取った方が大きいこと、８行目は$h^{-1}$がLipschitz連続であること、９行目は一様ノルムの定義、最後の行は$\gamma < \frac{1}{L_h L_{h^{-1}}}$であることを使った。

関数 h(z)

　命題１を満たすような関数$h$を考えなければならない。論文では

h(z)= \rm{sgn}(z)\left(\sqrt{|z|+1}-1 \right) +\epsilon z　,　\epsilon=10^{-2}

としている。$\rm{sgn}$は符号関数である。また関数$h^{-1}$は、

h^{-1}(z)= \rm{sgn}(z) \left\{\left(\frac{\sqrt{1+4\epsilon(|z|+1+\epsilon)}-1 }{2\epsilon}\right)^2  -1 \right\}

とする（らしい）。

　最初に、関数$h(x)$と関数$h^{-1}(x)$は、微分可能で単調増加であることを説明する。
実際に、関数$h(x)$を微分してみる。$x\neq0$の場合、$\mathrm{d}|x|/\mathrm{d}x = \rm{sgn}(x)$ を使えば

\begin{align}
\frac{\mathrm{d}h}{\mathrm{d}x}&=\frac{(\rm{sgn}(x))^2}{\sqrt{|x|+1}} +\epsilon \\
& =  \frac{1}{\sqrt{|x|+1}} +\epsilon 
\end{align}

となる。$x=0$の場合、$z \to 0^+$の極限を考えて

\begin{align}
\lim_{z \to 0^+} \frac{h(x+z)-h(x)}{z}\Biggr|_{x=0}&=\lim_{z \to 0^+}\frac{\sqrt{z+1}-1  +\epsilon z}{z}\\
&=\lim_{z \to 0^+}\frac{1}{2\sqrt{z+1}}+\epsilon = \frac{1}{2}+\epsilon
\end{align}

となる。２行目はロピタルの定理を使った。
$z \to 0^-$の極限を考えて

\begin{align}
\lim_{z \to 0^-} \frac{h(x+z)-h(x)}{z}\Biggr|_{x=0}&=\lim_{z \to 0^-}\frac{-\sqrt{-z+1}-1  +\epsilon z}{z}\\
&=\lim_{z \to 0^-}\frac{1}{2\sqrt{-z+1}}+\epsilon = \frac{1}{2}+\epsilon
\end{align}

となる。以上から、$h(x)$は微分可能である。さらに明らかに

\begin{align}
\frac{\mathrm{d}h}{\mathrm{d}x}& =  \frac{1}{\sqrt{|x|+1}} +\epsilon > 0
\end{align}

なので、単調増加である。

　同様に関数$h^{-1}(x)$を微分してみると

\begin{align}
\frac{\mathrm{d}h^{-1}}{\mathrm{d}x}&=\left(\frac{\sqrt{1+4\epsilon(|x|+1+\epsilon)}-1 }{2\epsilon}   \right)\frac{4\epsilon (\rm{sgn}(x))^2}{\sqrt{1+4\epsilon(|x|+1+\epsilon)}}\\
& = \frac{1}{\epsilon}\left(1-\frac{1}{\sqrt{1+4\epsilon(|x|+1+\epsilon)}}  \right) 
\end{align}

となる。$z \to 0^+$の極限を考えて

\begin{align}
\lim_{z \to 0^+} \frac{h^{-1}(x+z)-h^{-1}(x)}{z}\Biggr|_{x=0}&=\lim_{z \to 0^+}\frac{1}{z} \left(\left(\frac{\sqrt{1+4\epsilon(z+1+\epsilon)}-1 }{2\epsilon}\right)^2  -1  \right)  \\
&=\lim_{z \to 0^+}\frac{1}{\epsilon}\left(1-\frac{1}{\sqrt{1+4\epsilon(z+1+\epsilon)}}  \right) = \frac{2\epsilon}{2\epsilon+1}
\end{align}

$z \to 0^-$の極限も同様に

\begin{align}
\lim_{z \to 0^-} \frac{h^{-1}(x+z)-h^{-1}(x)}{z}\Biggr|_{x=0}& = \frac{2\epsilon}{2\epsilon+1}
\end{align}

となる。以上から、$h^{-1}(x)$は微分可能である。さらに明らかに

\begin{align}
\frac{\mathrm{d}h^{-1}}{\mathrm{d}x}&=\frac{1}{\epsilon}\left(1-\frac{1}{\sqrt{1+4\epsilon(|x|+1+\epsilon)}}  \right) >0
\end{align}

なので、単調増加である。

　次に、関数$h(x)$と関数$h^{-1}(x)$のLipschitz連続について説明する。Lipschitz連続については、Lipschitz定数$L_h,L_{h^{-1}}$が求まること（存在する事）を言う。

　$x,y\in\mathrm{R}$かつ$x<y$とすると、平均値の定理を用いれば

\begin{align}
\frac{|h(x)-h(y)|}{|x-y|}&\leq \sup_{\xi\in(x,y)} \left|\frac{\mathrm{d}}{\mathrm{d}x}h(\xi)\right| \\
&\leq \sup_{\xi\in\mathrm{R}} \left|\frac{\mathrm{d}}{\mathrm{d}x}h(\xi)\right| \\
&\leq \sup_{\xi\in\mathrm{R}} \left|\frac{1}{\sqrt{|\xi|+1}} +\epsilon\right| \\
&\leq \frac{1}{2}+\epsilon
\end{align}

と計算され、Lipschitz定数 $L_h=1/2+\epsilon$ が求まる。関数$h^{-1}(x)$も同様に

\begin{align}
\frac{|h^{-1}(x)-h^{-1}(y)|}{|x-y|}&\leq \sup_{\xi\in(x,y)} \left|\frac{\mathrm{d}}{\mathrm{d}x}h^{-1}(\xi)\right| \\
&\leq \sup_{\xi\in\mathrm{R}} \left|\frac{\mathrm{d}}{\mathrm{d}x}h^{-1}(\xi)\right| \\
&\leq \sup_{\xi\in\mathrm{R}} \left|\frac{1}{\epsilon}\left(1-\frac{1}{\sqrt{1+4\epsilon(|\xi|+1+\epsilon)}}  \right)\right| \\
&\leq \frac{1}{\epsilon}
\end{align}

と計算され、Lipschitz定数 $L_{h^{-1}}=1/\epsilon$ が求まる。
　
　以上より次の命題2が成り立つことが分かる。ただし（３）については確認していない。

命題 2

　$\epsilon > 0$とし

h(x)= \rm{sgn}(x)\left(\sqrt{|x|+1}-1 \right) +\epsilon x

とすると
(１).　$h(x)$は単調増加。
(２).　$h(x)$はLipschitz連続で、Lipschitz定数は $L_h=1/2+\epsilon$である。
(３).　$h(x)$の逆関数は、

h^{-1}(z)= \rm{sgn}(z) \left\{\left(\frac{\sqrt{1+4\epsilon(|z|+1+\epsilon)}-1 }{2\epsilon}\right)^2  -1 \right\}

である。
(４).　$h^{-1}(x)$は単調増加。
(５).　$h^{-1}(x)$はLipschitz連続で、Lipschitz定数は $L_{h^{-1}}=1/\epsilon$である。

最後に

割引率$\gamma$は、

\gamma < \frac{1}{L_h L_{h^{-1}}}

を満たす必要があるが、$\epsilon=10^{-2}$とすると

\frac{1}{L_h L_{h^{-1}}}=\left(\frac{2\epsilon+1}{2\epsilon}\right)^{-1}=\frac{0.02}{1+0.02}

となってしまい$\gamma$を小さく取らないといけなくなる。

　論文では、
Proposition A.2 shows that the transformed operator is a contraction, the discount factor γ we use in practice is higher than $\frac{1}{L_h L_{h^{-1}}}$.
と書かれているので気にしなくても良いだろう。。。

実装はまだしていないので、結果は機会があれば。