多目的エージェント「JAT」公開：1つのモデルでゲームもロボット制御もこなす

公開 2024年4月22日拡散指数 33

汎用エージェントプロジェクト「JAT」がオープンソースで公開された。Gatoの再現と改良を目指したもので、モデルとデータセットを同時提供。単一のTransformerで157のタスクに対応し、専門エージェント比65.8%の性能を達成している。

JAT（Jack of All Trades）は、強化学習の意思決定、ロボット制御、基本的な画像・言語処理を単一ネットワークで実行できるオープンソースの汎用Transformerエージェントだ。

DeepMindのGatoを再現・拡張したプロジェクトであり、専門環境ごとにモデルを用意することなく、1つの重みで複数のタスクをこなす。

できること

JATは157種類の訓練タスクに対応している。

Atariのレトロゲームをプレイするだけでなく、MuJoCoによる連続制御シミュレーション、Meta-Worldでのロボットアーム操作、BabyAIでの指示に従った移動タスクなどをこなせる。

さらにテキストや画像の入力にも対応しており、初歩的な視覚・言語タスクも同じインターフェースから処理できる。

JATのベースにはEleutherAIのGPT-Neoが使われており、時系列の意思決定データを効率よく扱う独自の埋め込み機構を備えている。

観測値、報酬、行動の各埋め込みを交互に配置し、因果関係に沿って次の行動を予測する。観測値が画像ならCNN、連続値なら線形層、離散値なら線形射影を用い、テキストにはGPT-2トークナイザ、画像にはViTエンコーダを併用する。

行動の予測だけでなく将来の観測結果を補助的に予測させる損失設計を取り入れており、適切な重み付けを行うことで学習効率が向上したという。

評価タスク全体において、JATは専門エージェント比で平均65.8%の性能を記録した。BabyAIでは99.0%、MuJoCoでは84.8%、Meta-Worldでは65.5%に達し、Atari 57では専門家の14.1%、人間比で37.6%のスコアを出し、21本のゲームで人間を上回った。

モデルの重み、ソースコード、各タスクの専門エージェントはすべてHugging Face Hubで公開されている。

さらに、数十万件のエキスパート軌道データを収録した「JAT dataset」も公開されており、Pythonのdatasetsライブラリから直接読み込んで利用できる。

これまでGatoのような汎用意思決定モデルは論文発表にとどまり、外部の研究者が重みやデータセットを検証する手段が限られていた。

JATはモデル本体だけでなく、高品質な軌道データと専門方策までオープンにした点で価値が高い。汎用エージェントや身体性AIの研究に取り組む開発者にとって、信頼できるオープンな基準点となる。

現状は単純な行動模倣（Behavioral Cloning）で学習されているため、専門家の性能を追い越すことは難しいが、ゲームからロボット制御まで網羅する基盤として今後の発展が期待できる。