この記事は自作している強化学習フレームワークの解説記事です。

・フレームワークの記事

・GitHub

前：MuZero
MuZeroの別の派生：Stochastic MuZero

EfficientZero（2021年）

強化学習は多くの成功を収めてきましたがサンプル効率の悪さが課題としてありました。
例えばAlphaZeroでは訓練に2100万回のゲームをプレイする必要があります。
しかし人間は1日に約5回ほどしかプレイできないので同じ経験を積むのに11500年かかってしまいます。

EfficientZero は MuZero をベースに以下のアプローチを実施し、サンプル効率を上げたアルゴリズムです。

自己教師あり環境モデル（a self-supervised environment model）の学習
エンドツーエンドの学習による累積誤差(compounding error)の軽減
オフポリシー問題の対策

性能は以下で、Arari 100k において従来のSoTA（多分SPRかな？）と比べて平均値176%、中央値163%上がり、初めて人間レベルを超えた手法となります。

※Atari100kは従来のAtariベンチマークの200Mステップ（=200,000,000）を100kステップ（=100,000）に減らしたベンチマーク

論文とGitHubは以下です。

・EfficientZero
https://arxiv.org/abs/2111.00210
https://github.com/YeWR/EfficientZero

MuZeroとEfficientZeroの違い

MuZeroを書いたのもかなり昔なので復習もかねて違いを図示して見ました。
以下はMuZeroのアーキテクチャです。

左側が unroll steps（学習時のステップ）で、右側が simulation steps（MCTSで探索する仮想ステップ）です。

流れとしてはまず左下の盤面 $o_t$ をRepresentationで内部表現 $s_t$ に変換し、Predictionで評価値 $v_t$ と方策 $p_t$ を予測します。
次に $s_t$ をルートノードとしたMCTSで方策 $\pi_t$ を計算し、そこから次のアクション $a_t$ をだします。
その後、$s_t$ と $a_t$ を元にDynamicsモデルで1ステップ進め、報酬 $r_t$ と次の状態 $s_{t+1}$ を予測するという流れです。

EfficientZeroでは以下になります。

赤い部分①②③+LSTMが変更箇所で、順に見ていきたいと思います。

①自己教師あり学習における一貫性損失（Self-Supervised Consistency Loss）

MuZeroの環境モデル（ダイナミクスモデル）は報酬関数、価値関数、方策関数から学習されますが、どれも環境を学習する情報としては不足していました。
報酬関数はスカラー値の学習で疎な環境だとほとんど情報がありません。
価値関数はブートストラップ（更新に自分自身の値を使用）で学習されるのでノイズが多く、方策関数は探索フェーズの学習です。

EfficientZeroでは次の状態も教師データに含める事で、環境モデルとしての学習情報を取り入れこれを解決します。
学習にはSimSiam（論文リンク）自己教師ありフレームワークを採用しています。

SimSiamは1枚の画像を2種類の方法でランダムに変形し、変形後の画像同士を学習させる事で画像の特徴量を取得する自己教師あり学習となります。

自己教師あり学習と強化学習を組み合わせた他の研究としてSPRがあります。
SPRとEfficientZeroの違いは、SPRはBYOLをベースにしているのに対し、EfficientZeroはSimSiamをベースにしている点で差異があります。

SimSiamによる自己教師あり学習の概要は以下です。（論文より）

現在の状態 $O_t$ と次の状態 $O_{t+1}$ は representation を通して内部表現 $s$ に変換されます。
現在の状態 $s_t$ はアクション $a_t$ と next-state を通して、 $\hat{s_{t+1}}$ に変換され、 $\hat{s_{t+1}}$ と $s_{t+1}$ を学習することで内部表現を獲得します。

②エンドツーエンドの学習

「偶発的不確実性（aleatoric uncertainty）」は不確実性の中でもどれだけ学習しても改善する見込みのない不確実性の事です。（環境で決められたランダムな挙動など）
逆にモデルを学習することで改善できる不確実性を「知識的不確実性（epistemic uncertainty）」といいます。

