AVGに適格度トレースを導入する

\displaystyle \newcommand{\bs}{\boldsymbol}

　先日、Action Value Gradient (AVG)を試して手元で動くところまで確認しました。

　リプレイバッファを使わずバッチサイズ1のオンライン強化学習で、Humanoid-v5での報酬が伸びていることが確認できています。

　一方、2Mステップほど回してようやくある程度の性能になるというように、サンプル効率がやや悪いようには思われます。もとの論文でも「7 Conclusion」のLimitations and Future Workで触れられており、適格度トレースは対応策の一つなのではないかと挙げられているため、今回はこれに挑戦してみました。

実装

　適格度トレースはTD(λ)を後方観測として実装するもので、トレースベクトル \bs{z} を以下のように構成します。

\displaystyle \begin{align} \bs{z} _ 0 &:= \bs{0} \\ \bs{z} _ t &:= \gamma \lambda \bs{z} _ {t - 1} + \nabla \hat{v}(S _ t, \bs{w} _ t) \end{align}

かつ、1ステップTD誤差

\displaystyle \begin{align} \delta _ t := R _ {t+1} + \gamma \hat{v}(S _ {t+1}, \bs{w} _ t) - \hat{v}(S _ t, \bs{w} _ t) \end{align}

を考えたときに、更新が以下のようになります。

\displaystyle \begin{align} \Delta ^ B _ t := \alpha \delta _ t \bs{z} _ t \end{align}

　ちゃんと考えたわけではありませんが行動価値についても同じようになるだろうと信じて、そのまま実装すると以下のようになると思われます。（全体はこちら。以下は主要部分を抜粋）

class AVG:
    def __init__(self, cfg: argparse.Namespace) -> None:
        ...

        # コンストラクタの中でトレースベクトルを準備
        with torch.no_grad():
            self.eligibility_traces_q = [
                torch.zeros_like(p, requires_grad=False) for p in self.Q.parameters()
            ]


    def update(...) -> None:
        ...

        q = self.Q(obs, action.detach())  # N.B: Gradient should NOT pass through action here
        with torch.no_grad():
            next_action, action_info = self.actor(next_obs)
            next_lprob = action_info["lprob"]
            q2 = self.Q(next_obs, next_action)
            target_V = q2 - self.alpha_lr * next_lprob

        reward = self.symlog(reward)
        delta = reward + (1 - done) * self.gamma * target_V - q

        ...

        self.qopt.zero_grad()
        if self.use_eligibility_trace:
            q.backward()
            with torch.no_grad():
                for p, et in zip(self.Q.parameters(), self.eligibility_traces_q):
                    et.mul_(self.et_lambda * self.gamma).add_(p.grad.data)
                    p.grad.data = -2.0 * delta * et
        else:
            qloss = delta**2
            qloss.backward()
        self.qopt.step()

    def reset_eligibility_traces(self) -> None:
        for et in self.eligibility_traces_q:
            et.zero_()

　use_eligibility_trace というフラグで使うか使わないかを切り替えられるようにしています。use_eligibility_trace のオン・オフによらず、Optimizerの step() でパラメータを更新をすることは統一したかったため、gradを直接更新するという形での実装にしました。

　理屈として、TD(λ)において \lambda=0 とした場合、1ステップTDと一致します。この実装においてもuse_eligibility_trace=Trueかつself.et_lambda=0としたときには、use_eligibility_trace=Falseと数値的に同じこと（同じシード値を使ったときに1回更新した後のネットワークパラメータが完全に一致すること）を確認しました。

結果

　適格度トレースを使わない1ステップTDおよび、適格度トレースを使い \lambdaを0.0(これは1ステップTDと一致するはずです), 0.1, 0.2, 0.4, 0.8などで実験しました。

　いくつかのシードで実験してみましたが、どれもあまり適格度トレースが有効である結果とはなりませんでした。

シード値 結果
46
47
48

　よく考えると \lambda = 0.0のものが1ステップTDと長期的には完全一致していないのはおかしいことに思われます。更新1回では厳密に一致していたはずですが、どこかのタイミングでずれ始める要因がランダム性や演算誤差にあるのかもしれません。

分析

　適格度トレースを入れてもサンプル効率が改善しない理由を知るために、1ステップTDでの学習について、TD誤差を記録してプロットしました。

　これを見ると、TD誤差自体はそこまで大きくないので価値関数はある程度学習できているのではないかと思います。つまり、サンプル効率が悪い原因は価値関数の部分ではないのかもしれません。

　一方、学習が進むと常にTD誤差が負になっているということから、次ステップでのQ値が非最適なもの、つまり次の行動として非最適な行動を選んでいるのではないかと読み取れます。

　総合的に見て、価値関数よりも方策の学習に難があるのではないかと考えられます。

　実際、コード中では lprob として得られる行動の対数確率密度をプロットすると、とても小さくなっていることがわかります。

　エントロピー正則化が強力に効いていると予想されますが、一方で、エントロピー正則化の係数 alpha_lr をちょっと小さくすると急に lprob が大きくなって同じ行動ばかりになってしまうようなので、ここの調整はかなりシビアであるようです。

まとめ

　AVGに適格度トレースを導入してみましたが、あまり効果はないようでした。現状のサンプル効率の悪さは価値関数よりも方策由来なのではないかと思われます。次回では方策を改善する方法について検討します。

From Past to Future: Rethinking Eligibility Tracesを読んだメモ

è«–æ–‡

概要

後方観測的実装によるTD(λ)法を、ニューラルネットワークによる関数近似器と同時に使うといくつか問題が生じます。この論文では特に、非線形関数近似による方策評価では、TD(λ)によって維持されるトレースベクトルが古くなり、状態価値の更新が後方観測の意図する方向と一致しなくなる可能性があるという問題を主題にします。現在注目しているタイムステップでのトレースベクトルを式変形していくと、過去方向への価値関数のような定式化に至ることを確認します。実験により、人工的なタスクにおいてはTD(λ)より性能が良い結果が出ています。

双方向価値関数の定義まで

タプル (S, A, P, R, d _ 0, \gamma) として定義される一般的なマルコフ決定仮定(MDP)を考えます。正確な定義は論文を参照してください。

