Wikipediaはオフライン強化学習を補助できるか? オフライン強化学習に言語タスクでの事前学習を導入！

Offline Reinforcement Learning 2023年10月11日

3つの要点
✔️ オフライン強化学習において, 大規模データを集める困難さを解決するために, ドメインの異なる言語タスクでの事前学習を提案
✔️ 言語タスクで事前学習したモデルの性能をオフライン強化学習時に転移させるための技術を提案
✔️ 実験的に収束スピードと性能の面で既存手法を上回ることを確認

Can Wikipedia Help Offline Reinforcement Learning?
writtenby Machel Reid, Yutaro Yamada, Shixiang Shane Gu
(Submitted on 28 Jan 2022 (v1), last revised 24 Jul 2022 (this version, v3))
Comments: Published on arxiv.
Subjects: Machine Learning (cs.LG); Artificial Intelligence (cs.AI); Computation and Language (cs.CL)

code：

本記事で使用している画像は論文中のもの、紹介スライドのもの、またはそれを参考に作成したものを使用しております。

背景

近年オフライン強化学習では, Decision Transformerのように, 強化学習をシークエンスモデリングタスクとして定式化して, 自己回帰的な機構を備えたモデルによって方策を実現する枠組みが成功を収めています. 一方でそうした手法は一から訓練すると収束のスピードが遅いという課題を抱えていました. 本論文では, 言語タスクや視覚タスクのような, 異なるドメインでシークエンスモデルを後続の強化学習のために事前学習する手法を開発し, その有効性を実験的に検証しました.

手法

まず, シークエンスモデリングとしての強化学習の基本的な要素について説明します. ここでは, 以下の軌道と呼ばれる状態と行動と未来のリターンの系列 $t$がデータとして与えられることを仮定します.

$$t = (\hat{R}_1, s_1, a_1, \dots, \hat{R}_N, s_N, a_N)$$
ここで,$s_i, a_i$は時刻$i$での状態と行動, $\hat{R}_i$は $\hat{R}_i = \sum_{t=i}^{N}r_t$です. それぞれの時刻の$s_i, a_i, \hat{R}_i$をトークンだと捉えると言語モデルの訓練と同じ枠組みが使えることがわかると思います.
本論文では, 言語や視覚タスクなど, 強化学習とはドメインの異なる問題でTransformerを事前学習することを目指します[図1]. 彼らはその際に, 言語や視覚タスクのための表現と, 強化学習によって獲得される表現の乖離が問題となると考え, 事前学習で得られた表現ベクトルと, 強化学習で得られる表現ベクトルの乖離をなくすために2つのテクニックを提案しています.

1. Similarity loss between Language representation and offline RL representation

ここで, 事前学習されたTransformerの語彙サイズを$V$として, それぞれの埋め込みベクトルを$E_1, \dots, E_j, \dots, E_V, \quad \forall j, E_j \in \mathbb{R}^d$とします. また, 各系列で状態, 行動, 未来のリターンを別々に埋め込んだ表現ベクトルを$I_1, \dots, I_i, \dots, I_{3N}, \quad \forall i, I_i \in \mathbb{R}^d$とします. オフライン強化学習のための表現ベクトル$I_1, \dots, I_{3N}$が既に得られた言語の表現ベクトルと近くなるように以下の損失を導入します.

$$\mathcal{L}_{\mathrm{cos}} = - \sum_{i=0}^{N} \max_j \mathcal{C}(I_i, E_j)$$

ここで, $\mathcal{C}$はコサイン距離です. この損失によって, 強化学習で得られる表現が言語タスクで獲得されたものと乖離しないように学習を進めて, 事前学習で高い性能を達成しているTransformerのシークエンスモデルとしての性能を援用することを期待しています.

2. Language model co-training

