SWIM: Learning What Matters from Behaviors for JEPA World Models / SWIM：因行知重，基于逆动力学行为加权的 JEPA 具身世界模型

Paper: 2610.36662
Authors: Jiaming Hu, Shi Bo, Yan Zheng, Tian Wang, Florian Dubost
Categories: cs.RO, cs.AI, cs.LG

The Gap

Joint-Embedding Predictive Architectures (JEPAs, pioneered by Yann LeCun) have emerged as the leading conceptual alternative to generative pixel-level world models (like Sora or Dreamer). By predicting abstract feature embeddings $\mathbf{s}_{t+1} = f(\mathbf{s}_t, \mathbf{a}_t)$ rather than rendering full $H \times W \times C$ video frames, JEPAs avoid wasting capacity on unpredictable pixel noise (e.g., flickering background leaves, rippling water, lighting shifts).

However, applying JEPAs to real-world embodied robotics and physical manipulation reveals a core objective alignment failure:

The Visual Salience Bias:
- Self-supervised visual objectives (like masked latent prediction or contrastive losses) align with information-theoretic visual variance, not task utility.
- A large, colorful robotic arm or dynamic background shadows generate massive latent loss signals, whereas a millimeter-scale displacement of a small screw or key contact point generates negligible gradient signal.
The Curse of Uniform State Supervision:
- To counteract this, researchers often supervise JEPAs using simulator ground-truth state vectors $\mathbf{x} \in \mathbb{R}^D$ .
- However, standard practice applies uniform Euclidean weighting across all normalized state coordinates: $d(\mathbf{x}, \mathbf{x}') = \sum_{i=1}^D (x_i - x_i')^2$
- Redundant coordinates (e.g., multi-sensor joint readouts) overcount the exact same degrees of freedom, while completely irrelevant variables (e.g., tabletop texture coordinates) drown out the critical task dimensions (e.g., gripper-object latching).
The Manual Feature Engineering Trap:
- Hand-crafting task-specific state feature weights for every single manipulation task defeats the purpose of autonomous generalist world models.

   THE TASK-ALIGNMENT FAILURE IN STANDARD JEPA WORLD MODELS

   Observation (Visual Space):
   ┌──────────────────────────────────────────────┐
   │ [Huge Moving Robot Base] -> Dominates Latent Loss (High Salience, Zero Task Value)
   │ [Dynamic Shadow Shifts] -> Dominates Latent Loss (High Salience, Zero Task Value)
   │ [0.5mm Keyhole Alignment]-> Negligible Latent Loss (Low Salience, 100% Task Critical!)
   └──────────────────────────────────────────────┘
                          │
                          ▼
   Standard JEPA Latent Space:
   - Plans trajectories that minimize visual arm movement rather than inserting the key!

The Increment

One sentence: SWIM (State Weighting with Inverse-Dynamics Model) automatically infers the physical task importance of each state dimension by training an inverse-dynamics model to predict executed actions from observed state displacements, using these behavioral weights to align JEPA latent representations without any planning-time overhead.

Core Mechanism

The guiding insight of SWIM is rooted in classical control theory: a physical state feature matters if and only if knowing its displacement is necessary to explain or choose the agent’s actions. If changing coordinate $x_j$ requires a totally different action, $x_j$ is critical; if varying $x_k$ leaves the required action unchanged, $x_k$ is task-irrelevant noise.

   SWIM ARCHITECTURE & TRAINING PIPELINE

   [Offline Trajectory Dataset] D = {(s_t, a_t, s_(t+1))}
                    │
                    ▼
   [Stage 1: Inverse-Dynamics Action Prediction]
   Displacement: Δx = x_(t+1) - x_t
   Weighted State Features: w ⊙ Δx
                    │
                    ▼
   Predict Action: a^_t = MLP(w ⊙ Δx)
   Loss: L_ID = || a^_t - a_t ||^2  ==>  Backprop optimizes Feature Weights: w* ∈ R^D
                    │
                    ▼ Learned Optimal Weights: w* (Freeze weights!)
   [Stage 2: Metric-Aligned JEPA Latent Representation Training]
   Ground-Truth Weighted State Metric:
        d_SWIM(x, x') = || w* ⊙ (x - x') ||_2
                    │
                    ▼ Contrastive / Alignment Objective
   Force JEPA Latent Distance d_Z(z, z') ∝ d_SWIM(x, x')
                    │
                    ▼
   [Inference / Online Planning Phase]
   - Inverse-dynamics model and simulator state are COMPLETELY DISCARDED!
   - Pure Vision Input: Robot plans action rollouts in the calibrated JEPA latent space.

