世界モデルにポケモンをプレイさせる方法

世界モデルにポケモンをプレイさせる方法 人工知能 世界モデル 人工知能の分野で最も魅力的な研究の一つは、世界モデルに関するものだと考えています。世界モデルには多くの種類があり（そしてその言葉自体が少し過負荷になっていますが）、世界モデルにおける最もエキサイティングなアーキテクチャの一つは、ヤン・ルカン氏のJEPA（Joint Embedding Predictive Architecture）です。多くのバリエーションがあり、特に私の目を引いたのはLeWorldModel1でした。これは私のRTX 3080 Tiでローカルにトレーニングできるほど小さく、以前の多くのJEPAアーキテクチャよりもシンプルな設計でした。 サンフランシスコはポケモン関連の記念品で覆われていました–1026匹のポケットモンスターが登場する大きな広告が、マーケットストリートの地下鉄路線やバス停に貼られていました。SFは今年、ポケモンワールドチャンピオンシップの開催地であり、多くの風変わりで喜びに満ちたポケモン使いが、SFダウンタウンの穏やかな丘やコンクリートの金融街をさまよっていました。おそらく、すべての広告が私を潜在的に操作したのかもしれませんが、私たちの世界モデルにとって良い「世界」は、すべてを始めた1996年のゲーム、ポケットモンスター 赤であると決めました。 2 ポケモンのメカニクス ポケットモンスター 赤は、世界を探索し、ポケモンと呼ばれる生き物を集め、それらをバトルで使用するゲームです。 方向ボタンを使って歩き回ったり、メニューを移動したりします。Aボタンは物と対話したり、会話を進めたり、選択を確認したりします。Bボタンは通常、キャンセルしたり、戻ったりします。スターターを選択するということは、モンスターボール（ポケモンを保持する容器）に近づき、選択を確認する会話を通過することを意味します。 ゲームの序盤、プレイヤーはオーキド博士の研究室にいます。そこで博士は、最初のポケモンとしてフシギダネ、ヒトカゲ、ゼニガメのいずれかを提供します。これらは3匹の「スターター」で、研究室のテーブルにあるモンスターボールの中にそれぞれ待っています。 上記のGIFは、モデルトレーニングの最終結果です。モデルは、ゼニガメを選択する一連のボタン操作を計画しました。モデルは簡単なテストでは有望に見えましたが、予想以上に多くの困難や後退がありました。 計画 目標は比較的単純でした：オーキド博士の孫を倒すこと。これを複数のステップに分解し、次のように考えました。 オーキド博士の研究室に行く オーキド博士との会話を完了する スターターポケモンを選択する 研究室を出ようとする オーキド博士の孫と対戦して倒す 最初の実験としては野心的すぎるかもしれないとすぐに気づいたので、オーキド博士の研究室のセーブ状態からスターターを選択することに絞り込みました。3匹のスターターのいずれかを取得できれば、成功とみなされます。 研究室のセーブされた位置から、モデルにAボタンを12回押させるだけで十分です。しかし、モデルはそれを学習できるでしょうか？モデルは単に目的なくさまようかもしれませんし、オーキド博士の研究室の中で永遠の苦痛に陥るかもしれません。あるいは、モデルは数回のAボタン操作の後にBボタンを押して、目標にほぼ到達したときに努力をキャンセルするかもしれません。 そもそも世界モデルとは何か？ 世界モデルとは、現在の状態または観測から始まり、何らかのアクションが発生してその状態を変更し、新しい観測を生み出すモデルです。現実的には、状態と、それによって生成される新しい状態との間に、何らかの相関、できれば因果関係があるはずです。世界モデルの目標は、この相関を学習することです。現在の観測がプレイヤーがモンスターボールの前にいるオーキド博士の研究室のスクリーンショットであり、アクションが左ボタンを押すことで、プレイヤーを左に1タイル移動させ、新しい最終スクリーンショットを生成するとしましょう。目標は、世界モデルが左ボタンを押すとプレイヤーが左に移動するという直感をある程度発達させることです。 $$ o_{t+1}\ ext{≈}F(o_t,a_t). $$ ここで、\(o_t\)は現在のスクリーンショット、\(a_t\)は押されたボタン、そして\(F\)は世界モデルが次のスクリーンショット\(o_{t+1}\)を予測するために学習する関数です。 しかし、目標は一般化することです。モデルは、オーキド博士の研究室で左ボタンを押すとプレイヤーが左に移動することを学習するだけでなく、どこでも左ボタンを押すとプレイヤーが左に移動することを学習する必要があります。 3 余談ですが、この設定は強化学習に似ているように聞こえるかもしれません。しかし、決定的に重要なのは、世界モデルは報酬なしで状態遷移を理解することを学習するということです（これについては後述します）。 世界モデルはどのように予測を学習するか スクリーンショットから埋め込みへ したがって、私たちが本当に興味を持っているのが環境の状態遷移とダイナミクスであり、私たちが持っているのが観測（この場合はゲームのスクリーンショットの形式）だけであるなら、モデルはスクリーンショットを予測することを学習するのでしょうか？ いいえ、モデルが実際に学習するのは、潜在空間、つまり埋め込み空間内での予測です。まず、エンコーダーが必要です。これは、スクリーンショットが与えられると、埋め込みを生成します。 埋め込みとは何か？ スクリーンショット エンコーダー→0.24−1.070.63⋮0.18192個の数値 埋め込みは、入力を数値のベクトルとして学習された表現です。エンコーダーは、画像、文、または音をそのようなベクトルに変換でき、別のモデルが比較、予測、または入力として使用できるようにします。 エントリは手動でラベル付けされた特徴ではありません。1つの座標が「色」または「位置」を意味する必要はありません。情報は多くのエントリに分散される可能性があり、元の入力のいくつかの詳細は完全に失われる可能性があります。埋め込みが有用かどうかは、エンコーダーがどの区別を保持することを学習したかに依存します。 現在の観測\(o_t\)を表す埋め込みを使用して、次に予測モデルを用意できます。これは、アクション\(a_t\)が与えられると、将来の埋め込みを予測しようとします。この将来の予測埋め込みを使用して、将来の観測\(o_{t+1}\)によって生成された実際の埋め込みと比較しようとすることができます。これは実際にはかなり単純です。新しいスクリーンショットを同じエンコーダーに渡すだけで、実際の将来の埋め込みが得られます。目標は次のようになります。 $$ \\\\hat z_{t+1}\ ext{≈}z_{t+1}. $$ ここで、\\\hat z_{t+1}\\(は予測された次の埋め込みであり、\(z_{t+1}\)は実際に続いたスクリーンショットの埋め込みです。 次のポケモンのスクリーンショットの埋め込みを予測する スクリーンショットを上から読み取ります。現在のスクリーンショットと次のスクリーンショットは同じエンコーダーを通過します。現在の埋め込みとAボタンが予測子に供給されます。平均二乗誤差は、その予測埋め込みと実際の次のスクリーンショットの埋め込みを比較します。色付きのブロックは、ベクトルを模式的に表しています。 プレスA 予測子 予測埋め込み MSE 現在の埋め込み ターゲット埋め込み 同じ重み エンコーダー エンコーダー 現在のスクリーンショット 次のスクリーンショット 1ステップについてこれを書き出すと、次のようになります。 $$ z_t = E(x_t), ext{ }\\\hat z_{t+1} = P(z_t,a_t), ext{ }z_{t+1} = E(x_{t+1}). $$ ここで、\(x_t\)は現在のスクリーンショット、\(a_t\)はアクション、そして\(E\)と\(P\)はエンコーダーと予測子です。重要なのは、予測された埋め込み\\\hat z_{t+1}\\(はスクリーンショットに変換されることはなく、\(z_{t+1}\)と直接比較できることです。 予測損失と崩壊 私たちの予測とグラウンドトゥルースの埋め込みとの違いを学習するために、信頼できる損失関数である平均二乗誤差を使用します。 $$ \\\mathcal L_{\ ext{pred}}}=\rac{1}{D}\\\ ext{∑}_{d=1}^{D}(\\\hat z_{t+1,d}-z_{t+1,d})^2, ext{ }D=192. $$ これは少し複雑に見えますが、距離関数として想像できます。基本的に、2つの埋め込みベクトルの間の距離を測定しようとしています。 すべては十分に単純に見えます。この関数を最小化すれば、ポケモンをプレイできる世界モデルが得られるはずですよね？ しかし、それは決して単純ではなく、実際には危険な退化ケースが発生する可能性があります。エンコーダーと予測子の両方をトレーニングしていることを思い出してください。そして、エンコーダーは各トレーニング例で2回出現します。 エンコーダーは現在のスクリーンショットを\(z_t=E(x_t)\)に変換します。 予測子は\(z_t\)とボタン\(a_t\)を受け取り、\\\hat z_{t+1}=P(z_t,a_t)\\(を推測します。 同じエンコーダーが実際の次のスクリーンショットを\(z_{t+1}=E(x_{t+1})\)に変換します。 トレーニングは\\\hat z_{t+1}\\(と\(z_{t+1}\)の間の距離を縮小し、予測子とエンコーダーの両方を更新します。 最悪の場合、エンコーダーはすべてのスクリーンショットを次のように埋め込み始める可能性があります。