MuZeroはLSTMか？

#はじめに
この記事は妄想です。論文やMuZero解説にMuZeroがLSTMと見なせると説明している方はいませんが、そのように考えることも出来るのではと思った理由を書いてみます。

#MuZeroの基本
MuZeroはモデルは下記の図に示される3種類のモデルがあります。
・モデルhは入力oから潜在層sへの変換
・モデルfは潜在層sから方針(次の手)Pと評価値Vへの変換
・モデルgは潜在層sと行動aから次の潜在層sと報酬rへの変換


AlphaZeroでは入力oから方針Pと評価値Vへの変換でした。これはMuZeroの入力oから潜在層sへの変換のモデルhと潜在層sから方針(次の手)Pと評価値Vへの変換のモデルfの繋がったものです。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/8aa74885-9cc9-d9ed-2632-74e2e3ef6278.png)　　　![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/2b42333f-d88e-2a2d-fc54-76a821dc8d1f.png)

実装を確認するとモデルhはrepresentation_network、モデルfはprediction_network、モデルgはdynamics_networkと呼ばれており、モデルgの外側で潜在層sと行動aはtorch.catによって結合されている。

```python:
    def representation(self, observation):
        encoded_state = self.representation_network(observation)
        ...
        return encoded_state_normalized

    def prediction(self, encoded_state):
        policy, value = self.prediction_network(encoded_state)
        return policy, value

    def dynamics(self, encoded_state, action):
        ...
        x = torch.cat((encoded_state, action_one_hot), dim=1)
        next_encoded_state, reward = self.dynamics_network(x)
        ...
        return next_encoded_state_normalized, reward
```
損失関数はvalue, reward, policyの損失の三種類で教師データのtargetはゲームをして(モンテカルロ木探索して)得られる。

```python:
    def loss_function(value, reward, policy_logits, target_value, target_reward, target_policy):
        # Cross-entropy seems to have a better convergence than MSE
        value_loss = (-target_value * torch.nn.LogSoftmax(dim=1)(value)).sum(1)
        reward_loss = (-target_reward * torch.nn.LogSoftmax(dim=1)(reward)).sum(1)
        policy_loss = (-target_policy * torch.nn.LogSoftmax(dim=1)(policy_logits)).sum(1)
        return value_loss, reward_loss, policy_loss
```
#強化学習におけるRNN
強化学習においてゲームの種類として考えると、将棋やチェスは悪い手を打てばすぐに即死するから報酬までの道のりは短いゲームである。一方、Atariのゲーム(たとえばパックマンとか)では行える行動が囲碁などより少なく(Atariでは方向キーと1ボタンしかない)一手だけ悪い手を打ってもなかなか死亡しないため、報酬までの道のりが長いゲームと見なされる(多分)。中間報酬を設定すれば報酬までの道のりを相対的に短くできるが、それを設定するにはルールの知識が必要である。報酬までの道のりが長いと現在の打つ手が長期的にみて悪手かどうかという判断がしづらいのではと考えられる。
このような報酬までの道のりが長い場合においてRNNが使われるというアイデアは[DRQN](https://arxiv.org/pdf/1507.06527.pdf)や[R2D2](https://openreview.net/pdf?id=r1lyTjAqYX)で既に存在します。強化学習にRNNを使うというアイデアは目新しいものではありません。



#MuZeroはRNNか？
最初に自分がMuZeroの図を見た時に感じたことは内部記憶を持つRNNに似ているように見えるという事でした。考えてみればMuZeroの前のAlphaZeroは次の手の予測を行い、それを仮に数手先まで打ってみて状態の評価を返し学習しますが、これは時系列データの逐一予測のようなものです。
この時、AlphaZeroでは一手毎の状態ごとに畳み込みhを計算する必要がありました。（モンテカルロ木探索によって同じ状態の演算は省略可能ですが）
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/ff3d2487-522d-7c8c-be3a-9e11f5dfff36.png)**：AlphaZero
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/c2d266ce-73a1-10e5-f904-85f4d633e609.png)：MuZero
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/9b60cc7f-c3df-8907-8ba9-3eaf04af3476.png)：RNN**