Inverse-Dynamics Feature Weighting:
- An inverse-dynamics model $g_\theta$ predicts action $\mathbf{a}_t$ from state $\mathbf{x}_t$ and future displacement $\Delta \mathbf{x} = \mathbf{x}_{t+k} - \mathbf{x}_t$ : $\hat{\mathbf{a}}_t = g_\theta(\mathbf{w} \odot \Delta \mathbf{x})$
- Here, $\mathbf{w} \in \mathbb{R}^D_{\ge 0}$ is a learnable importance vector. Coordinates that are causally coupled with control actions receive high weights $w_i$ , while passive or redundant dimensions are naturally regularized to near-zero.
Behavior-Calibrated Metric Alignment:
- The learned weights define an anisotropic metric space $\mathcal{M}_{\mathbf{w}^*}$ over states: $d_{\mathbf{w}^*}(\mathbf{x}_A, \mathbf{x}_B) = \sqrt{\sum_{i=1}^D w_i^* (x_{A,i} - x_{B,i})^2}$
- During JEPA representation training, the latent distance between visual encodings $z = E(I)$ is explicitly supervised to preserve this behavior-weighted distance: $\mathcal{L}_{\text{metric}} = \left| \|z_A - z_B\|_2 - d_{\mathbf{w}^*}(\mathbf{x}_A, \mathbf{x}_B) \right|^2$
Zero Planning-Time Overhead:
- The inverse-dynamics model and state annotations are used solely during offline pretraining.
- At test time, the robot operates purely from image observations, planning trajectories using Model Predictive Path Integral (MPPI) or CEM in the aligned JEPA latent space with zero extra FLOPs.

The Evidence

The authors benchmarked SWIM against standard visual JEPAs and uniform state-weighted world models across five challenging robotic manipulation environments in RoboSuite and MetaWorld (Peg Insertion, Door Opening, Drawer Closing, Nut Assembly, Block Stacking):

Manipulation Task	Unsupervised JEPA	Uniform State JEPA	Hand-Engineered State	SWIM (Ours)	Relative Gain over Uniform
Peg Insertion (Tight Clearance)	32.4%	58.1%	76.5%	84.2%	+26.1%
Door Opening (Continuous Latch)	48.6%	69.3%	78.4%	86.7%	+17.4%
Nut Assembly (Screw Fit)	28.5%	51.2%	68.0%	79.5%	+28.3%
Drawer Closing (Contact Dynamics)	62.1%	81.4%	87.2%	92.1%	+10.7%
Block Stacking (Multi-Object)	41.0%	64.8%	74.6%	81.3%	+16.5%
Overall Average Success Rate	42.5%	65.0%	76.9%	84.8%	+19.8% Average

Empirical highlights:

Consistent Superiority: SWIM achieves an 84.8% average planning success rate, beating uniform state weighting (65.0%) by nearly 20 percentage points across all tasks.
Surpassing Human Feature Engineering: In all five environments, SWIM outperformed manually tuned heuristic weights (84.8% vs 76.9%), discovering subtle non-linear kinematic couplings between gripper pitch and contact normal forces that human engineers overlooked.
Robustness to Distractor Objects: When irrelevant moving visual distractors (colored spheres bouncing on the table) were introduced, standard visual JEPA success collapsed from 42.5% to 11.2%, while SWIM maintained 82.9% success because distractor coordinates had zero weight in the inverse-dynamics action objective.

Key Takeaways

Behaviors Define Physical Relevance: Not all pixels or state variables are created equal. Inverse dynamics provides a mathematically principled, automated filter to separate causal control variables from passive noise.
Latent Metric Alignment Bridges Vision and Action: Supervising latent distances with behavior-weighted state metrics grounds JEPA world models in task reality without requiring pixel generative decoders.
Clean Decoupling of Training and Inference: Leveraging rich simulator state during training to calibrate vision models while keeping test-time deployment strictly image-based offers an ideal paradigm for scalable robot learning.

论文: 2610.36662
作者: Jiaming Hu, Shi Bo, Yan Zheng, Tian Wang, Florian Dubost
分类: cs.RO, cs.AI, cs.LG

核心痛点

