我想复制 LeWorldModel，这是一个小型的联合嵌入预测结构（JEPA），能够从像素和动作中学习世界动态。原始论文将其用于 Push-T 中的无奖励规划。但由于我热爱视频游戏，同时希望更深入地了解 LeCun 的 JEPA 结构，我决定从头开始编写整个架构，并在超级马里奥兄弟上进行训练。模型通过了我最初认为重要的每个测试。它能够推广到保留的剧集，使用动作，并且在五步预测未来时比强基线表现更好。原始的无奖励规划可以使马里奥朝着附近的图像目标移动，并在目标的两到五个像素内完成。 :D 一度看起来这个模型学会了玩游戏。然后我将目标移动到更远的关卡... 马里奥无法可靠地跳过第一个主要障碍或朝着单个远离的目标图像导航。这个模型学会了预测游戏，但这并不意味着它学会了如何在游戏中取得进展。 D: 这篇文章既是技术性的 walkthrough，也是一个事后见之明，描述我构建的内容，我是如何测试的，我犯的错误，以及逐步揭示真实问题的实验。（许多教训在事后看来似乎很明显 T^T）整个架构在介绍每个方程之前，先看看整个机器是很有帮助的：让我们从绿色路径开始。每个训练样本包含四个马里奥画面。视觉编码器将每个画面压缩为一个称为隐变量（zz）的192维表示：zt=Eθ(xt),zt∈R192。你可以把隐变量看作模型对截图的私人描述。红色路径包含控制器输入。每对观察之间相隔五个模拟器帧，每帧包含六种可能的按钮状态：frames: [batch, 4, 3, 224, 224] actions: [batch, 4, 5, 6] # 左，右，上，下，A，B 动作编码器将每个5 × 6的按钮序列压缩为另一个192维向量。帧和动作的隐变量然后被输入到因果预测器中。它的任务是回答：根据先前框架的样子和按下的按钮，下一帧的隐变量应该是什么样的？预测器包含六个变换器块，其中每个帧将关注于先前的帧。但是，我们该如何注入动作呢？动作通过自适应层归一化零（AdaLN-Zero）进入这些变换器块。不是简单地将动作向量附加到帧向量，AdaLN-Zero 将每个动作转换为三种控制：Shift：为帧特征添加依赖于动作的偏移量 Scale：调整特定特征的大小 Gate：控制变换器更新当前状态的强度。这些如何影响变换器块？通常我们正常的注意力会给每个帧其先前上下文，然后通过馈送前馈网络（MLP）合成该信息，但 AdaLN 根据动作修改两个阶段。例如，跳跃动作可能会放大与垂直运动相关的隐变量特征，并缩小对预测跳跃影响较小的特征。Shift 将归一化特征移向依赖于动作的不同基线。最后，gate 决定注意力或 MLP 更新对预测状态的影响程度。这些控制是分别为注意力和 MLP 分支产生的，使得这个块总共得到了六个值：每个分支的 shift、scale 和 gate。“零”只是意味着它们的权重从零开始，因此预测器在没有随机动作效果的情况下开始，并在训练过程中逐步学习打开哪些门。现在，在训练期间，经过六个变换器块后，一个小的投影头生成三个预测的未来隐变量：z^1, z^2, z^3。它们与真实的三个下一帧产生的隐变量进行比较：Lpred=MSE⁡([z^1,z^2,z^3],[z_1,z_2,z_3])。现在我们只想将这个损失降低到 0... 但显然，模型有一个简单的作弊方式，我们可以将所有隐变量向量做成一样！预测将变得完美，因为马里奥，一个管道，一个新世界看起来都是相同的（表示坍缩）。所以为了防止它作弊，我们使用 SIGReg 1！它通过鼓励真实帧隐变量保持多样性和信息性来防止这种坍缩。因此我们新的损失函数变为：L=Lpred+0.1LSIGReg。所以这就是整个架构！现在进入实际结果。但它真的学会了吗？LeMario 在来自 32 个马里奥关卡的 280 个剧集中的 737,134 帧上进行了训练。为了验证模型是否真的学习了，我们不能仅仅关注更低的损失。更低的损失不足以证明它学会了动态，因为附近的帧往往看起来太相似，以至于预测“什么都不做”可能是有效的。