どこから見てもメンダコ

軟体動物門頭足綱八腕類メンダコ科

深層分布強化学習 ②QR-DQN

QR-DQNをtensorflow2で実装します。
元論文: [1710.10044] Distributional Reinforcement Learning with Quantile Regression

前記事:
horomary.hatenablog.com

参考:

https://physai.sciencesconf.org/data/pages/distributional_RL_Remi_Munos.pdf

Going beyond average for reinforcement learning | DeepMind

Quantile regression - Wikipedia


はじめに

DeepMindDQNに代表される典型的なQ学習においては、状態行動価値Q(s, a)の期待値関数近似します。

一方、前記事で実装を紹介したCategorical DQN ([1707.06887] A Distributional Perspective on Reinforcement Learning)は、状態行動価値Q(s, a)を明示的に確率分布Z(s, a)としてモデル化することを提案し、これにより大きくパフォーマンスが向上することを当時のatari環境のSotAという結果で示しました。

本記事で紹介するQR-DQNはCategoricalDQNの直接の後継手法*1です。Categorical DQNでは価値分布をそのままカテゴリ分布で近似しようとしたのに対し、QR-DQNは状態行動価値分布Z(s, a)の分位点を近似するというアプローチによりCategorical DQNの残した多くの課題を解決しました。


Categorical DQNの分布モデル

分布強化学習でモデル化したい真の(ground truth?)状態行動価値分布Z(s, a)は連続分布であるはずですが、連続分布は大変扱いづらいのでCategorical DQNではその名の通りZ(s, a)をカテゴリカル分布で近似します。Categorical DQN論文ではカテゴリカル分布のビン数=51の場合がatari環境でもっとも性能が良かったので、この場合をとくにC51と呼称しています。*2

f:id:horomary:20210328225227p:plain:w600
カテゴリ分布によるZ(s, a)のモデル化

状態行動価値分布Z(s,a)へのベルマンオペレータの適用は下図のように行います。rewardによって分布が水平スライドし、割引率によって分布が縮むようなイメージです。※見た目にわかりやすいようにreward=7, 割引率γ=0.6という極端な値で作図していることに留意ください。

f:id:horomary:20210328230353p:plain:w600
分布ベルマン方程式

状態行動価値分布をCategorical分布で近似するC51のアプローチはいくつかの大きな問題を抱えています。

1つはベルマンオペレータの適用によって分布のビン幅がずれることです。上図でも元の分布Z(s,a)のビン幅である赤破線からTZ(s, a)のビン幅はずれてしまっていることがわかります。よってCategorical DQNではこのずれたビン幅を無理に再割り当てして修正する処理*3が必要なのですが、この処理の実装がかなり煩雑&やや重い*4です。

別の問題はカテゴリカル分布では有限領域しか扱えないため、分布の最大値/最小値の設定が非常に重要なハイパーパラメータになってしまうことです。 この問題は学習初期と学習終盤で報酬のスケールが大きく変化するような場合には顕著な問題となります *5。また、最大/最小幅を大きくとった場合はカテゴリカル分布の性質上ビンの数を十分に多くしないと細かな分布の形状を捉えにくいという問題も生じます。

さらにCategorical DQNの最大の問題は、Categorical DQN論文で証明された"p-Wasserstein距離を分布間の距離尺度に設定するとベルマンオペレータが縮小写像である"という理論とCategorical分布のKL距離をロス関数とする実装にギャップがあることです。大雑把には、確率的勾配降下法でWasserstein距離をロス関数にすると biased gradient になるので、言っていることとやっていることが違うのだけどKL距離をロスにするヒューリスティックな実装にしたよ、という感じです。( 前記事を参照)


QR-DQNの分布モデル

Categorical DQNではZ(s,a)をそのままカテゴリカル分布で近似しましたが、QR-DQNではZ(s,a)の累積分布関数Fを近似します。※Z(s,a)とその累積分布関数Fは1対1変換であるのでどちらを近似してもよいことに留意。

f:id:horomary:20210329004825p:plain:w500
Z(s,a)とその累積分布関数F

ここで、QR-DQNのポイントは累積分布関数Fそのものではなく、Fの逆関数をカテゴリカル分布で近似することです。

f:id:horomary:20210329010227p:plain:w500
各ビンは分位点と解釈する

したがって、Categorical DQNでは各ビンの値はZ(s,a)がある状態行動価値θをとる確率でしたが、QR-DQNでは各ビンの値はZ(s,a)の τ%分位点 (Quantile)の値となります。あえてZ(s,a)の累積分布関数の逆関数をカテゴリカル分布で近似することにより、前述したCategorical DQNの問題点を解消することができます。

まず、Categorical-DQNではx軸のカテゴリカル分布でZ(s,a)を近似していましたが、ベルマンオペレータの適用によってビン幅がずれるため煩雑なビンの再割り当て処理(projection)が必要でした。一方、QR-DQNではカテゴリカル分布で価値分布の累積分布関数をy軸にそってモデル化する(つまり累積分布関数の逆関数を近似)ためビン幅ずれ問題に煩わされることは無くなりました(下図)。

f:id:horomary:20210331230243p:plain:w500
Z(s, a)とTZ(s, a)でquantileは当然変わらない

また、カテゴリ分布の最大値/最小値の設定に悩まなくてよくなりました。なぜならば累積分布関数の逆関数は0-1の有限区間で定義される関数であるためです。

さらに、Categorical DQN論文の最大の残課題は理論的にはWasserstein距離を最小化したいのだけれども、Wasserstein距離をそのままSGDのロス関数にするとBiased gradientとなってしまうので仕方なく分布間のKL距離を最小していたことです( 前記事を参照)。

そこでQR-DQNではターゲット分布の分位点を予測することが1-Wasserstein距離を最小化することを示し、このためにSGDのロス関数に 分位点回帰を使用することを提案しました。これにより直接Wasserstein距離をロス関数として使用することを回避してWasserstein距離を最小化できます

f:id:horomary:20210401001211p:plain:w500
論文Fig.2より:分位点を予測することが1-Wasserstein距離を最小化するになることの視覚的な説明


分位点回帰

分位点回帰とそのロス関数を簡単に説明します。

分布 \displaystyle{ Z
}の70%分位点を予測することを考えます。この分布 \displaystyle{ Z
}の 10%, 30%, 50%, 70%, 90% 分位点を  \displaystyle{ \hat{Z}
} = [-1.23, -0.29, 0. , 0.29, 1.23] とします。*6

f:id:horomary:20210401232756p:plain:w500
ターゲット分布Z

70%分位点の予測値をθと置くと、論文より分位点ロスは下式となります。
 \displaystyle{ \delta_{u \lt 0}} \displaystyle{ (Z - \theta) \lt 0 } のとき1、そうでなければ0という意味です。

f:id:horomary:20210401232954p:plain:w500
分位点ロス関数

この分位点ロスの視覚的な説明が下図です。ポイントは分布 \displaystyle{ \hat{Z} } のすべてのサンプルについて計算した分位点ロス(赤破線で表示)の平均が最終的な分位点ロスであることです。直感的には、予測値θより大きい値との距離総和と予測値θより小さい値との距離総和を予測したい分位点に応じてバランスしているという感じです。

f:id:horomary:20210402001048p:plain:w500
70%分位点(τ=0.7)を予測したい場合


分位点Huberloss

この分位点ロスをニューラルネットのロス関数にそのまま使うとu=0付近で滑らかでないため学習が不安定化するらしく、論文ではQuantile HuberLossを提案しています。と言っても |u|≦1のときは \displaystyle{ \rho_{\tau}(u) = 0.5u^{2}(\tau - \delta_{u \lt 0}) }、|u|>1のときは \displaystyle{ \rho_{\tau}(u) = (u-0.5)(\tau - \delta_{u \lt 0}) } とただのHuberLossに分位点重みがかかるだけのなので特に難しくはありません。


QR-DQNの実装

Breakout (ブロック崩し)環境向けにQR-DQNを実装します。 ネットワーク構造とネットワーク更新以外はオリジナルのDQNと完全に同じです。

horomary.hatenablog.com


QRネットワークの実装

ネットワーク構造自体はCategorical DQNとまったく同じです。構造は同じですが解釈が違うだけです。

アクション選択もCategorical DQNの場合と同様に価値分布Z(s, a)の平均値が最も大きいactionを選択します。ここで、分位の刻み幅を均等にとっている場合は、E[Z(s,a)]は分位点の単純平均と一致することに留意しましょう。


分位点ロスによるネットワーク更新


