JEDI: Joint-Embedding Demonstration Compression for In-Context Imitation Learning

Jiazhao Liang; Zhewen He; Shuaihang Yuan; Hao Huang; Yi Fang

cs.ROOct 29, 2026

JEDI: Joint-Embedding Demonstration Compression for In-Context Imitation Learning

Authors: Jiazhao Liang, Zhewen He, Shuaihang Yuan, Hao Huang, Yi Fang

Organizations: New York University Tandon School of Engineering, Brooklyn, NY 11201, USA · NYUAD Center for Artificial Intelligence and Robotics (CAIR) and Embodied AI and Robotics (AIR) Lab, New York University Abu Dhabi, Abu Dhabi 129188, UAE

Abstract

In-context imitation learning (ICIL) adapts robot policies through test-time demonstrations, but long multimodal trajectories constrain how many examples fit in context. We ask what a compressed demonstration should preserve. We propose JEDI, which compresses demonstrations into latent slots trained to predict the query episode's next-stage keyframe in a frozen video-JEPA space. Task and process counterfactuals help the slots preserve execution details, while the predictor selects demonstrations from a bank. Post-training with keypoint tokens enables human-hand conditioning on robots. Experiments on LIBERO and real robots show that JEDI supports many-shot context; LIBERO objective controls show improved success and demonstration following over action-only compression. These findings support query-future prediction as a training objective for preserving execution-relevant information in compressed demonstrations.

Figures & tables

Figure 2: Overview of JEDI. (1) A compressor maps each demonstration to $k$ latent slots, which are cached in a demonstration bank. (2) The slots are trained to predict the next keyframe of the query episode in the embedding space of a frozen V-JEPA 2.1 encoder, with task and process negatives. (3) The predictor ranks the demonstrations in the bank, and only the policy runs in the control loop. (4) Post-training on 10 tasks uses equal shares of human-hand and dual-arm robot demonstrations with the same objective; unseen tasks are executed from human-hand context without test-time weight updates.

Method	LIBERO (%, $\uparrow$ )			RLBench (%, $\uparrow$ )
Method	Spatial	Object	Goal	RLBench (%, $\uparrow$ )
ICRT ( Fu et al., 2025 )	58.5 $\pm 3.4$	44.6 $\pm 4.6$	24.6 $\pm 2.1$	49.3 $\pm 3.8$
RoboSSM ( Yoo et al., 2026 )	66.8 $\pm 3.1$	54.3 $\pm 4.3$	37.3 $\pm 3.6$	54.4 $\pm 2.8$
Instant Policy ( Vosylius and Johns, 2025 )	83.1 $\pm 3.2$	69.1 $\pm 2.9$	47.1 $\pm 2.8$	77.1 $\pm 3.3$
ContextFlow ( Ding et al., 2026 )	85.6 $\pm 2.8$	76.4 $\pm 3.2$	56.5 $\pm 2.2$	73.6 $\pm 1.3$
JEDI (ours)	92.7 $\pm 1.6$	83.6 $\pm 2.9$	63.0 $\pm 2.3$	82.1 $\pm 2.0$

Table 1: Success rate on unseen manipulation configurations in LIBERO-Spatial, LIBERO-Object, LIBERO-Goal and RLBench. The first five rows use one demonstration (

N{=}1

); input adapters are distinguished in Appendix E . Values are mean

\pm

standard deviation (%) over 3 training seeds. Each LIBERO seed uses 4 held-out tasks and 100 rollouts per task (1,200 trials per suite across seeds; Appendix F ). The best

N{=}1

mean in each column is in bold.

	SR $\uparrow$			Order			Height
Compressor objective	$N{=}1$	$N{=}8$	CFR $\uparrow$	match $\uparrow$	$\text{DD}_{\text{task}}$ $\uparrow$	$\text{DD}_{\text{proc}}$ $\uparrow$	err. $\downarrow$
Action-only (matched)	67.1	64.2	34.1	52.2	–	–	4.1
+ $\mathcal{L}_{\text{JEPA}}$	73.4	78.1	46.3	56.2	1.11	1.02	2.5
+ $\mathcal{L}_{\text{JEPA}}$ + $\mathcal{L}_{\text{cf}}$ (task negatives)	72.2	76.6	69.7	57.4	2.06	1.09	2.6
Label contrast (no predictor)	68.7	68.7	51.5	73.5	–	–	3.5
JEDI (task + process negatives)	79.8	80.2	74.9	84.1	2.21	1.49	2.0

Table 3: Core compressor-objective ablations (same backbone, data and schedule). SR, CFR and order-match are in %; height error is in cm. SR uses instructions and averages the three LIBERO suites of Table 1 ;

\text{DD}_{\text{task}}

\text{DD}_{\text{proc}}

