分布強化学習 （distributional reinforcement learning ）の概念を深層強化学習 へ導入したCategorical DQN （C51）をtensorflow2で実装します。
why restrict ourselves to the mean?
― [1707.06887] A Distributional Perspective on Reinforcement Learning
参考資料
Deep Reinforcement Learning Hands-On | Packt
https://physai.sciencesconf.org/data/pages/distributional_RL_Remi_Munos.pdf
『強化学習』（森村 哲郎）｜講談社BOOK倶楽部
前提手法：Deep-Q-Network (DQN )
horomary.hatenablog.com
はじめに
[1707.06887] A Distributional Perspective on Reinforcement Learning
従来の深層強化学習 （Q学習）では期待割引報酬和を最大化するため に、状態行動価値Q(s, a)の期待値を関数近似 します。Q関数がうまく近似できているならば、Q関数に従って期待値の大きな行動(action)を選択しつづければ報酬和も最大化されます。
Q関数の期待値が高い行動を選択しつづけること自体は問題ありません。しかし、状態行動価値Q(s, a)の期待値を直接関数近似 するのではなく、状態行動価値Q(s, a)の確率分布を関数近似 しそこから期待値を算出する方が、Q関数が環境をよく表現できるので学習が安定化するのではないか？ というのが論文の要旨です。
この論文ではまずベルマン方程式を分布の概念を導入した分布版ベルマン方程式を提案しました。
さらに、価値分布をカテゴリ分布で表現したQ学習をatari 環境向けに実装し、分布強化学習 の有用性を当時のatari 環境におけるSOTAという結果で証明しました。atari 環境においてはカテゴリ分布のビン数が51に設定されたことから論文内ではこの手法を指してC51 と呼称しています。
論文より
分布強化学習 (distributional reinforcement learning )自体は古くから存在するアイデア だったのですが、
この手法が発表されたことにより深層学習×分布強化学習 の研究が活発化し、
QR-DQN ,
IQN ,
FQF など多くの後継手法が発表されることとなりました。
※単語が似ていて紛らわしいのですが分布強化学習 (distributional reinforcement learning )と分散強化学習 (distributed reinforcement learning )は全く異なる概念です。java とjavascript くらい違います。
rayで実装する分散強化学習 ①A3C（非同期Advantage Actor-Critic） - どこから見てもメンダコ
※論文では状態をxで表記していますがここではsで表記します。また、表記の簡単のために状態遷移は決定論 的とします。
通常のベルマン方程式
通常のQ学習では下式のようにベルマンオペレータTを定義し、ベルマンエラー
Q(s, a)はスカラ値なのでベルマンエラー 教師あり学習 と同様にMSE（平均二乗誤差）*1 なんかを最小化するだけです。
分布版ベルマン方程式
分布版ベルマン方程式では、状態行動価値を期待リターンの確率分布Z(s, a)で表現したベルマンオペレータTを定義します。
遷移先状態s'での行動a'は
通常のベルマン方程式も分布版ベルマン方程式も見た目上はあまり変わりませんが、通常のベルマンオペレータは”スカラ値Q(s', a') に割引率を掛けて即時報 酬を足す”、という非常にわかりやすい処理であるのに対して、分布版ベルマンオペレータでは”確率分布Z(s', a') に割引率を掛けて即時報 酬を足す”、というやや想像しにくい四則演算を行います。
論文Fig.1を改変して掲載
論文内fig.1 はZがカテゴリ分布として表現されている場合での分布版ベルマンオペレータ適用のイメージです。①割引率を掛けることで分布が縮み、②即時報 酬が足されることで分布が横に平行移動する 、 という処理が可視化されています。（③のずれたビンを直す処理については後述）
同じことをPython でやってみます。
コードを見ればとくに難しいことはありませんね。カテゴリ分布のビンの幅が変わっている ことに注意してください（後述）。
ベルマンエラーの計算（分布版）
ベルマンオペレータ 2つの分布間の距離を適切に定義 する必要があります。
2つの分布間の距離尺度 といえばKL-divergenceやWasserstein-metricなどいろいろありますが、C51ではカテゴリ分布でZを表現しているので深層学習ではお馴染みのcategorical cross entropy を２つの分布間の距離として、これを最小化するようにネットワークを更新していきます。
しかし、上述したようにベルマンオペレータ適用後のビンの幅が変わってしまっているため 、そのままではクロスエントロピー を計算することができません。（灰線がZ(s, a)のビン幅、赤船がTZ(s, a)のビン幅）
そこで、
初見では意味不明な数式ですがコードを見るとそれほど難しいことは言っていないことに気が付きます。
ビン幅を再割り当てしたエントロピー を計算するだけです。
Categorical DQN （C51）のtensorflow2実装
上述した分布版ベルマンオペレータの適用とビンの再割り当て処理(projection)を除くとほぼDQN と同じなので実装の要点だけ掲載。
実装全体はgithub へ：
github.com
カテゴリ分布Q関数
ネットワーク構造は基本DQN ですが、最終Dense層が"action_space×カテゴリ分布のビン数"のunitsを出力したうえでこれを(バッチサイズ, action_space, n_atoms)にreshapeし、各action軸でsoftmaxをとることで確率分布を表現します。
ちなみに論文ではカテゴリ分布のビンのことをatoms、およびそのビンに割り当てられた確率をsupportと呼称しています。英語だとそういう表現なんですね。
ネットワーク更新
SGD するだけです。
クロスエントロピー の計算はlogをとる前に微小量を加えておかないと打切り誤差でLog(0)を計算してnanを吐くことがあるので注意。（実際2Mstepくらい進んだあたりでnan吐いて学習失敗した）
分布版ベルマンオペレータの適用
上述のコードとやってることは同じですが、ミニバッチ単位での処理のためにややコードの見通しが悪くなっています。
またdoneのときの処理、つまり
BreakoutDeterminisitc-v4 での結果
4000エピソードちょいでおよそ400万stepを学習した結果です。サチる みたいです。
マシンはGCP のn1-standard-4(4-vCPU, 15GBメモリ) + GPU K80 のプリエンティブルVM インスタンス を使って24時間学習しました。プリエンティブルVM は激安な代わりに24時間で自動シャットダウンされます。費用はざっくり計400円くらい？
付録
分布ベルマン方程式の理論的保証
本文では価値分布をカテゴリ分布で表現し、categorical cross entropy（※論文ではKL divergenceのクロスエントロピー 項と言っている）を分布間の距離尺度とする分布版ベルマンオペレータを天下り 的に正しいもののように紹介しましたが、実際はこのような方法の理論的保証はされておらず”APPROXIMATE DISTRIBUTIONAL LEARNING”と明記されています。
[1707.06887] A Distributional Perspective on Reinforcement Learning の前半は分布ベルマン分布の理論的保証に割かれているのですが、ここで保証されているのは固定方策かつZ(s, a)を連続な分布で表現した場合にZ(s, a)とTZ(s, a)の累積分 布関数の逆関数 のLp-Wasserstein距離を分布間の距離尺度としたときに分布ベルマンオペレータが縮小写像 となる（収束する）ということです。
続報でもこのあたりの議論が深められています。[1902.03149] Distributional reinforcement learning with linear function approximation
しかし、分布版ベルマンオペレータの収束性の話は当然通常のベルマンオペレータの収束性の議論の理解が前提なので、そもそもこのあたりの知識に自信がない人（私です）は、強化学習 の青い本『強化学習』（森村 哲郎）｜講談社BOOK倶楽部 ）2章 プランニング
こちらの記事も大変参考になりました。Thank you!Wasserstein Distance, Contraction Mapping, and Modern RL Theory | by Kowshik chilamkurthy | Medium
経験分布からのWasserstein LossはBiased
上述の通り、論文ではWasserstein距離を分布間距離尺度としたときに分布ベルマン方程式が収束するということを証明したのに、CategoircalDQNの実装ではカテゴリ分布のKL距離を採用しています。これは経験分布のWasserstein距離はBiasedであるためです。
これについてはLemma 7, Proposition 5で説明されていますが、私は理解に時間がかかったので同じような方のために解説しておきます。
具体例として２つの分布PとQの間のwasserstein距離を考えます。分布Pは表（=1）と裏（=0）がどちらも0.5の確率で出るコインの分布（ベルヌーイ分布の確率質量関数）とします。一方、分布Qは表（=1）がpの確率で、裏（=0）が 1-p の確率で出る歪んだコインの分布とします。
この分布Pと分布Qの1-wasserstein距離 d(P, Q)= |0.5 - p| であることは視覚的にも容易にわかります（下図）。
PとQのwst距離の視覚的理解
分布Pは無限のコイントス 試行によって得られる真の分布であることに注意してください。もし分布Pが経験分布である場合、極端にはコイントス 試行１回行って表が出たという結果からの経験分布の場合、経験分布Pと分布QのWST距離は 1-p となります(下図)。同様に裏が出た場合のWST距離はpとなります。
経験分布Pと分布QのWST距離
ここで、PのコインについてN回の試行を行い、各サンプルから得られるPの経験分布と分布QのWST距離のサンプル平均を考えます。表が出る確率は0.5でありこの場合のWST距離は1-p, 裏が出る確率も0.5でありこの場合のWST距離はpであるので、このサンプル平均 E[d(Pi, Q)]は 0.5×p + 0.5×(1-p)= 0.5 となります。上述の通り、分布間の真のWST距離は |0.5 - p| であるので、よってp=0 or 1の場合を除いて経験分布からのWasserstein距離にバイアスがかかることが理解できます。
SGD においてミニバッチを構成する各サンプルはまさに報酬分布R(s)および遷移分布P(s' | s, a)からの１回のサンプリングから得られた経験分布であるため、各サンプルのWST距離のミニバッチ平均もまた同様にBiasedとなります。
なぜ分布の推定が学習を安定化させるか？
atari2600環境で当時のSOTAを達成し単体のDQN 拡張手法としては最良の結果を示したCategorical DQN ですが、状態行動価値分布を推定することでなぜ学習が安定するのかについてはdiscussion項で考察されているものの明確な結論は示されていません。
強化学習 の青い本『強化学習』（森村 哲郎）｜講談社BOOK倶楽部 ）8.2.1.4 カテゴリDQN 法 でも、
性能改善の要因として、マルチタスク 学習のように目的タスク（期待リターン推定）の学習と並行して関連する多数のタスク（リターン分布推定）を学習することによる効果などが考察されていますが、いまだ議論の段階で、今後さらなる解析が期待されます。
と、述べられています。
他にはこちらの記事も参考になりました。https://clarelyle.com/posts/2019-02-08-aaai.html
マウスの脳は分布強化学習 を行うか？
最近出たDeepmind とハーバード大 のNature論文。マウスのドーパミン 細胞における神経活動の測定から分布強化学習 っぽさを示す結果が得られたらしい。興味ある方はどうぞ。
www.nature.com
horomary.hatenablog.com
