https://horomary.hatenablog.com/entry/2024/11/03/140544

たった２時間のゲームプレイで人間相当性能に到達可能なサンプル効率の高い強化学習手法 ”Bigger, Better, Faster”を実装します。

背景：強化学習実用の課題は劣悪なサンプル効率

レトロゲームで人間並みのパフォーマンスを実現したDQN (Deep Q-Network) が登場してからわずか10年間ほどで深層強化学習は驚くべき発展を遂げましたが、一方で深層強化学習の実世界応用の成功例は一部の例外を除き*1まだまだ限られています。

強化学習の実用のための最大の課題はサンプル効率（学習効率）です。

Atari環境（レトロゲームベンチマーク）において最先端の手法であるMuZeroは人間をはるかに超えたパフォーマンスを発揮しますが、その超人性能は実ゲームプレイ時間に換算しておよそ4000時間もの試行錯誤に支えられています。しかし、ゲームやシミュレータならともかく実世界のビジネスにおいて4000時間の多数の大失敗を含む試行錯誤が可能な状況はほとんどなく、これが実用上の大きな課題となっています。

一方で人間はわずか十数分の試行錯誤で一定のパフォーマンスに到達することが可能*2であり、そのサンプル効率（学習効率）は最先端の強化学習よりもはるかに良好です。もし、このような人間レベルの高いサンプル効率（学習効率）を備えた強化学習手法があるならば、強化学習実用の場が大きく広がることが期待できます。

画像は https://deepmind.google/discover/blog/agent57-outperforming-the-human-atari-benchmark/ より

強化学習におけるサンプル効率向上アプローチ

評価指標： Atari-100Kベンチマーク

強化学習分野において、Atari環境では200M環境ステップ（＝4000時間のゲームプレイ相当）完了時の性能でアルゴリズムを評価することが慣例となっていますが、しかし前述のように4000時間の試行錯誤はあまりにも非現実的な設定であるため、より実用を意識したアルゴリズムにはAtari-100K（100K環境ステップ＝２時間のゲームプレイ相当）が評価に使われるようになっています。

2024年現在、Atari-100Kにおける有力なアプローチは大きく２つに分類できます。

①リセット法によるリプレイ率増大
②環境シミュレータ（世界モデル）のデータ駆動構築

BBF論文より主要な高サンプル効率アルゴリズムの性能比較（IQM=1.0で人間相当）

①リセット法によるリプレイ率の増大

該当手法： SR-SPR(2022), BBF: Bigger, Better, Faster(2023)

Q学習のようなオフ方策強化学習アルゴリズムにおいてサンプル効率を上げたいならばリプレイ率（経験再生の頻度）を上げることが手っ取り早い解決策のように思いますが、実際には学習が極めて不安定になるためうまくいきません。これはリプレイ率を上げすぎることにより初期段階で価値関数の過学習が発生し同じ行動ばかり選択するようになるため、試行錯誤によってより良い経験を取得することができなくなり性能向上が頭打ちになるためです。*3

この問題の解決策としてSR-SPRでは価値関数を定期的にソフトリセットすることで過学習を防ぐという荒業でリプレイ率の増大を実現する「リセット法」を提案しました。本稿ではSPRの後継手法である BBF(Bigger, Better, Faster) の実装を紹介します。

②環境シミュレータ（世界モデル）のデータ駆動構築

該当手法： DIAMOND(2024), EfficientZeroV2(2024), DreamerV3(2023) など多数

もしも環境の高精度シミュレータが存在するならばノーコストで試行錯誤ができるのでそもそもサンプル効率が論点になりません、このために教師あり学習により環境シミュレータ、いわゆる世界モデル（World Models）を構築します。このアプローチは実空間に復元可能なシミュレータを構築する方法と、潜在変数空間のみのシミュレータを構築する方法の大きく２流派が存在します。

前者の”実空間に復元可能なシミュレータ”を構築する方法は、動画予測というタスクの困難さにより従来それほどうまくいっていなかったのですが、近年登場した拡散モデルの恩恵を受けた高性能なアルゴリズムが出現し始めています。

後者の"潜在変数空間のみのシミュレータ"を構築する方法は、シミュレータと言いつつ人間には何が起こっているのかまったくわからないという説明可能性上の弱みはあるものの高い性能を発揮することが分かっています。ここに分類される代表的なアルゴリズムとしてはMuzeroのサンプル効率改善版である EfficientZero や、その後継の EfficientZeroV2 などが挙げられます。

Bigger Better Faster: BBF (2023)

