LeWAM: End-to-End World and Action Modeling with JEPAs
作者: Minghao Fu, Tavis Siebert, Eryk Halicki, Randall Balestriero
分类: cs.RO, cs.AI
发布日期: 2026-10-08
备注: Project website: https://le-wam.github.io/
💡 一句话要点
提出LeWAM以解决控制任务中的世界与动作建模问题
🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture)
关键词: 联合嵌入预测 控制策略 动作规划 机器人控制 高效学习
📋 核心要点
- 现有的JEPA方法在控制任务中通常需要分阶段训练,导致效率低下和性能限制。
- LeWAM通过联合训练潜在状态预测和动作流匹配,简化了模型训练流程,提升了控制能力。
- 实验结果显示,LeWAM在多个任务上显著提高了成功率和规划速度,表现优于传统方法。
📝 摘要(中文)
联合嵌入预测架构(JEPA)能够在不进行像素重构的情况下学习紧凑的潜在世界模型。现有JEPA管道通常在添加控制策略或求解器之前,先训练视觉表示和动作条件动态。本文提出LeWAM,一种端到端的JEPA控制模型,具有两个主要特点:首先,采用潜在状态预测和动作流匹配的联合训练目标;其次,对潜在状态施加各向同性高斯约束,使得单一模型能够支持直接策略执行和通过评估采样动作进行高效规划。与之前的JEPA在视觉规划中的表现相比,LeWAM在接触丰富任务上的成功率从28.6%提升至89.7%,在长时间任务上的成功率从10.9%提升至46.2%,并且规划速度提升了32.3倍。在仿真和物理机器人行为克隆任务中,LeWAM的表现也与现有方法相当。可解释性分析表明,LeWAM有效学习了与控制相关的表示和动作分布,从而有效指导测试时的规划和控制。
🔬 方法详解
问题定义:本文旨在解决现有JEPA方法在控制任务中效率低下和性能不足的问题,特别是在复杂环境下的动作规划和执行。现有方法通常需要分阶段训练,导致模型的整体性能受限。
核心思路:LeWAM的核心思路是通过联合训练潜在状态预测和动作流匹配,构建一个端到端的控制模型。这种设计使得模型能够在同一框架内进行直接策略执行和高效规划。
技术框架:LeWAM的整体架构包括两个主要模块:潜在状态预测模块和动作流匹配模块。潜在状态预测模块负责学习环境的潜在表示,而动作流匹配模块则用于优化动作选择和执行。
关键创新:LeWAM的主要创新在于引入了各向同性高斯约束,使得潜在状态的分布更加平滑,从而提高了模型的泛化能力和控制精度。这一创新与传统JEPA方法的分阶段训练方式形成鲜明对比。
关键设计:在模型设计中,LeWAM采用了特定的损失函数来平衡潜在状态预测和动作流匹配的训练目标,同时在网络结构上进行了优化,以支持高效的动作评估和规划。
🖼️ 关键图片
📊 实验亮点
LeWAM在接触丰富任务中的成功率从28.6%提升至89.7%,在长时间任务中的成功率从10.9%提升至46.2%。此外,LeWAM的规划速度提升了32.3倍,显示出其在效率和效果上的显著优势。
🎯 应用场景
LeWAM的研究成果在机器人控制、自动驾驶、智能制造等领域具有广泛的应用潜力。通过高效的世界和动作建模,LeWAM能够提升机器人在复杂环境中的自主决策能力,推动智能系统的进一步发展。
📄 摘要(原文)
Joint-embedding predictive architectures (JEPAs) learn compact latent world models without pixel reconstruction. Existing JEPA pipelines typically train visual representations and action-conditioned dynamics before adding a control policy or solver. We introduce LeWAM, an end-to-end JEPA for control with: (1) a joint training objective over next latent state prediction and action flow matching, and (2) an isotropic Gaussian constraint on the latent state, enabling a single model to support direct policy execution and efficient planning by evaluating sampled actions with its learned latent dynamics. Compared with prior JEPAs in visual planning, LeWAM improves success rate from 28.6% to 89.7% on contact-rich tasks and from 10.9% to 46.2% on long-horizon tasks, and plans 32.3 times faster. On behavior cloning tasks in simulation and physical robots, LeWAM is also competitive with established methods. Interpretability analyses show that LeWAM learns control-relevant representations and action distributions effectively guiding test-time planning and control.