LeMario: スーパーマリオブラザーズでJEPA世界モデルを訓練する — 技術解説と事後検証

Benjamin Bai氏のLeMarioプロジェクトは、LeWorldModel（小型のJoint-Embedding Predictive Architecture、JEPA）をゼロから再現し、スーパーマリオブラザーズで学習させたものです。詳細な技術的解説では、完全なアーキテクチャ、学習設定、そして成功と失敗を率直に検証したポストモーテムが紹介されています。

アーキテクチャ概要

モデルは4連続のマリオフレームを処理します。ビジョンエンコーダは各フレームを192次元の潜在ベクトルに圧縮します: z_t = E_θ(x_t), z_t ∈ R^192。アクション（5エミュレータフレーム分の左、右、上、下、A、Bボタンの状態）は別途アクションエンコーダで192次元ベクトルに符号化されます。

フレームとアクションの潜在ベクトルは、6つのトランスフォーマーブロックを持つ因果予測器に入力されます。アクションはAdaptive LayerNorm Zero（AdaLN-Zero）を介して注入され、アテンションとMLPの各ブランチにシフト、スケール、ゲート制御を生成します。ゲートは更新が予測状態に与える強さを制御します。重みはゼロから始まり、予測器は徐々に学習します。

予測器は3つの未来の潜在ベクトルを出力します: ẑ₁, ẑ₂, ẑ₃。これらは実際の次の3フレームの潜在ベクトルとMSE損失で比較されます: L_pred = MSE([ẑ₁,ẑ₂,ẑ₃], [z₁,z₂,z₃])。表現の崩壊を防ぐため、SIGReg正則化が追加されます: L = L_pred + 0.1 * L_SIGReg。

成功した点

• モデルは未見のエピソードに汎化した。
• アクション情報を効果的に活用した。
• 強力なベースラインより優れた5ステップ先予測を達成。
• 報酬なしの生の計画で、近距離の画像目標に対してマリオを2～5ピクセル以内に移動できた。

失敗した点

• 目標がステージの遠くにある場合、マリオは最初の大きな障害を確実に越えられなかった。
• 単一の遠距離目標画像に向かってナビゲートできなかった。
• 結論: ゲームを予測することは、進歩を学習することとは異なる。モデルはピクセルレベルのダイナミクスを学習したが、長期的な計画や目標指向の行動は獲得できなかった。

重要な教訓

Bai氏は、ほとんどの教訓は後から見れば明白だと指摘しています。短期的予測と長期的計画のギャップは重要であり、標準的な予測損失では捉えられません。記事では、この限界を徐々に明らかにした実験について詳述しています。