環境モデルによる報酬の予測は十分に学習した後でも大きな予測誤差がある事が分かりました。
これは元の環境による偶発的不確実性によって引き起こされている問題で、特に遠い未来を予測する場合に大きな問題になります。

モデルベース学習ではこの将来の予測誤差が大きくなる事を「状態エイリアス問題(state aliasing problem)」といいます。

私の所感ですが、MuZeroは次の遷移に対して確率を考慮していません（決定的な遷移を想定し、MDPな環境を想定していない）
なのでこれは確率的な遷移によって引き起こされる問題を言っている気がします。

図（論文より）はAtariのPongで右のプレイヤーが失点する様子を表しています。（画像は、左が0step、真ん中が10step後、右が20step後の画像）

ここで0stepの画像から将来を予測し、どの行動でいつ失点するかを正確に予測することは困難です。
しかし、右のプレイヤーが動かないとしたらどうでしょうか。
十分なstep経過すれば失点することは簡単に想像できるかと思います。

人間はポイントを失う正確なステップは予測しません。
しかし、より長い未来を想像することで、確実に起こりうる事象を予測します。
この直感に着想を得て、価値をエンドツーエンドで予測する手法を提案します。

UCT

UCTの知識が必要なので復習しておきます。（MuZeroのPUCTの事で、EfficentZeroではUCTと言っています）
UCTは以下でした。

$Q(s,a)$が行動価値、$P(s,a)$が方策、$N(s,a)$が訪問回数、$N(s,b)$が$a$の兄弟ノードの訪問回数を表します。

ここで行動価値は以下で計算されます。

$r$は予測した報酬です。
ここで $\sum_{i=0}^{k-1}\gamma^i r_{t+i}$ を value prefix と名付けます。

MuZeroでは報酬 $r$ を学習してから value prefix を計算していましたが、EfficientZero では value prefix 自体を直接学習します。
これにより独立した各ステップの報酬を合計するより、ある一連の流れを踏まえた報酬の合計を予測できるのでより正確な値を予測できます。

また value prefix は入力数が可変になるので EfficientZero ではLSTMを採用して対応しています。

③環境モデルを用いたオフポリシーの修正

MuZeroは経験を取得する時の方策と学習時の方策が異なるオフポリシーなアルゴリズムです。
オフポリシーアルゴリズムでは経験を使いまわせる利点がある一方、オフポリシー特有の問題もあり MuZero でも発生します。

・オフポリシー問題とは
簡単に言うと、過去と現在の方策の違いにより獲得する報酬が異なり、その結果計算される価値の期待値が正確ではなくなるという問題です。

Rainbowで使われているMulti-step learningでも同じ問題を抱えており、その対策について過去記事を書いているので興味があれば見てみてください。

対策ですが、価値を学習するための教師データの $z$ は過去の方策で値を計算しているので古い結果となります。

$$
z_t = \sum_{i=0}^{k-1} \gamma^i u_{t+1} + \gamma^k v_{t+k}
$$

EffecientZeroでは環境モデルがあるので、これを用いて新しい方策で直接値を計算することで解決します。
（モデルフリー系のアルゴリズムではできない方法です）

具体的には現在の方策でMCTSを実施し、その結果を利用します。

$$
z_t = \sum_{i=0}^{l-1} \gamma^i u_{t+1} + \gamma^l v^{MCTS}_{t+l}
$$

ここで $l$ はステップ数になり、$l <= k$ とし、サンプリングされた軌跡が古いほど$l$は小さくします。

Reanalyze

上記ですが、ステップ数 $l$ 以外はMuzeroのReanalyzeと同じです。
$l$ は以下で計算されます。

各数値の意味は以下です。

$k$ : TD steps と呼ばれる経験収集時のステップ数（$k=5$）
$T_{current}$ : 現在の学習回数
$T_{s_t}$ : データ収集時の学習回数
$T_{total}$ : 学習予定の全学習回数（$T_{total}=100k$）
$\tau$ : 係数で0.3

