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

# LeMario: 在JEPA模型上训练的超级马里奥兄弟 来源: https://www.benjamin-bai.com/projects/lemario 我原本想复现LeWorldModel（https://arxiv.org/abs/2603.19312），这是一种小型的联合嵌入预测架构（JEPA），它能从像素和动作中学习世界动态。原始论文将其用于Push-T任务的无奖励规划。但因为我热爱电子游戏，又同时想更深入了解LeCun的JEPA架构，我决定从头编写整个架构，并在《超级马里奥兄弟》上训练它。 该模型通过了所有我最初认为重要的测试。它能泛化到未见过的片段，利用动作信息，并且比强基线更好地预测五步后的未来。原始的无奖励规划能将马里奥移向附近的图像目标，并最终停在目标的两到五个像素内。:D 有那么一刻，看起来模型学会了玩游戏。但当我将目标移到关卡更远处时……马里奥无法可靠地跳过第一个主要障碍，也无法朝着单个遥远的目标图像前进。 模型学会了预测游戏，但这并不意味着它学会了如何在游戏中取得进展。T^T 这篇博文既是技术讲解也是事后分析，包括我构建了什么、如何测试、犯下的错误，以及逐步揭示真正问题的实验。（大多数教训事后看来都很明显 T^T） ### 整体架构 在逐一介绍每个公式之前，先整体看看整个机器： LeMario JEPA架构图，包括视觉编码器、动作编码器、因果预测器和SIGReg目标 让我们从绿色路径开始。每个训练样本包含四帧马里奥画面。**视觉编码器**将每一帧压缩成一个192维的表示，称为**潜在向量**（zz）： zt = Eθ(xt)，zt ∈ R^192 你可以把潜在向量看作模型对一幅截图的私有描述。 红色路径包含控制器输入。每对观测之间间隔五帧模拟器帧，每一帧包含六个可能的按键状态： ``` frames: [batch, 4, 3, 224, 224] actions: [batch, 4, 5, 6] # 左、右、上、下、A、B ``` **动作编码器**将每个 `5 × 6` 的按键序列压缩成另一个192维向量。 帧和动作的潜在向量随后被送入**因果预测器**。它的任务是回答： > 给定之前的画面内容和按下的按键，下一帧的潜在向量应该是什么？ 预测器包含六个Transformer模块。其中每一帧都会关注之前的帧。但动作如何注入呢？ 动作通过**自适应层归一化零初始化**（AdaLN-Zero）进入这些Transformer模块。 图示展示AdaLN-Zero如何将动作相关的偏移、缩放和门控注入Transformer的注意力和MLP分支 AdaLN-Zero不是简单地将动作向量附加到帧向量上，而是将每个动作转化为三种控制： - **偏移**：为帧特征添加一个与动作相关的偏移量 - **缩放**：调高或调低特定特征。 - **门控**：控制Transformer更新当前状态的强度。 这些如何影响Transformer模块？通常我们的普通注意力会给每一帧提供之前的上下文，然后通过前馈网络（MLP）综合这些信息，但AdaLN根据动作修改这两个阶段。 例如，一个跳跃动作可能会**缩放**与垂直运动相关的潜在特征，并降低对预测跳跃不太重要的特征。**偏移**将归一化特征移动到不同的、与动作相关的基线。最后，**门控**决定注意力或MLP更新应对预测状态产生多大影响。 这些控制分别针对注意力分支和MLP分支产生，因此整个模块共有六个值：每个分支一个偏移、一个缩放和一个门控。 “零初始化”只是意味着它们的权重从零开始，因此预测器最初不会产生随机的动作效果，并在训练过程中逐渐学习打开哪些门控。 在训练期间，经过六个Transformer模块后，一个小型投影头产生三个预测的未来潜在向量： z^1，z^2，z^3 这些与三个真实下一帧产生的潜在向量进行比较： Lpred = MSE ( [z^1，z^2，z^3]，[z1，z2，z3] ) 现在我们只想让这个损失降到0……但显然模型有一个简单的作弊方法：它只需让所有潜在向量相同即可！预测会变得完美，因为马里奥、水管、新世界看起来都一模一样（表示坍缩）。 因此，为了防止作弊，我们使用**SIGReg**¹（https://www.benjamin-bai.com/projects/lemario#sigreg）！它通过鼓励真实帧的潜在向量保持多样性和信息性来防止坍缩。所以我们的新损失函数变成： L = Lpred + 0.1 LSIGReg 这就是整个架构！现在来看实际结果。 ### 但它真的学到了吗？ LeMario在来自32个马里奥关卡的280个片段共737,134帧上进行了训练。为了验证模型是否学到东西，我们不能只看更低的损失。更低的损失并不足以证明它学到了动态，因为邻近帧通常看起来非常相似，预测“什么都没变”是一个很强的基线。 在未见过的片段上，我将LeMario与持久性基线以及真实帧历史结合打乱动作的情况进行了比较： | 方法 | 一步误差 | 五步误差 | |------|----------|----------| | LeMario | **0.013773** | **0.077717** | | 预测无变化 | 0.014472 | 0.142473 | | 打乱动作 | 0.016555 | 0.114648 | 打乱动作使一步误差提高了20.2%。在五步递归预测中，LeMario比持久性基线好45.5%，而打乱动作则差47.5%。预测得越远，按键的作用越明显。 > LeMario已经学会了基于玩家动作的短视界马里奥动态！ ### 让它触碰控制器 现在进入有趣的部分。一旦模型能够想象未来，我们就可以搜索这些未来，让它选择马里奥应该做什么。 ### 在想象中搜索 为了把只能根据给定动作预测下一帧的模型，变成能预测多个动作和未来多步的模型，我使用了交叉熵方法（CEM）！ 给定当前图像xt和目标图像xg，编码器产生zt和zg。CEM随后： 1. 采样数百个动作序列。 2. 通过LeMario向前推演每个序列。 3. 根据预测的最终潜在向量与zg的接近程度打分。 4. 保留最佳候选。 5. 在它们周围重新采样并重复。 图示交叉熵方法通过想象的LeMario未来进行规划 CEM找到了预测目标距离远低于随机候选的动作序列。我有了一个能想象的模型、一个能搜索其想象的优化器，以及一个无需奖励工程的目标图像。这太棒了。 ### 然后马里奥几乎没动 原始JEPA和CEM的动画展示，马里奥几乎未能向附近目标移动 我开始设置一个很小的目标。马里奥从`x=40`出发；目标画面显示他在`x=72`。原始JEPA+CEM最终停在`x=44`。 通俗地说，它表现得很差。 此时我不知道哪个部分出了问题。也许是预测器错了，也许是CEM坏了，或者是编码器完全忽略了马里奥。我必须以最简单的问题开始：那192个数字里是否包含了马里奥的位置？ ### 192个数字里面有什么？ 潜在向量只有192个数。我需要一种方法来询问其中包含什么信息，同时不改变编码器。 ### 它遗忘马里奥了吗？ 我冻结了JEPA，训练了一个小型的**探针**²（https://www.benjamin-bai.com/projects/lemario#the-position-probe），从潜在向量中恢复马里奥的模拟器坐标。因为编码器不能改变，探针恢复出的任何位置都必须是JEPA已经学到的信息。 ``` 水平位置：MAE = 9.30 px，R² = 0.997 垂直位置：MAE = 21.62 px，R² = 0.188 ``` 探针起作用了！马里奥的水平位置几乎完美可恢复。垂直状态较弱，但编码器显然学到了关于玩家的有用信息。 ### 那个“修复”一切的探针 我临时用探针预测的水平位置来评分CEM想象的未来。LeMario仍然想象未来，CEM仍然选择动作，探针只改变了未来排名的方式。 对于一个目标在`x=72`，探针评分的CEM将马里奥从`x=40`移动到`x=71`。通过局部重新规划，他后来达到了`x=176`，目标在`x=177`。 这是第一个有效的展开！JEPA模型可以想象有用的水平运动，而探针可以找到它。 这让我相信潜在规划在理论上是可行的。如果马里奥的位置已经在表示中，并且学到的动态能够移动它，也许第一个目标只是与起始帧太相似了。 ### 尝试半个关卡外的目标 规划器的任务是找到连接两个嵌入的动作。如果起始帧和附近目标已经具有相似的嵌入，那么几乎什么都不做看起来就像成功。于是我移除监督探针，再次使用原始潜在规划，但这次使用可达的目标图像，大约在Worlds 1-1、2-1和3-1的中途。 以关卡3-1为例：这一次马里奥移动得更远。不再是`x=40`到`x=44`，三次运行大约达到`x=290–307`。他仍然在第一个有意义的危险附近死亡，但原始潜在规划不再是什么都不做了。 CEM完全按照我的要求做了：寻找连接两个嵌入的动作。一个更远的目标显然在潜在空间中创造了足够的压力，让马里奥移动起来。 但马里奥距离目标还有1,442个世界像素，而编码器将他最终场景的潜在距离分配为仅`0.164`。CEM预测为`0.153`，所以预测器并没有疯狂地产生幻觉。编码器本身认为错误场景相当接近（正如我们从上面起始帧和目标帧看到的那样）。 马里奥的滚动摄像机解释了原因。两个遥远的位置尽管在游戏中是两个不同的地点，但看起来可能非常相似！所以模型实际上在认真工作，它相信自己找到了一条通往“目标”的路径，但实际上只是困在了一个看起来与目标嵌入相似的位置。 ### 好吧，让目标更小 于是我推测，如果我们将目标拆分成更小的检查点，那么帧之间的差异会更大，所以我将人类运行的记录拆分成多个中间图像目标。这有所帮助，原始潜在规划达到了`x=314`，这是项目中无探针取得的最大进展。 但更短的目标并没有真正解决问题。 马里奥在2像素内到达了第一个目标。对于第二个目标，他超调到了`x=283`，然后修正回来，停在`x=239`，与参考点相差5像素。 从视觉上看，马里奥到达了正确位置。但基准测试仍然将该检查点标记为失败，因为最终嵌入不够接近目标嵌入（可能由于HUD或其他小细节）。 这表明距离只是问题的一部分。更小的目标帮助马里奥走得更远，但潜在距离仍然是判断他是否到达的脆弱方式。 下一个目标要求马里奥跳跃，规划器再次失败。这与之前的探针结果一致：水平位置表示良好，而垂直位置则弱得多。 > 更小的目标有助于规划视界。但这并没有使潜在距离成为进展的可靠度量。 至此，这些失败看起来不再互不相关。它们指向三个主要问题。 ### 预测状态不是控制状态 编码器因能表示有助于预测未来图像的任何信息而获得奖励。摄像机位置、敌方状态、动画、计时器状态都可能有用。 控制器需要不同的东西，一个距离对应于可控进展的状态。 ### CEM搜索模型的弱点 CEM完全按照我的要求去做：找到LeMario认为能到达目标的行为。它无法识别模型何时出错。 这使得模型的弱点变得明显。它把视觉上相似的位置当作同一个地方，并且难以推理垂直运动。 ### 马里奥改变了Push-T背后的假设 Push-T使用了专家轨迹中的附近目标、固定摄像机、平滑运动，以及视觉相似性与进展一致的场景。它的模型在20,000个专家片段上训练了十个周期。 LeMario在280个片段上训练了一个周期，这些片段分布在32个关卡中。附近目标变成了半个关卡之外。固定摄像机变成了滚动摄像机。平滑运动变成了动量、跳跃、坑、敌人、动画和死亡。 我复制了架构，同时改变了几个使其规划规则起作用的条件，却没有意识到这些条件对方法的成功同样至关重要。 ### 我的总结 尽管LeMario没有完美地学会玩《超级马里奥兄弟》，但它仍然能够捕捉到短视界的、基于动作的动态。 回顾过去，很多错误事后看来都很明显。我过于强调复现架构本身，而忽略了数据集、环境、评估和底层假设同样重要。 我还学到，意外的结果并不标志着实验的结束。当马里奥几乎不动时，我必须将系统的每个组成部分（预测器、表示、规划器和评估指标）隔离出来，并单独测试。每一次失败和后续实验都加深了我对问题和模型的理解。 对于未来的研究项目，我想更早地验证核心假设，在信任任何指标之前建立简单的基线，并围绕我真正关心的行为设计评估，而不是仅仅优化训练损失。 LeMario没有成为我最初想象的那个马里奥智能体。但它让我很好地了解了世界模型研究领域。 ### 附录 ### 1 SIGReg SIGReg阻止编码器将所有帧映射到同一个潜在向量。它检查真实帧的潜在向量是否保持多样性和近似高斯分布。 ``` 192维JEPA潜在向量 → 投影到1,024个随机方向 → 在每个方向上的17个点评估 → 与标准高斯分布比较 → 取误差的平均值 ``` SIGReg在批次中对每个时间步独立应用。 ### 2 位置探针 探针对JEPA进行了微调。我冻结了编码器，收集了60条World 1-1轨迹，并将完整轨迹分割为训练集和验证集。这样得到了3,203个训练潜在向量和859个未见过的潜在向量。 MLP设计得故意很小： ``` 192维JEPA潜在向量 → Linear(192, 128) → GELU → Linear(128, 2) → 预测 (x, y) ```