ある方策 \pi が固定されたとき、その方策に従って得られる収益 G _ t := R _ t + \gamma R _ {t+1} + \gamma ^ 2 R _ {t+2} + \dots, \mathrm{where} \, R _ t = R(S _ t, A _ t) を計算することを方策評価と言います。

この収益の期待値をパラメータ \theta を持った関数 v _ \theta(s) = \mathbb{E} _ \pi \lbrack G _ t | S _ t = s \rbrack によって推定します。

1ステップTD誤差を \delta _ t = R _ t + \gamma v _ {\theta _ t}(S _ {t+1}) - v _ {\theta _t}(S _ t) 、学習率を \alpha 、トレースベクトルの初期値を e _ {-1} = 0 として適格度トレースでは以下のように更新します。

\displaystyle \begin{align} e _ t &= \gamma \lambda e _ {t - 1} + \frac{\partial v _ \theta(S _ i)}{\partial \theta} \Bigg| _ {\theta = \theta _ t}\\ \Delta \theta _ t &= \alpha \delta _ t e _ t \end{align}

e _ tは一つの式にまとめると

\displaystyle \begin{align} e _ t = \sum _ {i = 0} ^ t (\lambda \gamma) ^ {t - i} \frac{\partial v _ \theta (S _ i)}{\partial \theta} \Bigg| _ {\theta = \theta _ i} \end{align}

となります。

ここで重要なことは、式の右端が \theta = \theta _ i であり、トレースベクトルはそのときどきのタイムステップでのパラメータを使った勾配を集約しているということです。つまり、たとえば現在のタイムステップから見ると、古い時点での価値関数の勾配がベクトルに貯められています。価値関数は逐次的に更新されているので、古い価値関数を更新するべき方向が、現在の価値関数を更新するべき方向であるかはわかりません。

これは価値関数が線形である場合には問題になりません。（※ちゃんと理解したわけではありませんが、線形なので結局順番を入れ替えられて問題にならなくなるということだと思っています。）ニューラルネットワークなどの非線形関数を用いて価値関数を近似し始めた段階で生じる問題となります。

ここで、現在のパラメータで過去の勾配を集約する \tilde{e} _ t を考えてみることとします。（一番右端の添字だけが i から t に変わります）

\displaystyle \begin{align} \tilde{e} _ t = \sum _ {i = 0} ^ t (\lambda \gamma) ^ {t - i} \frac{\partial v _ \theta (S _ i)}{\partial \theta} \Bigg| _ {\theta = \theta _ t} \end{align}

こちらの方が求めたい勾配にはなりますが、一方で、もちろん計算が大変になります。毎回のステップで、過去のステップを全て戻って計算することは非現実的です。

ここで話は変わりますが、上記の適格度トレースは単一エピソードについてのものです。 Van Hasselt et al.(2020)は、期待トレースベクトル z(s) を考えています。

\displaystyle \begin{align} z(s) := \mathbb{E} \lbrack e _ t | S _ t = s \rbrack \end{align}

これのトレースベクトルを、先程の現時刻でのパラメータで考える \tilde{e} _ t に置き換え、式変形を行います。

\displaystyle \begin{align} \mathbb{E} \lbrack \tilde{e} _ t | S _ t = s \rbrack &= \mathbb{E} \left\lbrack \sum _ {i = 0} ^ t (\lambda \gamma) ^ {t - i} \frac{\partial v _ \theta (S _ i)}{\partial \theta} \Bigg| _ {\theta = \theta _ t} | S _ t = s \right\rbrack \\ &= \frac{\partial}{\partial \theta} \mathbb{E} \left\lbrack \sum _ {i = 0} ^ t (\lambda \gamma) ^ {t - i} v _ \theta (S _ i) | S _ t = s \right\rbrack \Bigg| _ {\theta = \theta _ t} \end{align}

この式変形は、期待値と勾配の線型性から来ています。

この式の中身を関数 f で置きます。

\displaystyle \begin{align} f(\theta, s) = \mathbb{E} \left\lbrack \sum _ {i = 0} ^ t (\lambda \gamma) ^ {t - i} v _ \theta (S _ i) | S _ t = s \right\rbrack \end{align}

これを意味の面から解釈すると、状態 S _ t に至るまでの価値を \gamma\lambda で減衰させながら過去方向に集約していったものだとみなせます。

過去方向への収益を \overleftarrow{G _ t} := \sum _ {i=1} ^ t (\gamma\lambda) ^ i R _ {t - i} として、関数 f が v _ {\theta ^ \ast} = v ^ \pi であるような最適なパラメータ \theta ^ \ast においてどうなるかを確認します。

Lemma 3.1.

\displaystyle \begin{align} f(\theta ^ \ast, s) &= \mathbb{E} \left\lbrack \sum _ {i = 0} ^ t (\lambda \gamma) ^ {t - i} v _ {\theta ^ *} (S _ i) | S _ t = s \right\rbrack \\ &= \mathbb{E} \left\lbrack \sum _ {i = 0} ^ t (\lambda \gamma) ^ {t - i} v ^ \pi (S _ i) | S _ t = s \right\rbrack \\ &= \frac{1}{1 - \gamma^2 \lambda} \mathbb{E} \left\lbrack G _ t + \overleftarrow{G _ t} - \gamma (\lambda \gamma) ^ {t+1} G _ 0 | S _ t = s \right\rbrack \end{align}

となり、収益によって記述できます。

収益と同様に、価値関数も定めることができます。

\displaystyle \begin{align} \overleftarrow{{v} ^ \pi} (s) := \mathbb{E} \left\lbrack \sum _ {i = 1} ^ t (\lambda \gamma) ^ {i} R _ {t - i} | S _ t = s \right\rbrack \end{align}

これを逆向き価値関数と呼ぶことにします。区別を明確にするため、一般的な価値関数を前向き価値関数 \overrightarrow{v} = v とします。

ここで両者の合計として新しい価値関数の概念を定義します。

\displaystyle \begin{align} \overleftrightarrow{v} (s _ t) := \overleftarrow{v}(s _ t) + \overrightarrow{v}(s _ t) \end{align}

ここで、Lamma 3.1.で示した式からは - \gamma (\lambda \gamma) ^ {t+1} G _ 0 の項が無くなっていますが、これは t \to \infty で0になるためです。