再探索時は、MCTSのディリクレノイズは再度サンプルされ、方策と価値の更新は方策を99%、価値は100%で置き換えられます。

また、再探索の時間ですが再度探索するので倍になります。
しかしこれはMuzeroのReanalyzeと同じで別プロセスで計算すれば無視できるとの事です。

その他の改善

平均Q値メカニズム

Q値ですが、未訪問のノードは仮で0を置いていました。
これは最悪の状態を表しているので平均Q値メカニズム（a mean Q value mechanism）で改善します。（Elf OpenGoで実装されている内容）

計算式は以下です。

訪問回数が0のノードは「親ノード＋1回以上実行された別アクションのノード」のQ値の平均を使います。
親ノードも含んでいるのは多分別アクションもすべて0の場合に分母が0になるのを防ぐためだと思います。

Q値の正規化

MuZeroではQ値を最小値最大値を使用して0～1に正規化していました。
ただこれは、最小値と最大値の差が小さい場合に値が大きくスケールし、過大評価される可能性があります。
これに閾値を設けて対処します。（ソフト最小値最大値更新; soft minimum-maximum updates）

ここで $\epsilon$ は min-max の最低限の値を保証する値で 0.01 になります。

EfficientZero V2（2024年）

EfficientZero V2（EZ-V2）はEfficientZeroに連続アクションや離散アクション、画像入力や低次元入力等、ざまざまな環境でも動作するように拡張した手法です。

性能は以下で DreamerV3（青）や各ドメインのSOTA（緑）を超えた性能との事。

EfficientZeroとの違いは以下です。

MCTSでの探索をガンベル探索に変更し、シミュレーション回数を大幅に削減
ガンベル探索にガウス分布が使えるように拡張し、連続行動空間に対応
Search-based Value Estimation
アクションのEmbedding化
収集されたサンプルの優先度を計算ではなく最大値に変更
ネットワークとハイパーパラメータの調整

本記事ではこの中の上3つを解説します。

論文とGitHubは以下です。

・EfficientZeroV2
https://arxiv.org/abs/2403.00564
https://github.com/Shengjiewang-Jason/EfficientZeroV2

Gumbel-Top-k Trick

（Gumbel-Top-k Trick 論文：https://arxiv.org/abs/1903.06059 ）

Gumbel-Top-k Trick の前にまず Gumbel-Max trick（ガンベル最大トリック）について話します。
Gumbel-Max trick は、カテゴリカル分布からsoftmaxを通すことなく値をサンプリングする手法です。
（以前書いた記事「離散空間におけるReparameterizationTrick」もよかったらどうぞ）

Gumbel-Max trick は一番確率が高いものしかサンプリングできませんでした。
これを上位k個でサンプリングできるようにしたのが Gumbel-Top-k Trick です。

まず上位1個（Gumbel-Max trick）で値をサンプルする方法は以下です。

$\text{logits}(a)$ がsoftmaxを通す前の値で、$g(a)$ がガンベル分布に従った乱数です。
logits+ガンベル乱数の値が最大値となる確率が、カテゴリカル分布の確率と一致します。

そしてこれはk個に拡張できるようです。

選ばれた値を除いてもう一度サンプリングするだけですね。
上位k個の最大値がカテゴリカル分布の確率と一致する事を示したのが Gumbel-Top-k Trick となります。

Gumbel search と木探索

ガンベル探索は離散空間をサンプリングする手法なのでこれを連続空間に拡張します。

提案手法は任意の状態$s$に対して以下の操作をします。

現在の方策で$K$個のアクションをサンプリング（主に活用側のアクション）
別の分布で$K$個のアクションをサンプリング（主に探索側のアクション）

ここで別の分布はガウス分布（現在の分布）をより平坦化させた（分散を増やした）分布となります。
これらの行動集合 $A_S=[A_{S1},A_{S2}]$ に対してバンディットアルゴリズムを適用することで連続空間でも探索を可能としています。

