在超级马里奥游戏中训练JEPA世界模型

背景

近日，一篇名为《LeMario: Training a JEPA World Model on Super Mario Bros》的技术文章在Hacker News上引发关注。作者尝试复现并扩展一个基于联合嵌入预测架构（Joint-Embedding Predictive Architecture，JEPA）的世界模型，将其应用于经典游戏《超级马里奥兄弟》。该工作源于对LeCun提出的JEPA架构的深入学习兴趣，以及将视频游戏作为模型验证平台的热情。作者从头编写了整个架构，并使用马里奥游戏的帧序列和按钮动作进行训练，最终揭示了模型“能预测却不会玩”的深层问题。

核心内容

整体架构

LeMario 的架构由三个主要组件构成：视觉编码器、动作编码器和因果预测器。

视觉编码器（绿色路径）：每个训练样本包含四帧马里奥游戏画面。视觉编码器将每一帧压缩成一个192维的向量，称为“潜变量”（latent）。这个潜变量可以理解为模型对屏幕截图的一种私有描述。

动作编码器（红色路径）：控制器输入包含六个按钮（左、右、上、下、A、B）。每对观测帧之间间隔5个模拟器帧，因此动作数据形状为 [batch, 4, 5, 6]。动作编码器将每个5×6的按钮序列压缩成另一个192维向量。

因果预测器：帧潜变量和动作潜变量被送入因果预测器。预测器的任务是：给定之前的帧画面和按下的按钮，预测下一帧的潜变量应该是什么。预测器包含六个Transformer块，每个帧会关注之前的帧。动作信息通过自适应层归一化零初始化（Adaptive LayerNorm Zero，AdaLN-Zero）注入Transformer块。

AdaLN-Zero将每个动作向量转化为三种控制信号：

• 偏移（Shift）：为帧特征添加与动作相关的偏移量。
• 缩放（Scale）：放大或缩小特定特征。
• 门控（Gate）：控制Transformer更新当前状态的强度。

这些控制信号分别作用于注意力分支和MLP分支，每个分支各有shift、scale、gate三个值，共六个值。“Zero”表示权重初始化为零，使预测器从无随机动作影响开始，逐步学习训练中应打开哪些门控。

经过六个Transformer块后，一个小型投影头产生三个预测的未来潜变量。这些预测潜变量与真实下一帧的潜变量进行比较，计算损失。为防止模型通过使所有潜变量完全相同来“作弊”（即表示崩溃），作者使用了SIGReg1正则化，鼓励真实帧的潜变量保持多样性和信息量。

训练与测试

LeMario在32个马里奥关卡的280个片段、共737,134帧上训练。为了验证模型是否真正学到了动力系统，作者不仅看损失下降，还对比了“持久性基线”（假设画面不变）和打乱动作后的表现。

• 打乱动作使一步预测误差增加20.2%。
• 在五步递归预测中，LeMario比持久性基线好45.5%，打乱动作则差47.5%。
• 预测越远，按钮动作越重要。

结论：LeMario学会了基于玩家动作的短时程马里奥动力学。

实际操控测试

为了让模型能规划未来，作者使用交叉熵方法（Cross-Entropy Method，CEM）进行搜索。给定当前图像和目标图像，编码器生成潜变量后，CEM进行以下迭代：

• 采样数百个动作序列。
• 通过LeMario将每个序列向前滚动。
• 根据预测的最终潜变量与目标潜变量的接近程度评分。
• 保留最佳候选，围绕它们重新采样并重复。

CEM找到了远低于随机候选的预测目标距离。但实际测试时，当目标从x=40移动到x=72，模型只能将马里奥移动到x=44——几乎没动。

诊断：潜变量中有什么？

作者冻结JEPA，训练一个小的探针（probe）从潜变量中恢复马里奥在模拟器中的坐标。结果：

• 水平位置：MAE = 9.30像素，R² = 0.997（几乎完美恢复）
• 垂直位置：MAE = 21.62像素，R² = 0.188（很弱）

这说明编码器确实学到了关于玩家的有用信息，但垂直状态信息较弱。

用探针“修复”规划

作者临时改用探针预测的水平位置来评分CEM的想象未来。对于目标x=72，探针评分的CEM将马里奥从x=40移动到x=71。通过局部重规划，最终达到x=176（目标x=177）。这是第一次成功执行！证实了JEPA模型能够想象有用的水平运动，且探针能够找到它。

更远的目标失败

当目标设在半个关卡之外时，模型无法可靠地跳过第一个主要障碍或导航到一个单一的远处目标图像。模型学会了预测游戏，但并没有学会如何通过游戏取得进展。

关键要点

• LeMario 使用 JEPA 架构从像素和动作中学习世界动力学，预测未来帧的潜变量，而非直接预测像素。
• 训练中采用 SIGReg1 正则化防止表示崩溃，确保潜变量具有多样性。
• 在短时程预测（五步以内）上，LeMario 显著优于持久性基线，且动作信息对预测至关重要。
• 直接使用潜变量目标进行规划（CEM）在短距离目标上表现很差，仅移动了4像素。
• 通过探针从潜变量中提取马里奥的水平位置，发现编码器很好地保留了水平位置信息，但垂直位置信息很弱。
• 用探针替代潜变量评分后，CEM 成功实现了短距离导航，但远距离导航仍失败。
• 核心教训：模型学会了预测局部画面的变化，但缺乏对关卡全局进展的理解，无法规划复杂动作序列（如跳跃障碍物）。

意义与影响

LeMario 项目是一次对 JEPA 架构的深入实践，揭示了“预测”与“智能规划”之间的关键差距。虽然模型能准确预测短时程的帧变化，并在简单目标上实现导航，但面对需要多步策略、克服障碍的复杂任务时，它表现出了本质上的局限性。

这一结果对世界模型研究具有警示意义：仅仅学会预测下一个帧是不够的，模型需要具备更高级的抽象表示来理解长期目标、障碍物的几何结构以及动作的因果链。作者在文中坦言，这些教训在事后看来显而易见，但亲自构建、测试并失败的过程，恰恰是理解 JEPA 能力边界的最佳方式。

此外，探针实验表明，即使编码器保留了丰富的位置信息，规划器（CEM）在潜空间中的搜索方式可能无法有效利用这些信息。未来工作可以考虑改进规划算法（如使用更强大的搜索策略或分层规划）或增强编码器对垂直运动等关键信息的编码能力。

LeMario 虽然不是一篇正式的论文，但其技术细节和诚实的失败记录为后续研究提供了宝贵参考：如何让世界模型真正“学会玩”，而不仅仅是“学会预测”。