PSG-JEPA：仅靠前向预测够吗？

01 Motivation

近期 JEPA-based 世界模型（如 LeWM）从观测序列中学习 action-conditioned 的潜动力学。但作者指出：前向预测目标并不显式监督"从单个 latent 解码机器人物理状态、从 latent 对解码状态变化"这件事，于是留下一个 "robot-centric identifiability gap"——视觉与动作之间的对齐关系没有被编码进表征，下游 planner / policy 只能自己从头去学，进而限制规划与策略的性能。

"control requires relating visual observations to the robot's state and to how actions change it" —— 因此把 vision–action 对齐放到表征层本身，是控制的基础。

02 Method

PSG-JEPA 在标准前向 JEPA 损失之外并联两条 physical state grounding 目标，用来"塑形"潜空间。共享 encoder 对每帧观测 o_t+i 产出 z_t+i=E_φ(o_t+i)；两条目标各挂一个轻量 head，只在训练时使用、训练后丢弃，因此推理结构与算力与基线完全一致。

① Static state grounding —— 单帧锚到本体状态

把每个 latent 拉向机器人本体感知状态（joint angles、gripper state、end-effector pose）。损失为
ℒ_static = (1/T) Σ_i=0..T−1 MSE(H_s(z_t+i), s_t+i)。 这保证"从单个 latent 就能线性/MLP 读出机器人物理状态"。

② Dynamic transition grounding —— 潜向量对锚到多步关节角变化

把 latent 对 (z_t+i, z_t+i+k) 拉向对应的 multi-horizon 关节角变化，横跨所有 horizon k ∈ {1,…,T−1}：
ℒ_dynamic = (1/|𝒦|) Σ_k∈𝒦 (1/|ℐ_k|) Σ_i∈ℐk MSE(Δq̂_t+i,k, Δq_t+i,k)。 相比只在相邻帧上做 one-step inverse-dynamics，这里的 multi-horizon 端点对提供了更丰富的"状态变化"监督。

③ 组合损失

总损失为 ℒ_PSG = ℒ_JEPA + λ_g(ℒ_static + ℒ_dynamic)，其中 λ_g = 0.1。state head H_s 与 transition head H_Δ 均在训练后丢弃，推理架构和计算成本保持不变。

03 Experiments

在三个层级评测：(1) 通过 probing 检验 latent 的物理可辨识性；(2) 在冻结 latent 上做 goal-conditioned planning；(3) 在仿真与真机上做 policy learning。基线包括 LeWM（同 backbone 前向预测）、LeWMActionIDM（相邻对加一步 inverse-dynamics 目标）、DINOv2（预训练视觉参照）。

Q1 · 物理可辨识性探针（OGBench-Cube，Pearson r，linear/MLP）

Table 1 单帧状态探针：PSG-JEPA 在 EE-yaw 上 r 达 0.94/0.98，而 LeWM 仅 0.08/0.08、LeWMActionIDM 仅 0.11/0.10——前向预测/单步 IDM 几乎"看不见"末端偏航。

Table 2 相邻对的 transition 探针：JointVel 0.75/0.75、GripVel 0.69/0.76、Action 0.80/0.86，均为最优或持平最优。

Q2 · 冻结 latent 上的 goal-conditioned planning（成功率 %，3 seeds 均值±std）

PSG-JEPA 在小 planner 训练预算和小数据下优势最大：OGBench-Cube 全量数据、GC-IDM 仅训 5 epoch 时即达 95.0±0.7，而 LeWM 80.7±1.9；训到 100 epoch 时达 98.7±1.2。

Q3 · 策略学习：仿真 + 真机

LIBERO-Goal（10 任务，每 seed 50 rollout，3 seeds）：PSG-JEPA 85.3±3.9%，高于 LeWM 77.7±0.5%、LeWMActionIDM 82.6±2.2%、DINOv2 80.1±5.3%。真机 Dual-arm Cobot（Mobile ALOHA 设计，每任务 50 trials）三项任务全部超越 LeWM，平均 79.3% vs 60.0%。

Ablations（Table 5，Cube 规划 SR@5ep 与 LIBERO-Goal）

逐一移除组件均导致下降，说明"每个 grounding 组件都有贡献"：去掉 transition head 时规划仅 81.3±2.4%、去掉 state head 时探针 r 掉回 0.69；把 multi-horizon 换成仅相邻对（adjacent/endpoint）规划为 93.5/93.7%——均不及完整版的 95.0±0.7% 与策略 85.3±3.9%。

04 Limitations