今回こそ、強化学習に革命を起こせるのか？Transformerのチャレンジ
3つの要点
✔️ 強化学習を同時確率分布の系列モデリングを自己回帰モデルで扱うDecision Transformer
✔️ GPTのCausal Transformerをシンプルに適応したアーキテクチャ
✔️ Offline RLの設定でModel-free手法のSOTAを達成
Decision Transformer: Reinforcement Learning via Sequence Modeling
written by Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, Igor Mordatch
(Submitted on 2 Jun 2021 (v1), last revised 24 Jun 2021 (this version, v2))
Comments: Accepted by arxiv.
Subjects: Machine Learning (cs.LG); Artificial Intelligence (cs.AI)
code：
はじめに：TransformerがRLでも革命を起こせるのか？
Zero-shot言語生成モデルGPT-3や分布外な画像生成モデルDALL-Eをはじめ、最近の研究ではTransformerが高次な意味的概念を学習できることを示しています。NLP生まれのTransformerがCVの世界で大活躍していることをヒントに、そのTransformerが強化学習の世界にも革命を起こせるのか？今回はその可能性を秘めている研究を紹介します。
有望なTransformerを強化学習への応用を試みる研究は前からありました。それらと異なり、この研究では強化学習の扱い方を変えて状態、行動と報酬の同時確率分布を自己回帰モデルとして扱うDecision Transformerを提案しました。Decision Transformerは事前学習+ファインチューニングという組み合わせを目指すOffline RLの問題設定において、既存のModel-freeのSOTA手法と同等な精度を出しています。
ここでは、Decision Transformerを直感的に理解するためにトイデータを用いた説明を行います。
Figure 2(左）に示した有効グラフの最短パスを見つけるタスクを考えてみましょう。強化学習の問題として捉えると、エージェントがゴールノードにたどり着いたら報酬を0とし、そうでなければ-１を与えます。ある時刻以降に得られる収益(return-to-go)の考え方を用いると、Figure 2（中）のようなデータセットがランダム方策で得られると思います。例えば、あるノードに-2が付くのは、その軌道データにおいて2ステップ後にゴールノードに辿り着くことを意味します。ここで状態、行動、未来の収益を一つのトークンとした系列データに対して、言語生成モデルGPTを用いて次のトークンを予測する自己回帰モデルで学習してみました。評価時は現時刻の状態と目的報酬と過去の情報から次に行くべきノード（行動）を予測させます。その結果、学習用データセットから状態と目的報酬のペアが関連しているトークンが該当する行動を予測できました。またFigure 2(右)のように学習データを単に模倣するのではなく、よりも優れた方策の学習が可能になることが分かります。
提案手法：Decision Transformer
提案手法のアーキテクチャは極めてシンプルです。次のような軌道の系列データが与えられる設定を考えます。時刻tの状態と行動とそれ以降に得られる収益(Return-go-to)を一つのトークンとし、KつのトークンをGPTのCausal Transformerに入力する仕組みとなっています。トークンは、状態Sに対しては畳込み層(CNN)、収益(Return-to-go)や行動(Action)を線形レイヤーを通してconcat結合して得られます。
Return-to-go：
訓練時はまず、オフラインデータセットから系列長=Kのミニバッチをサンプリングします。そして、時刻tより以前の情報で得られたトークンをGPTで提案したCausal Transformerに入力し、自己回帰モデルとして次の行動を予測するように学習します。つまり、時刻Tの状態とそれ以降に得られる収益(Return-go-to)と過去の情報を入力して時刻Tの行動を予測し、予測した行動を用いて次の時刻の予測を行います。
最後に、離散的な行動を取る場合はクロスエントロピー損失関数で、連続的な行動をとる場合は平均二乗誤差の損失関数を用いて学習させます。
ただし、評価時はコンテキスト長Kと目標収益をハイパーパラメータとして与えます。具体的な流れはAlgorithm 1に示します。
Offline RL ベンチマークでの実験と評価
比較手法はOffline RL専用のアルゴリズムと模倣学習(Behavior Cloning)を選びました。提案手法Decision TransformerはModel-freeな手法であるため、Model-freeのSOTAモデル Conservative Q-Learning(CQL)をはじめBEAR、BRACと比較実験を行いました
そして、実験タスクは高次元な画像観測かつ長期的なクレジット分配を求める行動空間が離散的なAtariと、きめ細かい連続的な制御を必要とするOpenAI Gymとしました。
Figure 3からDecision Transformerはどちらの実験においても既存のSOTA手法と同等な精度を出しており、提案手法の有効性が示せました。
Atari
DQN方策によって取得した500万トランザクションの軌道からなるデータセットの1%を訓練に用いました。またプロゲーマーの点数を100とし、ランダム方策を0とした評価を導入しました。
４つのAtariタスク（Breakout, Qbert, Pong, and Seaquest）で評価実験を行った結果(Table 1)、３のタスクでSOTAのCQLと同程度な精度を出しており、その他の手法より全てのタスクで精度が上手回りました。
OpenAI Gym
このセッションでは、D4RLベンチマークに加えて2Dのreacher環境を加えて、実験を行いました。2D reacherタスクはゴールが与えられてそれ達成するタスクであるため、報酬は離散的である点はその他タスクとは異なっています。次に説明する３種類のデータセットを用いました。
1. Medium：スコアで見ると、SACを用いて専門家レベルの1/3を達した「中レベル」方策によって集めた100万 timesteps のデータセット
2. Medium-Replay：SACで「中レベル」に達するまでに、バッファに貯めた全部のデータセット
3. Medium-Expert：「中レベル」の方策と「専門家レベル」の方策によってそれぞれ集めた100万timestepsのデータセット
Table 2図に示されている通り、Model-freeのSOTAモデルCQLと最近提案されている有望な手法と全部比較し、提案手法DTがほとんどのタスクにおいて精度が良かったことがわかります。
OpenAI Gymの方がAtariよりも提案手法の有効性をわかりやすく示せました。
考察
論文ではDecision Transformerの性質に関して多くの実験を行なって調べました。詳細を知りたい方は論文をご覧ください。
ここでは提案手法Decision Transformerは模倣学習とどう異なるか、目的収益とコンテキスト長Kが結果にどう影響するのかといった特性に関する分析と、これらの特性によって長系列のタスクにてもたらす優位性の分析と二つの視点から議論の内容を紹介します。
1. Decision Transformerは、一部のデータセットに対して特定の収益を得る設定の模倣学習とどう違うのか？
データセットをエピソード単位で収益高い順に並べ、上位のX%のデータのみを用いて模倣学習したPercentile Behavior Cloning(%BC)を使って、提案手法がBCとの違いについて分析しました。結果を見ると、Decision Transformerは10%BCと同等な精度が得られました。これは良い方策を学習する情報は上位10%のデータセットに集約されていることを示しています。
一方で、replay bufferの1%のデータのみを用いたAtariでの実験では、%BCは精度が低く、Decision Transformerの方が良いことがわかりました。従って、少数データの場合Decision Transformerは、効率よく最適なデータを選び出して模倣学習を行なっていることが推測できます。
2. Decision Transformerは目的収益と状態から行動を予測するが、その目的収益が結果にどう影響を与えるのか？
Decision Transformerがこれ以降に得られる収益の意味を理解しているかを目的報酬の設定を調整して評価実験を行いました。
Figure 4に示したように、ほとんどのタスクにおいては与えた目的収益と実際に得られた収益は相関しており、Decision Transformerは目的収益によって適切な行動を選択できることがわかります。
さらに、AtariのSeaquestタスクにおいては、データセットにある最も高い収益よりも高い収益が得られたことは、Decision Transformerには探索する能力を備えていることを意味します。
3. Decision Transformerのコンテキスト長Kが結果にどう影響を与えるのか？
フレームトラッキングを用いる一般的な強化学習においては一つ前の状態(K=1)があれば十分と考えられてきました。しかし、Decision Transformerにおいては過去の情報にアクセスする重要さを実験で示し、Atari Gameには過去の情報が役立つことが推測できます。
4.Decision Transformerが長期のクレジット分配問題とスパース報酬の設定にうまく対応できるか？
系列が長い問題設定にどう対応するかは強化学習の大きな課題となっている。ここでは、鍵のある部屋にエージェントが行動を取った後に、空の部屋を経て、ドアのある部屋に入るグリッドベースのKey-to-Door環境を考えます。一つ目の部屋で鍵を取得しなければ、最後のヘアにあるドアを開けられずタスクは失敗します。時間の経過によって長期間のクレジット分配ができるかを評価するために、空の部屋を用意しています。
この実験では、ランダムに動く方策によって収集したデータセットを用いてDecision Transformerの学習を行いました。コンテキストの長さをエピソードの長さ(K=episode length)としました。成功したデータのみを模倣する%BCとDecision Transformerだけが効率よく方策の学習することができたことがわかります。また、TD Learning(CQL)はQ-valuesを長い時間ステップの伝搬ができず、低い精度を出しました。
ここでは正確に状態を捉えられるかを評価するために、Decision Transformerを行動とともに報酬を予測するように修正し、最初から報酬を与えずに予測してもらいました。Figure 5の結果(左)からDecision Transformerは部屋が変わる時に予測する報酬を更新していることがわかります。また注意機構の重み（Attention weight）は鍵を取ったとドアを開けた際に励起され、正確に機能していることが読み取れます。
まとめ
この記事では、言語系列モデルのアイデアを強化学習に適応するモチベーションから生まれたDecision Transformerの紹介を行いました。言語モデルのアーキテクチャをほとんど変えないDecision Transformerは、スタンダートのOffline RLベンチマークにおいてModel-freeのSOTA手法と同等程度なパフォーマンスが出せました。
シンプルな手法で結果も良いのですが、Decision Transformerの１番の貢献は自然言語処理の技術進化の恩恵を強化学習でも受けられる可能性を示したことだと考えています。特に大規模な事前学習とファインチューニングの組み合わせを目指すOffline RLに、Decision Transformerのアーキテクチャが自然にフィットしていることから、強化学習に大きな改革を起こせる研究となることを期待しております。
この記事に関するカテゴリー