联合嵌入预测架构（JEPA，由 Yann LeCun 等人倡导）被广泛视作替代 Sora、Dreamer 等生成式像素世界模型的最具前景的学术方向。通过在抽象潜在表征空间中预测下一时刻的状态 $\mathbf{s}_{t+1} = f(\mathbf{s}_t, \mathbf{a}_t)$ ，而非吃力不讨好地渲染整张高清像素视频帧，JEPA 完美避开了背景落叶、水波反光、复杂光影变化等无法预测的高频视觉噪声。

然而，将纯视觉 JEPA 应用于物理具身机器人（Embodied Robotics）精密操作时，暴露了致命的**“任务目标对齐失效”**：

视觉信息量不等于任务价值：
- 现有的自监督视觉表征目标（如掩码特征预测或对比学习）本质是在捕捉“视觉像素信息量方差”，而非“具身操作任务效用”。
- 巨大且颜色鲜艳的机械臂底座运动、随光线晃动的桌面投影会引发巨额的潜在表征损失，而一把钥匙旋转 1 毫米、插销细微对准的决定性接触点，其产生的梯度信号几乎微不可察。
均匀状态监督的平庸陷阱：
- 为缓解纯视觉偏见，研究人员常借助仿真器中记录的物理状态向量 $\mathbf{x} \in \mathbb{R}^D$ 辅助监督表征。
- 但业内普遍采用全维度均匀欧式加权： $d(\mathbf{x}, \mathbf{x}') = \sum_{i=1}^D (x_i - x_i')^2$
- 这种平均主义导致严重失真：多传感器重复记录的冗余坐标被严重过分放大，而无关的桌面纹理参数则掩盖了决定成败的夹爪-物体咬合临界状态。
人工调参的泛化困局：
- 为每个机器人任务手工编写特征加权规则不仅极其耗时繁琐，更违背了通用具身大模型自主进化的初衷。

   传统 JEPA 具身世界模型的“注意力错位”

   观测视角 (视觉像素空间):
   ┌──────────────────────────────────────────────┐
   │ [巨大机械臂躯干平移] -> 占据 80% 潜在损失 (视觉显著，但对抓取毫无意义)
   │ [房间背景光影斑驳]   -> 占据 15% 潜在损失 (视觉噪声，任务完全无关)
   │ [0.5毫米插销孔径对准] -> 仅占 1% 潜在损失 (肉眼难见，但决定任务生死！)
   └──────────────────────────────────────────────┘
                          │
                          ▼
   传统 JEPA 潜在空间规划出的轨迹：
   - 极其擅长“让机械臂在空中摆出优雅姿态”，却屡屡插不进钥匙！

创新方案

一句话总结：SWIM（State Weighting with Inverse-Dynamics Model，基于逆动力学状态加权）创新性地以“动作预测”为因果试金石，利用离线轨迹中状态位移对实际操作动作的解释能力自动沉淀物理特征权重，以此诱导 JEPA 潜在距离度量，在零推理开销下大幅增强机器人物理规划能力。

核心机制

SWIM 的核心洞察源自经典控制理论的深刻直觉：如果某个物理坐标的位移必须用不同的控制动作才能解释，那么该坐标才是真正决定任务走向的关键变量；反之，若某个坐标无论怎么变都不改变机器人的执行动作，它就是纯粹的被动游离噪声。

   SWIM 算法架构与两阶段训练流

   [离线交互轨迹数据集] D = {(s_t, a_t, s_(t+1))}
                    │
                    ▼
   [第 1 阶段: 逆动力学动作因果预测]
   观测状态物理位移: Δx = x_(t+1) - x_t
   特征加权投影: w ⊙ Δx
                    │
                    ▼
   逆动力学网络预测动作: a^_t = MLP(w ⊙ Δx)
   损失函数: L_ID = || a^_t - a_t ||^2  ==>  反向传播自动收敛特征权重: w* ∈ R^D
                    │
                    ▼ 固化已学习的最优物理权重向量: w* (冻结参数)
   [第 2 阶段: 行为度量对齐的 JEPA 表征学习]
   加权物理流形真实距离度量:
        d_SWIM(x, x') = || w* ⊙ (x - x') ||_2
                    │
                    ▼ 潜在表征距离对齐损失
   约束 JEPA 视觉编码空间距离 d_Z(z, z') 正比于 d_SWIM(x, x')
                    │
                    ▼
   [最终部署与在线 MPC 轨迹规划]
   - 逆动力学网络与状态真值标签被彻底丢弃！
   - 机器人纯依靠单目摄像头画面，在校准后的 JEPA 潜在空间中极速搜索最优动作序列。

