皆さんはDecision Transformerをご存じですか？
これは、2021年に登場したTransformerを用いた強化学習モデルです。AtariやOpenAI Gymなどをはじめ、多数のモデルフリーのオフライン強化学習タスクにおいてベールラインと同等か、それ以上の性能を達成することができたモデルです。
Transformerというと、自然言語処理をはじめ、多種多様な領域で大きな成果をあげているニューラルネットワークのアーキテクチャですから、それを深層強化学習の方策ネットワークに用いた研究なのかな？と思われるかもしれませんが、そうではありません。Decision Transformerでは、強化学習を系列問題として解きます。他の強化学習手法は価値の最大化問題として方策獲得を目指しますが、Decision Transformerではそれを踏襲していないため、それらの概念は不必要となっています。
強化学習について勉強したことがある方であれば、そのゴリゴリの数学で理解が難しい分野の上に成り立っているものであることはご理解いただけると思いますが、Decision Transformerはそれを崩すような衝撃的なものでした。このようなDecision Transformerについて本記事では分かりやすく解説していきます。
Transformerの躍進
Transformerの概要
Transformerは、2017年にGoogleがAttention is All You Needという論文でニューラル機械翻訳として提案した手法です。ニューラル機械翻訳というのはニューラルネットワークを用いて翻訳タスクを実現しようとする自然言語処理タスクの１つで、従来一般的に使用されてきた再帰的ニューラルネットワーク（Recurrent Neural Network: RNN）の欠点を解決する形で登場しました。その課題というのは、位置依存性に弱い点と、並列化による計算の高速化ができないというものです。特に、性能面で言うと位置依存性に弱いというのは、「マクロな情報は捉えられる一方で、ミクロな情報は捉えられない」ということになるため、正確な機械翻訳を実現するうえで大きな課題となっていました。それを解決することができる手法として提案されたTransformerは本当に大きなブレイクスルーだと思います。
それ以降、BERTやGPTなどTransformerを用いたモデルが大きな成果を上げています。話題のChatGPTもTransformerを用いて構築されたGPT-3.5やGPT-4が使用されています。
※GPT-4のアーキテクチャは公開されていませんが、Generative Pre-trained Transformerの略であるGPTシリーズであることには変わりないため、Transformerをベースとしたモデルであることは間違いないでしょう。
Transformerは自然言語処理分野で登場したアーキテクチャですが、画像や音声など、他の分野でも高性能を達成できることが明らかになり、現在は多種多様な分野で使用されています。
Transformerについては以下の記事で詳しく解説していますので、詳細を知りたい方はぜひそちらもご参照ください。
Transformerの特徴
Transformerが提案された背景課題と、それを解決できる手法であることはご理解いただけたと思いますが、Transformerには更に驚くべき特徴があります。それは、スケーリング則（Scaling Law）と呼ばれる性質を持つ点です。
スケーリング則とは、主にモデルのパラメータ数やデータセットのサイズを大きくすると、モデルの性能が指数関数的に高くなるという法則です。そして、その性能には上限が無いという特徴があります。上限が無いという特徴から、モデルの大規模化による性能向上の流れが加速し、GPT-1, GPT-2, GPT-3と大規模化が進みました。ChatGPTのGPT-3.5の性能の高さもTransformerのスケーリング則のおかげです。
余談ですが、パラメータ数の大規模化は実質的に計算コストの増大を意味するため、データセットのサイズを大きくする方に舵をきるとことが増えてきている印象です。
このように大規模化して性能を向上させると、今までは獲得が難しかったような高次的な意味概念を獲得できるという現象がみられています。これにより、ゼロショットで汎化できたり、分布外の出力に対応できたりするようになりました。
このTransformerの特徴はとても興味深く、これを強化学習にも取り入れたいという流れは必然ではないでしょうか？
強化学習
強化学習の概要
強化学習は報酬に基づいて最適な行動獲得を目指す機械学習領域です。教師あり学習のように明確な答えが有るわけではないけど、なんとなく良し悪しは判断することができるという、人間のもつ曖昧さをゴリゴリの数学のもとで成り立っている機械学習で扱えるようにした点が強化学習の凄さです。
具体的にはエージェントと環境の２つをマルコフ決定過程に基づいてモデル化します。マルコフ決定過程では、状態、行動、即時報酬の３つを考え、次の状態は、現在の状態と行動のみに基づいて決定されるということを考えます。そこで、現在にとって良い行動とは、将来に得られる報酬から逆算して決められるべきだという考えのもと、価値関数を定義し、ベルマン方程式をもとに再帰的に学習します。強化学習では、ベルマン方程式を近似的に解こうとし、MC法、TD法などがあります。
オンライン強化学習とオフライン強化学習
強化学習は、試行錯誤をもとに学習を行います。すなわち、環境とエージェントが相互作用できることが前提になっているのです。学習時に環境とエージェントの相互作用を必要とする強化学習手法をオンライン強化学習と言います。
しかし、必ず環境とエージェントが相互作用できるとは限りません。例えばビジネスや医療を考えてみると、試行錯誤による失敗はあまり許容できません。既に先人の知恵があるなら、それを学んで実践すべきだと言えるでしょう。そこで、既にあるデータのみを用い、環境とエージェントの相互作用を伴わない強化学習手法であるオフライン強化学習が注目され始めています。
オフライン強化学習の課題
強化学習というのは、基本的に試行錯誤により足りない情報は適宜収集し学習する方法をとります。しかし、オフライン強化学習では追加でデータを収集することができません。すなわち、データが存在しない部分については価値を適切な値に設定することができないのです。この場合、価値の過大評価などが起こることになります。このような問題を分布シフトと呼びます。
本題であるDecision Transformerはオフライン強化学習手法の１つですが、強化学習問題を系列問題として扱うため、上で述べたような分布シフトが発生せず、注目を浴びています。
Decision Transformer
それではDecision Transformerについて説明します。Decision Transformerは以下のような構造をしています。
Decision Transformerで使用されているTransformerは、2017年に提案された元祖Transformerを改良したCausal Transformerを用いたGPTアーキテクチャが使用されています。
ここで扱う軌道\(\tau\)を
$$
\tau = \{\hat{R}_1, s_1, a_1, \cdots,\hat{R}_{t-1}, s_{t-1}, a_{t-1},\hat{R}_t, s_t, a_t ,\cdots, \hat{R}_T, s_T, a_T\}
$$
とします。ここで、報酬\(\hat{R}_t\)は以下のように表されるものとします。
$$
\hat{R}_t = \sum_{t'=t}^T r_{t'}
$$
Decision Transformerの入力は報酬・状態・行動、出力は行動です。報酬はスカラー、状態や行動はベクトルであり、扱うデータは画像やモータの角度など多種多様です。その為、報酬・状態・行動ごとに別々の埋め込み層を適用しTransformerへ入力できるベクトルへサイズへと変換します。そして、位置埋め込みベクトルを足し合わせたものをTransformerへの入力とします。
学習では\(\tau\)を入力して、\(\{\cdots,a_{t-1}, a_t,\cdots\}\)を出力するようにパラメータを最適化します。ただし、先読みを防止するため、\(a_t\)を予測するときは、\(\{\hat{R}_1, s_1, a_1,\cdots,\hat{R}_{t-1}, s_{t-1}, a_{t-1},\hat{R}_t, s_t \}\)のみを使用するようにマスク操作をします。
学習するときは、\(\hat{R}_t\)はデータセットに依存して自動的に決定されますが、評価するときは報酬部分は私たちが決めてあげる必要があります。これはDecision Transformerの大きな特徴ですが、強化学習を報酬による条件付き系列問題として解いているので、指定する報酬の値が小さいと、質の低いとされる行動が出力され、逆に大きいと、質の高いとされる行動が出力されることになります。
価値を定義して方策を学習するタイプの強化学習手法の場合、価値が最大となる行動の確率分布しか分かりませんが、Decision Transformerでは、それ以外の行動についても報酬の予測値を変えてあげることで出力させることができます。
成果
Decision TransformerのパフォーマンスをTD学習と、BC（模倣学習の一種）で比較したものを以下に示します。評価する際に使用したタスクはAtari、OpenAI Gym、Key-To-Doorです。最もインパクトのデカいのはKey-To-Doorの比較です。Key-To-Doorとは、鍵を使用しないと開かないドアがあり、それを超えないとゴールにたどり着けないタスクです。鍵を取りに行くか行かないかで、ドアに向かう価値が変化するため、従来の機械学習手法では扱うことが困難なものとなっています。一方で系列問題として解いたDecision Transformerはパフォーマンスに圧倒的な差を付けました。
なぜ、Decision TransformerはKey-To-Doorタスクにおいて大きな差をつけて高精度を達成することができたのでしょうか。Decision Transformerで用いられているTransformerは自然言語処理において大きな成果を出せたモデルです。自然言語は任意の単語や文章が遠く離れた場所の単語や文章の意味に影響を与えるという複雑性を持ちます。強化学習における行動も類似した複雑性を持つと考えられます。しかし、TD法など価値を用いる強化学習手法は、割引報酬和を考えるため、近視眼的な行動を獲得しがちであり、一方で、Decision Transformerはその流れを踏襲せず、Transformerの得意とする系列問題として扱ったため高精度を達成することができたのでしょう。
まとめ
最後に簡単にまとめます。Decision Transformerとは、Transformerを用いて強化学習問題を系列問題として扱う手法です。AtariやOpenAI Gymなどをはじめ、多数のモデルフリーのオフライン強化学習タスクにおいてベールラインと同等か、それ以上の性能を達成することができました。
強化学習もTransformerのスケーリング則による高性能化が可能であることを示しました。
参考文献
[1] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, Igor Mordatch. "Decision Transformer: Reinforcement Learning via Sequence Modeling, " arXiv preprint arXiv:2106. 01345, 2021.