本論文では, オフライン強化学習のフェーズでも言語タスクでの学習も続けるということを行う. このことによって, オフライン強化学習中も, より直接的に言語タスクなどのシークエンスモデルングタスクの恩恵を受けられることを期待している. 最終的な目的関数は以下のようになります.
$$\mathcal{L} = \mathcal{L}_{\mathrm{MSE}} + \lambda_1 \mathcal{L}_{\mathrm{cos}} + \lambda_2 \mathcal{L}_{\mathrm{LM}}$$
ここで, $\mathcal{L}_{\mathrm{MSE}}$はDecision TransformerのようなTransformerを用いたオフライン強化学習のための損失関数, $\mathcal{L}_{\mathrm{LM}}$は言語タスクのための損失, $\lambda_1, \lambda_2$はハイパーパラメタです.

実験設定

以下に本実験の設定を説明します. 本実験では, 言語タスク, 画像認識タスクでの事前学習の有効性を, オフライン強化学習のベンチマークを用いて検証しています. 以下に, 事前学習に用いるモデル, 比較対象の強化学習ベースライン, 性能の検証に用いたタスクについて簡単に説明します.

事前学習モデル

言語タスク: 1. GPT-2-small 2. Decision Transformerとの公平な比較のために, 同じパラメタ数のTransformerをWikipedia-103データセットで訓練したモデル. 彼らはこれをChibiT (小型の言語モデル)と呼んでいます.
画像認識タスク: 1. CLIP (Contrastive Language–Image Pre-training). CLIPは, text-encoderとimage-decoderからなり, キャプションと画像のマッチングを予測するよう訓練されたモデル. それぞれのencoderはTransformerからなります. 2. ImageGPT. ImageGPTはGPTと同じアーキテクチャを持ち, 言語タスクではなくピクセルの予測タスクによって訓練されています.

強化学習ベースライン

本研究では強化学習のベースラインとして, Transformerを事前学習なしで用いるDecision Transformer (DT), Transformerを用いないオフライン強化学習手法としてCQL, TD3-BC, BRAC, AWRを用いています.

タスク

本実験では, モデルをAtari, Open AI Gym Mujocoそれぞれに対して, オフライン強化学習専用のデータセットD4RLを用いて評価しています. D4RLは, それぞれのタスクについて, クオリティの異なる行動方策によって集めたデータを提供しています.

Atariタスク: Breakout, Qbert, Pong, Seaquestの4つのゲームを利用し、ベースラインモデルと比較しています.
OpenAI Gymタスク: HalfCheetah, Walker2d, Hopperの3つのタスクでモデルを訓練し、評価しています.

結果と解析

図2からわかるように, 言語タスクで事前学習されたChibiT, GPT2は事前学習なしのDTに比べて高い性能をはっきしていることがわかります. また, オフライン強化学習手法のCQLやTD3-BCと比べても同等かより高い性能を達成していることがわかります. この結果から, まず性能面で言語タスクによる事前学習が有効であることを見ることができます. 以下により詳しい解析の説明を行なっていきます.

収束速度

彼らは, 事前学習なしのDTと言語タスクにより事前学習したモデルの収束スピードを比較しています. ここで, 収束は, 平均リターンと最大リターンの差が正規化したスコアで2以内となることであると定義されています.図3を見てわかるように, 言語タスクにより事前学習されたモデルChibiTとGPT2は収束のスピードがDTに比べて2倍以上早いことがわかります.

Vision vs. Language

図2を見ると, 画像認識タスクで事前学習を行なったCLIP, iGPTいずれも言語タスクで事前学習を行なったモデルに比べて性能が低いことがわかります. 特に性能が著しく低いのは, 画像認識タスクのみで事前学習されたiGPTです. 著者らはこれを、言語モデリングとトラジェクトリモデリングの間の基本的な類似点に起因するものとし, 仮説検証のために注意機構のパターンを可視化しています. [図4]. 図4から, iGPTのattention patternは, DTのattention patternと著しく異なり, 解釈可能性にも乏しいことがわかります. こうした追加実験と合わせて, 強化学習においては, Transformerの事前学習タスクとして言語タスクの方がより有益であるということが示唆されています.

まとめ

いかがだったでしょうか. 今回はTransformerを用いたオフライン強化学習の性能と収束速度を, 言語タスクを用いた事前学習によって向上させられることを示した論文を紹介しました. 一見全く性質の異なる強化学習と言語タスクの間にシナジーが生まれるというのは非常に面白いですね. 今後もTransformerを用いた強化学習から目が離せませんね.