やってることは上述の分位点回帰の説明と同じです。しかし、上述の例では70%分位だけを計算していましたがQR-DQNでは設定されたすべての分位についてそれぞれ分位点ロスを計算する必要があるのでけっこう煩雑です。そこで、やってことがわかりやすいようにbatchsize=1の場合を下記に示しておきます。

import numpy as np
import tensorflow as tf

N = 5  #:分位の分割数
quantiles = np.array([0.1, 0.3, 0.5, 0.7, 0.9], dtype=np.float32)

target_quantile_values = np.array([23, 35, 42, 56, 76], dtype=np.float32).reshape(1, -1)
quantile_values = np.array([20, 32, 45, 50, 70], dtype=np.float32).reshape(1, -1)

target_quantile_values = tf.repeat(target_quantile_values, N, axis=0)
quantile_values = tf.repeat(quantile_values.reshape(-1, 1), N, axis=1)

td_error = target_quantile_values - quantile_values
indicator = tf.where(td_error < 0, 1., 0.)

#: k=1.0の場合のhuberloss
huberloss = tf.where(tf.abs(td_error) < 1.0, 
                     0.5 * tf.square(td_error), 
                     tf.abs(td_error) - 0.5)
quantiles = tf.repeat(quantiles.reshape(-1, 1), 5, axis=1)
quantile_weights = tf.abs(quantiles - indicator)

quantile_huberloss = quantile_weights * huberloss
total_quantile_huberloss = tf.reduce_mean(quantile_huberloss, axis=1, keepdims=True)
loss = tf.reduce_sum(total_quantile_huberloss, axis=0)


Breakoutでの学習結果

BreakoutDeterministic-v4環境(ブロック崩し)において、GCPのn1-standard-4(4-vCPU, 15GBメモリ) + GPU K80 のプリエンティブルVMインスタンスを使って24時間学習した結果十分なパフォーマンスを確認できました。

Breakoutはatariの中では比較的単純な環境であることを考慮して、Adamの学習率は論文より高め設定のlr=0.00025(論文記載はlr=0.00005) & 分位点の刻み数Nを論文より小さめ設定のN=50(論文記載は分位点の刻み数N=200)にしています。

f:id:horomary:20210401004514p:plain:w500

f:id:horomary:20210401004556p:plain:w500

コード全文: github.com


次:FQF

そのうち


*1:Bellemareさんが著者リストに入ってる

*2:Categorical 51でC51。もしDistributional 51でD51命名されてたとしてもやっぱり蒸気機関車

*3: 論文ではprojectionと呼称

*4:とくにバッチサイズ大きいと処理が重い。このあたりの煩雑さがパフォーマンスは優秀なのにApeX-DQNではハブられた理由なのではないかと邪推している

*5: atari環境ではreward clippingが有効なのであまり問題になりません

*6:分かりやすさのため分位を明示しているが、確率密度に従ってサンプリングされていれば分位が分かっている必要はない

rayで実装する分散強化学習 ③Ape-X DQN

深層強化学習における超大規模分散並列化の有用性を示したApeX-DQN(Distributed Prioritized Experience Replay)をtensorflow2とrayで実装します。手法の構成要素自体はRainbowとだいたい同じであるため、本記事の焦点は分散並列学習の実装です。

f:id:horomary:20210227174718p:plain:w1000

rayで実装する分散強化学習
Pythonの分散並列処理ライブラリRayの使い方 - どこから見てもメンダコ
rayで実装する分散強化学習 ①A3C(非同期Advantage Actor-Critic) - どこから見てもメンダコ
rayで実装する分散強化学習 ②A2C(Advantage Actor-Critic) - どこから見てもメンダコ


前提手法:
DQNの進化史 ①DeepMindのDQN - どこから見てもメンダコ
DQNの進化史 ②Double-DQN, Dueling-network, Noisy-network - どこから見てもメンダコ
DQNの進化史 ③Prioritized experience replay, Multi-step learning, Categorical DQN - どこから見てもメンダコ
DQNの進化史 ④Rainbowの実装 - どこから見てもメンダコ


はじめに

[1803.00933] Distributed Prioritized Experience Replay
Distributed Prioritized Experience Replay | OpenReview

Distributed Prioritized Experience Replay、あるいはApe-X*1はその名の通り 優先度付き経験再生を大規模分散並列学習に対応させた手法です。Ape-XはDDPGにもDQNにも適用可能な手法ですが、後者に適用された場合にはApe-X DQNと呼称されます。Ape-X DQNはオフポリシー強化学習の大規模分散並列化は訓練時間*2を短縮するだけでなく、パフォーマンスの向上にも寄与することを当時のatari環境における圧倒的SotAで示しました。

f:id:horomary:20210227130719p:plain:w500
並列化しているので訓練時間が短縮されるのは当然であり、パフォーマンス向上が著しいことのインパクトが強い

DQNはオフポリシー手法であるので、収集した遷移情報(経験)を何度でも再学習してよいはずなのですが、Ape-Xはそのような循環式の経験再生よりも、源泉かけ流し的な贅沢な経験再生の方がパフォーマンスが良くなるということを示しました。この発見が以降の強化学習手法における大規模分散並列学習トレンドを加速していくこととなります*3


Ape-X DQN の概要

Ape-Xの基本コンセプトは遷移情報を収集するプロセス、遷移情報を蓄積を担うプロセス、および勾配計算してネットワークを更新するプロセスを完全に分離することによる効率化です。この分散学習アーキテクチャはFig.1に示されており、Learner, Actor, Replay という3つの主要な役割があることが分かります。

f:id:horomary:20210227144313p:plain:w600
apexの分散並列アーキテクチャ

Learnerの役割

Leanerプロセスには1CPU, 1GPUが割り当てられます。

Leanerの役割はReplayから供給されるミニバッチでひたすらにQネットワークを更新しつづけること、およびActorからのネットワーク重み同期要求に応じることです。学習速度を最大化する(≒GPU稼働率を最大化する)ためにReplayからのミニバッチの供給を途切れさせないことが重要となります。

Actorの役割

各Actorプロセスには1CPU, 0GPUが割り当てられます。このActorプロセスは論文では最大360並列実行されています。

Actorの役割は、遷移情報の収集各遷移の初期優先度の算出です。ActorはQネットワークを持ち自律的にrolloutを行います。100step程のrolloutを行った後に勾配計算は行わず集めた遷移情報をそのままRepalyプロセスに送信します。ただし、遷移情報の送信時にはローカルQネットワークでの推論により初期優先度(∝TD誤差)を算出しておきます。

オリジナルの優先度付き経験再生では、各遷移の初期優先度には最大値を割り当てることで必ず一回は経験が再生されるようにしていましたが、経験再生される速度よりも経験の供給速度の方が圧倒的に速いApe-Xアーキテクチャでそれを行うと直近の遷移ばかりが再生されることになりReplayが意味をなさないため、Actorプロセスで初期優先度を計算することにより遷移情報をふるいに掛けています。換言するとApe-Xアーキテクチャでは一度も再生されないまま消えゆく遷移情報もあるということで、当然サンプル効率は劣悪です*4

ActorがローカルQネットワークを保持して自律的にrolloutを行うという点ではA3Cアーキテクチャと同じですが、Ape-XではA3Cと異なりActorは勾配計算せずReplayへ遷移情報を送信するだけです。これは、勾配情報だとグローバルQネットワークへの反映遅れに気を使う必要がありますが、遷移情報ならばReplayへの反映が多少遅れても問題ないので大規模分散学習にて扱いやすいためです。また、勾配計算するならActorにもGPUが無いと厳しいですが推論だけならCPUだけでもそれほど苦しくないという利点もあります。

また、分散並列化されたActorはすべて異なる探索率εが割り当てられる、というのも重要なポイントです。従来のシングルActorのDQNでは、高い探索率εで学習を開始しゆっくりとεを下げていくアニーリング方式によって探索と活用のトレードオフをバランスしていました。これに対してApe-Xではさまざまな探索率のActorが存在するので自然に多様な経験を収集することができます。このような異なる方策(探索率)を持った並列Actorでのrolloutは、DQNがoff-policyであること生かしたテクニックと言えます。

f:id:horomary:20210227171147p:plain:w300
並列Actorはすべて異なる探索率εを割り当てられる

Replayの役割

Replayの役割はActorからの遷移情報受け取りLeanerに供給するミニバッチの作成、およびLeanerからの更新優先度情報の受け取りです。実態はただの優先度付きReplayBufferなのですが、ActorともLeanerともやり取りしなければいけないため一番忙しいプロセスです。


Rainbowからの継承要素