average the last 200 training steps. The final line tests height transfer with order-only negative labels. Best in bold; further controls are in Table 9 .

Method	Canonical	Paraphrase	Conflict
JEDI	70.0	71.7	78.8
Action-only	60.0	70.0	45.0
Instruction-only	65.0	55.0	–

Table 4: Instructions stating the object order on LIBERO-Order (success, %;

N{=}8

; 20/60/80 trials per column). Conflict: opposite-order context. Instruction-only: context masked. See Appendix I .

Method	R1	R2	R3
Retrieval-1	$13/25$	$4/10$	—
Action-only, $N{=}1$	$15/25$	$4/10$	—
JEDI, $N{=}1$	$15/25$	$6/10$	—
JEDI, $N{=}8$	$21/25$	$\mathbf{8/10}$	$10/15$
Robot ( $N{=}8$ )	$\mathbf{23/25}$	$\mathbf{8/10}$	$\mathbf{11/15}$
Mixed ( $N{=}4{+}4$ )	$22/25$	$\mathbf{8/10}$	$10/15$

Table 5: Real-robot success (successes/trials; 5 per task). R1: unseen manipulation (5 tasks); R2: instruction–demonstration ambiguity (2); R3: demonstration–demonstration ambiguity (3). Best in bold. —: R3 needs several demonstrations. Hardware and tasks in Appendix H .

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Item	Setting
Model and optimization
Slots per demonstration $k$	8
Compressor	joint perceiver, $L_{c}{=}4$ layers
Demonstration subsampling	8 frames; 128 state/action steps
Vision encoder	SigLIP So400m/14 (from $\pi_{0}$ )
query side	fine-tuned

Appendix

Table 7: Hyperparameters and benchmarks.

Method (no instruction)	$N$	Spatial	Object	Goal	Avg.
$\pi_{0}$ , no demonstration	0	52.8	15.4	9.1	25.8
Action-only (matched)	1	84.6	55.2	34.1	58.0
Action-only (matched)	16	81.9	46.5	20.8	49.7
JEDI	1	80.3	63.0	74.9	72.7
JEDI	16	87.2	73.8	61.2	74.1

Appendix

Table 8: Success without instructions (%; protocol of Table 1 ). Avg.: mean over suites. Best per column in bold.

	SR $\uparrow$			Order			Height
Compressor objective	$N{=}1$	$N{=}8$	CFR $\uparrow$	match $\uparrow$	$\text{DD}_{\text{task}}$ $\uparrow$	$\text{DD}_{\text{proc}}$ $\uparrow$	err. $\downarrow$
+ Reconstruction	66.8	63.2	36.4	52.2	–	–	3.8
+ Trace target	71.4	76.1	46.3	60.5	1.31	1.10	2.5
+ Prediction without slots	64.2	66.6	39.7	55.4	–	–	3.3
JEDI, gradient through $o_{t}$	77.5	83.9	60.1	84.9	2.58	1.54	1.4

Appendix

Table 9: Additional compressor controls under the settings and metrics of Table 3 . SR, CFR and order-match in %; height error in cm. Best per column in bold.

Objective	Labels used	Order	Height
Action-only (matched)	none	52.2	55.1
Label contrast	order only	68.3	52.2
Label contrast	height only	51.8	75.4
Label contrast	order + height	73.5	76.8
JEDI	task only	57.4	63.9
JEDI	order only	76.1	70.5

Appendix

Table 10: Prediction versus label supervision. Order-match and height-match rates (%) on LIBERO-Order using the indicated labels for negatives (JEDI) or contrastive pairs (Label contrast). Off-diagonal entries evaluate variation omitted from those labels; positive support/query pairs remain process-matched. Best in bold.

Variant	$N{=}1$	$N{=}8$	CFR
	SR
JEDI (default)	79.8	80.2	74.9
DINOv2 target	76.4	74.3	62.7
Final-frame target	68.2	76.4	65.2
One slot per demonstration	70.9	66.0	51.3
No instruction dropout	74.3	78.2	45.1

Appendix

Table 11: Key ablations (%). Best per column in bold.

Method	$N$	Prefill (ms)	Hz	Mem. (GB)
Raw context	1	178	3	12
Raw context	2	425	2	13
ContextFlow	1	26	13	8
JEDI	1	28	13	8
JEDI	4	20	10	8
JEDI	8	44	10	10

Appendix

Table 12: Inference cost on one A100. Prefill: one-time cost of encoding the support set; Hz: control frequency with cached slots. Bank workloads use 8 pseudo-query steps per demonstration, batched across pair/time inputs. The final block gives measured bank-construction times.

Attention mass on relevant slots (%)
	Distractors $N_{d}$
