🔍

論文:強化学習でもスケーリング則がある?

に公開

1000 Layer Networks for Self-Supervised RL: Scaling Depth Can Enable New Goal-Reaching Capabilities

主張

  • 自己教師ありの対照強化学習 (Contrastive RL) において、ネットワークを深くすると劇的に性能が向上した -> 強化学習にもスケーリング則がある。
  • ただし、滑らかに能力が向上するのではなく閾値を超えると急激に向上する。
  • 一般的に強化学習では浅い MLP を用いるが、残差接続と LayerNorm と Swish を組み合わせることで critic 側では最大 1024 層 (actor 側では 512 層) まで安定した学習ができた

ゴール条件付き強化学習

ゴール条件付きMDP

\mathcal{M}_g = (S, A, p_0, p, p_g, r_g, \gamma)

を考える

  • S は状態の集合、A は行動の集合で、各時刻 t における状態 s_t と行動 a_ts_t \in S, a_t \in A
  • 初期状態 s_0 は分布 p_0 に従ってサンプリングされる
  • p は遷移関数 p(s_{t+1} | s_t, a_t)
  • ゴール g \in \mathcal{G}p_g からサンプルし、 f: \mathcal{S} \rightarrow \mathcal{G} で状態とゴール空間を関連付ける

報酬は次のステップでゴール状態に到達する確率密度で定義。r_g(s_t, a_t) \stackrel{\mathrm{def}}{=} (1 - \gamma) p(s_{t+1} = g | s_t, a_t)
ここから割引報酬和を計算すると、そのポリシーがゴール状態 g をいつか訪れる割引付き確率となる。ただし、実験では sparse reward(ゴール近傍で1)という設定を主に使っている

対照強化学習 (Contrastive RL)

Criticの学習
Criticは状態 s と行動 a の埋め込みを得る \phi(s, a) とゴール g の埋め込みを得る \psi(g) からなる。
(s_i, a_i, g_j) の組について、i = j のケース (同一軌跡の未来の状態) を正例、i \neq j のケース (他軌跡のゴール状態) を負例として、以下のような損失関数 \mathcal{L} を最小化するように学習する。結果として、(s, a, g) が到達可能なゴールに対応している組であるほどスコア f_{\phi, \psi}(s, a, g) が大きくなるように \phi\psi が学習される。

\begin{gathered} f_{\phi, \psi}(s, a, g) = ||\phi(s,a) - \psi(g)||_2 \\ \mathcal{L}(\phi, \psi) = - \mathbb{E} \left[\sum_i \log{\frac{\exp{f_{\phi, \psi}(s_i, a_i, g_i)}}{\sum_{j=1}^{K} \exp{f_{\phi, \psi}(s_i, a_i, g_j)}}} \right] \end{gathered}

Actorの学習
Actorはゴール条件付きポリシー \pi_{\theta}(a | s, g) で、単純にスコア f が大きい行動を選ぶよう学習する。

\max_{\theta} \mathbb{E} \left[f_{\phi, \psi}(s, a, g) \right]

実験結果


どのタスクでも Depth を増やすことで成績は向上する。Depthが深くなるにつれて線形に向上するというよりは閾値を超えた時に大きく向上する。


閾値を超えた時に大きく向上するためには残差接続が必要。


Depth が浅い内はバッチサイズを大きくしてもあまり意味がないが、閾値を超えるとバッチサイズを大きくするほど成績が向上する。


同じデータを使って学習するなら、Depth の深いネットワークのほうが性能が良い
Deep Collector と Shallow Collector の Collector の結果を比較すると、Collector (リプレイデータを収集するネットワーク) 自体も深い方が性能が良い、つまり Depth が深いネットワークが探索したほうが良い探索ができている。


SAC、SAC+HER、TD3+HER、GCBC、GCSL では Depth を深くしてもほとんど性能の向上が見られない。

知見・考察

  • この論文では対照強化学習を対象として実験を行っている。TD法をベースとした方法では深くしても性能は飽和・悪化する。
  • ただし、対照強化学習という設定に固有の話というよりは「自己教師あり学習」「分類問題として学習信号を用いる」の 2 点が重要なのではないかと考えているようだ。
  • 対照強化学習だけでは駄目で、オンラインでの探索を組み合わせないと Depth を深くすることの恩恵はない。

Discussion