Ape-X DQNは優先度つき経験再生の後継手法というよりは、DQNの改良トリック全部盛り手法である Rainbow + 大規模分散並列学習 と表現するほうが正確でしょう。実際に、Rainbowが採用していた6つのDQN改良トリック(Double Q-learning, Dueling network, Noisy-network, Prioritized Experience Replay, Multi-step learning, Categorical DQN)のうち、Ape-X DQNではNoisy-networksとCategorical DQN (C51) 以外はすべて採用しています。

上述の通りApe-Xでは各Actorに異なる探索率εを割り当てることにより(計算パワーの力で)探索と活用のバランスをとるので、同じく探索戦略であるNoisy-networkを除外することはごく自然です。一方で、Rainbowに採用された6つDQN改良トリックのうち単体でもっともパフォーマンスの高いC51を除外するのは明らかに違和感がありますが、OpenReviewでの回答を見る限りでは単に実装の煩雑さを嫌っただけのようです*5

Q1: on using all Rainbow components and on using multiple learners.

These are both interesting directions which we agree may help to boost performance even further. For this paper, we felt that adding extra components would distract from the finding that it is possible to improve results significantly by scaling up, even with a relatively simple algorithm. (https://openreview.net/forum?id=H1Dy---0Z)

horomary.hatenablog.com

大規模並列Actorの効果検証

Actorを増やすほどReplayBuffer内の繊維状の入れ替わりサイクルが短くなるため、より最近に収集された遷移情報が再生されやすくなります。また、ある経験が再生される回数が少なくなり源泉かけ流しに近くなっていきます。これはある意味でon-policy学習のやり方でQネットワークを訓練していると表現できます。

on-policyっぽくDQNの訓練を行うことがパフォーマンスの向上の理由ならば、actorの数を大規模並列化せずともApe-Xと同等のパフォーマンスが得られるはずです。論文ではこれについての検証実験を行っており、Fig.6はactorの並列数(n)=32に固定したうえで、ある遷移情報が再生される回数(k)を変化させるとパフォーマンスにどう影響するのかを示しています。同じactor数(n=32)では再生回数kの違いは大差ないことがわかります。さらにactorの並列数(n)=256の場合はn=32の場合と比較してパフォーマンスに大きな差をつけています。

この結果から論文では経験再生の新陳代謝の速さに由来するon-policyっぽい学習だけでなく、大規模並列actorによって生成される多様な経験がパフォーマンスに寄与していると結論付けています。

f:id:horomary:20210227175440p:plain:w600
経験のrecencyの検証実験(Fig. 6)と探索率εの多様さの検証実験(Fig.7)

ついでにFig.7では各Actorへの探索率εの割り当ての多様性の効果について検証しています。各Actorにすべて異なる探索率を割り当てた時と、割り当てる探索率を6つに減らした時にどうパフォーマンスが変わるかの検証実験です。前者はepsilons = np.linspace(0.01, 0.4. num_actors)で後者はepsilons = np.linspace(0.01, 0.4. 6)という感じのイメージ*6だと思います。結果は直感通りで、パフォーマンスにそこまでの大差なしとのこと。


CartPole環境での簡易実装

まずは分散並列アーキテクチャを理解するために、CartPole環境で優先度付き経験再生以外のDQN改良トリックを除外したシンプルな実装を示します。

分散学習の実装

Ape-X アーキテクチャで重要なのはLearner(=GPU)を休ませないことです。このためにLeanerには16セットのミニバッチを渡し、Learnerがネットワークの更新をしている裏でReplayはせっせとActorから経験を受け取っていきます。*7 この流れはrayを使うことですっきり記述できます。

Pythonの分散並列処理ライブラリRayの使い方 - どこから見てもメンダコ


実装上のポイントは39行目のfinished_learner, _ = ray.wait(wip_learner, timeout=0)です。timeout=0を指定したray.waitは実行時点で対象プロセスが未完了である場合、finished_learnerとして空リストを返すためLearnerプロセスの終了判定が可能です。このLeanerプロセス終了判定をActor to Replayでの遷移情報送付が1回行われるごとに実行することで、Learnerプロセスが空き次第すぐに次のminibatchを渡すという疑似的な割り込み処理を実装することができます。

この実装は45,46行目でReplayプロセスが行う、次のミニバッチセットの作成と優先度の更新処理がLeanerプロセスに比べて十分に速いことを前提としていることに注意してください。もしReplayプロセスが遅すぎる、あるいはLeanerプロセスに渡すミニバッチの数が少ないためにLeanerプロセスの完了が早すぎる場合にはActorからの遷移情報送付が滞ってしまいます。このため、一度のネットワーク更新ごとに何回Actorからの遷移情報送付が行われているかはしっかりチェックしておきましょう。

Actorへ最新の重みを渡すためにray.putを使用していることに留意してください。ray.putは大きめのデータ、この場合はメインQ関数の重みを多数のリモートActorに配布する処理を効率化してくれます。

Ray Core Walkthrough — Ray v2.0.0.dev0

ApeXのコアとなるコードはこれだけです。 rayのおかげでシンプルに実装できていると思います。以下では各プロセスの詳細実装を紹介しますが、DQNおよび優先つき経験再生を理解できていればとくに難しいことは無いはずです。


Actorの実装

Actorプロセスは一定stepのrolloutを行って遷移情報するだけなので実装はごく単純です。A3Cと同様にApeXではActorがローカルQネットワーク(LearnerのQ関数のコピー)を持ち自律的にrolloutを行うので、Actor.rolloutではまず初めにleanerのQ関数と重みの同期を行います*8。その後、100step分のrolloutを行い、収集した遷移情報とそれらについてのTD誤差(初期優先度の計算に使用)を返します。


Replayの実装

Replayは単なる優先度付き経験再生です。Ape-XではReplayプロセスが行うミニバッチ作成(Replay.sample_minibatch)の速度パフォーマンスが求められるので、高速な重み付きサンプリングができるSumTree構造で優先度を保存しています。他の留意点として、オリジナルの優先度付き経験再生では、Importance Sampling weights(もどき)のハイパラであるβをアニーリングしていましたが、Ape-Xでは固定値になっています。

horomary.hatenablog.com


Learnerの実装

LearnerはReplayから受け取ったミニバッチでひたすらネットワーク更新するだけのプロセスです。16セットのミニバッチを消費したら最新の重み、および更新された優先度をメインプロセスへ返却します。優先度付き経験再生を理解していれば何も難しいことはありません。

horomary.hatenablog.com


Qネットワーク

特筆することは何もありませんが一応載せておきます。


学習結果:CartPole-v0

CartPole-v0のハイスコア200点に到達するまでおよそ20秒!

f:id:horomary:20210301223207p:plain:w400
1cycle(横軸)ごとにLeanerは16セットのミニバッチでネットワーク更新

(もちろんマシンスペックに依存しますが)8並列Actorの環境ではLearnerが16セットのミニバッチ(各batch_size=32)を消化する間にActorからReplayへの遷移情報送信が25回程度行われていました。


Atari環境(Breakout)での実装

さて、CartPoleでの簡易実装がうまくいったのでBreakout(ブロック崩し)でDQN改良トリックまで含めて本格実装したコード例も掲載します。と言いたいところだったのですが、過去記事で紹介したRainbowの実装と丸被り & コードが長大なので結果だけ示します。改良トリックの詳細は過去記事で、実装全体はGithubでご確認ください。大筋は上に示したCartPoleと同じですが、Dueling-net, Double-DQN, Multi-step Learning、およびメモリ節約のために遷移情報をzlibで圧縮するコードが追加されています。

horomary.hatenablog.com

github.com

学習結果:BreakoutDeterministic-v4

リソースの都合上*9、actorは20並列にしかできませんでしたがそれでも分散学習の威力を実感できる結果となりました。探索率εの大きいactorが混じっていることによる正則化?効果のおかげかパフォーマンスが大崩れせず順調に学習が進みます。ただし、Breakout環境では探索率εの上限値が論文通りの0.4では小さすぎるためか学習の立ち上がりが悪く感じたのでεの上限は0.5に変更しています。

f:id:horomary:20210301215216p:plain:w500
20actorでトータル15時間学習(ε=0.01)

f:id:horomary:20210301215756p:plain:w500
Learnerのlossの経過


次:R2D2

そのうち。

*1:Asynchroneous Prioritized EXperience replay?

*2:wallclock time, 実世界での時間

*3:そして一般人や小規模ラボが参入しにくくなった

*4:サンプル効率についてはFIg.10を参照

*5:C51はネットワーク更新時のCPU処理も多いから設計が面倒になるのかも

*6:εの割り当ての具体的な数値は記述無し

*7:もし実装レベルまで論文を再現したいならtensorflow.Queueで実装する

*8:論文では400stepごとに重みを同期、とあるのでこの実装のように毎回同期はしない

*9:GCPへのリソース割り当て増加リクエストが通らなかった

Segment Tree(セグメント木)による重み付きランダムサンプリング

競技プログラミング界隈では一般教養であるらしいセグメント木のSum-tree構造で高速な重み付きサンプリングを実装します。


はじめに

強化学習の重要手法である優先度付き経験再生(Prioritized Experience Replay)では、重みづけされた100万の経験(遷移情報)からランダムにサンプリングしてミニバッチを作成する、という処理があります。このような重みづけサンプリングはnp.random.choiceの引数pに重み情報を与えることで楽に実装できます。コードの見通しが大変よくなるので過去記事ではこの方法での実装例を紹介しました。

しかし論文ではsum-treeデータ構造で実装すると速いと書いてあります。本記事ではせっかくなのでこちらの実装を試してみます。

[1511.05952] Prioritized Experience Replay

DQNの進化史 ③Prioritized experience replay, Multi-step learning, Categorical DQN - どこから見てもメンダコ

DQNの進化史 ④Rainbowの実装 - どこから見てもメンダコ


A. numpy.choiceによる重み付きランダムサンプリング

まずはベースラインとしてnumpy.random.choiceによる重み付きランダムサンプリングのパフォーマンスを見ます。要素数は100万で各要素には0-5の優先度が割り当てられます。DQNでのミニバッチサイズの32に従って1iterで32要素をサンプリングします。また、Breakout(ブロック崩し)環境ではそれなりに学習が進むと1 episodeで200回くらいはミニバッチ作成するので200iter繰り返します。

f:id:horomary:20210215220249p:plain:w500

結果は約3.2秒となりました。1episodeあたり3.2秒なら趣味で強化学習やるくらいなら許容できる程度ではありますがちょっと遅いですね。


B. 累積和による重み付きランダムサンプリング

つぎに愚直な実装として逆関数法による重み付きランダムサンプリングを実装します。累積密度関数が計算できる確率分布なら逆関数法を使うことでい一様乱数から目的の確率分布に従う乱数に変換できます。

逆関数法 - Wikipedia

逆関数法を用いた乱数生成の証明と例 | 高校数学の美しい物語

たとえば、4要素のリストにおいて各要素の優先度が [4, 2, 1, 3] のときは、0≦ z ≦ 4+2+1+3 = 10 の範囲で一様乱数を発生させ、累積和がzとなるのがどの要素のときかを調べることで優先度に従ったサンプリングを行うことができます。

f:id:horomary:20210215223407p:plain:w600
0≦z≦10で乱数を発生させて累積和がzに該当する要素を選択すれば優先度の大きさにサンプリング確率が従う

f:id:horomary:20210215224616p:plain:w500

たしかに優先度に従ってサンプリングできていることがわかります。では要素数が増えた時のパフォーマンスがどうなるかをnp.random.choiceと同じ条件で確かめます。

f:id:horomary:20210215231046p:plain:w500

結果は144秒、遅い! すべての要素に対して累積和チェックをしているので計算量がNになってしまっていることが原因です。

C. Sum-tree構造を活用した重み付きランダムサンプリング

上で重いのは累積和がzになるのはどの要素番号のときであるかを調べる処理です。これはSegment-tree(セグメント木)構造を使うことで高速に検索することができます。さきほどと同様に各要素の優先度が [4, 2, 1, 3] のときのSum-treeを構築すると下図のようになります。

f:id:horomary:20210215232341p:plain:w500
[4, 2, 1, 3]に対するSum-tree

たとえば累積和が6.5を超える要素番号を検索したいとしましょう。ルートノードである10の左子ノードが6なので、要素0, 1までの累積和が6であることがわかります。よって、要素番号2,3の区間における累積和が0.5(= 6.5 - 6)になる要素を探せばよいというわけです。そこでルートノードの右子ノード4に進みます。この子ノードを見ると1, 3なので左子ノードで累積和が0.5になることが分かります。左子ノードは実要素なのでここで探索終了となります。

実際に格納される要素数Nが  N = 2^{K} のとき、Sum-treeの深さ(階層?)はKになるので探索回数は要素数Nに対してlogNとなり効率的であることがわかります。

Sum-TreeのPython実装

競プロ界隈の人はわざわざ遅いpythonで実装とかしないかもしれませんが、強化学習で使う分にはそこそこのパフォーマンスが出ればよいのでpythonでSum-Treeを実装します。この実装は ray/segment_tree.py at master · ray-project/ray · GitHub から抽象化を削りシンプルに再実装したものです。

実装のポイント:
・格納される要素数がNのときSumtree全体の要素数は2N-1
・ルートノードのインデックス番号を1に設定すると左子ノードのインデックス=2×親のインデックス、右子ノードのインデックス=2×親のインデックス+1 となり便利なので、インデックス番号0を使わない長さ2NのリストでSum-treeを実装する
__setitem____getitem__ を活用してsumtreeであることを感じさせない使い勝手を実現する

f:id:horomary:20210215235032p:plain:w500
クラス外から見た要素番号(赤)と実体の要素番号(青)

f:id:horomary:20210216004207p:plain:w500

速度パフォーマンスの確認

numpyのときと同様にバッチサイズ32のミニバッチを200回作った時の速度を計測します。

f:id:horomary:20210216004530p:plain:w400

サンプリング時間だけ見ればnumpy.random.choiceより30倍程度は速いことがわかります。


おわりに

__setitem____getitem__ が一番輝くのはsegment tree説

DQNの進化史 ④Rainbowの実装

Deep-Q-Network (2013) 以降の深層強化学習(Q学習)の発展を、簡単な解説とtensorflow2での実装例と共に紹介していきます。今回はDQNの改良トリックを全部盛りにしたら強いんでは?という脳筋発想によって生まれた手法であるRainbowを実装します。


DQNシリーズ
DQNの進化史 ①DeepMindのDQN - どこから見てもメンダコ
DQNの進化史 ②Double-DQN, Dueling-network, Noisy-network - どこから見てもメンダコ
DQNの進化史 ③Prioritized experience replay, Multi-step learning, Categorical DQN - どこから見てもメンダコ
DQNの進化史 ④Rainbowの実装 - どこから見てもメンダコ


はじめに

[1710.02298] Rainbow: Combining Improvements in Deep Reinforcement Learning

2017年に発表されたRainbowは、それまで報告されてきたDQN改良トリックをすべて搭載したDQNの総まとめ的な手法です。具体的にはオリジナルのDQNに、Double Q-learning, Dueling-network, Noisy-network, Prioritized Experience Replay, Categorical DQN(C51, or Distributional DQN), Multi-step learningの6つの手法を全部盛りにすることにより当時のatari環境のSotAを更新しました。

f:id:horomary:20210211143435p:plain:w400
虹色の線がオシャレ(論文Fig.1)

手法自体に目新しいことは無いので本記事では実装レベルの解説をしていきます。


構成要素の寄与について

論文ではRainbowの各構成要素をひとつ抜いた時にどれだけパフォーマンスが下がるか、という実験をしています。これによると優先度付き経験再生(prior)、分布強化学習(distributional)、Multi-step learningの寄与が大きいようです。

f:id:horomary:20210211144800p:plain:w300
構成要素を一つ抜いたらどれだけパフォーマンスが下がるか(論文 Fig. 3)

ただし、あくまで要素の一つ抜きだけでありすべての組み合わせを試しているわけではないので、各要素の寄与はFigに示されているほど単純に見積もることはできません。実際にtensorflowのチュートリアルDQN C51/Rainbow  |  TensorFlow Agents)ではDistributional DQNとMulti-step learningの組み合わせだけでRainbowと同等のパフォーマンスが得られたという記述がされています。