基于逆动力学的行为特征加权：
- 构建轻量逆动力学模型 $g_\theta$ ，以加权位移 $\mathbf{w} \odot \Delta \mathbf{x}$ 为输入预测机器人当时实施的控制动作 $\mathbf{a}_t$ 。
- 可学习重要性向量 $\mathbf{w} \in \mathbb{R}^D_{\ge 0}$ 会在反向传播中自然分化：凡是与控制动作强因果耦合的维度（如夹爪指尖与目标的相对欧拉角）被赋予高权重，而机身无用晃动维度权重趋向于零。
行为流形诱导的各向异性度量对齐：
- 固化后的权重向量 $\mathbf{w}^*$ 构建了一个全新的物理加权度量空间： $d_{\mathbf{w}^*}(\mathbf{x}_A, \mathbf{x}_B) = \sqrt{\sum_{i=1}^D w_i^* (x_{A,i} - x_{B,i})^2}$
- 在训练视觉 JEPA 编码器 $z = E(I)$ 时，强行拉通潜在嵌入距离与行为物理距离的保真度： $\mathcal{L}_{\text{metric}} = \left| \|z_A - z_B\|_2 - d_{\mathbf{w}^*}(\mathbf{x}_A, \mathbf{x}_B) \right|^2$
推理规划期零开销（Zero Overhead）：
- 逆动力学网络与仿真器物理真值仅作为离线“脚手架”存在于训练期。
- 实际机器人推理时，直接以纯视觉单目画面输入已对齐的 JEPA 潜在空间，配合模型预测路径积分（MPPI）或交叉熵（CEM）算法进行轨迹推演，不耗费任何额外 FLOPs。

验证成效

作者在 RoboSuite 和 MetaWorld 包含的 5 个严苛机械臂精细操作场景（密差销轴插入、门锁把手扳动、抽屉阻尼推拉、螺母螺栓旋合、多积木高塔搭建）中进行了全面规划对决：

具身操作任务	纯无监督视觉 JEPA	均匀状态监督 JEPA	专家手工调优状态加权	SWIM (本研究)	相对均匀监督的净增益
紧密销轴插入 (Peg Insertion)	32.4%	58.1%	76.5%	84.2%	+26.1% 成功率
弹簧门把扳动 (Door Opening)	48.6%	69.3%	78.4%	86.7%	+17.4% 成功率
螺栓精准旋合 (Nut Assembly)	28.5%	51.2%	68.0%	79.5%	+28.3% 成功率
阻尼抽屉推拉 (Drawer Closing)	62.1%	81.4%	87.2%	92.1%	+10.7% 成功率
多积木稳定叠放 (Block Stacking)	41.0%	64.8%	74.6%	81.3%	+16.5% 成功率
5 大任务综合平均规划成功率	42.5%	65.0%	76.9%	84.8%	大幅领跑近 20 个百分点

核心评测突破：

全面碾压基线模型：SWIM 取得了 84.8% 的极高综合规划成功率，相比业内通行的均匀状态加权方案（65.0%）大幅净胜近 20 个百分点。
超越人类资深工程师经验：在全部 5 项任务中，SWIM 的自动因果加权表现全面击败了由人类工程师数周手动调配的专家启发式权重（84.8% vs 76.9%），敏锐捕捉到了诸如夹爪微小侧倾与接触面法向力之间的隐蔽非线性约束。
对抗动态视觉干扰坚如磐石：在桌面上人为引入不可预测的弹跳彩色小球与炫目光影晃动时，传统视觉 JEPA 的规划成功率直接从 42.5% 崩塌至 11.2%，而 SWIM 依旧稳定在 82.9%，因其逆动力学模块早已将这些视觉小球的权重降维湮灭。

关键启示

以行为因果定义物理价值：在物理世界的表征学习中，绝不能被无意义的几何色块牵着鼻子走；动作条件下的逆动力学因果关系，是筛选关键维度的最强大数学滤波器。
度量对齐打通纯视觉与纯动作的鸿沟：通过潜在距离对齐将行为流形注入表征空间，无需生成每一个像素，即可让非生成式世界模型拥有极致敏锐的物理感知力。
“富态离线训练，轻量纯视觉推理”是具身 AI 的黄金范式：在训练期最大化挖掘物理特权信息的监督价值，在部署期彻底卸载特权依赖保持纯视觉极速推演，为规模化具身模型落地树立了标杆。