ベルマン方程式の確認

まず前向き価値関数に関するベルマン方程式を確認します。

r ^ \pi (s) = \sum _ {a \in A} \pi (a|s) R(s, a) および \overrightarrow{P ^ \pi} (s'| s) = \sum _ {a \in A} \pi (a|s) P(s'|s, a) とすると

\displaystyle \begin{align} \overrightarrow{v ^ \pi} (s _ t) = r ^ \pi (s _ t) + \gamma \sum _ {s _ {t+1}} \overrightarrow{P ^ \pi}(s _ {t+1} | s _ t) \overrightarrow{v ^ \pi} (s _ {t+1}) \end{align}

となります。後ろ向き価値関数についても同様に

\displaystyle \begin{align} \overleftarrow{v ^ \pi} (s _ t) = \lambda \gamma \overleftarrow{r ^ \pi} (s _ t) + \lambda \gamma \sum _ {s _ {t-1}} \overleftarrow{P ^ \pi}(s _ {t-1} | s _ t) \overleftarrow{v ^ \pi} (s _ {t-1}) \end{align}

となります。これらから考えると

\displaystyle \begin{align} \overleftrightarrow{v ^ \pi} (s _ t) = \frac{1}{1 + \gamma^2 \lambda} \big( r ^ \pi(s _ t) (1 - \gamma ^ 2 \lambda) + \\ \gamma \sum _ {s _ {t+1}} \overrightarrow{P ^ \pi}(s _ {t+1} | s _ t) \overleftrightarrow{v ^ \pi} (s _ {t + 1}) + \\ \lambda \gamma \sum _ {s _ {t-1}} \overleftarrow{P ^ \pi}(s _ {t-1} | s _ t) \overleftrightarrow{v ^ \pi} (s _ {t - 1}) \big) \end{align}

となります。これは頑張って式変形をすると得られます。この式では、 \overleftrightarrow{v ^ \pi} (s _ t) が \overleftrightarrow{v ^ \pi} (s _ {t + 1}) と \overleftrightarrow{v ^ \pi} (s _ {t - 1}) から求まっていることがわかります。

ここからそれぞれのベルマンオペレータ