Although C51 and n-step updates are often combined with prioritized replay to form the core of the Rainbow agent, we saw no measurable improvement from implementing prioritized replay. Moreover, we find that when combining our C51 agent with n-step updates alone, our agent performs as well as other Rainbow agents on the sample of Atari environments we've tested.
(C51とn-step更新は、優先リプレイと組み合わされてRainbowエージェントのコアを形成することがよくありますが、優先リプレイを実装しても測定可能な改善は見られませんでした。さらに、C51エージェントをnステップ更新のみと組み合わせると、テストしたAtari環境のサンプルで他のRainbowエージェントと同様に機能することがわかりました。)


レーニングループ

レーニングループはDQNとほぼ同じです。


ハイパーパラメータはrainbow論文ではなく、 Categorical DQN論文(Distrubutional DQN)に従っていることに注意してください。これはBreakout環境ではDistributional DQN単体の方がパフォーマンスが良いためです。

f:id:horomary:20210211162307p:plain:w600
breakoutはrainbowより分布DQN単体の方が強い

Q-networkの実装(tensorflow2)

Qネットワークには Dueling-network, Categoical DQN (Distributional DQN), Noisy-network の3要素が導入されます。


ReplayBufferの実装

経験バッファには 優先度付き経験再生, Multi-step learning の2要素が導入されます。


