Technical Summary: Demo-JEPA

Problem Statement

Robotic imitation learning faces a fundamental bottleneck known as the embodiment gap: the physical and kinematic mismatch between a demonstrator (e.g., a human or a different robot) and the learner. Traditional approaches attempt to bridge this gap through action-level alignment, relying on shared action spaces, explicit action correspondences, heuristic retargeting, or large-scale multi-embodiment co-training. However, these methods are fundamentally limited for heterogeneous agents because identical semantic intents (e.g., "grasp an object") often require radically different motor commands, torque profiles, and joint configurations depending on the specific morphology. Consequently, learning cross-embodiment skills via low-level action alignment struggles to generalize across mismatched morphologies without massive amounts of aligned data.

Methodology: Demo-JEPA

The authors propose Demo-JEPA (Joint-Embedding Predictive Architecture), a framework that reframes cross-embodiment imitation not as action reproduction, but as latent goal-conditioned planning. The core hypothesis is that demonstrations should be viewed as implicit specifications of future goals (what state the demonstrator is trying to realize) rather than sequences of motor primitives (how they execute it).

The framework operates in three main stages within a shared predictive representation space:

1. Action-Conditioned World Model

The target agent utilizes a JEPA-based world model (specifically V-JEPA 2.1) to learn the dynamics of its own embodiment.

• Encoder ($E$): Maps visual observations ($o_t$) to latent states ($z_t$).
• Dynamics Predictor ($F_{wm}$): Models transitions conditioned on the current latent state, robot state ($s_t$), and action ($a_t$): $\hat{z}_{t+1} = F_{wm}(z_t, s_t, a_t)$.
• Planning: Given a target latent goal, the agent uses the Cross-Entropy Method (CEM) to optimize an action sequence that minimizes the distance between the predicted future latent state and the goal.

2. Dreamer Predictor (Embodiment-Aware Goal Inference)

This is the novel component responsible for translating source demonstrations into target-compatible goals.

• Input: The current target observation ($o^t_k$) and a pair of source frames ($o^s_k, o^s_{k+n}$) representing a temporal offset in the demonstration.
• Mechanism:
  • A shared JEPA encoder maps observations to latent representations ($z^t_k, z^s_k, z^s_{k+n}$).
  • Cross-Attention Modules: Two modules capture (1) cross-embodiment correspondence ($f_{emb}$) and (2) temporal evolution of the behavior ($f_{mot}$).
  • 3D Convolutional Fusion: The features are aggregated via a 3D convolution ($\phi$) to preserve spatiotemporal structure, followed by a Transformer predictor ($T$) to decode the fused features into a target-compatible latent goal ($\hat{z}^t_{goal}$).
• Training: The predictor is trained in Stage 1 using action-free cross-embodiment demonstrations to minimize latent reconstruction loss ($L_{pred} = \|\hat{z}^t_{goal} - z^t_{k+n}\|^2_2$). Temporal perturbation is applied to the source frames to enhance robustness.

3. Action Co-Training (Stage 2)

Once the Dreamer Predictor is frozen, the world model's dynamics predictor ($F_{wm}$) is fine-tuned. The objective is to align the world model's latent rollouts with the goal distribution provided by the Dreamer Predictor, optimizing a planning loss ($L_{plan}$).

Inference with Adaptive Goal Updating

During execution, the system does not advance the source reference frame at a fixed frequency. Instead, it computes the latent discrepancy ($D_k$) between the newly observed target state and the current goal. The system only advances to the next source reference pair if $D_k < \epsilon$ (the goal is reached). This adaptive mechanism prevents compounding errors caused by kinematic mismatches.

Key Contributions

1. Novel Formulation: The paper proposes reframing cross-embodiment imitation as latent goal-conditioned planning, effectively circumventing the bottleneck of explicit action alignment.
2. Demo-JEPA Framework: Introduction of a purely vision-driven framework featuring an embodiment-aware Dreamer Predictor. This component bridges the morphological gap by inferring target-compatible future latent states from source demonstrations without requiring action-level correspondences or shared action spaces.
3. Empirical Validation: Extensive evaluation in both simulated (RLBench) and real-world robotic manipulation tasks demonstrates that Demo-JEPA achieves superior generalization to unseen tasks and radically different embodiment configurations compared to prior baselines.

Experimental Results

The framework was evaluated across three suites with increasing distribution shifts: Behavior Grounding (seen tasks), Cross-Embodiment Bridging (unseen actions), and Zero-Shot Generalization (unseen configurations).

• Simulation (RLBench): Demo-JEPA matched specialized in-domain planners in behavior grounding but significantly outperformed baselines (VPP, XSkill) as distribution shift increased. In Zero-Shot Generalization, Demo-JEPA achieved a 0.36 average success rate, whereas VPP and XSkill achieved only 0.04 and 0.03, respectively.
• Real-World: In physical experiments (UR5e source $\to$ Franka target), Demo-JEPA achieved a 0.55 average success rate in Cross-Embodiment Bridging and 0.25 in Zero-Shot Generalization, outperforming VPP and XSkill which failed significantly in zero-shot scenarios.
• Ablation Studies:
  • Removing the 3D Convolutional fusion module degraded performance on complex, articulated tasks (e.g., "Change Channel"), highlighting the importance of explicit spatiotemporal modeling.
  • Comparisons with "Naive" (direct latent transfer) and "Oracle" (ground-truth target future) baselines confirmed that the Dreamer Predictor successfully converts heterogeneous demonstrations into target-compatible goals, approaching oracle performance without privileged information.
• Scaling: The study found that task diversity (exposure to diverse task semantics) is a more critical driver of transferable latent mappings than the volume of per-task trajectories.

Significance and Claims

The paper claims that Demo-JEPA successfully translates abstract task progression into an embodiment-compatible goal space by capturing the underlying predictive dynamics of a demonstration rather than memorizing surface-level statistics.

• Robustness: The method demonstrates that inferring latent goals is more effective than using demonstrations as fixed action sequences or skill prototypes, particularly when distribution shifts are severe.
• Flexibility: By avoiding action-level correspondence, the framework supports flexible imitation across heterogeneous embodiments using only visual demonstrations and the target agent's own interaction experience.
• Limitations: The authors modestly acknowledge that performance is currently constrained by the modeling bottlenecks of existing action-conditioned world models, which limits efficacy in highly complex, high-precision tasks. Additionally, the current pipeline still requires temporal or progress-aware alignment during training, and future work aims to develop fully unaligned learning strategies.