Method	0	4	8	16
Action-only (matched)	64.2	43.3	42.9	42.8
JEDI	80.2	79.0	78.1	76.1
JEDI, no distractor mixing	80.6	78.2	74.7	74.4
Retrieval-1 (from the mixed pool)	77.5	74.5	74.3	73.1

Appendix

Table 13: Distractor demonstrations. Success (%) with

N{=}8

relevant demonstrations and

N_{d}

distractors, and attention on relevant slots. Retrieval-1 selects one demonstration from this pool.

Method	R1 subset	R2 subset	R3 subset
JEDI, $N{=}8$	$15/20$	$8/10$	$5/10$
Action-only, $N{=}8$	$11/20$	$5/10$	$2/10$
Instruction-only	$10/20$	$4/10$	$4/10$
No keypoint tokens	$11/20$	$5/10$	$5/10$
No predictive loss	$9/20$	$6/10$	$3/10$

Appendix

Table 14: Matched real-robot context and post-training controls (successes/trials; 10 per condition). Subsets use the R1/R2/R3 predicates of Table 5 .

Figure 4: Matched context-size comparison on LIBERO-Object. Identical instructions, two held-out tasks and 20 paired trials per point.

Method	No instruction	Instruction only
JEDI	65.0	25.0
Action-only	60.0	–
ContextFlow	35.0	–

Appendix

Table 15: Input controls on LIBERO-Object (success, %; 20 paired trials per cell). Instruction-only retains the JEDI policy with its demonstration context masked. The order-instruction controls are in Table 4 .

ICL for Robotics

In-context learning (ICL) for robotics lets robots adapt from demonstrations and interaction history without updating model weights.

See all 43 papers

Demo-JEPA: Joint-Embedding Predictive Architecture for One-shot Cross-Embodiment Imitation
May 20, 2026Jingyang He, Guangrun Li, Jieyu Zhang +3Robot Skill Learning World Models for Robotics
State Key Laboratory of Multimedia Information Processing, School of Computer Science, Peking University · Beijing Innovation Center of Humanoid Robotics · University of Washington
In-context Robot Learning Made Simple: A Democratized Recipe for Manipulation Tasks
Sep 29, 2026Minxing Li, Minghao Han, Weizhi Zhao +10ICL for Robotics Robotic Manipulation
NLPR, Institute of Automation, Chinese Academy of Sciences (CASIA) · Amap, Alibaba Group
ContextFlow: In-Context Flow Matching for Robot Manipulation
Sep 6, 2026Jian Ding, Xianjie Dai, Roei Herzig +5ICL for Robotics Robot Imitation Learning
KAUST, Thuwal, Saudi Arabia · University of California, Berkeley, USA · Zurich, Switzerland

May 20, 2026J/K move · Enter open · S save

Demo-JEPA: Joint-Embedding Predictive Architecture for One-shot Cross-Embodiment Imitation

Jingyang He, Guangrun Li, Jieyu Zhang, Chengkai Hou+2

State Key Laboratory of Multimedia Information Processing, School of Computer Science, Peking University · Beijing Innovation Center of Humanoid Robotics · University of Washington

Robot Skill Learning World Models for Robotics Joint-Embedding Predictive Architecture Generalization in Robotic Manipulation Cross-Embodiment Robot Learning Imitation Learning+1

Open paper AskPDF arXiv

Robotic imitation learning is often treated as reproducing demonstrated actions, but actions are inherently embodiment-specific. When demonstrations come from humans or robots with different morphology, kinematics, or action spaces, this action-centric view requires shared action spaces, heuristic retargeting, or large-scale multi-embodiment co-training. We instead view demonstrations as implicit specifications of future goals: the target agent should infer what state the demonstrator is trying to realize, rather than how the demonstrator executes it. We propose Demo-JEPA, a cross-embodiment imitation framework that decouples demonstration intent from embodiment-specific execution. Built on a JEPA-based world model, Demo-JEPA translates source visual demonstrations into target-compatible future latent trajectories in a shared predictive representation space. The target agent then uses these latent trajectories as subgoals and realizes them through planning under its own learned forward dynamics. Because Demo-JEPA avoids action-level correspondence and requires only visual demonstrations plus the target agent's own interaction experience, it supports flexible imitation across heterogeneous embodiments. Experiments on RLBench and real-world manipulation tasks show that Demo-JEPA matches specialized in-domain planners and generalizes to unseen tasks and embodiment configurations where prior methods fail.

Abstract

ICL for Robotics

Explore similar work

Demo-JEPA: Joint-Embedding Predictive Architecture for One-shot Cross-Embodiment Imitation

In-context Robot Learning Made Simple: A Democratized Recipe for Manipulation Tasks

ContextFlow: In-Context Flow Matching for Robot Manipulation