ネットワーク更新(tensorflow2)

ネットワーク更新に絡んでくるのは、Double Q-learning、 優先度付き経験再生、Categorical DQN(Distributional DQN)の3要素です。

見ればわかりますがDistributional DQNの存在がTD誤差の計算をやたらと煩雑にしています。この詳細は過去記事を参照ください。

horomary.hatenablog.com


Breakoutでの学習結果

Breakout(ブロック崩し)を GCPのn1-standard-4(4-vCPU, 15GBメモリ) + GPU K80 のプリエンティブルVMインスタンス*1で24時間学習させました。1Mstep未満で40点取れてるので動作確認としては十分なスコアだと思います。(※rainbow論文では200Mstepを学習)

f:id:horomary:20210211172902p:plain:w500

いろいろ試してみましたが、どうもBreakout環境はNoisy-networkとの相性が悪い印象を受けました。また、優先度つき経験再生を雑に実装*2しているため処理速度がかなり遅く24時間で1Mstepしか進行していません。

実装全体はgithubへ:
github.com


次:Ape-X DQN

C-51がいないおかげでRainbowよりApe-Xのほうが実装が楽に感じる。

horomary.hatenablog.com

*1:24時間でシャットダウンされる代わりに激安なインスタンス

*2:低スぺ対応のためにzlibでオブジェクト圧縮したり重みつきサンプリングをnp.random.choiceで実装したり

DQNの進化史 ③優先度付き経験再生, Multi-step learning, C51

Deep-Q-Network以降の深層強化学習(というか深層Q学習)の発展を、簡単な解説とtensorflow2での実装例と共に紹介していきます。今回は経験再生の改良である優先度付き経験再生(Prioritized experience replay)、方策勾配法ではよく使われるMulti-step learning, そして深層分布強化学習の有用性を示したCategorical DQN を紹介します。

DQNシリーズ
DQNの進化史 ①DeepMindのDQN - どこから見てもメンダコ
DQNの進化史 ②Double-DQN, Dueling-network, Noisy-network - どこから見てもメンダコ
DQNの進化史 ③Prioritized experience replay, Multi-step learning, Categorical DQN - どこから見てもメンダコ
DQNの進化史 ④Rainbowの実装 - どこから見てもメンダコ

前提手法:DQN horomary.hatenablog.com


Prioritized experience replay(2015)

[1511.05952] Prioritized Experience Replay

オリジナルのDQNではReplayBufferに蓄積した遷移情報からのランダム選択によってミニバッチを作成します。しかし、遷移情報をランダムに選択するのでは思いがけず上手くいったような貴重なイベント(遷移)を学習する効率が悪いですね。そこで、Prioritized experience replay(優先度つき経験再生)ではその名の通り、意外性の高い遷移を優先してReplayBufferからサンプリングします。

具体的には、TD誤差δの大きい遷移情報ほど意外性が高いと見なし、TD誤差δの絶対値の大きさに応じてサンプリングされる確率に重みをつけます。*1

TD誤差δ:
 \displaystyle{
\delta = r_t  + \gamma \max_{a'} Q_{target}(s_{t+1}, a') - Q(s_t, a_t)
}


バッファ内のi番目の遷移情報がサンプリングされる確率P(i):
 \displaystyle{
P(i) = \frac{(|\delta_{i}| + \epsilon)^{\alpha}}{\sum_{k}^{N} (|\delta_{k}| + \epsilon)^{\alpha}  }
}

P(i)内のαはハイパーパラメータ(0≦α≦1)です。α=0のとき P(i) = 1/N になりランダムサンプリングと同一となります。また、εはサンプリング確率が完全にゼロになってしまうことを防ぐための適当な微小量です。

このような重みづけサンプリングは速度パフォーマンスを気にしないならばnp.random.choiceで楽に実装できます。速度を気にするならSegmentTree*2で実装しましょう。簡単のためにここでは前者の実装例を示します。

#: Nは蓄積されている遷移情報の総数
 probs = priorities / priorities.sum()
 indices = np.random.choice(np.arange(N), p=probs, replace=False, size=batch_size)

sumtreeでの高速な優先度付きサンプリングについては別記事を参照ください。

horomary.hatenablog.com


補正TD誤差によるネットワーク更新

優先度に応じてサンプリングすると同じ遷移情報を執拗に学習することになり学習の安定性を損なう恐れがあるため、Q関数の更新時には遷移から計算されるTD誤差にサンプリング確率に応じた補正を行います*3

補正重み:
 \displaystyle{
w_i =  \Big(\frac{1}{N} \cdot \frac{1}{P(i)} \Big)^{\beta}
}

βは補正の強さを決めるハイパーパラメータ(0≦β≦1)です。β=0のときw=1で補正無しとなりβ=1のとき完全な補正となります(優先度付きのサンプリング確率がP(i)であるため)。βはオリジナルDQNにおける探索率εと同様、学習中にアニーリングを行います。atari環境ではβ=0.4から始めて学習終了時にβ=1.0になるように線形に増加させるのが良いようです。

補正重みはミニバッチ平均をとる前のTD誤差に適用します

ネットワークの更新のために計算したTD誤差で各遷移の優先度を更新するのも忘れないようにしましょう。


リプレイバッファの実装:

atari環境で100万ステップを愚直に蓄積するとメモリ消費が大変なことになるのでzlibで圧縮しています。

【追記:2021/02/15】
パフォーマンス検討の結果、この実装例で処理速度のボトルネックになっているのはnp.random.choiceではなく、listからndarrayへの変換処理が入る probs = np.array(self.priorities) / sum(self.priorities)でした。よって、self.prioritieslistではなくnp.ndarrayで実装するとget_minibatchの速度パフォーマンスが10倍くらい改善されます。


Multi-step learning

Multi-step learning というアイデアは新しいものではありませんが、Rainbow (2017) やApe-X-DQNの構成要素となっているので紹介しておきます。

まず、オリジナルのDQNでは下式のように1step分の遷移情報を使用してTD誤差を計算します。

1step-TD誤差
 \displaystyle{
\delta = r_t  + \gamma \max_{a'} Q_{target}(s_{t+1}, a') - Q(s_t, a_t)
}

これに対してMulti-step learningではその名の通り、1stepじゃなくNstepの遷移情報でTD誤差を計算します。ここでは具体的にRainbow(2017)でも採用されている3stepTD誤差の式を示します。

3step-TD誤差
 \displaystyle{
\delta = r_t  + \gamma ^{1} r_{t+1} +  \gamma ^{2} r_{t+2}  + \gamma ^{3} \max_{a'} Q_{target}(s_{t+3}, a') - Q(s_t, a_t)
}

これだけです、とくに難しいことはないですね。

Multistep-learningを採用することで遅延報酬が伝搬しやすくなると考えられます。例えば、Breakout(ブロック崩し)ではボールを弾く(行動選択)タイミングに対して、ブロックを崩す(報酬を得る)のは何フレームか後です。このような行動選択に遅延して生じる報酬の因果関係を学習しやすくなることが期待できます。とはいえ、あまり長く先まで見すぎるとそれはもはやモンテカルロ推定でありバイアスが大きくなりすぎるのでほどほどにします。

補足:
atari環境では問題なく機能しますがMultistep-learningを採用したDQNは厳密にはoff-policyではないことに留意。

Off-policy n-step learning with DQN - Data Science Stack Exchange

実装例

Multistep-learningのコンセプトは単純なのですがどう実装するか、というかどの部分に実装するかは迷うところです。とりあえずreplay_bufferにこの役割を担わせることにしました。こうするとオリジナルDQNからのコード変更が少ない気がします。


Categorical DQN(2017)

[1707.06887] A Distributional Perspective on Reinforcement Learning

Rainbow以前の単体でもっとも強力なDQN拡張手法はおそらくこのCategorical DQN (C51)でしょう。*4 Dueling-networkや前述の優先度付き経験再生など他のDQN拡張手法を使用せずに単体で当時のatari環境のSotAを奪いました。

f:id:horomary:20210208220529p:plain:w500
A Distributional Perspective on Reinforcement Learn より

通常のQ学習では状態行動価値Q(s, a)の期待値を近似するのに対して、Categorical DQNでは状態行動価値Q(s, a)の分布を近似することを狙います。状態行動価値Qの期待値でなく分布を近似することでなぜ性能が大幅向上するかは実はよくわかっていないものの、その圧倒的な性能で深層分布強化学習という分野を切り開き、D4PGQR-DQNIQN など多くの後継手法を生みました。

