LeMario：在超级马里奥兄弟上训练 JEPA 世界模型

← 返回项目 LeMario：超级马里奥兄弟在 JEPA 模型上进行训练 Benjamin Bai · 2026 年 7 月 · LeWorldModel 论文 · GitHub · X 我想重现 LeWorldModel，这是一个小型联合嵌入预测架构 (JEPA)，可以从像素和动作中学习世界动态。原始论文将其用于 Push-T 中的无奖励规划。但是，由于我喜欢视频游戏，同时想更深入地了解 LeCun 的 JEPA 架构，所以我决定从头开始编写整个架构，并在《超级马里奥兄弟》上对其进行训练。该模型通过了我最初认为重要的所有测试。它推广到了持续的事件，使用了行动，并比强大的基线更好地预测了五步未来。原始的无奖励计划可以让马里奥朝着附近的图像目标前进，并在目标的两到五个像素内完成。 :D 有一瞬间，模型看起来像是已经学会了玩。然后我将目标进一步移入关卡……马里奥无法可靠地跳过第一个主要障碍或导航到单个遥远的目标图像。该模型已经学会了预测比赛，但这并不意味着它已经学会了如何在比赛中取得进展。 D：这篇文章既是技术演练，也是事后剖析，我构建了什么，如何测试它，我犯过的错误，以及逐渐暴露真正问题的实验。 （事后看来，大多数教训似乎很明显 T^T ） 整个架构 在分别介绍每个方程之前，它有助于立即了解整个机器：让我们从绿色路径开始。每个训练样本包含四个 Mario 帧。视觉编码器将每一帧压缩为 192 个数字表示，称为潜在 ( z z z )： z t = E θ ( x t ) , z t ∈ R 192 z_t = E_\theta(x_t), \qquad z_t \in \mathbb{R}^{192} z t ​ = E θ ​ ( x t ​ ) , z t ​ε R 192 您可以将潜在特征视为模型对屏幕截图的私人描述。红色路径包含控制器输入。每对观察结果由五个模拟器帧分隔，每个帧包含六种可能的按钮状态： 帧：[batch, 4, 3, 224, 224] 动作：[batch, 4, 5, 6] # Left, Right, Up, Down, A, B 动作编码器将每个 5 × 6 按钮序列压缩为另一个 192 个数字的向量。然后，框架和动作潜伏被输送到因果预测器中。它的工作是回答：给定前一帧的样子以及按下了哪些按钮，下一帧的潜在帧应该是什么样子？预测器包含六个变压器块。每个帧都会关注前面的帧。但是，我们如何注入行动呢？这些动作通过 Adaptive LayerNorm Zero (AdaLN-Zero) 进入这些变压器块。 AdaLN-Zero 不是简单地将动作向量附加到帧向量，而是将每个动作转换为三种控件： Shift：向帧特征添加依赖于动作的偏移量 Scale：向上或向下转动特定特征。 Gate：控制变压器更新当前状态的强度。现在这些对变压器块有何影响？通常我们的正常注意力会给每个帧提供其先前的上下文，然后将其通过前馈 (MLP) 来合成该信息，但是 AdaLN 根据操作修改这两个阶段。例如，跳跃动作可能会放大与垂直运动相关的潜在特征，并缩小对于预测跳跃不太重要的特征。转变将标准化特征移向不同的依赖于动作的基线。最后，门决定注意力或 MLP 更新对预测状态的影响程度。这些控制是针对注意力和 MLP 分支分别生成的，总共为块提供了六个值：每个分支的移位、缩放和门。 “零”只是意味着它们的权重从零开始，因此预测器开始时没有随机动作效果，并逐渐学习在训练期间打开哪些门。现在，在六个 Transformer 块之后的训练过程中，一个小投影头会产生三个预测的未来潜伏： z ^ 1 , z ^ 2 , z ^ 3 \hat z_1,\hat z_2,\hat z_3 z ^ 1 ​ , z ^ 2 ​ , z ^ 3 ​ 这些与三个真实下一帧产生的潜伏进行比较： L pred = MSE ⁡ ( [ z ^ 1 , z ^ 2 , z ^ 3 ] , [ z 1 , z 2 , z 3 ] ) \mathcal L_{\text{pred}} = \operatorname{MSE} \left( [\hat z_1,\hat z_2,\hat z_3], [z_1,z_2,z_3] \right) L pred ​ = MSE ( [ z ^ 1​ , z ^ 2 ​ , z ^ 3 ​ ] , [ z 1 ​ , z 2 ​ , z 3 ​ ] ) 现在我们只想将这个损失降低到 0...但显然模型有一种简单的方法可以作弊，我们可以使所有潜在向量相同！预测将变得完美，因为马里奥、一根烟斗、一个新世界看起来都一样（表征崩溃）。所以为了防止它作弊，我们使用 SIGReg 1 ！它通过鼓励潜在的真实框架保持多样性和信息量来防止这种崩溃。因此我们的新损失函数变为： L = L pred + 0.1 L SIGReg \mathcal L = \mathcal L_{\text{pred}} + 0.1\mathcal L_{\text{SIGReg}} L = L pred ​ + 0.1 L SIGReg ​