\displaystyle \begin{align} \overleftarrow{\mathcal{T}}(\overleftarrow{v _ i} (s)) &:= \lambda \gamma \overleftarrow{r ^ \pi} (s) + \lambda \gamma \sum _ {s'} \overleftarrow{P ^ \pi}(s' | s) \overleftarrow{v _ i} (s') \\ \overleftrightarrow{\mathcal{T}}(\overleftrightarrow{v _ i} (s')) &:= \frac{1}{1 + \gamma^2 \lambda} \big( r ^ \pi(s') (1 - \gamma ^ 2 \lambda) + \\ & \gamma \sum _ {s''} \overrightarrow{P ^ \pi}(s'' | s') \overleftrightarrow{v _ i} (s'') + \\ & \lambda \gamma \sum _ {s} \overleftarrow{P ^ \pi}(s | s') \overleftrightarrow{v _ i} (s) \big) \end{align}

が導かれます。ここで、 \pi から導かれるマルコフ連鎖にエルゴード性があると仮定すると、縮小写像になるため、それぞれの価値関数は収束します。

これに基づき、価値関数を \phi, \psi でパラメータ化しているとして、それぞれのTD誤差を

\displaystyle \begin{align} \overleftarrow{\delta _ t} &:= \lambda \gamma R _ {t - 1} + \lambda \gamma \overleftarrow{v} _ {\phi _ t} (S _ {t - 1}) - \overleftarrow{v} _ {\phi _ t} (S _ {t}) \\ \overleftrightarrow{\delta _ t} &:= \frac{1}{1 + \gamma^2 \lambda} \big( R _ t (1 - \gamma ^ 2 \lambda) + \gamma \overleftrightarrow{v} _ {\psi _ t} (S _ {t + 1}) + \lambda \gamma \overleftrightarrow{v} _ {\psi _ t} (S _ {t - 1}) - \overleftrightarrow{v} _ {\psi _ t} (S _ {t}) \big) \end{align}

と定め、更新が

\displaystyle \begin{align} \Delta \phi _ t = \alpha \overleftarrow{\delta _ t} \frac{\partial \overleftarrow{v} _ {\phi _ t} (S _ {t})}{\partial \phi} \Bigg| _ {\phi = \phi _ t} \\ \Delta \psi _ t = \alpha \overleftrightarrow{\delta _ t} \frac{\partial \overleftrightarrow{v} _ {\psi _ t} (S _ {t})}{\partial \psi} \Bigg| _ {\psi = \psi _ t} \\ \end{align}

となります。

実験

モデル化

改めて、定義した価値関数については \overleftrightarrow{v} (s _ t) := \overleftarrow{v}(s _ t) + \overrightarrow{v}(s _ t) なので、これをニューラルネットワークでモデル化する際もいくつかのバリエーションが考えられます。どのようなモデル化が良いのかは不明であるため、3通りを試すことにします。

左からBiTD-FR, BiTD-BiR, BiTD-FBiと呼ぶことにします。

実験ドメイン

今回導入した価値関数は、入力が似ているのに状態価値が大きく異なるような状況で有効に働くことが予想されます。よって、以下のような人工的なChain環境を考えます。

隣接状態への移動で正負交互に大きな報酬を得ることとします。これにより隣接状態について、上部に表示されているように解析的な価値は飛び飛びに変わります。下段に表示されているように表現は似たベクトルとして定義することにします。

方策として左右に50%ずつの確率で移動するものを固定し、この方策について方策評価を行う能力を検証します。

結果

結果は以下のようになります。

提案手法であるBiTD-FR, BiTD-BiR, BiTD-FBiはどれもTD(λ)より良い結果となっています。パラメータλに対する振る舞いとしても、λが0.4から0.6あたりの大きさで最も良い結果となっています。

また、価値関数のモデル化としてはBiTD-FRとBiTD-FBiがよく、BiTD-BiRはやや悪めでした。良いものの共通点としては前向き価値を直接出力するモデルになっているので、その点は重要であると考えられます。

所感

まだ非常に人工的なタスクにおける評価のみにとどまっているため、実践的な設定でどこまで有効かはわかりませんが、ある程度有望なのではないかとも感じます。

完全に証明を追いきったわけではないのですが、理屈の概要としてはそこまで奇妙なトリックに基づくようなものでもないので、損にならない補助タスクとしては入ってくる余地が大いにあるのではないかと思います。

2024年このブログの振り返り

取り組みへの振り返り

　今年は、週に1回はなにかしらの更新をするというルールでやっていた。1,2時間程度で書ける雑なもので良いからとにかく1記事は書くつもりだったが、8月中旬の引っ越しの週では1回だけ落としてしまった。本当に1時間分も余裕が無かったということはないので、これは完全に手抜きでしかなかった。

　それ以外でも、明らかにしょうもない記事だなと自分で感じるものは多々ある。この習慣が無いよりは良かったと思うが、まだ改善する余地もあるだろう。来年はまた一歩進めるようにしたい。

　相変わらずなにをやっていくかフラフラしていた1年だった。以下、書いてきたものを月ごとに振り返る。それぞれの時期で興味があったことを示しているはずだ。

1月

　1月は基本的には強化学習のAtari-100kベンチマークにおけるSOTA手法を動かしたりしていた。このベンチマークは比較的計算資源貧者に優しい部類のものだと思うので、できれば今後も追いかけてみたい。

2月

　2月は、先月に引き続き強化学習を触ったり、突然世界モデル側を見てFSQと戯れたり、AHC030にかなり時間をかけたりしていた。AHCに対するモチベーションは安定しておらず、すごくやるときと参加すらしないときがある。予めスケジュールを把握しておけばモチベーション上げやすい気がするので、多少気をつけてみたい。

3月

　3月は業務で使いそうな技術をつまみ食いしている。Gaussian Splattingもガウス過程もLevenberg-Marquardt法も、今どれだけ使っているかというと微妙。

4月

　4月は急にMambaが気になって読んだ。この後出るMamba2とかも含めて、追いかけて良かったと思う。

5月

　5月も技術をいろいろつまみ食いしたりAHCやったり、謎の動きである。最適化手法を行きあたりばったりに学んだのは、結局あまり身になっていないようにも感じる。

6月

　6月は、途中でMamba2が混入しているが、それ以外はずっとstreet-gaussians-nsを触っている。結構手を入れたので、半年後に手順を再現する必要が生まれたときもそこそこスムーズにできた。結局なににも結びついてないところは残念であるが、これは一つの成果ではあるかもしれない。

7月

　7月も相変わらずstreet-gaussians-nsと、さらにglimにも触っており、Localization/Mappingの人間らしい動きではあるかもしれない。とはいえ、2024年末の時点でLocalization/Mappingに対する興味がどの程度あるかというと、正直かなり薄まっているとは感じる。

8月

　8月は一番ひどい月だった。全然時間をかけていない。そういう月もある。

9月

　9月は完全に迷走して『哲学探求（鬼界彰夫訳）』とか読み出している（実際は8月後半からかなり時間をかけて読んだ覚えがあるが）。いろいろ考えた結果、拡散モデル〜Flow Matchingあたりで気を取り直し、最終的にはMineRLやろうというところに辿り着いているのは良かったかもしれない。ここからしばらく触っていくことになる。

10月

　10月はFlow Matchingによる世界モデルをMineRLで実装するというところに向けて実装を進めていた。ある程度有望であるとは思うが、やっぱりこれだけでは、と思うところもある。

11月

　11月もMineRLをやっている。そして、Streaming Deep Reinforcement Learning Finally Worksという強烈な論文も見つけた月だった。適格度トレースも学び直して、オンライン強化学習への興味が再燃している。手元でも再現してみたい……、なんで2ヶ月も経ってまだなにもできていないのだろう。

12月

　12月もなかなかに話題が散らかってしまった。しかし、その中でも内発報酬のアイデアは自分でも悪くないと思うし、Action Value Gradientも理解と動かしてみることの両方をちゃんとやれたのは良かった。

内容についての総括

　やはり自分は基本的に強化学習に興味があるのだと思う。特にオンライン強化学習。そして計算資源は多くないので必然的にドメインは絞られる。あるいは、計算機がなくても理論面については進められるはず。自ずとやっていくことは収斂していくだろうとも感じるし、積極的に一本芯を通したいとも思う。来年もまた同じような振り返り記事を作ったときに、テーマが浮かび上がってくるような記事のリストになっていると良い。

Action Value Gradientを読み、試す

è«–æ–‡

Action Value Gradientという手法を提案している論文があり、NeurIPS 2024に通っている。

https://openreview.net/forum?id=DX5GUwMFFb

基本的にはSoft Actor-Critic(SAC)を拡張した手法であり、リプレイバッファやバッチ更新を行わないオンライン強化学習の設定で有望な結果を示している。

主な追加点は

  • ActorとCriticにpenultimate normalizationを入れる
  • 観測を正規化する
  • TD誤差も正規化する

であり、逆にいくらかのSACの要素は取り払われている。

AVGのOpenReviewにおいて、SACとの比較が表で簡潔に述べられている。

Q関数の冗長化やターゲットネットワークの作成など、強化学習の学習を安定させるための工夫が除去されて、リプレイバッファも排除されている。

Soft Actor-Critic

まずSACについておさらいする。表の中で外された部分には触れず、コアのアルゴリズム部分についてまとめる。

参考資料

内容

Clean RLのドキュメントによると、SACに関する論文としては結構たくさんある。基本的にはDeep Deterministic Policy Gradient(DDPG)の拡張として確率的な方策になり、エントロピー正則化が入ってQ関数もそれに伴って変わる、という流れになる。

OpenAI Spinning Upによると、off-policy手法であり、TD3と類似していてClipped Double-Q トリックを使っている。

SACの基本バージョンでは連続行動空間でのみ使える。

確率分布 P に対してエントロピーは H(P) = \mathbb{E} _ {x \sim P} \lbrack - \log P(x) \rbrack である。求めたい最適方策 \pi ^ * を以下のように定める。

\displaystyle \begin{align} \pi ^ * = \mathrm{arg} \max _ \pi \mathbb{E} _ {\tau \sim \pi} \left\lbrack \sum _ {t=0} ^ \infty \gamma ^ t \left( R(s _ t, a _ t, s _ {t + 1}) + \alpha H(\pi(\cdot|s _ t)) \right) \right\rbrack \end{align}

通常の収益に対してエントロピーが追加されている。ここで \alpha \gt 0 は調整のためのパラメータである。 \alpha は学習中に変えていくこともあれば、固定することもある。

この条件のもとでは価値関数もそれぞれ変わる。

\displaystyle \begin{align} V ^ \pi(s) &= \mathbb{E} _ {\tau \sim \pi} \left\lbrack \sum _ {t=0} ^ \infty \gamma ^ t \left( R(s _ t, a _ t, s _ {t + 1}) + \alpha H(\pi(\cdot|s _ t)) \right) | s _ 0 = s \right\rbrack \\ Q ^ \pi (s, a) &= \mathbb{E} _ {\tau \sim \pi} \left\lbrack \sum _ {t=0} ^ \infty \gamma ^ t R(s _ t, a _ t, s _ {t + 1}) + \alpha \sum _ {t=1} ^ \infty \gamma ^ t H(\pi(\cdot|s _ t)) | s _ 0 = s, a _ 0 = a \right\rbrack \end{align}

価値関数の学習

ベルマン方程式を考えると

\displaystyle \begin{align} Q ^ \pi (s, a) &= \mathbb{E} _ {s' \sim P, a' \sim \pi} \left\lbrack R(s, a, s') + \gamma \left( Q ^ \pi (s', a') + \alpha H(\pi(\cdot|s')) \right) \right\rbrack \\ &= \mathbb{E} _ {s' \sim P, a' \sim \pi} \left\lbrack R(s, a, s') + \gamma \left( Q ^ \pi (s', a') - \alpha \log \pi(a'|s') \right) \right\rbrack \end{align}

となる。エントロピー項を入れるといっても、本当に方策のエントロピーそのものを計算するというのではなく、次の行動での対数確率を組み込む形になる。上記式からTD誤差を計算して価値関数を更新する。

方策の学習

最大化したいものは以下である。

\displaystyle \begin{align} V ^ \pi (s) &= \mathbb{E} _ {\alpha \sim \pi} \lbrack Q ^ \pi (s, a) \rbrack + \alpha H(\pi(\cdot|s)) \\ &= \mathbb{E} _ {\alpha \sim \pi} \lbrack Q ^ \pi (s, a) - \alpha \log \pi(a|s) \rbrack \end{align}

方策の最適化にReparameterization trickを使う。これはランダムなノイズ \xi に対して確定的な変換を入れることで行動を振り出すものである。

\displaystyle \begin{align} \tilde{a} _ \theta (s, \xi) = \tanh(\mu _ \theta(s) + \sigma _ \theta(s) \odot \xi) \end{align}

Reparameterization trickにより、行動に対する期待値（これは分布が方策のパラメータに依存する問題を持つ）をノイズに対する期待値に書き換えることができる。

\displaystyle \begin{align} \mathbb{E} _ {a \sim \pi _ \theta} \lbrack Q ^ {\pi _ \theta}(s, a) - \alpha \log \pi _ \theta(a | s) \rbrack = \\ \mathbb{E} _ {\xi \sim \mathcal{N}} \lbrack Q ^ {\pi _ \theta}(s, \tilde{a} _ \theta(s, \xi)) - \alpha \log \pi _ \theta(\tilde{\alpha} _ \theta(s, \xi) | s) \rbrack \end{align}

元のSACでは二つのQ関数で小さい方を使ったりするが、結局AVGでは削除されているパートになるため割愛する。

Action Value Gradient

ここまでのSACの損失関数を整理するとAVGの実装となる。

https://github.com/gauthamvasan/avg/blob/main/avg.py

Q lossの方は

    #### Q loss
    q = self.Q(obs, action.detach())    # N.B: Gradient should NOT pass through action here
    with torch.no_grad():
        next_action, action_info = self.actor(next_obs)
        next_lprob = action_info['lprob']
        q2 = self.Q(next_obs, next_action)
        target_V = q2 - self.alpha * next_lprob

    delta = reward + (1 - done) *  self.gamma * target_V - q
    qloss = delta ** 2
    ####

Policy lossの方は

    ploss = self.alpha * lprob - self.Q(obs, action) # N.B: USE reparametrized action

となる。

※ 一見、plossのものがQ関数にも勾配が繋がってしまっているように見えるが、この損失に対応するオプティマイザがpolicyに関するパラメータしか対象に取っていないので問題ないのだろうか。

追加された要素としては

Observation Normalization

学習を始めてから得てきた観測の値から、逐次的に平均・分散を求め、それを使って正規化する。そこではWelfordのアルゴリズムを使っている。

実装としてはGymnasiumのNormalizeObservationを使えば1行で実現できる。

https://gymnasium.farama.org/v0.29.0/_modules/gymnasium/wrappers/normalize/

Penultimate Normalization

最終層の手前の層について、ノルムが1になるようにニューラルネットワークの活性値を正規化する。

\displaystyle \begin{align} \hat{\psi} _ \theta(S) = \frac{\psi _ \theta(S)}{||\psi _ \theta (S) || _ 2} \end{align}

Scaling Temporal Difference Errors

TD誤差自体を正規化することも考える。Return-based Scaling: Yet Another Normalisation Trick for Deep RL で提案されているように

\displaystyle \begin{align} \sigma _ \delta ^ 2 &:= \mathbb{V}[R + \mathbb{V}[\gamma]\mathbb{E}[G ^ 2] \\ \bar{\delta} _ t &:= \delta _ t / \sigma _ \delta \end{align} ]

として正規化する。それぞれの平均・分散は、やはり逐次計算手法で求める。

avg.py のコードではこのTD誤差の正規化は見当たらない。

https://github.com/gauthamvasan/avg/blob/main/avg.py

incremental_rl/td_error_scaler.py でクラス自体は存在しているので、適切に利用すれば効果が見られると思われる。

https://github.com/gauthamvasan/avg/blob/main/incremental_rl/td_error_scaler.py

実験

実際にAVGのコードを動かした。もとのコードのままなので実験ドメインはHumanoidである。

3時間ほど回して2Mステップほど進み、以下のようにエピソードごとの収益が着実に伸びていることが確認できた。

最後あたりでの挙動を動画で見ると以下の通り。

すぐ倒れずにある程度耐えていることがわかる。

まとめ

AVGを動かし、適切に動作することが確認できた。

GymnasiumのNormalizeRewardを使うと、TD誤差ではなく報酬段階でのスケーリングはできる。当面はそちらで \gamma が定数であるならそれで問題ないように見える。もう少し理解を深めて、違う環境に対しても適用してみたい。

AVGの問題点は、論文中でも指摘されているように、サンプル効率の悪さなので、そこを改善するやり方を探してみたい。

VLAなど

è«–æ–‡

　最近いろいろ考えていると、結局、計算資源貧者としては、既存のものをFine-Tuningするということになるのだろうなぁと感じる。

　そうなるとVision-and-Languageモデルから派生して行動を取れるようにする、Vision-Language-Actionモデルとして使うのが一番実践的になるのだろう。

　そのあたりについていくらか論文を見てみた。

è«–æ–‡1 : [2406.09246] OpenVLA: An Open-Source Vision-Language-Action Model

　7Bパラメータのモデルでロボットマニピュレーションタスクに調整する。VLMの初期値としてはPrismatic-7B VLMを使う。

　データとしてはOpen-X Embodiment（70個体のロボットによる2M軌跡のデータ）から選別、また多少の追加などをして970kの軌跡で学習させる。

　学習は結局、データに対するトークン予測になる。もちろん行動は連続的なので、1%から99%の間の値を等間隔に切って256ビンの離散化をする。

　Llama tokenizerでは、100個のspecial tokensが準備されているが、256には足りない。なので、Llamaで使われる頻度が低い方から256トークンを上書きする（[7]）と同じやり方を採用する。

　Low-rank adaptation(LoRA)を使う。必要に応じて量子化もする。

　画像サイズとしては224x224とし、Vision EncoderもFine-Tuningした方が良い。

　学習率は2e-5で固定する。warmupは特に効果がなかった。

計算資源

  • 学習
    • 64基のA100GPUで14日学習
      • A100単機で換算すると21,500 時間
    • バッチサイズ2048
  • 推論
    • 15GBのGPUメモリ消費（bfloat16を使った場合）
    • 4090GPUで6Hz

所感

　やはりこれでも計算資源はかなり必要になってしまう。この量のA100を集められることはない。

　そして学習自体もトークン予測に終始しているため、強化学習をやっていきたいとは思う。

è«–æ–‡2 : [2405.10292] Fine-Tuning Large Vision-Language Models as Decision-Making Agents via Reinforcement Learning

　こちらはActionを直接出力するわけではないのでVision-Language-Actionモデルではないが、テキストとして行動を出させて、それを行動にマッピングする関数を挟んでいる。

　そして、そのテキストを出力する確率から行動の確率を計算する。これがわかれば強化学習が回せるので、PPOで最適化する。

　タスクとしては以下のようなものをやっている。

　簡単なものではたとえば画像的に表示されたテキストを見て数字を合わせるようにするなど。

　モデルとしては llava-v1.6-mistral-7b から始めている。

　基本的に4GPUを使っており、それぞれで128trajectoryを集めて512バッチサイズとして学習を行う。

　PPOの実装としてはpytorch-a2c-ppo-acktr-gail を使っている。

所感

　こちらでも 8 A100s DGX machineが前提になっており、なかなか真似できそうではない。もう一段階小さいものでないと使えないだろうか。3Bあるいは1Bくらいまでで探したくなる。

　ざっと探すと [2407.07726] PaliGemma: A versatile 3B VLM for transfer などになるのだろうか。 RaviNaik/Llava-Phi2 · Hugging Faceもある。このようなものから多少なりともなんとかならないかというところをやっていく必要がある。

世界モデルと内発報酬

　結局MineRLでどういう方向のやり方を試していきたいのだろうと考えたときに、どうも「世界モデルを使った内発報酬」としてまとめるのが良いのではないかと感じられてきた。今回は具体的な実装というよりも、アイデア出しと調査という感じになった。

アイデア

　世界モデルを使った内発報酬の設計について考えていた結果、「以前では予測できなかった状態について、今は予測できるようになっているとき、報酬を与える」ことで上手くいくのではないかという考えに至った。以下、その思考の流れを示す。

　内発報酬は、多くの場合、新規な状態に対して報酬を与えることを大方針とする。これと世界モデルを組み合わせようと思ったとき、最も単純なやり方としては、次状態を予測する世界モデルを学習して予測誤差を報酬とすることが考えられる。

　一方、これだとNoisy TVといった設定が問題になる。Noisy TVとは、常にランダムな画像を映し続けるテレビが環境の中にある場合のことで、単に「世界モデルの予測誤差」を報酬としている場合、このNoisy TVの前に居座り続けると常に予測誤差が大きくなるため、それが望ましいものとなってしまう。

　Noisy TVでなくとも、たとえばMinecraftで単純に予測誤差を報酬としていると、とにかくカメラ方向をぐちゃぐちゃに動かしまくると予測は難しくなると思われるので、そういう行動が強化されるのではないかと予想される（実証したわけではない）。

　Noisy TVや異常に難しい行動に共通することは、どちらも根源的に学習が困難であるというところだと思う。つまり、内発報酬は学習可能であるのにまだ未学習なものに対してのみ与えられるべきなのではないか。

　学習可能性とは要するに学習が進んだ将来での予測誤差であり、理想的には「現在の予測誤差と将来の予測誤差」を比較したいことになるのだが、もちろん将来の予測誤差を今得ることはできないので、両者を同じだけ時間を過去方向にスライドさせて「過去の予測誤差と現在の予測誤差」で代用するということになる。

　こういった思考を経て、過去の時点での予測誤差と、現時点での予測誤差の比較によって報酬を決めるという考えに至る。

　実装の想定としては、現在の世界モデルのパラメータ \theta と、それを適当な時刻間隔でEMA(exponentially moving average)で取ったモデル \theta _ \mathrm{EMA}のそれぞれで状態 s _ {t - 1}、行動 a _ {t - 1}から次状態 s _ {t}を予測して損失の差分を取る。 fを遷移予測の世界モデルとし、損失を \mathcal{L}(\hat{s}, s) = || \hat{s} - s || ^ 2 などと定めて

\begin{align} R _ \mathrm{intrinsic} = \mathcal{L}(f(s _ {t - 1}, a _ {t - 1}; \theta _ \mathrm{EMA}), s _ t) - \mathcal{L}(f(s _ {t - 1}, a _ {t - 1}; \theta), s _ t) \end{align}

論文調査

　きっと似たアイデアがあるはずと思って論文を調べている。まずはキーワード検索と、最近のarXivから内発報酬に関連するものを調べた。

　しかし、あまり考えているものに近そうなものはなかったので、最近あった強化学習の教科書的なもの（Kevin Murphy, Reinforcement Learning: An Overview）の「5.2.4 Intrinsic reward」から探した。

　これの末尾のJürgen Schmidhuber氏のPDFでかなり明確に、1ページ目の右下に

Intrinsic rewards measuring the model’s improvements (first derivative of learning progress) due to the learning algorithm (thus measuring the degree of subjective surprise or fun),

とある。

　ここからはこの論文を読みつつ、引用している論文などを探していきたい。

簡単な実験

　世界モデルのEMAで想定通りの挙動になるかを確認した。学習ステップ40回中に1回、0.9999で更新を行った。

ステップ GT 現パラメータでの推論 EMAパラメータでの推論
10,000
25,000
50,000
100,000

　概ね想定通りに学習にラグを作れている。

その他

　必要な技術パーツとして考えると、もっと強化学習の実装部分に詳しくならないとどうしようもない。それの練習としてはMinecraftは難しすぎると思われるので、違うドメインのもので練習した方が良さそうではある。

MineRL学習 その9 バッチ教師あり学習へ戻る

　前回まで一通りMineRLで各種学習が（性能はともかく）回ることは確認できた。

　おおむね

  1. バッチでの教師あり学習
  2. バッチサイズ1、シャッフルなしでのstream教師あり学習
  3. 連続データでのstream強化学習

という3種類のものでそれぞれ検証していくことができそう。

　今回からはまた(1)の「バッチでの教師あり学習」に戻って、ここからはちゃんと性能をある程度求めていく学習を行うこととする。

データ生成やり直し

　教師あり学習をやるに際して、今回またデータを作り直すこととする。

工夫1 : 行動を10フレームごと固定する

　行動を毎ステップランダムに選択すると、長押しでのブロック破壊に到達できる確率がとても低い。これをある程度低減するために、行動選択をするステップを毎回ではなく何回かに一回とし、行動を選んでいないときは前回と同じ行動を取り続けるようにする。今回は10フレームごとに行動選択をするようにした。

　これはDreamerV3あるいはその元の論文で、より賢いものが行われている。（これらの論文では、目の前のブロックを壊せるだけの長さで長押しするという行動が用意されているらしい）

工夫2 : インベントリ開く行動を行わない。

　また、今回は自身の動きに応じた適切な予測ができるかどうかを確認するため、インベントリを開く行動は取らせない状態で行動させた。

学習

　データを18,000フレーム x 10セット分だけ生成し、学習を行った。

　パラメータ

  • 画像サイズ : 64x64
  • 学習率 : 1e-4
  • 重み減衰 : 0.0
  • 入力系列長 : 16

　50,000ステップの学習を3回行った。

　学習後のものについて予測に基づいた画像を生成した。以下の結果は16ステップ分のGTを入れて、次の1ステップ分を予測するということを繰り返している。

番号 GT 予測
0000

　カメラ方向の変化に対してある程度尤もらしい推測ができていそうである。

所感

　ある程度時間を取ってちゃんと学習をさせるとある程度のものができそうなことはわかってきた。もう少し回してより鮮明に得られないかを試してみたい。

DiffusionDriveを読んだメモ

è«–æ–‡
\displaystyle \newcommand{\bs}{\boldsymbol}

https://arxiv.org/abs/2411.15139

概要

自動運転のPolicy（軌跡出力）として拡散モデルを使うことを考える。

既存手法と合わせて分類すると以下のようになる。

  • (a) 一つの軌跡を回帰する(Transfuser, UniAD, VAD)
  • (b) 予め用意した軌跡候補から選択する(VADv2)
  • (c) ガウス分布からの拡散モデル
  • (d) 【提案手法】アンカーありのガウス分布からの切り捨て拡散モデル

最も単純に拡散モデルを適用した場合、(c)のTransfuserDPと名付けるものとなる。これの問題点には以下の二つがある。

  • (1) 生成に20ステップほどかかってしまい遅い
  • (2) 複数のガウス分布からサンプリングしても一つの軌跡になってしまいがち

この論文では、これらの問題に対して

  • (1) 生成の元となるガウス分布はアンカーを中心とした複数のガウス分布からのサンプリングにする
  • (2) 生成ステップを途中で切り上げることで高速化する

という工夫を入れる。

手法詳細

センサーデータを入力として、軌跡（位置の系列） \tau = {(x _ t, y _ y)} ^ {T _ f} _ {t = 1} ( T _ f は予測する長さ )を出力する問題を考える。

今回考える拡散モデルでは、センサーデータから得られる表現を条件として、軌跡について、最初はノイズ的なものから始めて、適切な軌跡になるように逆拡散過程をたどっていく。

提案手法では、まず N _ \mathrm{anchor} 個のアンカーを用意する。これはK-Meansで学習データから構築する。一つのアンカーは軌跡である。 \bs{a} _ k = { (x _ t, y _ t) } ^ {T _ f} _ {t = 1} これとノイズの混合を拡散モデルで扱う。

\displaystyle \begin{align} \tau _ k ^ i = \sqrt{\bar{\alpha} ^ i} \bs{a} _ k + \sqrt{1 - \bar{\alpha} ^ i} \bs{\epsilon}, \qquad \epsilon \sim \mathcal{N}(0, \bs{I}) \end{align}

であり、拡散ステップ i は i\in\lbrack1, T _ \mathrm{trunc}\rbrack として、 T _ \mathrm{trunc} は元の拡散ステップの最大値よりも十分に小さいものを使う。学習において、このノイズ付き軌跡を用いて、ノイズが除去された軌跡と、そのスコアを求めるようにする。

\displaystyle \begin{align} \{\hat{s} _ k, \hat{\tau} _ k\} ^ {N _ \mathrm{anchor}} _ {k=1} = f _ \theta (\{\tau ^ i _ k \} ^ {N _ \mathrm{anchor}} _ {k=1}, z) \end{align}

z は条件付けのための情報である。

スコアの教師として、GTの軌跡に最も近いアンカーにのみ y _ k = 1 を、それ以外に y _ k = 0 を与える。

工夫を入れない拡散モデルであるTransfuserDPでは、出力が結局一つの軌跡にまとまってしまうモード崩壊が起きる。例としてTransfuserDPについて20個のノイズから20ステップで生成を行うと以下の上半分のようになり、普通のTransfuserDPでは20個のノイズが一つの軌跡に潰れてしまうことがわかる。下半分が表現しているように、提案手法のDiffusionDriveのようにアンカーを元にした生成を行うと、Top-10あたりには、ある程度尤もらしい軌跡が残るようになる。

後に実験の項目でこれを定量的に評価するために、以下のような指標を準備する。

\displaystyle \begin{align} \mathcal{D} = 1 - \frac{1}{N} \sum _ {i=1} ^ N \frac{ \mathrm{Area}(\tau _ i \cap \bigcup _ {j=1} ^ N \tau _ j) }{\mathrm{Area}(\tau _ i \cup \bigcup _ {j=1} ^ N \tau _ j)} \end{align}

要するにすべての軌跡が作るエリアと各軌跡のIOUの平均を1から引いたもの。mIoUが高いということは多様性がないことを意味するので、1から引いたこの指標は高いほど多様性が高い。

実験

NAVSIMデータの概要

評価スコア

  • PDMスコア
    • 各スコアの重み付け和
    • NC : no at-fault collisions
    • DAC : drivable area compliance
    • TTC : time-to-collision
    • Comf : comfort
    • EP : ego progress

アーキテクチャ

  • 正面3つのカメラを1024x256として利用
  • LiDARはBEVとして利用
  • navtrainã‚’100エポック
    • 8個の4090GPUs
    • バッチサイズ512
  • テスト時
    • 8点の軌跡で4秒分を予測

結果

比較

他の手法と比較した結果が以下である。

提案手法は、PDMSのスコアが最も高く、かつAnchorの数も少なく抑えられている。

段階的な性能変化

Transfuserから段階的に提案手法まで発展していく過程を表にすると以下のようになる。

読み取れることは

  • 2行目 : 単純にTransfuserDPにすると推論速度が大幅に上がってしまい遅い
  • 3行目 : TruncatedDiffusionなどの工夫を入れると高速になる
  • 4行目 : さらにネットワークアーキテクチャを工夫することで一番良い性能となる

RTX 4090で45FPSで動作するほど高速である。

各種パラメータの比較

ネットワークアーキテクチャの違いや、アンカーの数などを比較すると以下のようになる。

所感

Diffusion Policyがかなり盛り上がってきているのを感じる。拡散モデルの、データから多峰的な分布を学習できるという点が良いとされている。速度が主に問題であるとされるが、この論文の手法のように改善できるとするものがでてきている。

データから学習していく設定ではDiffusion Policyはかなり有望に思える。一方、個人的な興味としては強化学習まで結びつけたい。拡散モデルをスコアを基準にしたランジュバン・モンテカルロ法として捉えると、Q関数のスコアがそれに近いという論文がある。(https://openreview.net/forum?id=CKqiQosLKc) （レビューで理論面での指摘が入っており、撤回されているが）ちょっと内容を見てみたい。

MineRL学習 その8 ShortcutModelの実装

MineRL

　前回は、オンラインで学習できそうなことを確認した。

　しかし、結局毎ステップで完全に次状態予測をしなければ、そこでの損失を行動学習に利用できない。（フローマッチングの学習自体は次状態の予測とは異なる）

　次状態の生成に何回も推論をする必要があると大変である。よってその高速化がしたくなる。

　以下の論文ではShortcutモデルとして学習の仕方を工夫すると高速化が達成できると述べられている。

　これを実装してみた。

　学習した結果は以下の通り。

Shortcut損失なし Shortcut損失あり

　Shortcut損失なし（重みを0にした）場合ではShortcut損失が大きくなっており、重みを1にした場合はほぼ維持するような形になった。

　最初のステップでとても小さいことが少し気になる。なにかしらバグらせているかもしれない。

　NFE（推論時のニューラルネットワークを通す回数）を変更して生成の結果を確認した。。左から「GT」「NFE=64」「NFE=16」「NFE4」「NFE4」縦はタイムステップ。

Shortcut損失なし Shortcut損失あり

　Shortcut損失のない段階から、ほぼNFE=4までは生成結果が変わらない。Shortcut損失を入れた改善というのはあまり見られない。

　学習時間は

  • Shortcut損失なし : 1時間27分
  • Shortcut損失あり : 2時間26分

であった。

　ちょっと狙った効果が出てない、あるいは最初からそこそこ得られている？ みたいなところがよくわからないので、次回はここの調査を進めつつ、オンライン学習でもこの再構成損失を利用するようにしたい。

MineRL学習 その7 オンライン学習

MineRL

　前回はバッチサイズが1、およびデータのシャッフルなしで学習できることを確認した。記事の最後でstep学習にしたときバグる問題があると言っていたが、これは単純なミスですぐに修正できた。

　今回はこれMineRLとのオンライン相互作用から学習できるように実装を行った。

　学習結果は以下の通り。

　学習最後付近での再構成結果（左が実際の画像・右が予測画像）

　学習時間としては10万ステップの学習にちょうど4時間ほどであった。

　環境が終わってしまったとき、どうもresetするだけでは2回目以降の動作が上手くできないという問題は残っている。それへの対処として、環境が終わる制限時間を変更することができた。MineRLの以下の部分を修正すると良い。

　次回はなんとか適格度トレースも共存させて、ActorCriticとして学習できるようにしたい。現状のところ報酬はまともに得られていないと思うので、環境からの報酬ではなにも学習はできないだろう。せっかく世界モデル的なものがあるので、それの誤差などを用いてなんらかの内発報酬という形にしてなにかできないかということを考えたい。