Categorical DQNは単体でもっとも強力なDQN拡張手法であると同時に、単体でもっとも実装が煩雑なDQN拡張手法でもあります。難解ではありませんが煩雑です。 ゆえに解説するとわりと長くなるので実装は別記事を参照ください。

horomary.hatenablog.com


次:Rainbow

horomary.hatenablog.com

*1:TD誤差の絶対値ではなくTD誤差の大きさのリプレイバッファ内順位に応じて重みづけるrank-baseの方法も提案されているがatari環境では絶対値ベースの方法がパフォーマンスがよいらしい

*2:競プロの一般教養らしい. rllibの実装がわかりやすい ray/segment_tree.py at master · ray-project/ray · GitHub

*3:勾配クリップと似たような効果かなと思っている

*4:Distributional DQNと呼称されることもあります。

DQNの進化史 ②Double-DQN, Dueling-network, Noisy-network

Deep-Q-Network (2013) 以降の深層強化学習(Q学習)の発展を、簡単な解説とtensorflow2での実装例と共に紹介していきます。今回はオリジナルのDQNを微修正するだけで実装可能な改良手法である、Double DQN , Dueling-network, そしてNoisy-network の3つを紹介します。

DQNシリーズ
DQNの進化史 ①DeepMindのDQN - どこから見てもメンダコ
DQNの進化史 ②Double-DQN, Dueling-network, Noisy-network - どこから見てもメンダコ
DQNの進化史 ③Prioritized experience replay, Multi-step learning, Categorical DQN - どこから見てもメンダコ
DQNの進化史 ④Rainbowの実装 - どこから見てもメンダコ

前提手法:DQN horomary.hatenablog.com


Double DQN(2015)

[1509.06461] Deep Reinforcement Learning with Double Q-learning

オリジナルDQNのベルマンオペレータでは行動選択もQ(s, a)の評価もtarget-Q-networkが行います。

DQN
 \displaystyle{
\mathcal{T}Q(s, a) = r  + \gamma \max_{a'} Q_{target}(s', a')
}

 \displaystyle{
a' = \text{arg} \max_{a'} Q_{target}(s', a')
}

Double-DQN論文では、ベルマンオペレータ内の \displaystyle{
\max_{a'} Q(s', a')
}において、行動選択するネットワークとQ値を評価するネットワークが同一である場合に推定Q値が過大評価される傾向があることを指摘しまし*1た。さらに、その低減策として行動決定をQ-networkによって行い、Q(s, a)の評価はtarget-Q-network によって行うDouble-Q-learingの適用を提案しました。

Double DQN
 \displaystyle{
\mathcal{T}Q(s, a) = r  + \gamma \max_{a'} Q_{target}(s', a')
}

 \displaystyle{
a' = \text{arg} \max_{a'} Q(s', a')
}


実装例

実装はとても単純です。


Q値の過大評価の直感的理解

論文を乱暴に要約すると、 \displaystyle{Q(s, a)} にランダムノイズが乗る場合には \displaystyle{\max Q(s, a)}でQ値の過大評価が起こるよ、ということを説明しています。

直感的な理解のために、5次元アクション空間での状態sにおける  \displaystyle{ Q^{*}(s, a) = } [  \displaystyle{0, 0, 0, 0, 0 } ] の場合、つまりどの行動を取ってもQ(s, a) = 0 である極端な状態*2を考えましょう。各アクションaについての  Q^{*}(s, a) の推測値に一切のノイズが無い場合は当然ながら、

 \displaystyle{
E \left[{ \underset{a}{\max} Q^{*}(s, a) }\right] = 0
}

となります。しかし、報酬やダイナミクスの確率的要因によって生じる正規分布ノイズε が各状態行動価値の推測値に乗る場合、
つまり  \displaystyle{ Q^{*}(s, a) = } [  \displaystyle{ 0+\epsilon_0, 0+\epsilon_1, 0+\epsilon_2, 0+\epsilon_3, 0+\epsilon_4 } ] である場合は、

 \displaystyle{
E \left[{ \underset{a}{\max} Q^{*}(s, a) }\right] \neq 0
}

となり真のQ値からずれます。このことをjupyter notebookでの簡単な実験で確かめます。

f:id:horomary:20210205010153p:plain:w500

真のQ値は0であるので確かに過大評価が生じています。この過大評価は行動選択を、平均値が同じだが別のノイズが乗ったネットワーク(target-Q-networkに対するQ-network)で行うことによって低減することができる、というのがDouble DQNの直感的理解です。

2つのQ-networkでノイズを低減するというDouble DQNの考え方は連続値制御のためのオフポリシー手法であるTD3(DDPGの後継手法)でもclipped-double-Q-learningという名称で採用されています。

horomary.hatenablog.com


Dueling Network(2016)

[1511.06581] Dueling Network Architectures for Deep Reinforcement Learning

Dueling Networkは強化学習では珍しくネットワーク構造を工夫した手法です。Dueling-networkアーキテクチャではQネットワークが出力する"状態価値 V(s)" と "アドバンテージ A(s, a)" の和をとったものをQ(s, a)とします。*3

f:id:horomary:20200609221814p:plain:w600
Dueling-network論文の図に注釈をつけたもの

アドバンテージA(s, a)とは状態行動価値Q(s, a)から状態価値V(s)を引いた値であり、 \displaystyle{
 Q(s, a)  =  V(s) + A(s, a)
} で定義されます。

状態価値V(s)とアドバンテージA(s, a)を分けて学習するアーキテクチャは、どのような行動をとっても価値がほとんど変わらないような状態において学習を促進すると考えられます。たとえば下図のボンバーマンのような状態を考えてみます。

f:id:horomary:20200609215031p:plain:w300
どの行動を選択しても価値がほとんど変わらない状態の例

どんな行動をとっても結果が変わらないことは明らかであるので、Q(s, a) ≒ V(s) です。V(s)は行動選択に依存しない価値であるのでDueling-networkアーキテクチャではどのアクションaを選択してもQ(s, a)の近似性能が大きく改善します。一方、Q(s, a)を直接出力するDQNアーキテクチャではこのような明らかに絶望的な状態sでもすべての行動aを試さなければよい近似性能が得られません。


実装例

学習の安定性のために、アドバンテージA(s)は平均値でスケーリングしてからV(s)に足していることに注意してください。*4


Noisy Network(2017)

[1706.10295] Noisy Networks for Exploration

探索と活用のトレードオフのバランスをどうとるかは強化学習における最重要問題の一つです。オリジナルのDQNでは探索促進のためにε-Greedy法を採用しました。ε-Greedy法は行動の決定時にεの確率で完全にランダムなアクションを採用し、1-ε の確率で貪欲方策でアクション決定を行うというごく単純な方策です。

ε-Greedy法の問題は探索率εの適切な設定が難しいことです。学習の序盤では探索率εを高めにとる必要がありますし、ゲームに習熟してきたらεを低めにする必要があります*5。そこで、DQNでは最初の100万ステップで探索率εを1.0から0.1まで線形に減少させ、それ以降はε=0.1で固定するというアニーリング方式でεの値を調整します。

epsilon_scheduler = (lambda steps: max(1.0 - 0.9 * steps / 1000000, 0.1))


しかし、ε-Greedy方策では探索率εのアニーリングのスケジューリング関数がエージェントのパフォーマンスを大きく左右する重要なハイパーパラメータとなってしまうため大変扱いづらいです。この問題へのアプローチのひとつが Noisy networkであり、この手法ではε-Greedy法のようにアクション選択に直接確率的ノイズを乗せるのではなく、Q-netwrok自体がノイズを発生することによってアクション選択が確率的になるためε-Greedy方策を使わなくても探索が促進されます。

このようなNoisyなネットワークはDense層の重みwとバイアスbが毎回の推論ごとに正規分布からサンプリングされるようにすることで実現します。 すなわち、 \displaystyle{
 w \sim \mathcal{N}(\mu^{w}, \sigma^{w}),  \ b \sim \mathcal{N}(\mu^{b}, \sigma^{b})
} であり、正規分布パラメータが学習すべきパラメータです。

ただし、本当に確率的サンプリングしてしまうと誤差逆伝搬できなくなるので実際は正規分布パラメータのμとσをtrainable parameterとし、標準正規分布ノイズを外から与えます。これは変分オートエンコーダなどでもおなじみのreparameterization trickですね。

f:id:horomary:20210205215824p:plain:w600
Noisy-network論文より

Factorized Gaussian Noisy Network

パラメータに確率的摂動を与えるNoisyNetworkでは、推論するたびにDense層のパラメータと同じ数の正規分布ノイズをサンプリングしなければいけません。例えばDQNではconv層から抜けた直後のDense layer では入力ユニット数が約3000で出力ユニット数が512なので、wのパラメータ数 + bのパラメータ数 = 3000×512 + 512 ≒ 150万 になります。致命的なほどではありませんがちょっと嫌ですね。

そこでNoisy-networkではFactorized Gaussian noise というトリックで、入力ユニット+出力ユニット数分の正規分布ノイズから疑似的にパラメータ数と同数の正規分布ノイズを生成することで計算量を減らす方法を提案しています。正直なところ tf.random.normalを使えばこのトリック使わなくてもあんま処理速度変わらなくない?とは思いましたがせっかくなのでこの方法で実装します。

実装例

DQNのすべてのDense LayerをNoisyDenseに置き換えればNoisy-networkの実装完了です。


次:優先度つき経験再生, Multistep-learning, C51

horomary.hatenablog.com

*1:Q値の過大評価自体は以前から指摘されていたがこの論文がきれいにまとめた

*2:マリオなら落下死が確定した状態とかね

*3:ただし、Dueling-networkはあくまでV(s)とA(s, a)っぽいものを学習出来たらいいな、ということを狙ったネットワーク構造というだけであり本当にV(s), A(s, a)を学習している保証はないことに留意ください。

*4:論文では他に A = A - max(A)とする方法も提案されているが平均を引く方が性能が良いらしい。

*5:CartPole環境程度ならεを固定値にしてもあんま問題ありません

DQNの進化史 ①DeepMindのDQN

DeepMindDQNからR2D2くらいまでの深層強化学習(Q学習)の発展の歴史を、簡単な解説とtensorflow2での実装例と共に紹介していきます。

まずは深層強化学習の新たな時代を切り開いたDeepMindDQN(2013)です。論文からはわかりにくいatari環境向けの実装上のテクニックDQNを構成する各要素が後継手法でどのように改良されていったかのレビューに焦点を置いてBreakout(ブロック崩し)向けにtensorflow2での実装例を紹介します。


DQNシリーズ
DQNの進化史 ①DeepMindのDQN - どこから見てもメンダコ
DQNの進化史 ②Double-DQN, Dueling-network, Noisy-network - どこから見てもメンダコ
DQNの進化史 ③Prioritized experience replay, Multi-step learning, Categorical DQN - どこから見てもメンダコ
DQNの進化史 ④Rainbowの実装 - どこから見てもメンダコ


はじめに

[1312.5602] Playing Atari with Deep Reinforcement Learning

Human-level control through deep reinforcement learning | Nature

Human-level control through deep reinforcement learning (pdf)

DeepMind社によって発表された Deep-Q-Network (arxiv 2013, nature 2015) は、"強化学習に深層CNNを導入することでレトロゲームを人間レベルでプレイできるようになる!" という結果でもって世界に衝撃を与え、深層強化学習の新時代を切り開いたまさにエポックメイキングと言うべき手法です。発表は2013年ながら最新の強化学習手法でもDQNをベースにしたものは多く、深層強化学習の初学者は必ず学ぶことになる手法です。

しかしDQNは初学者が必ず学ぶことになる手法、であるにも関わらずそこそこ実装難易度が高いです。これはDQNのパフォーマンスがハイパーパラメータに非常にsensitiveである上に、論文を読むだけでは分かりづらい実装上の細かいテクニックが多く存在するためです。

そこで、本記事ではatari環境を攻略するための実装上のテクニックおよびDQNの各構成要素がその後の後継手法でどう発展していったかに焦点を置いてtensorflow2での実装例を紹介していきます。

※注意: DQNそのものの説明はすでに良記事がたくさんありますので基本的なアルゴリズム解説は最低限です。


レーニングループの実装

ネットワーク更新以外のエージェントの実装は以下のようになります。

それではトレーニングループの各要素を確認していきましょう。

BreakoutDeterministic-v4

Breakoutはいわゆるブロック崩し環境です。gymには似た名前の環境(Breakout-v0, Breakout-v4とか)が多数実装されていますが、基本的には BreakoutDeterministic-v4を使ってください。このブロック崩し環境では必ず指示した通りの行動が実行され、高すぎるフレームレートを間引くため毎回4フレームスキップします。これ以外の環境だと指示した通りに動かなかったりフレームスキップ数がランダム化されたりと環境の遷移が確率的になりクソゲーでは?)、攻略難易度が劇的上昇します。

