Jack of All Trades (JAT) 多目的トランスフォーマーエージェント

Hugging Face は、Jack of All Trades (JAT) を発表しました。これは、単一のネットワークで幅広いビジョン・言語および意思決定タスクを実行することを目的とした汎用エージェントプロジェクトです。このプロジェクトは Gato アーキテクチャのオープンな再現として、包括的なエキスパート RL エージェントのセット、専門的なトレーニングデータセット、そしてマルチモーダルトランスフォーマーモデルを提供します。

JAT データセットとエキスパートポリシー

JAT プロジェクトは、汎用エージェントのトレーニング専用に設計された最初のデータセットを提供し、数十万件のエキスパート軌跡を含んでいます。これらの軌跡は、最先端のエキスパートエージェントを複数の多様な環境でトレーニングすることにより生成されました。含まれる環境は以下の通りです：

  • Atari: クラシックなアーケードゲーム。
  • BabyAI: シンプルなナビゲーション環境（BabyAI ボットを利用）。
  • Meta-World: ロボット操作タスク。
  • MuJoCo: 物理ベースの連続制御。

統一されたユーザーインターフェースを実現するために、JAT データセットは Wikipedia、Oscar、OK-VQA、Conceptual-Captions からのテキストデータも組み込んでいます。

JAT エージェントアーキテクチャ

JAT は EleutherAI の GPT-Neo 実装上に構築されています。このアーキテクチャは、観測埋め込み、行動埋め込み、対応する報酬を交互に配置することで、シーケンシャルな意思決定タスクに対応できるよう設計されています。

埋め込みとエンコーディングのメカニズム

JAT はデータモダリティに応じて異なるエンコーディング戦略を使用します：

  • Images (e.g., Atari): 畳み込みニューラルネットワーク (CNN) で処理。
  • Continuous Vectors: 線形層で処理。
  • Discrete Values: 線形射影層で処理。
  • Text: GPT-2 の戦略でトークン化。
  • General Images: ViT 型エンコーダで処理。

予測と損失関数

モデルは因果マスクを用いて自己回帰的に次の埋め込みを予測し、観測を1ステップずらすことで、エージェントは過去の観測と行動に基づいて次の行動を予測しなければなりません。損失は各モダリティごとに別々に計算されます：

  • MSE Loss: 画像と連続値に使用。
  • Cross-Entropy Loss: 離散値に使用。

最終的な損失はシーケンス内の各要素の損失の平均です。

パフォーマンス結果

JAT は 157 のトレーニングタスクで評価され、4 つの主要ドメインにおけるエキスパートエージェントと比較して平均 65.8% のパフォーマンスを達成しました。

ドメイン別パフォーマンス

  • BabyAI: エキスパートスコアの 99.0% を達成し、1 つのタスクでのみエキスパートの 50% を超えられませんでした。
  • MuJoCo: エキスパートスコアの 84.8%。
  • Meta-World: エキスパートスコアの 65.5%。
  • Atari 57: エキスパートスコアの 14.1%（人間パフォーマンスの 37.6% に相当）で、21 ゲームで人間のパフォーマンスを上回りました。

モデルは NLP と CV タスクにおいて基本的な能力を示すものの、最大の強みは単一のネットワークでこれら多様な RL ドメインにおけるエキスパートのパフォーマンスを模倣できる点にあります。

観測予測の影響

JAT プロジェクトで行われた研究では、エージェントに将来の観測を予測させる（報酬の最大化に加えて）ことが学習を改善するかどうかを調査しました。観測損失と行動損失のバランスを取るために重みパラメータ $\kappa$ を使用した結果、$\kappa = 0.005$ で「最適点」が見つかりました。

この特定の重みでは、観測を予測する学習がエージェントの学習効率を向上させました。しかし、$\kappa$ が高すぎる（例：0.5）場合、観測予測の目的が学習プロセスを妨げました。これは、補助的な目的が適切にバランスを取れば汎用エージェントに有益であることを示唆しています。

今後の研究課題

JAT エージェントをさらに進化させるために、Hugging Face は改善すべき主な 3 つの領域を特定しました：

  1. Data Quality: より多様なエキスパートエージェントからの軌跡を収集してデータセットを拡張し、バイアスを低減する。
  2. Offline RL: 基本的な行動模倣から脱却し、エージェントがサブ最適軌跡を活用してエキスパートを上回る可能性を持たせる。
  3. Sampling Strategies: 均一サンプリングではなく、より困難なタスクに焦点を当てる動的なマルチタスクサンプリング戦略を実装する。

Sources