(GROOVE X Advent Calendar 2024 2日目の記事です）

こんにちは！ふるまいチームのエンジニア、市川です！

普段はLOVOTの意思決定エンジンを作っています。趣味や仕事で論文を読むこともあります。今後、読んだ論文をInside of LOVOTで紹介できればと思っています（紹介する技術がLOVOTで使われているわけではないです）。

今回紹介するのは「STEVE-1」です。これはMinecraftというゲーム上で、エージェントに指示を出して行動させる研究です。

arxiv.org

前提知識

この論文を理解するためには、いくつかの重要な概念を押さえておく必要があります。

CLIP（Contrastive Language-Image Pre-training）

[2103.00020] Learning Transferable Visual Models From Natural Language Supervision

出典：Alec Radford et al., "Learning Transferable Visual Models From Natural Language Supervision", arXiv:2103.00020.

CLIPは、画像・テキストのペアデータを大量に用いて訓練されたモデルです。画像とテキストを共通の潜在空間（同じ意味のモノが近くの座標になる事が期待される空間）に埋め込むことができます。言うなれば、同じ座標系で画像やテキストの位置を求めることができるため、画像とテキスト間の類似度（距離）を計算できます。例えば画像と"dog"というテキストとの類似度を調べれば、"dog"が写っているかどうかを推定できます。後述のunCLIPの通り、画像生成のタスクにも利用されています。

MineCLIP

[2206.08853] MineDojo: Building Open-Ended Embodied Agents with Internet-Scale Knowledge

出典：Linxi Fan et al., "MineDojo: Building Open-Ended Embodied Agents with Internet-Scale Knowledge", arXiv:2206.08853.

MineCLIPは、CLIPをMinecraft用に特化させたようなモデルです。 Minecraftのゲームプレイ動画とそれに対応するキャプションのペアデータを用いて訓練されており、動画とテキストを共通の潜在空間に埋め込むことができます。そして、埋め込まれた動画とテキストの位置関係から類似度を計算することができます。これにより、例えば「Shear sheep to obtain wool（羊毛を得るために羊の毛を刈って）」というテキストと「羊に近づいている様子」の動画の類似度は高くなります。「羊の毛を刈っている様子」の動画の類似度は更に高く、逆に「羊から離れている様子」の動画の類似度は低くなります。つまり、テキストで目標を与えた時に、それがどれくらい達成できているかを類似度によって評価できるわけです。これは、強化学習における報酬の評価にも使うことができます。

unCLIP

[2204.06125] Hierarchical Text-Conditional Image Generation with CLIP Latents

出典：Aditya Ramesh et al., "Hierarchical Text-Conditional Image Generation with CLIP Latents", arXiv:2204.06125.

unCLIPは、テキストキャプションから画像を生成する手法です。図は、点線から上がCLIPの学習プロセス、下がunCLIPの生成プロセスを示しています。 unCLIPは、priorとdecoderで構成されます。

prior：CLIPで作られた「テキスト埋め込み」と「ノイズの乗った画像埋め込み」を受取り、ガウシアン拡散モデルによりテキストを条件としてノイズを取って、ノイズの無い画像埋め込みを出力する
decoder：埋め込み表現から画像を生成する（CLIPの画像埋め込みを反転するように作られたモデル）

これにより、入力テキストと画像を条件に、新たな画像を生成する事ができます（decoderでテキストの埋め込み表現だけから画像を生成することもできます）。

VPT (Video PreTraining)

[2206.11795] Video PreTraining (VPT): Learning to Act by Watching Unlabeled Online Videos

出典：Bowen Baker et al., "Video PreTraining (VPT): Learning to Act by Watching Unlabeled Online Videos", arXiv:2206.11795.

VPTは、画像からエージェントの行動を出力するような基盤モデルの事前学習手法です。大量のラベルなし動画で事前学習させることを目的とします。論文ではMinecraftのゲームプレイに適用されています。

IDM（Inverse Dynamics Model）

まず、IDMと呼ばれる、動画を入力して行動を出力するモデルを作ります。 Inverse Dynamicsと呼ばれているのは、環境の変化を引き起こした行動を推測するようなタスクであるためです。 Minecraftの場合は、プレイ動画（連続したフレーム）を入力して、各フレームにおいてプレイヤーが行った低レベル制御（マウスとキーボード操作）を出力させます。このモデルの学習のために、人間プレイヤーの操作とゲーム動画を記録したデータセットが必要になります。つまり教師有り学習が必要なのですが、基盤モデルに学習させる大量の動画に比べると遥かに小規模なデータセットで良いようです。

VPTによる基盤モデルの事前学習

作成したIDMを用いて、大量のラベルなし動画に行動のラベル（疑似ラベル）を付けます。それを教師データとして基盤モデルを学習させます。ここで作るのは、画像を入力し、人間がとる可能性の高い行動を出力するようなモデルです。 Minecraftの場合は、ゲーム画面を受け取り、人間がしそうな低レベル制御（マウスとキーボード操作）を出力させます。これによって事前学習されたモデルは、強化学習などの手法によって下流タスクにファインチューニングすることができます。

STEVE-1

ここからが本題となるSTEVE-1の紹介です。

STEVE-1は、テキスト指示を受け取って行動を出力するようなモデルです。上述のMineCLIPやVPTの基盤モデルと、unCLIPのようなPrior modelを組み合わせた階層的な構造となっています。 Prior modelにはCVAE（Conditional Variational Autoencoder）が採用されています。

推論について

推論時の流れは下記の通りです。

指示の埋め込み
- テキスト指示を受け取り、MineCLIPのエンコーダで埋め込みベクトルにします
埋め込まれた視覚目標の生成
- テキストの埋め込みベクトルをPrior model（CVAE）に入力し、対応する画像の埋め込みベクトルを生成します
- 画像の埋め込みベクトルはテキスト指示に対応しているため、行動で目指すべき視覚目標（簡単に言えば指示を実行できているゲーム画面）であると考えられます
行動の生成
- 現在のゲーム画面と視覚目標を方策モデル（VPTモデル）に入力し、低レベル制御（マウスとキーボード操作）を出力します
- 方策モデルはResNetやTransformer-XLで構成され、下記のように計算されます
  - ゲーム画面をResNetに入力し、画像の特徴を出力します
  - 視覚目標を線形変換します
  - この2つを足したものをTransformer-XLに渡し、低レベル制御を出力します
行動の実行
- 算出された低レベル制御を実行します

学習について

STEVE-1の学習工程では、方策モデル（VPTモデル）のファインチューニングと、Prior model（CVAE）の学習を行います。

方策モデルのファインチューニング

事前学習されたVPTのモデルを基盤とし、視覚目標を達成するようにファインチューニングします。効率化のため、人間がラベルを作る代わりに、教師データを自動生成します。具体的には、行動した後に過去のタイムステップをランダムに選択し、その視覚を更に過去のタイムステップの目標として設定します。

Classifier-Free Guidance

方策モデルはしばしば指示を無視し、事前学習で学習した行動に従ってしまうことが確認されています。これを軽減するために、画像生成モデルで使用されているテクニックであるClassifier-Free Guidanceを適用します。具体的には、方策モデルにより、視覚目標を条件としたロジット（行動の確率分布）と、無条件のロジットを同時に計算します。そして、 $\displaystyle \lambda$ パラメータを用いて2つのロジットを組み合わせて、2つのモデル間のトレードオフを行います。

$\displaystyle {\rm logits} = (1 + \lambda) ({\rm conditional \ logits}) - \lambda \ ({\rm unconditional \ logits})$

$\displaystyle \lambda$ の値をある程度大きく設定することで、視覚目標を条件としたロジットの影響が大きくなり、パフォーマンスが向上する事が確認されています。無条件ロジットを出力可能にするため、学習時には視覚目標を10%の確率でドロップアウトさせます（推論時にはドロップアウトさせません）。

Prior model（CVAE）の学習

Prior modelはテキストの埋め込みベクトルから視覚目標を作るモデルでした。学習データとして、テキスト指示と、それに対応する短い動画（16フレーム）のペアを用意します。本研究では2000件のデータが手動で作成されています。更にGPT-3.5-turboを用いて8000件のテキスト指示を追加し、計10000件のテキスト指示それぞれに対してゲームプレイの動画から最も類似度の高いセグメントを検索し、ペアデータとしています。類似度の判定にはMineCLIPが用いられます。

主な結果

本研究では短期的タスク、長期的タスクなど、様々な評価が行われています。詳しくは論文を見てください。また、下記でデモ動画を見ることもできます。

https://sites.google.com/view/steve-1

ここでは短期的タスク、長期的タスクの結果について簡単に紹介します。

短期的タスク

短期的な11のタスク（木を切る、穴を掘るなど）でテストされました。これらはゲームプレイの最初の2.5分以内に達成可能で、完了までに複数の手順を必要としないタスクです。

上図の左(a)では、VPTエージェントと、STEVE-1エージェントの結果が比較されています。 STEVE-1はテキスト指示だけでなく、方策モデルに直接視覚目標を与えることもできるため、その両方の結果が掲載されています。評価項目として、初期地点からの移動距離、取得できたSeed、Wooden Log、Dirtの最大アイテム数があります。従来のVPTエージェントのパフォーマンスを大きく上回っている事が分かります。

上図の右(b)では、STEVE-1エージェントのテスト中に視覚目標をどれだけ達成できたかを表しています。与えられた指示と、評価項目の対応表となっており、色は視覚目標と埋め込まれたゲーム画面の最短距離です（小さいほど目標を達成できたと推定できます）。対角線の数字はその要素のランクです（0であれば行の最小要素）。

下記はテキスト指示によって動いている様子です。

youtu.be

長期的タスク

塔の建設、木製の板の作成のような、複数の手順を必要とする、より長期的なタスクも実験されています。目標を直接プロンプトする方法と、タスクをいくつかのサブタスクに分解してプロンプトを順番に与える方法の2種類が試されています。結果は下図の通りです。

下図左下はオレンジ色の線が直接プロンプトした結果、青線がサブタスクに分解した結果です。直接プロンプトすると成功率がほぼゼロになっています。サブタスクに分解した場合は50%程度の確率で塔を建設し、70%の確率で木製の板を作成できています。このように、サブタスクに分解することで一定の成功率向上は見られますが、長期的なタスクにはまだ課題があります。

下記は塔の建設の様子です。

youtu.be

まとめ

STEVE-1について紹介させて頂きました。 MineCLIP、VPT、unCLIPを組み合わせた手法でした。短期的タスクでは良い成績を残していましたが、長期的なタスクではまだ課題が見られました。しかし面白い研究結果であり、Optimus-1のような後続の研究に繋がっていきます（これも今後紹介できればと思います）。

arxiv.org