LeMario: Super Mario BrosをJEPAワールドモデルで学習させる

プロジェクトに戻る LeMario: Super Mario BrosをJEPAモデルで学習させる Benjamin Bai · 2026年7月 · LeWorldModel論文 · GitHub LeWorldModel、つまりピクセルとアクションから世界のダイナミクスを学習する小規模なJoint-Embedding Predictive Architecture (JEPA) を再現したいと思いました。元の論文では、これを報酬なしのプランニングにPush-Tで使用していました。しかし、私はビデオゲームが好きで、同時にLeCunのJEPAアーキテクチャについてより深く学びたかったので、アーキテクチャ全体をゼロから書き、Super Mario Brosで学習させることにしました。 モデルは、当初重要だと考えていたすべてのテストをパスしました。未知のエピソードに汎化し、アクションを使用し、強力なベースラインよりも優れた5ステップ先の予測を行いました。生の報酬なしプランニングは、マリオを近くの画像目標に向かって動かし、ターゲットから2〜5ピクセル以内で完了することができました。:D 一瞬、モデルはプレイを学習したように見えました。しかし、目標をレベルのさらに奥に移動させると... マリオは最初の大きな障害物を飛び越えたり、単一の遠い目標画像に向かってナビゲートしたりすることが信頼性高くできませんでした。モデルはゲームを予測することを学習しましたが、それはゲームを進む方法を学習したことを意味しませんでした。D: この記事は、技術的な解説とポストアナルシスの両方です。私が構築したもの、テスト方法、犯した間違い、そして徐々に本当の問題を露呈させた実験について説明します。（ほとんどの教訓は、後から考えると明白なものばかりです T^T ） アーキテクチャ全体 各方程式を個別に紹介する前に、機械全体を一度に見ると役立ちます。 緑色のパスから始めましょう。各トレーニングサンプルには4つのマリオのフレームが含まれています。ビジョンエンコーダーは、各フレームを「潜在変数（z）」と呼ばれる192個の数値の表現に圧縮します。 zt=Eθ(xt),zt∈R192 潜在変数とは、モデルがスクリーンショットを独自に記述したものだと考えることができます。 赤いパスはコントローラー入力を含んでいます。各観測ペアは5つのエミュレーターフレームで区切られており、各フレームには6つの可能なボタン状態があります。 フレーム: [バッチ, 4, 3, 224, 224] アクション: [バッチ, 4, 5, 6] # 左, 右, 上, 下, A, B アクションエンコーダーは、各5x6のボタンシーケンスを別の192個の数値ベクトルに圧縮します。次に、フレームとアクションの潜在変数が因果予測器に渡されます。その仕事は、次の質問に答えることです。 前のフレームがどのように見え、どのボタンが押されたかを考えると、次のフレームの潜在変数はどのようになるべきか？ 予測器には6つのトランスフォーマーブロックが含まれています。各フレームは前のフレームにアテンションを向けます。しかし、アクションをどのように注入するのでしょうか？ アクションは、Adaptive LayerNorm Zero (AdaLN-Zero) を介してこれらのトランスフォーマーブロックに入力されます。単にアクションベクトルをフレームベクトルにアタッチするのではなく、AdaLN-Zeroは各アクションを3種類の制御に変換します。 シフト: アクションに依存したオフセットをフレーム特徴量に加算します。 スケール: 特定の特徴量を増減させます。 ゲート: トランスフォーマーが現在の状態をどれだけ強く更新するかを制御します。 これらがトランスフォーマーブロックにどのように影響するのでしょうか？通常、通常の注意機構は、各フレームに前のコンテキストを与え、それをフィードフォワード（MLP）を通して合成しますが、AdaLNはアクションに応じて両方のステージを変更します。 例えば、ジャンプアクションは、垂直方向の動きに関連する潜在特徴量をスケールアップし、ジャンプの予測にあまり重要でない特徴量をスケールダウンする可能性があります。シフトは、正規化された特徴量を、アクションに依存した異なるベースラインに移動させます。最後に、ゲートはアテンションまたはMLPの更新が予測状態にどれだけ強く影響するかを決定します。 これらの制御は、アテンションブランチとMLPブランチそれぞれに個別に生成され、ブロックには合計6つの値（各ブランチにシフト、スケール、ゲート）が生成されます。「Zero」は、その重みがゼロから始まることを意味するため、予測器はランダムなアクション効果なしで開始し、トレーニング中にどのゲートを開くかを徐々に学習します。 トレーニング中、6つのトランスフォーマーブロックの後、小さなプロジェクションヘッドが3つの予測された未来の潜在変数を生成します。 z^1, z^2, z^3 これらは、3つの実際の次のフレームによって生成された潜在変数と比較されます。 Lpred = MSE([z^1, z^2, z^3], [z1, z2, z3]) この損失を0に減らしたい... しかし、モデルが簡単に不正行為をする方法が1つあります。すべての潜在変数を同じにすることができます！マリオ、パイプ、新しい世界がすべて同じように見えるため、予測は完璧になります（表現の崩壊）。 そのため、不正行為を防ぐためにSIGReg1を使用します！これは、実際のフレームの潜在変数が多様で情報豊富であり続けることを奨励することで、この崩壊を防ぎます。したがって、新しい損失関数は次のようになります。 L = Lpred + 0.1LSIGReg これがアーキテクチャ全体です！それでは、実際の結果に進みましょう。 実際に学習したのか？ LeMarioは、32のマリオレベルにわたる280エピソードから737,134フレームで学習しました。モデルが学習したことを確認するために、損失が低いかどうかだけを見ることはできませんでした。損失が低いだけでは、ダイナミクスを学習したことを証明するには十分ではありませんでした。近くのフレームはしばしば非常に似ているため、「何も変化しない」と予測することは強力なベースラインとなります。 未知のエピソードで、LeMarioをそのパーシステンスベースラインや、シャッフルされたアクションとペアになった実際のフレーム履歴と比較しました。 手法 | 1ステップ誤差 | 5ステップ誤差 ---|---|--- LeMario | 0.013773 | 0.077717 変化なしと予測 | 0.014472 | 0.142473 アクションをシャッフル | 0.016555 | 0.114648 アクションのシャッフルは、1ステップ誤差を20.2%増加させました。5回の再帰ステップでは、LeMarioはパーシステンスを45.5%上回りましたが、シャッフルされたアクションは47.5%悪化しました。遠くまで予測するほど、ボタンが重要になりました。 LeMarioは、プレイヤーのアクションに条件付けられた短期的なマリオのダイナミクスを学習していました！ コントローラーに触れさせてみる さて、楽しい部分です。モデルが未来を想像できるようになったら、それらの未来を検索し、マリオに何をすべきかを選択させることができます。 想像を通して検索 次に、アクションが与えられたときに1フレーム先を予測するだけのモデルを、複数のアクションと未来のステップを予測できるものに変えるために、クロスエントロピー法を使用しました！ 現在の画像xtと目標画像xgが与えられると、エンコーダーはztとzgを生成します。CEMは次のことを行います。 数百のアクションシーケンスをサンプリングします。 各シーケンスをLeMarioを通してロールアウトします。 予測された最終潜在変数をzgと比較してスコアリングします。 最良の候補を保持します。 それらの周りで再サンプリングし、繰り返します。 CEMは、ランダムな候補よりもはるかに低い予測目標距離を持つアクションシーケンスを見つけました。私は、想像できるモデル、その想像を検索できるオプティマイザー、そして報酬エンジニアリングを必要としない目標画像を持っていました。それは素晴らしかったです。 しかし、マリオはほとんど動きませんでした 私は小さな目標から始めました。マリオはx=40から開始し、目標フレームは彼がx=72にいることを示していました。生のJEPA+CEMはx=44で終了しました。 平易な言葉で言うと、それはひどかったです。 この時点で、どの部分が失敗したのかわかりませんでした。予測器が間違っていたのかもしれない、CEMが壊れていたのかもしれない、あるいはエンコーダーがマリオを完全に無視していたのかもしれません。最も単純な質問から始める必要がありました。その192個の数値にはマリオの位置が含まれていたのでしょうか？ 192個の数値の中に何が入っているのか？ 潜在変数はわずか192個の数値です。エンコーダーを変更せずに、どのような情報が含まれているかを尋ねる方法が必要でした。 マリオを忘れたのか？ JEPAをフリーズし、潜在変数からマリオのエミュレーター座標を回復するための小さなプローブ2をトレーニングしました。エンコーダーは変更できなかったため、プローブが回復した位置は、JEPAがすでに学習した情報でなければなりませんでした。 水平位置: MAE = 9.30 px, R² = 0.997 垂直位置: MAE = 21.62 px, R² = 0.188 プローブは機能しました！マリオの水平位置はほぼ完璧に回復可能でした。垂直状態ははるかに弱かったですが、エンコーダーはプレイヤーに関する有用な情報を明らかに学習していました。 すべてを「修正」したプローブ プローブによって予測された水平位置を使用して、CEMの想像した未来を一時的にスコアリングしました。LeMarioはまだ未来を想像し、CEMはまだアクションを選択しました。プローブは、未来のランク付け方法を変更しただけでした。