またこの方法は、以下の式より方策の改善が保証されているとの事です。

詳細は論文に任せますが、概要を言うと $q(s, a_S^*)$ が新しく得られるQ値で、それが現在得られているQ値 $\mathbb{E}_{a \sim p_t}[q(s, a)]$ より同じまたは大きくなるというものです。

またこのアクションのサンプリング手法ですがルートノードのみで行い、子ノードでは現在の方策のみでサンプリングします。
これによりQ値推定のための分散を減らす効果と、より深く探索することで類似アクションの冗長なシミュレーションを回避する効果があるとの事です。

離散空間における学習用の方策

連続空間における方策は対数尤度を最大化します。（MCTSの結果得たアクションに近づける）

離散空間での方策は GumbelMuzero とほぼ同じで少しだけ変更してあります。
（GumbelMuzeroの論文: https://openreview.net/pdf?id=bERaNdoegnO）

内容は以下で、アクションが訪問済みか未訪問で completedQ の値が変わります。

$\sigma$はスケーリング関数（後述）となります。

訪問済みの場合

$$completedQ(a) = r(s,a) + \gamma v(s')$$

$r$は報酬で$\gamma$は割引率、$v(s')$は次の状態の状態価値です。（論文上は$completedQ(a)$ではなく$q(a)$と表記）

未訪問の場合

$$ completedQ(a) = \sum_{a'} \pi(a')q(a') $$

$a'$は訪問済みのアクションとなります。

スケーリング関数 $\sigma$

スケーリング関数は以下です。

$N$ は訪問回数で、$c_{visit}=50$、$c_{scale}=0.1$ となります。

2種類の方策

木探索の結果、以下2つの方策を学習に使います。

上記で計算した方策
MCTSの最良のアクション

最良アクションも採用することで特に行動数が多い環境に対して学習が早くなるようです。

図は論文より方策の学習を直感的に表した概要図で、左下の黄色が今いる地点、真ん中の星が目標です。
最良アクションは直接赤色の地点に行きますが、上の計算結果は全アクションの平均損失になるので茶色で示された箇所に移動します。
これはアクション数が多い場合に最適解に早く到達する事を促します。

Search-based Value Estimation（SVE）

上記EfficientZeroの「③環境モデルを用いたオフポリシーの修正」ではオフポリシー問題を軽減するために適応型ブートストラップ法（adaptive step bootstrapping method）を提案しました。
しかしこれは古い方策で得た報酬を使うので性能が低下する可能性があります。

EZ-V2では現在の方策と環境モデルから新しく価値推定を行います。
これを探索ベースの価値推定（ Search-Based Value Estimation; SVE）と呼びます。

計算にかかる時間ですが、これもMuzeroのReanalyzeと同じで別プロセスで計算できるので問題ないとの事。

実装

Reanalyzeはフレームワーク上での実装はしていません。
（実装自体はできるのですが、複雑になりすぎるので一旦保留にしています）
ですのでSVEも実装はしていません。

実装に関してはReanalyzeを除けばそこまで難しい箇所もないかと思います。

また学習結果ですが、体感ですがかなり少ないステップ数で学習できた印象です。
ただ、１回のステップがかなりかかる…。
MCTSのシミュレーション回数を減らせば早くなりますが精度とのトレードオフですね。

フレームワーク上の実装ではMCTSは並行して実施していません。（1回のシミュレーションを逐次に）
これはそもそもActorが並行されればシミュレーションも並行されるのでシミュレーションを更に並行する効果はないかなと思ってやっていません。
また、本家の実装ではMCTSはC++で実装されていましたね。

最後に

AlphaZeroからですが、WorkerとTrainerを逐次で学習すると学習が遅すぎるので経験の収集と学習は別プロセスで実行するのが必須な気がします。
またアルゴリズムに関しては、サンプル効率の高さはかなり実感できたのでMCTSの可能性をすごく感じた内容でした。