Difference between Breakout-v0, Breakout-v4 and BreakoutDeterministic-v4? · Issue #1280 · openai/gym · GitHub


POMDPへの対応

Q学習はMDP(マルコフ決定過程)を前提としています。MDPとは乱暴に言うなら適切なアクション判断に必要な情報はすべて現在の状態の観測に含まれている、という仮定が成立するような系です。Breakout環境では現在の状態の観測とは現在の画面(frame)にあたります。しかし1フレームだけではアクション決定には情報がまったく不十分です。

f:id:horomary:20210125000147p:plain:w400
ボールは上と下のどちらに進んでいるでしょうか?

適切なアクション選択のためには、もっと過去の情報を加味してアクションを決定する必要があります。このような系をPOMDP(部分観測マルコフ決定過程)と言います。そこでDQNでは直近4フレームを現在の状態と定義することで、本来POMDPであるatari環境をMDPっぽい系にすることに成功しQ学習の適用を可能にしました。

直近4フレームを現在の状態として疑似MDPとするDQNのアプローチはBreakoutのような反射神経要素の強いゲームでは十分に成功しましたが、一部のゲームでは直近4フレームでもなお情報が不十分であることが明らかでした。このようなPOMDP環境に対応するアプローチの一つがRNNの導入です。 このアプローチは当初([1507.06527] Deep Recurrent Q-Learning for Partially Observable MDPs)は思っていたほどうまくいかなかったのですが、2018年発表のR2D2Recurrent Experience Replay in Distributed Reinforcement Learning (ICLR 2019) )でついに大きな成功を収めました。

直近4フレームの保持

さて、DQNの疑似MDPアプローチのためには直近4フレームを保持しつづける必要があります。ここで便利なのがPython標準ライブラリcollections.dequeです。dequelistと同様に使えますが、指定された最大要素数maxlen)を超える要素が追加された場合は"First in First out"でもっとも古い要素を削除するので直近Nフレームを保持する処理をスッキリ書くことができます。
ちなみに、ゲーム開始直後は当然4フレームに満たないので初期フレーム×4を直近4フレームとすることに注意です。

フレームの前処理

CNNに入力するために各フレームは適切な前処理が必要です。具体的には余分な部分(スコア表記など)をトリミングし、二値化し、正方形にリサイズし、正規化します。ごく基本的な画像処理なのでPILでもtensorflow.imageでも自分の好きなライブラリで書きましょう。ただしopencvはなんか処理が重いので使わない方がいいです。


ε-Greedy方策によるアクション決定

局所最適解に陥らないために探索と活用(Exploration vs. Exploitation)のバランスをうまくとることは強化学習において重要な課題の一つです。 しかし、もっともシンプルなQ学習では常にQ関数を最大化するアクションを採用する決定論的方策である”貪欲方策 (Greedy policy)”を使用するので、探索力の弱さという問題を抱えることとなります。

貪欲方策
 \displaystyle{
a = \underset{a} {\text{argmax }} Q(s, a)
}

貪欲方策では学習初期にたまたま価値の高くなった行動をとり続けてしまいます。そこで、DQNではこのQ学習の探索力の弱さを補うためにε-Greedy法を採用します。この方法はごく単純で、行動の決定時にεの確率で完全にランダムなアクションを採用し、1-ε の確率で貪欲方策でアクション決定を行うことで探索力を高めようというものです。

ε-Greedy方策
(1-ε) の確率で:貪欲方策  \displaystyle{
a = \underset{a} {\text{argmax }} Q(s, a)
}
εの確率で: ランダムなアクション \displaystyle{ a} を選択

ε-Greedy法の問題は探索率εの適切な設定が難しいことです。学習の序盤では探索率εを高めにとる必要がありますし、ゲームに習熟してきたらεを低めにする必要があります*1。そこで、DQNでは最初の100万ステップで探索率εを1.0から0.1まで線形に減少させ、それ以降はε=0.1で固定するというアニーリング方式を採用しています(下図)。

epsilon_scheduler = (lambda steps: max(1.0 - 0.9 * steps / 1000000, 0.1))

f:id:horomary:20210125134709p:plain:w400
探索率εのアニーリング