RNNではない時系列データの逐一予測では出力の予想を行い、その出力を再度入力にいれてループが行われます。
・入力=>畳み込み=>潜在層=>出力=>入力=>畳み込み=>潜在層=>出力=>入力=>…

一方、時系列データのRNNでは畳み込み（または全結合）を行った潜在層を直接次に渡すことで「入力から潜在層の変換」や「潜在層から出力への変換」を一度だけで済ませる経路が存在します。
・入力=>畳み込み=>潜在層=>潜在層=>潜在層=>出力

結果的に報酬までの手数が多いゲームでは勾配消失を起こしにくいRNNモデルが有利です。

#評価値は長期記憶か？
そして、もう一つMuZeroは**モデルf**によって**方針P**と**評価値V**を返しますがこの**評価値V**はLSTMにおける長期記憶として働くように見なせないかという事です。一般にMuzeroの**モデルg**に**潜在層s0**と**行動a1**を与えると**潜在層s1**と**報酬r1**を返します。**評価値V0**と**報酬r1**の合計が次の**評価値V1**と見なせるとすれば、これがLSTMにおける長期記憶のように振舞うのかもしれません。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/61522fc5-d80a-1e8b-c8c8-517b54228d63.png)
評価値Vの更新には報酬rが足されます。

```python:
            value += (
                reward
                if game_history.to_play_history[index]
                == game_history.to_play_history[index + i]
                else -reward
            ) * self.config.discount ** i
```


#関係図
入力O、潜在層S、評価値V、方針P、行動a、報酬r、訪問回数Nの関係を図に示すと以下の様になる。
報酬rの経路がLSTMの短期記憶(潜在層S)から長期記憶(評価値V)の情報に見える。一方、長期記憶(評価値V)から短期記憶(潜在層S)へのフィードバックが見当たらない。MuZeroモデルは評価値V1=>方針P0の経路に相当するものはないように見える。(本当にLSTMなら長期記憶に依存した行動を行えるはず)
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/801ef3c9-9fce-80a4-f3b2-13da0bbe2f36.png)
モンテカルロ木探索は行動後の評価値V(と方針Pと訪問回数N)から行動を決定するが、それはMuZeroモデルの構造ではない。(例えば下記式)

```math
a=argmax[Q(s,a)+P(s,a)\frac{\sqrt{N(s)}}{1+N(s,a)}(c_1+log(\frac{1+N(s)+c_2}{c_2}))]
```
このため繰り返しゲームにおいて評価値に依存する方針、つまりプレイヤーがトータルで勝っている時と負けている時で行動パターンが変わるような表現はできないのではと思われる。(勝っているなら慎重に、負けているなら大胆にという方針の変更は出来ない)


#MuZeroの欠点
Atari57の中央や平均値での性能では一見素晴らしい性能を見せているMuZeroですが、困難なゲームにおいては性能が出ません。MuZeroはAtari57の57種のゲーム中51種は人間以上のスコアを取っている一方、6種では人間以下、特に下位5%のゲームでは実はR2D2以下の性能しか出ていません。
何故、困難なのかに関してはよく分かりませんが(報酬までのゲームの長さが長いためか、ゲーム情報が不完全なためか、ゲーム結果が運に左右されるためか)、ゲーム終了まで途中に獲得したスコアが貰えないゲームに対して内部報酬を予測するようにしたAgent57はこの困難なゲームにおいても学習できるらしいです。もっとも、Agent57はMuZeroよりはR2D2やNGUに近いのかもしれません。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/299928/d064dc4f-0c17-2d13-802c-7527d0e389ba.png)

#疑問
強化学習に詳しくない素人考えでは選択する行動に「投了」とか「スキップ(ゲーム終了するまで単一入力繰り返し)」を定義して追加すれば報酬までの道のりが長いゲーム途中でも簡単に中間報酬が得られるのではと思うのだが、そういう強化学習解説は見ない。こうすると朝三暮四のように目先の利益を追ってしまい、トータルの報酬を重視できないためだろうか。


#まとめ
MuZeroの評価値と潜在層がLSTMの長期記憶と短期記憶に相当するのではないかと思った。
ただし、強化学習にRNNが使われること自体はDRQNやR2D2のモデルで既に使われる為、目新しい訳ではないと思われる。また、厳密にはLSTMと見なすには一部の経路は足りていない。