はじめに

強化学習において内部報酬系の論文を調べていたらBYOL-Exploreというものがあったので見てみました。
（今回も解説のみで実装はありません）

他の手法との比較は以下です。

Atariから探索が難しい10ゲームを比較しており、各図がそれぞれのゲームの学習状況を表しています。
比較はBYOL-Explore, RND, ICM, RLでしており概要は以下です。(リンクは論文)

ICM(2018): [今の状態,アクション]から予測した次の状態と観測した次の状態の誤差を内部報酬として探索を促す、内部報酬による探索の原点的な手法
RND(2018): 探索率をランダムネットワークと今の状態の出力差で表現し、内部報酬として探索を促す手法
RL(2020): 引用先をみたらSPRの事でした。これは自己教師あり表現学習法で精度を上げた手法です（内部報酬の探索はなし）
BYOL-Explore(2022): この記事で紹介する論文です

ゲームとしては右上の Montezuma’s Revenge が有名ですね。
RND(ピンク)に比べてBYOL-Explore(青)が圧倒的なスコアを稼いでいます。

過去に書いた関係ある記事
- SND（RNDの派生手法）
- SPR
その他の強化学習の記事や自作強化学習フレームワーク
- フレームワークの記事: https://qiita.com/pocokhc/items/a2f1ba993c79fdbd4b4d
- GitHub: https://github.com/pocokhc/simple_distributed_rl

BYOL-Explore

かなりシンプルなアイデアに惹かれたので手法のみの解説となります。
ベースはSPRで、BYOLの拡張がされたQネットワークが使われます。

※図はSPRの論文より

OnlineEndoerは今の状態、TargetEncoderは次の状態を潜在空間$z$に変換します。
今の潜在空間$z_t$は遷移モデル（図のconv. transition model）で次の潜在空間$\hat{z_t}$に変換されます。
次の潜在空間$\hat{z_t}$とTarget側から計算された$\tilde{z_t}$を用いてコサイン類似度損失（図のCosine Similarily Loss）が計算されます。

BYOL-Exploreではこのコサイン類似度損失をそのまま内部報酬として扱います。
手法としては以上で、後は細かい話です。

世界モデル（Latent-Predictive World Model）

遷移モデルですが、SPRではConv層で表現していましたが、BYOL-Explore ではRNNを取り入れて以下のように表現します。

上部はTargetEncoderの内容でSPRと同じです。（EMAにより同期する）
下部のRNNですが、2種類あります。

・閉ループRNN（Closed-loop RNN cell）

これは過去の履歴を保持するRNNとなります。
入力は実測値を使い、今の状態 $f_\theta(o_t)$、前の行動 $a_{t-1}$、前の隠れ状態 $b_{t-1}$ から成ります。

・開ループRNN（Opend-loop RNN cell）

未来を予測するRNNとなります。
入力は閉ループRNNで作られた履歴 $b_t$ と任意のアクション $a_t$ となります。

内部報酬の正規化

訓練中の内部報酬は不確実性が時間とともに変化します。（"the non-stationarity of the uncertainties"と書かれてるけどうまく訳せなかった…）
この問題に対応するためにRNDと同じ手法で正規化します。
正規化方法は標準化（平均0、標準偏差1）で標準化に使う平均と標準偏差の値はサンプルで得た実測値をEMAで計算したものになります。

ChatGPTが作ったコード例は以下です。

class RunningNorm(nn.Module):
    """観測や報酬のRMS正規化を行う簡易クラス。

    Note:
    - 数値安定性のために分散の下限を設定。
    - 学習時のみ統計を更新し、評価時は固定。
    """

    def __init__(self, eps: float = 1e-6, momentum: float = 0.01) -> None:
        super().__init__()
        self.register_buffer("mean", torch.zeros(1))
        self.register_buffer("var", torch.ones(1))
        self.eps = eps
        self.momentum = momentum
        self.initialized = False

    @torch.no_grad()
    def update(self, x: torch.Tensor) -> None:
        if x.numel() == 0:
            return
        m = x.mean()
        v = x.var(unbiased=False)
        if not self.initialized:
            self.mean.copy_(m)
            self.var.copy_(v.clamp_min(self.eps))
            self.initialized = True
        else:
            self.mean.lerp_(m, self.momentum)
            self.var.lerp_(v, self.momentum)

    def forward(self, x: torch.Tensor, update: bool = True) -> torch.Tensor:
        """入力xを(z-score風に)正規化して返す。

        Args:
        x: 任意shapeのテンソル
        update: Trueなら統計を更新
        """
        if self.training and update:
            with torch.no_grad():
                self.update(x)
        std = (self.var + self.eps).sqrt()
        return (x - self.mean) / std

内部報酬の優先順位付け

オプションで報酬の優先順位付けを行うこともできます。
これは不確実性が高い報酬のみを学習に使う仕組みです。
これによりエージェントはモデルがまだよく分かっていない領域を重点的に探索することができます。

数式としては以下となり、バッチ内の平均値より高い報酬のみ使い、低い報酬は使わない形になります。

$$
r_{i,t}^j = \max\left(\frac{\ell_t^j}{\sigma_r} - \mu_{\ell / \sigma_r}, 0 \right)
$$

$\frac{\ell_t^j}{\sigma_r}$：正規化された内部報酬
$\mu_{\ell / \sigma_r}$：バッチ内の内部報酬の平均値（閾値）

BYOL-Explore がなぜ意味ある内部表現を獲得できるのか

論文内の"Intuition on why BYOL-Explore learns a meaningful representation."の箇所です。
BYOL（Bootstrap Your Own Latent）の直感と似ているとの事。

初期段階ではターゲットネットワークはランダムに初期化されています。
なのでオンライン側は未来のランダムな特徴を学習しつつ、未来の予測に役立つ情報も学習していきます。
未来の予測に役立つ情報はEMAでターゲットネットワークにもゆっくり伝播されます。
これにより、ターゲット側の表現も徐々に安定して改善され、オンライン側の品質も向上していきます。
このループによりオンライン表現の質が徐々に向上していくとの事です。

終わりに

名前から最初はGo-Explore系列の手法かな？と思ったら全然違いました。

SPRが既存の強化学習のモデルを壊すことなく追加できる形式で、さらにそのまま内部報酬まで追加できると個人的には興味深い内容でした。
また、性能もRNDよりもとても高いというのもすごいです。
ただ、元がBYOL（表現学習）なので内部報酬の直感的解釈がRNDより難しく感じています。