準最適なデータによるオフライン強化学習に関する研究がNeurIPS2023に採択されました。

本研究は、機械学習分野のトップ会議である、Thirty-seventh Conference on Neural Information Processing Systems(NeurIPS 2023)に採択されました(ポスター発表、採択率26.1%)。

Yueh-Hua Wu¹², Xiaolong Wang², Masashi Hamaya¹, “Elastic decision transformer” [paper] [project page] [code]. ¹OMRON SINIC X, ²UC San Diego

第一著者であるYueh-Hua WuさんはUC San Diegoの博士課程の学生で、本研究は弊社インターンシップで行われた成果となります。

背景

強化学習は環境と相互作用しながらエージェントを学習させる方法であり、ロボティクスや推薦システムに応用されます。これまで、強化学習はオンラインで環境と相互作用することが一般的でした。しかし、例えばロボットの強化学習においては、試行錯誤の中でロボットが壊れる可能性があるなど、データ収集が困難である、あるいは、できない場合も想定されます。このような問題に対し、近年では事前に取得したデータのみで学習する、オフライン強化学習という方法が提案されています。

特に、Transformerを使用したオフライン強化学習である、Decision Transformer(DT)が注目されています。DTは、過去の状態、行動、報酬の教示データの系列を入力とし、将来の行動を出力するモデルを、教師あり学習によって学習します。DTはConservative Q-Learningなどのオフライン強化学習より高い性能を示しましたが、“trajectory stitching”に問題があると指摘されています。Trajectroy stitchingとは、理想的なオフライン強化学習に求められる性質であり、準最適な教示データをつなぎ合わせ、最適に近い軌道を生成することを指します。

本研究は、Elastic Decision Transformer (EDT)という、可変長の入力により、stitchingの性能を向上させる方法を提案しました。例えば、短期的な入力を与えることで、過去の失敗データを適切に破棄することが可能となります。また、短期的な入力を与えると、出力の分散が大きくなりますが、探索を促進し、より改善された軌道を生成することが可能となります。一方で、すでに軌道が最適に近い場合、長期的な入力を与え、出力の安定性、一貫性を担保することができます。

Elastic Decision Transformer

以下にstitchingの簡単な例を示し、EDTの動機を説明します。横軸は時間、縦軸は価値(将来にわたる収益)を示します。今、二種類の軌道を持つデータセット(aとb)を考えます。基本的に、モデルは本来の軌道を予測する傾向にあります。例えば、aの状態にある場合は、次の時間はaの状態に遷移する可能性が高い傾向があります。Stitchingがないモデルを使用し、時間t-1でbの状態から開始した場合、より低い価値を持つ時間t+1におけるbの状態に遷移する可能性があります。この問題に対し、時間tの時点で、時間t-1の状態を入力しないことで、時間t+1でより価値の高いaの状態に遷移することが可能となります。逆に、時間t-1でaの状態から開始した時、aの状態を保持し、より高い価値を持つaの状態に遷移するようにします。

上記の動機をもとに、以下にEDTの概要図を示します。EDTは、可変長の状態を入力とし、どの長さの系列を考慮すると最も高い収益を得るかを計算します。そして、予想された収益をもとに、取るべき最適な行動を出力します。

実験

EDTの有効性を検証するために、シミュレーション実験を行いました。強化学習のベンチマークとして普及している、D4RLとAtariを使用しました。学習データとして、“medium-replay”という、エージェントが熟練者の1/3の性能に到達するまでの間に、再生バッファにあるすべての系列を保存したデータを使用しました。実験の結果、DTやImplicit Q-learningに比べてマルチタスク学習においてより高い性能を示しました。

YouTube Video

インターン募集

機械学習やロボティクスに興味がある方はぜひ応募お願いいたします。詳細は弊社のインターン募集のページからご覧いただけます。