本稿では「①リセット法によるリプレイ率の増大」アプローチの最新手法であるBBF (Bigger, Better, Faster) の簡単な解説とTF２による再現実装を行います。BBFはAtari-100K環境において当時のほぼSOTAという性能面での優秀さだけでなく、「実装がシンプル」かつ「GPU効率も良好」というお手軽さが魅力の手法です。

arxiv.org

手法解説

BBFは同じ著者による３連作（SPR→SR-SPR→BBF）の最新手法のため順に概要を解説していきます。

SPR(2020)

[2007.05929] Data-Efficient Reinforcement Learning with Self-Predictive Representations

SPR（Self-Predictive Representations）はQ学習の補助タスクとして自己教師あり表現学習を行うことで、「過学習の低減」と「より良い表現抽出」を可能にする手法です。この工夫によりNature DQN では0.25に設定されていたリプレイ率（=4環境ステップごとに１回ネットワーク更新）を1.0（=1環境ステップごとに１回ネットワーク更新）まで増大させることに成功しています。

自己教師あり学習の方法について、基本的にはネガティブサンプル不要な対照学習手法 BYOLに従いますが、CV一般向けのBYOLでは元画像とaugmentationした画像の類似度を最大化するのに対して、SPRのBYOLでは「現在状態」と「kステップ前の過去状態から将来予測によって得られた現在状態」の類似度を最大化します（下図）。

SR-SPR (2022)

Sample-Efficient Reinforcement Learning by Breaking the Replay Ratio Barrier | OpenReview

SR-SPRは、上述のSPRにネットワークを定期的にソフトリセットする「リセット法」を追加することにより、リプレイ率を増やすほど性能が増加する「リプレイ率のスケール則」を発見した手法です。より具体的にSR-SPRのリセット法とは、40K回のネットワーク更新ごとにEncoder CNN と transition model以外の層の重みを完全リセット（再初期化）する、Encoder CNNとtransition modelについては現在のweightとランダム初期化によって得られたweightで8:2で重みづけ和を取ることによりソフトリセットを行うというものです。リセット法が初期の過学習低減のための強制的な忘却機構として働くために効果的な表現抽出が行えることで性能が向上するのであると考えられます。

論文中に脳科学についての言及は一切ありませんが、個人的には人間の学習機構ともつながりがありそうだなと思ってます。全然倒せなくて行き詰っていたボスが一晩寝てもう一回やったらすんなりクリアできた、みたいな。

Bigger, Better, Faster (2023)

https://arxiv.org/pdf/2305.19452

Bigger, Better, Faster (BBF) は SR-SPRのハイパラ検討を詳細に行うことでIQM1.0 (人間相当性能)超えの大幅な性能向上に成功した手法です。SR-SPRから手法的にはほぼ変わっていないですがBBFの大きな貢献として、大きなモデルであればより強くリセットすることで「モデルサイズのスケール則」が成立することを発見したことがあります。

これによりCVやNLPなどの他分野では数年前から行われてきた「モデルサイズ増大による性能向上」の恩恵を強化学習分野でも受けることができるようになります。

BFF論文より：従来手法（SPR, SR-SPR）ではネットワークを大きくしても性能が上がらないがBBFではスケール則が成立

Tensorflow2による実装

実装全文： github.com

基本的にはDQN＋αであり実装上の難解な部分は少ないため、ここでは２つの重要なポイントだけを紹介します。

A. リセット法によるネットワーク摂動
B. SPRの自己教師あり表現学習向けtransition model

※なお、BBF論文ではベースDQNとしてRainnbowを採用していますが、実装がめんどいのでここではベースDQNとしてQR-DQNを採用してます。

horomary.hatenablog.com

A. リセット法によるネットワーク摂動

BBFのリセット法とは、40K回のネットワーク更新ごとにEncoder CNN と transition model以外の層の重みを完全リセット（再初期化）する、Encoder CNNとtransition modelについては現在のweightとランダム初期化によって得られたweightで5:5で重みづけ和を取ることによりソフトリセットを行うというものです。この実装はとてもシンプルです。

B. SPRの自己教師あり表現学習向けtransition model

前述のとおり、SPRのBYOLでは「現在状態」と「kステップ前の過去状態からの将来予測によって得られた現在状態」のコサイン類似度を最大化します。実装自体は基本的にただのBYOLなので難しくないのですが、将来予測のためのtransition modelの構造がややトリッキーなので実装を紹介します。Transition Modelでは過去状態S_tから行動履歴に従って時間発展させることでS_t+kを予測します。ちなみにMuZeroのtransition modelと同じ実装が採用されています。

※なお、コサイン類似度の最大化は単位ベクトルの二乗和誤差を取るのと同等なのでまじめにコサイン類似度を計算する必要はない。