このようにε-Greedy法では探索率εのアニーリングのスケジューリング関数がエージェントのパフォーマンスを大きく左右する重要なハイパーパラメータとなってしまうので大変扱いづらいという問題があります。この問題にアプローチした後継手法のひとつが Noisy network([1706.10295] Noisy Networks for Exploration, ICLR2018)です。Noisy networkではε-Greedy法のようにアクション選択に直接確率的ノイズを乗せるのではなく、Q-netwrok自体がノイズを発生することによって探索を促進します。

別アプローチとして分散強化学習手法 Ape-X( [1803.00933] Distributed Prioritized Experience Replay, ICLR2018)ではε-Greedy法を採用しながらも分散強化学習の強みを生かして並列化されたagentそれぞれに異なるεの値を割り当てるという、いい意味での力押しにより多様な探索を可能にしました。


Life losses as episode ends

Breakoutにおける1エピソードとは5つの残機をすべて失いGam eoverとなるまでです。当然、DQN論文に表記されているスコアもGame overになるまでに獲得したトータルスコアです。しかし、それでは残機を失うことが悪いことだとエージェントが学習しにくくなってしまうので、DQNでは残機が減ったら遷移情報(transition)ではエピソードが終了した扱いにするというトリックが使用されています。

最近の手法(R2D2)では、このトリックを使わないことによりatari環境の一部のゲーム(seaquest)で大幅に性能改善が見られることが報告されていますが、Breakoutについてはこのトリックを使用した方が学習の進みが良いです。


Experience Replay

強化学習はオンラインで遷移情報(経験)を収集するため、サンプル間の強烈な自己相関が発生し学習が不安定になるという問題があります。そこで、収集した経験をすぐに学習に使わずいったんバッファ(replay buffer)に蓄積し、学習時はバッファからランダムに経験を選択してミニバッチを作成する、という方法により自己相関を低減するのがExperience Replay *2 です。

Experience Replay はDQNの性能にもっとも大きく貢献している要素であることが論文のablation studyから分かります。

f:id:horomary:20210125230721p:plain:w700
Human-level control through deep reinforcement learning | Nature より

replay bufferの実装

遷移情報(経験)を蓄積するreplay bufferの実装自体は簡単です。もっともシンプルには遷移情報のタプル (state, action, reward, next_state, done) をリストに追加していくだけでもOKですが、ミニバッチ作成機能を持たせたいのでクラスで実装します。


このreplay bufferの実装は見通しがよいのですが、わかりやすさのためにメモリ効率を犠牲にしています。もしコードのわかりやすさを気にしなければメモリ消費は8分の1くらいで実装できるはずですが、強化学習においてデバッグのしやすさは本当に重要なので見通しを悪くしたくありません。

そこで、この実装では exp = zlib.compress(pickle.dumps(exp)) において遷移情報をpickle化したうえで圧縮することによりわかりやすさを保ったままメモリ消費を抑えています。DQNではreplay bufferには100万ステップ分の遷移情報を格納しますが、圧縮処理なしでこれをやるとRAMを60GBくらい消費する一方、圧縮すると処理が少し遅くなるかわりにメモリ消費は6GBくらいで済みます。

経験再生の改良

DQNでは100万ステップ分の遷移情報が蓄積されたリプレイバッファから毎更新ごとに32の遷移情報をランダムに選択してミニバッチを作成します。これではレアで高報酬な遷移を学習する効率が悪いことは明らかです。そこで後継手法として提案されたのが優先つき経験再生(ICLR2016, [1511.05952] Prioritized Experience Replay)です。この手法では遷移情報のTD誤差の大きさ(驚きの大きさ、あるいは意外性の高さ)に応じてリプレイバッファから選択される確率を高めることで、レアイベントを効率よく学習します。


Qネットワークの実装

DQNスタイルのQ-networkは状態sを入力としてアクション次元数のユニットを出力するようにします。ネットワーク構造自体はごく普通のCNNですが、初期のQ値が偏らないように重みの初期化スキームを適切に選択する必要があることに注意してください。 ここではhe_normalを採用しました。重みの初期化方法の選択はQ学習だけではなく方策勾配法でも同様に重要です。*3

画像認識分野では2012年のAlexNetから数年でResNetのような超大規模化が進んだのと対照的に、Qネットワーク自体の工夫や大規模化はあまり行われていません。パラメータ数が増えるとそれだけ必要なサンプル数が増えるからでしょうか?

Qネットワークの改良で明らかな成功を収めたのは、ネットワークの出力に近い部分をちょこっと変えた[1511.06581] Dueling Network Architectures for Deep Reinforcement Learning や RNNを導入したRecurrent Experience Replay in Distributed Reinforcement Learning (ICLR 2019) くらいじゃないでしょうか*4。他にはQ-networkの出力をスカラ値でなく確率分布に変更した C51([1707.06887] A Distributional Perspective on Reinforcement Learning)もQネットワークの改良と言えるかもしれません。


ネットワークの更新


reward clipping

DQNでは報酬を-1~+1の範囲にクリッピングし、報酬のスケールを揃えることで学習を円滑化します。ただし、Breakoutについてはもともと報酬のスケールがいい感じなのでreward clippingしなくてもあまりパフォーマンスに影響しません。

ベルマンエラーの計算

Qネットワーク更新のためのベルマンエラーを計算します。シンプルにはベルマンエラーのMSE(平均二乗誤差 )がQ-network更新のためのロス関数となります。が、実際にはMSEではなくhuber loss関数というものを使用します(後述)。

ベルマンエラー  \displaystyle{
= \mathcal{T}Q(s_t, a_t) - Q(s_t, a_t)
}

ここで、  \displaystyle{
\mathcal{T}Q(s_t, a_t) = r_t + \gamma \max_{a'} Q_{target}(s_{t+1}, a')
}


教師あり学習で言えばTQが教師ラベルでQが予測値なので、推測で推測を更新するという一見奇妙なことをしています。実際これをそのままニューラルネットワークでやると学習がものすごく不安定になります。そこで、DQNでは教師ラベルとなるTQは重みを凍結した過去のQ関数(target-Q-network)で算出することで学習の不安定化を低減しています。target-networkは定期的(DQNでは1万ステップごと)にメインのQ関数と重みを同期します。この同期頻度は頻繁にしすぎると学習が不安定化し、長くとりすぎると学習が進みにくくなるので重要なハイパーパラメータの一つとなっています。

実装について、TQ(s, a)の計算では  \displaystyle{ Q(s_{t+1}, \cdot) } の値を最大化するようなアクションa' (next_actions)を算出して、 \displaystyle{ Q(s_{t+1}, a') } を抽出する処理がありますが、tensorflowの仕様上onehot化したうえでreduce_sumをとるという少々ややこしいコードになっています。これがもしnumpyなら以下のようにすっきり書けるのですが、tensorflowではこのようなindexingはいまのところエラーになります。

n = next_qvalues.shape[0]  #: n==batch_size
max_next_qvalues = next_qvalues[np.arange(n), next_actions]

また、TQは教師ラベルなので勾配計算をしない=tf.GradientTapeの外で定義しておく必要があることに注意しましょう。tensorflow2で使えるtf.GradientTapeは普通の教師あり学習だとあまり有難みを感じませんが強化学習だとなかなか便利な記法です。


ロスと勾配計算

ロス関数について、前述した通りMSE(TQ - Q) でも学習はできますがDQNではMSEよりも外れ値に鈍感なhuber-loss関数を使用することで極端な勾配更新を防ぎます。huber-lossはtensorflow.keras.losses.Huber()として用意されているので自力実装しなくても大丈夫です。

ja.wikipedia.org

勾配計算について、DQN論文ではRMSpropを使用していますが、ここでは学習率0.00025のAdamを使用しました。


訓練の実行

24時間の学習でおよそ2.5Mステップでした。ブレークアウトは400点くらいが満点なのでいい感じに学習できてることがわかります。

f:id:horomary:20210125010821p:plain:w500

マシンはGCPのn1-standard-4(4-vCPU, 15GBメモリ) + GPU K80 のプリエンティブルVMインスタンス*5を使いました。24時間でざっくり400円くらいかと思います。

次: Double DQN, Dueling Network

horomary.hatenablog.com


*1:CartPole環境程度ならεを固定値にしてもあんま問題ありません

*2: Expereince Replayの提案自体はDQNではなくLong-Ji Lin. Reinforcement learning for robots using neural networks. Technical report, DTIC Document, 1993.

*3:[2006.05990] What Matters In On-Policy Reinforcement Learning? A Large-Scale Empirical Study

*4:私が知らないだけかも

*5:激安な代わりに最長24時間で停止するうえにGCPのリソース上の都合により突然シャットダウンされうることもあるVMインスタンス。プリエンティブルVMGPUは東京リージョンなどでは2021年1月現在使えないはずなので米国リージョンを選ぼう