No Gaussian Required: Contrastive Inverse Dynamics for JEPA World Models
作者: Jack Boylan, Chris Hokamp
分类: cs.LG, cs.AI
发布日期: 2026-08-18
备注: 17 pages, 5 figures. Code: https://github.com/jackboyla/action-contrastive-jepa
🔗 代码/项目: GITHUB
💡 一句话要点
提出对比逆动态方法以解决JEPA世界模型中的崩溃问题
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 世界模型 对比学习 逆动态 潜在表示 机器人控制 深度学习
📋 核心要点
- 现有的JEPA方法在学习世界模型时容易出现编码器崩溃的问题，导致模型性能下降。
- 本文提出的AC-MTM方法通过利用过渡数据中的动作信息，提供了一种新的反崩溃信号，避免了对目标网络和重建目标的依赖。
- 在多个标准任务中，AC-MTM的表现稳定且优于现有方法，特别是在复杂任务中取得了显著的成功率提升。
📝 摘要（中文）
联合嵌入预测架构（JEPA）通过预测未来嵌入来学习世界模型，但其目标存在常量编码器的平凡解，因此实际系统通常需要添加反崩溃机制。LeWM通过SIGReg正则化器强制潜在分布匹配各向同性高斯分布，从而稳定表示。本文提出的AC-MTM方法利用过渡数据本身提供反崩溃信号，保留LeWM的前向潜在预测目标，并增加一个训练专用的逆动态头，通过Action-NCE训练。实验表明，AC-MTM在多个标准任务上表现出色，尤其在多对象OGBench视觉场景任务中，成功率显著高于SIGReg，展示了其有效性和潜力。
🔬 方法详解
问题定义：本文旨在解决JEPA世界模型中编码器崩溃的问题，现有方法如SIGReg通过强制匹配高斯分布来防止崩溃，但这种方法存在局限性。
核心思路：AC-MTM方法的核心在于利用过渡数据中的动作信息作为反崩溃信号，通过训练逆动态头来增强模型的学习能力，而不依赖于固定的目标网络或重建目标。
技术框架：AC-MTM的整体架构包括前向潜在预测目标和逆动态头。前者与LeWM相同，而逆动态头则通过Action-NCE进行训练，识别产生每个潜在过渡的动作。
关键创新：AC-MTM的主要创新在于引入对比逆动态方法，提供了一种无分布假设的反崩溃信号，显著提高了模型的稳定性和性能。
关键设计：在设计中，AC-MTM采用了训练专用的逆动态头，使用Action-NCE损失函数，确保每个潜在过渡能够正确识别其对应的动作，避免了崩溃编码器的影响。实验中，模型在多个任务上均表现出色。
🖼️ 关键图片
📊 实验亮点
在四个标准像素控制任务中，AC-MTM的训练稳定性与SIGReg相当，而在更复杂的多对象OGBench视觉场景任务中，AC-MTM的成功率达到了80.0±2.0%，相比SIGReg的58.0±2.0%提升了20-24个百分点，展示了显著的性能优势。
🎯 应用场景
该研究的潜在应用领域包括机器人控制、自动驾驶和智能代理等，能够有效提升模型在复杂环境中的学习能力和决策性能。未来，该方法可能推动更高效的世界模型学习和更智能的自主系统发展。
📄 摘要（原文）
Joint-Embedding Predictive Architectures (JEPAs) learn world models by predicting future embeddings, but the objective admits a trivial solution of a constant encoder, so every practical system adds an anti-collapse mechanism (LeCun, 2022; Assran et al., 2023; Bardes et al., 2022; 2024). LeWorldModel (LeWM) prevents collapse with SIGReg, a regularizer that forces the latent distribution to match an isotropic Gaussian: the representation is stabilized by prescribing what it must look like, independently of the environment it models. We argue that the anti-collapse pressure can instead come from the transition data itself. Action-Contrastive Masked Transition Modeling (AC-MTM) keeps LeWM's forward latent-prediction objective and adds a training-only inverse-dynamics head trained with Action-NCE: each latent transition must identify the action that produced it among the other actions in the batch, a discrimination task that a collapsed encoder provably fails. The inverse branch is discarded after training, leaving test-time encoding, forward prediction, planning, and compute identical to LeWM. On four standard pixel-control tasks under a matched planning protocol, AC-MTM trains stably from scratch and matches SIGReg on average. On the harder multi-object OGBench Visual Scene task, results are consistent with the prescribed geometry becoming a bottleneck: AC-MTM reaches 80.0$\pm$2.0% success versus 58.0$\pm$2.0% for SIGReg, improving by 20-24 points in each training seed. A single 50-episode random-policy run gives a 52% baseline estimate. Contrastive inverse dynamics thus provides a distribution-free anti-collapse signal that requires no target network, stop-gradient, pretrained encoder, or reconstruction objective, and we characterize the action-space and observability assumptions under which it holds. We make our code available at https://github.com/jackboyla/action-contrastive-jepa