【論文瞬読】IDT：人間の意思決定を真似た革新的強化学習手法の紹介

こんにちは！株式会社AI Nestです。
今回は、強化学習（RL）の世界で注目を集めているIn-context Decision Transformer（IDT）という新しい手法について解説していきたいと思います。

タイトル：In-Context Decision Transformer: Reinforcement Learning via Hierarchical Chain-of-Thought
URL：https://arxiv.org/abs/2405.20692
所属：Institute of Artificial Intelligence, Jilin University, China; Key Laboratory of Symbolic Computation and Knowledge Engineering, Ministry of Education, Jilin University, China; Lehigh University, Bethlehem, PA, USA.
著者：Sili Huang, Jifeng Hu, Hechang Chen, Lichao Sun, Bo Yang

強化学習って？

まず、強化学習について簡単に説明します。強化学習とは、機械学習の一種で、試行錯誤を通じて学習していく方法です。例えるなら、ゲームの攻略法を自分で見つけていくようなものです。最初は下手でも、何度もプレイするうちに、だんだん上手くなっていく、そんなイメージですね。

従来のin-context RLの問題点

さて、IDTの話に入る前に、従来のin-context RLという手法が抱えていた問題点について触れておきましょう。この手法は、Transformerという自然言語処理などでよく使われる技術を使って、過去の経験から学習し、行動を決定します。

Figure1, 最小限の行動と高度な意思決定の試行錯誤比較。*はより良い結果を示す。 (a) 試行錯誤プロセスでは、記憶は経験から得られた最小限の行動で構成され、より良い行動を探すための文脈として機能する。 (b) 高次試行錯誤プロセスでは、記憶と探索は高度な意思決定を行う。1つの高次試行錯誤プロセスが複数の行動を制御するため、より少ないメモリで経験を保存し、より少ない計算コストでより良い意思決定を探索することができます。

ただ、この方法には計算コストという大きな課題がありました。過去の経験を長く記憶すればするほど、計算量が増えてしまい、処理が重くなってしまうのです。特に、長期的な視点で行動を決定する必要があるタスクでは、この問題が顕著でした。図1は、従来の試行錯誤とIDTの高レベルの試行錯誤を比較したものです。従来の試行錯誤では、過去の行動一つ一つを記憶し、そこからより良い行動を探索していました。一方、IDTの高レベルの試行錯誤では、高レベルの決定を記憶し、それを元に複数の行動を導き出すため、記憶すべき情報が大幅に減り、計算コストを削減できるのです。

IDTの登場！

そこで登場したのが、今回紹介するIDTです。IDTは、人間の意思決定の仕組みにヒントを得て、この計算コストの問題を解決しました。

人間の意思決定は、階層的な構造を持っていると言われています。例えば、旅行に行くときに、まず予算を決めて、次に交通手段を考え、最後に細かい行動を決める、といった感じです。IDTもこれと同じように、高レベルの決定と低レベルの行動という階層構造を採用しています。

高レベルの決定とは、長期的な目標や戦略に関わる決定で、低レベルの行動は、具体的な行動を指します。IDTは、まず高レベルの決定を行い、それに基づいて低レベルの行動を決定していくため、従来の手法に比べて計算コストを大幅に削減できるのです。

IDTの仕組み

Figure2, IDTのアーキテクチャは、高度な試行錯誤プロセスをシミュレートするために3つのモジュールに設計されています。まず、(1) Making Decisionsモジュールは、エピソード全体にわたるコンテキストを提供することで、高度な意思決定を予測します。エピソード全体にわたるコンテキストには、合計報酬の昇順に配置された複数の軌道が含まれます。次に、(2) Decisions to Goモジュールは、予測された高度な意思決定を条件として、cステップ先の行動を予測します。最後に、(3) Reviewing Decisions モジュールは、実行されたアクションを確認し、次のサイクルでの経験として活用します。Reviewing Decisions は、トレーニング時にはオフラインデータから高レベル決定の真のラベルをエンコードしますが、テスト時には実行されたアクションからエンコードします。

IDTは、3つのモジュールから構成されています。(図2)

1. Making Decisions: 過去の経験に基づいて、高レベルの決定を生成するモジュールです。

2. Decisions to Go: 生成された高レベルの決定を基に、具体的な行動を決定するモジュールです。

3. Reviewing Decisions: 実行した行動を振り返り、経験として学習するモジュールです。

これらのモジュールが連携することで、IDTは効率的に学習を進め、複雑なタスクにも対応できるようになります。

IDTのすごいところ

IDTのすごいところは、Grid WorldやD4RLといった様々なベンチマークで、従来の手法を上回る性能を達成したことです。特に、長期的な視点が必要なタスクで優れた結果を残しています。(図3)

Figure3, (a) テストと (b) トレーニング時間の結果。10,000 回の勾配更新ごとのトレーニング時間、Grid World における 50 エピソード、D4RL における 10 エピソードのテスト時間を報告します。ここでコンテキストサイズを測定するためにステップ数を使用していることに注意してください。ステップごとのトークン数はアルゴリズムによって異なる場合があります。AD の各ステップには、観察、行動、報酬、完了の 4 つのトークンが含まれます。IDTの「意思決定」モジュールとATには、さらに「リターン・トゥ・ゴー」トークンが含まれています。タスクの長さが長くなるにつれ、コンテキストの長さも指数関数的に増加せざるを得なくなり、計算コストが2乗的に増加します。これに対し、IDTは、一連の意思決定を高いレベルで構成しなおします。そのため、コンテキストは1エピソードの長さよりも短く、計算コストを大幅に削減できます。

また、IDTは外部からの指示も利用できます。例えば、ゲームの攻略動画を見せることで、より効率的に学習を進めることができるのです。

まとめ

今回は、強化学習における新しい手法、IDTを紹介しました。IDTは、人間の意思決定の仕組みにヒントを得て、計算コストの問題を解決し、複雑なタスクにも対応できる画期的な手法です。今後の発展が楽しみですね！

最後まで読んでいただき、ありがとうございました！