Inside-JEPA: From Prediction to Relational Grounding in Latent World Models

Hossein Resani; Javen Qinfeng Shi

cs.LGOct 22, 2026

Inside-JEPA: From Prediction to Relational Grounding in Latent World Models

Authors: Hossein Resani, Javen Qinfeng Shi

Organizations: Australian Institute for Machine Learning, Adelaide University · Responsible AI Research Centre

Abstract

Joint-embedding predictive architectures (JEPAs) learn world models in latent space, enabling efficient planning but leaving their states opaque. We introduce Inside-JEPA, which places relational channels for ordered entity pairs inside the JEPA latent state and shows that an inverse-dynamics loss can ground these channels without relation labels when its head reads them, either alone or as part of the full state. On Something-Something V2, latent prediction alone performs no better than a random relation head, while inverse-dynamics supervision that recovers actions from channel changes improves relational AUROC by .065-.103 across ten settings with annotated boxes, by .081 on RH20T robot episodes with held-out tasks, and by .050-.061 across five relations using off-the-shelf detection and tracking. An inverse head on entity latents decodes actions more accurately but barely improves relational grounding; shuffled labels, two label-free objectives, and action labels on only 5% of the pairs also fail to ground the channels. Given an action, predicted relational changes approach the quality of observed future relations (AUROC .599 vs .621). For one-step relational goals, a supervised readout of the 16 channels ranks the correct action among the top five of 65 in 43% of cases, compared with 24% for prediction-only latents and 17% for a frequency prior.

Figures & tables

Figure 2: Five arms, one difference. All arms share encoder, predictor, data, budget and loss weight and differ only in what the inverse head $h$ reads, and so in where its gradient enters the state (coloured outlines), or in its targets. PRED has no inverse head. INV-S (Inside-JEPA) reads the relational channels $s$ ; INV-Z reads the whole latent $z=(r,s)$ ; INV-R reads only the entity latents $r$ , so its gradient reaches $s$ only through the shared encoder; INV-SHUF is INV-S with shuffled targets $\pi(a_{t})$ for the inverse loss only. $P$ receives $a_{t}$ when action conditioning is enabled. The action inputs shown are omitted in the first four settings of Tab. 3 and in the action-free forecasting control.

Figure 3: How we test the relational channels (schematic). (a) Grounding: the AUROC of every relational channel

c

for every relation

\rho

is computed on validation pairs; Hungarian matching gives each relation one channel, and the mean matched AUROC is compared with that of a randomly initialised relation head (margin). (b) Forecasting relational change: from

z_{t}

P

forecasts

\hat{s}_{t+k}

with the true action

a_{t}

or a permuted one

\pi(a_{t})

; both are scored on the entries of

y_{t+k}

whose relation changes (outlined), and their difference is the action effect. (c) Planning: every candidate action is rolled through

P

, each forecast is scored by the likelihood of the relational goal

y^{\star}

under a readout, and we report how often the true action ranks in the top five.

Figure 4: Inverse dynamics grounds the relational channels across settings. Matched relational AUROC of PRED and Inside-JEPA ( INV-S ), with the gain in each setting; means over 3 seeds, or 5 where marked. tpl: templates; rel: relations.

	Grounding		Planning, top-5 of 65
Arm	Matched AUROC $\uparrow$	Margin $\uparrow$	Chan $\uparrow$	Inv $\uparrow$	Probe-Z $\uparrow$	Probe-S $\uparrow$
PRED	.593	+.013	.067	–	.244	.168
INV-SHUF	.539	$-$ .026	.136	–	.237	.184
INV-R	.608	+.020	.113	–	.311	.209
INV-Z	.670	+.083	.297	.482	.391	.380
INV-S	.658	+.080	.334	.374	.416	.425

Table 1: All arms on the main setup (65 templates, 13 relations, annotated boxes, frozen ResNet-18, 3 seeds; arms as in Fig. 2 ). Margin: over a seeded random head. Planning readouts, fitted on relation-labelled pairs: Chan , one calibrated channel per relation; Inv , the arm’s inverse head; Probe-Z , a probe on the whole forecast latent

z

; Probe-S , a probe on the 16 relational channels

s

only; prior .174, chance .077.

\uparrow

: higher is better; bold: best arm trained without relation labels; underline: second best. SUP-S is trained on the relation labels and shown only as a ceiling.

Figure 5: What drives grounding (main setup: 65 templates, 13 relations, annotated boxes). (a) Per-relation gain of INV-S over PRED (

\pm

2 SE, 3 seeds): contact, proximity and vertical relations ground; left_of , much_larger and the hand relations do not. (b) Nested fractions of the training data, all labelled (3 seeds): more data lifts PRED more than INV-S . (c) All data, with action labels on a fraction of the pairs (5 seeds): labels lift INV-S ; PRED has no inverse loss, so its solid line is a single value; dashed: the predictor also sees only the labelled pairs’ actions.

Loss on $s$	Matched AUROC $\uparrow$	Gain $\uparrow$	Effect $\uparrow$	Chan $\uparrow$	Probe-S $\uparrow$
None ( PRED )	.595	–	.003	.066	.162
Shuffled $a$ ( INV-SHUF )	.538	$-$ .057	.011	.122	.186
Arrow of time ( AOT-S )	.588	$-$ .007	.046	.190	.248
InfoNCE ( NCE-S )	.605	+.010	.006	.086	.150
$a$ on 1% (153)	.593	$-$ .002	.016	.114	.222
$a$ on 5% (701)	.603	+.008	.038	.185	.280

Table 2: Grounding needs correct action labels, in the thousands (65 templates, 13 relations, 5 seeds). Every row trains the same model with a different loss on the relational channels

s

; “

a

x

%”: inverse dynamics with action labels on a nested

x

% of the pairs (number of labelled pairs in brackets), while the predictor trains on all pairs with their true actions; hiding those actions from the predictor too gives the same gains ( Tab. 13 ). Gain over PRED ; Effect: action effect on forecasts of changing relations; planning top-5, prior .174.

\uparrow

: higher is better; bold: best; underline: second best.

Figure 6: Using the grounded channels. (a) Forecasting relational change (40 templates, 5 relations): AUROC on entries whose relation changes, when copying

s_{t}

, forecasting with a permuted or the true action, or observing

s_{t+k}

; only arms with grounded channels respond strongly to the action. (b) One-step planning from relational goals (top-5 of 65 templates, 13 relations) with probes on the forecast

r

s

z

; lines: frequency prior, chance and relation-supervised channels. (c) Relation-labelled pairs needed to fit the readouts.

Appendix figures & tables20 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 7: The thirteen relations , each with one example frame of a directed pair

a\to b

(rose, teal; hatched: a hand) that satisfies it. Boxes are normalised to

[0,1]

with

y

pointing down. Under each name is its rule as written in the code.

c

: box centre;

|\cdot|

: area;

d=\lVert c^{a}-c^{b}\rVert/\sqrt{2}

. A dashed ring marks the

d

threshold around

c^{a}

; dark grey marks

a\cap b

; the light-grey band in y_overlap is the shared vertical extent; dotted lines mark the compared centres or edges.

Setting	Tpl	Rel	PRED $\uparrow$	INV-S $\uparrow$	Gain $\uparrow$	Margin gain (2 SE)
First test	40	5	.536	.634	+.097	+.092 (.022)
All templates	65	5	.546	.633	+.088	+.084 (.021)
3 $\times$ epochs	40	5	.537	.639	+.103	+.089 (.026)
Fine-tuned backbone	40	5	.540	.608	+.068	+.080 (.034)
Action-cond. $P$	40	5	.535	.631	+.096	+.095 (.013)
all templates	65	5	.537	.638	+.102	+.091 (.040)

Appendix

Table 3: Grounding with annotated boxes , frozen ResNet-18 (fine-tuned in one row), 3 seeds. Every row passes C1 – C6 except 3

\times

epochs (120), which trained only PRED and INV-S and passes its own rule that the gap persists. Margin gain: INV-S ’s margin over the seeded random head minus PRED ’s. The first four rows use

P(z_{t})

; subsequent rows use

P(z_{t},u(a_{t})

Arm	Matched AUROC $\uparrow$	Random head	Margin $\uparrow$	Probe $\uparrow$	Flip rate (GT .212)	Own inverse top-1 $\uparrow$
PRED	.536	.536	+.000	.882	.242	–
INV-SHUF	.517	.535	$-$ .017	.856	.378	.025
INV-Z	.654	.535	+.119	.914	.324	.409
INV-S	.634	.542	+.092	.930	.302	.308

Appendix

Table 4: First test on real video (40 templates, 5 relations, compositional validation, 3 seeds). Random head and margin are seeded; inverse top-1 chance is .025.

\uparrow

: higher is better; bold: best; underline: second best.

Arm	Matched AUROC $\uparrow$	Margin $\uparrow$	Probe $\uparrow$	m3 forecast $\uparrow$	Copy	Obs.	Perm.	m4 $\uparrow$
PRED	.537	+.006	.896	.513	.506	.514	.512	+.001
INV-SHUF	.523	+.002	.856	.500	–	–	–	+.002
INV-R	.551	+.014	.919	.505	–	–	–	+.008
INV-Z	.635	+.097	.930	.593	.378	.620	.485	+.108
INV-S	.638	+.096	.938	.579	.400	.610	.506	+.074

Appendix

Table 5: Action-conditioned predictor on all 65 templates (5 relations, 3 seeds). m3: AUROC on changing entries with the true action; Copy: copying

s_{t}

; Obs.: observed future channels; Perm.: permuted action; m4: action effect. Seeded margins; C2 margin gain +.091 (2 SE .040).

\uparrow

: higher is better; bold: best; underline: second best.

Figure 8: Per-relation gain of INV-S over PRED in matched AUROC: annotated boxes on 40, 65 and all 174 templates, tracked objects and a frozen V-JEPA 2 encoder on 65 templates, and RH20T robot episodes. 3 seeds (V-JEPA 2: 5); a blank cell is a relation the setting did not evaluate.

Relation	Base rate	PRED	INV-S	INV-SHUF	INV-Z
left_of	.445	.55	.54	.51	.55
above	.443	.53	.53	.52	.55
overlap	.432	.51	.77	.53	.77
inside	.093	.55	.62	.51	.72
near	.745	.54	.71	.52	.67

Appendix

Table 6: Per-relation matched AUROC in the first test (validation base rates). The random head scores .52–.60 and the probe .80–.95 on every relation. Bold: best arm per relation.

Figure 9: Grounding gain by training-set size for the eleven non-hand relations, all thirteen and the two hand relations (65 templates, all pairs labelled, 3 seeds,

\pm

2 SE). Dashed: the C1 bar.

Data	Pairs	All 13 $\uparrow$	11 non-hand $\uparrow$	2 hand $\uparrow$	PRED	INV-S
12.5%	1,746	+.089	+.076	+.163	.531	.620
25%	3,490	+.095	+.101	+.063	.548	.643
50%	6,980	+.067	+.092	$-$ .067	.582	.650
75%	10,464	+.066	+.103	$-$ .137	.592	.658
100%	13,954	+.065	+.100	$-$ .128	.593	.658

Appendix

Table 7: Grounding gain by training-set size (65 templates, 13 relations, all pairs labelled, 3 seeds); the 12.5% and 75% rows were added later, under rules locked to test the hand/non-hand split.

Arm	IID binary AUROC $\uparrow$
Prediction only	.522 $\pm$ .018
+ inverse dynamics	.728 $\pm$ .009
+ inverse dynamics + unary inverse head	.743 $\pm$ .011
+ inverse dynamics, no residual bottleneck	.738 $\pm$ .005
+ inverse dynamics + object masking	.651 $\pm$ .023

Appendix

Table 8: Blocks2D ablations (mean

\pm

SD over 3 seeds; untrained head .641;).

Objects	Identity $\uparrow$	Found $\uparrow$	Outcome
DINOSAUR slots on DINOv2	.71	59%	+.018, fails
Slots re-ordered post hoc	.72	–	identity gate fails
SAVi-style tracking slots	.53	–	identity gate fails
SAM 2 automatic masks	.96	34%	coverage gate fails
GDINO + SAM 2, 1 frame	.94	64%	+.041, fails L1
GDINO + SAM 2, 3 frames	.87	79%	+.061, pass

Appendix

Table 9: Object sources tried. GDINO: Grounding DINO. Identity: the same object stays in the same slot or track over the clip; Found: annotated objects matched at IoU

\geq.5

\uparrow

: higher is better; bold: best.

Test	Tracked objects (annotated boxes)
Grounding, 40 tpl, 5 rel	+.061, pass (+.096)
no annotations at test	+.048, L1 fails (L2–L4 hold)
Grounding, 65 tpl, 5 rel	+.052, pass; post hoc 5 seeds +.050 (+.102)
no annotations at test, 5 seeds	+.047, L1 fails (L2–L4 hold)
Grounding, 65 tpl, 13 rel	+.043, below the bar; post hoc 5 seeds +.038 (+.065)
Forecast, true / permuted action, 5 seeds	.560 / .495 (.579 / .506)

Appendix

Table 10: Tracked objects : no box annotations in training; 3 seeds unless marked (marks refer to tracked objects). Values for annotated boxes in brackets. Planning: top-5 among 65 templates, prior .174. The forecasting action effect is 86% of the annotated-box one.

Figure 10: Where the gap to annotated boxes comes from , on one fixed set of validation pairs (gain of INV-S over PRED , 5 seeds,

\pm

SE). From top: all annotated boxes; annotated boxes for the objects the tracker (detection on one frame) found; its tracks with boxes snapped to the annotations; its tracks as run. Each pale segment is what one step costs. Right: the inverse head’s action accuracy. Seeds 3–4 were added after the locked three-seed reading.

Format	A $\to$ A	A $\to$ T	T $\to$ A	T $\to$ T
Objects first, hand last, hand flag	.103	.078	.028	.025
Hand first, no flag (tracker format)	.063	.051	.057	.052
Objects first, hand last, no flag	.116	.090	.052	.045
Hand first, hand flag	.060	.044	–	–
Random object order, hand last, no flag	.089	.068	–	–

Appendix

Table 11: Grounding gain by input format , training objects

\to

test objects (A annotated, T tracked; 65 templates, 5 relations). The A

\to

A and A

\to

T cells of the three no-flag rows and the T

\to

A cell of the tracker-format row use 5 seeds, the others 3.

Explanation	Test	Result
Captions name no direction	Real controls	left_of $-$ .031
Cameras disagree on image left/right	One camera per setup, 5 seeds	left_of $-$ .013
The action carries no left/right information	Mutual information per relation	second-most of five
The action is in the wrong frame	Controls mapped into each camera’s image axes	no gain in information (ratio .90); not trained
Too few channels	$K=32$ , on Something-Else	left_of +.014

Appendix

Table 12: Explanations for left_of not grounding. The first four were tested on RH20T and the capacity test on Something-Else (40 templates, 13 relations), where with

K=32

the extra channels go to relations that already ground.

Figure 11: Per-relation gain over PRED without action labels and with labels on part of the pairs (65 templates, 5 seeds). Left: label-free losses on the channels, arrow of time ( AOT-S ) and InfoNCE ( NCE-S ). Right: INV-S with action labels on 1–100% of the training pairs (locked design); the 100% column is INV-S itself.

	1% (153)	5% (701)	25% (3,490)	100%
Gain, locked $\uparrow$	$-$ .002 $\pm$ .011	+.008 $\pm$ .013	+.047 $\pm$ .009	+.066 $\pm$ .011
Gain, masked $\uparrow$	$-$ .002 $\pm$ .010	+.009 $\pm$ .015	+.047 $\pm$ .010	–
INV-S , masked	.596	.607	.645	–
PRED , masked	.597	.598	.598	–
Inverse top-1, locked $\uparrow$	.083	.122	.189	.318

Appendix

Table 13: Label fractions (65 templates, 13 relations, 5 seeds; gain

\pm

2 SE). No fraction below 100% passes C1 ; 25% misses it by .003 and holds C2 – C4 . The masked design hides the unlabelled pairs’ actions from the predictor too.

Arm	Dist $\uparrow$	Chan $\uparrow$	Inv $\uparrow$	Probe-Z $\uparrow$	Probe-S $\uparrow$	Probe-R $\uparrow$
PRED	.046	.067	–	.244	.168	.214
INV-SHUF	.118	.136	–	.237	.184	.213
INV-R	.075	.113	–	.311	.209	.301
INV-Z	.101	.297	.482	.391	.380	.329
INV-S	.146	.334	.374	.416	.425	.316
SUP-S	–	.394	–	–	.393	–

Appendix

Table 14: Planning from relational goals , top-5 among 65 templates, 13 relations, 3 seeds, readouts fitted on all labelled training pairs. Dist : squared distance of the matched channels; Chan : one calibrated channel per relation; Inv : the arm’s inverse head; Probe-Z , Probe-S , Probe-R : probes on the forecast

z

s

r

. Prior .174, chance .077.

\uparrow

: higher is better; bold: best arm without relation labels

Figure 12: The relational code is spread over the relational channels . (a) INV-S ’s readout AUROC for each relation from the one channel matched to it (light) and from a probe on all 16 relational channels (dark). (b) Planning top-5 with the one-channel readout Chan (light) and the channel probe Probe-S (dark), for INV-S and three label-free attempts to align relations with channels (the ICA rotation is applied after training); dashed: the frequency prior. INV-S is retrained here on the variants’ code path (5 seeds), so its values differ slightly from Tab. 14 .

Encoder input	PRED	INV-S $\uparrow$	Gain $\uparrow$	Probe-S $\uparrow$
ResNet-18, still frames	.593	.658	+.065	.425 (.244)
V-JEPA 2, still frames	.599	.648	+.050	.398 (.257)
still, 5 seeds	.601	.643	+.042	.390 (.252)
16 past frames	.625	.641	+.016	.399 (.308)
video start, still goal	.619	.640	+.021	.398 (.294)

Appendix

Table 15: Encoder (65 templates, 13 relations; matched AUROC). Planning: INV-S ’s channel probe, top-5, with PRED ’s whole-latent probe in brackets. The first two rows use 3 seeds, the others 5.

\uparrow

: higher is better

Figure 13: Share of the supervised ceiling that inverse dynamics recovers,

(\textsf{INV-S}-\textsf{PRED})/(\textsf{SUP-S}-\textsf{PRED})

in matched AUROC (light) or channel-probe AUROC (dark): 0 is PRED and 1 the label-trained ceiling (grey track). Light: one channel matched to each relation; dark: a probe on all 16 relational channels. Bars cut at

-

.5 continue to between

-

.61 and

-

1.97 (hand relations, whose ceiling gain over PRED is small). Ceiling 3 seeds, other arms 5.

Joint-Embedding Predictive Architecture

Joint-Embedding Predictive Architecture (JEPA) learns representations by predicting latent features of hidden or future data from observed context.

See all 245 papers

Is Forward Prediction Enough? Physical State Grounding for JEPA World Models
Aug 7, 2026Haodong Yan, Jiaguan Zhu, Mingyuan Jia +12Physical Reasoning World Models for Robotics
1The Hong Kong University of Science and Technology (Guangzhou), Guangzhou, China · 2COCO Matrix
D-JEPA: A Decision-Aligned Latent World Model
Sep 21, 2026Shuaijun Liu, Chengyu Wu, Qifu Wen +5World Models World Model-Based Planning
The Hong Kong University of Science and Technology (Guangzhou) · Boston University · Shanghai Jiao Tong University
Subspace-Decomposed JEPAs: Disentangling Progression and Content in Latent World Models
May 29, 2026Lucas Thil, Jesse Read, Rim Kaddah +1Disentangled Representation Learning Representation Learning
LIX, École Polytechnique Palaiseau, France · IRT SystemX Palaiseau, France · Safran Tech Chateaufort, France

Aug 7, 2026J/K move · Enter open · S save

Is Forward Prediction Enough? Physical State Grounding for JEPA World Models

Haodong Yan, Jiaguan Zhu, Mingyuan Jia, Ruiqing Yin+11

1The Hong Kong University of Science and Technology (Guangzhou), Guangzhou, China · 2COCO Matrix

Physical Reasoning World Models for Robotics World Model-Based Planning Joint-Embedding Predictive Architecture Action-Conditioned World Models State Representation Learning

Open paper AskPDF arXiv

Learning structured and control-relevant latent representations remains a key challenge for world models. Recent JEPA-based world models learn action-conditioned predictive latent dynamics from observation sequences. However, their forward-prediction objectives do not explicitly enforce reliable identifiability of robot-centric physical state from individual latents or state changes from latent pairs, which can limit downstream planning and policy performance. We propose PSG-JEPA, a physically grounded JEPA world model that shapes its latent space with two complementary grounding objectives beyond forward prediction: grounding individual latents in robot proprioceptive state, and grounding latent pairs in multi-horizon joint-angle changes. Both objectives are applied only during training, leaving the inference architecture and computational cost unchanged. To comprehensively evaluate PSG-JEPA, we conduct experiments at three levels: (1) latent identifiability via probing, (2) goal-conditioned planning on frozen latents, and (3) policy learning in simulation and on a real robot. Experiments demonstrate that our PSG-JEPA consistently outperforms state-of-the-art latent world-model baselines at all three levels.

Abstract

Joint-Embedding Predictive Architecture

Explore similar work

Is Forward Prediction Enough? Physical State Grounding for JEPA World Models

D-JEPA: A Decision-Aligned Latent World Model

Subspace-Decomposed JEPAs: Disentangling Progression and Content in Latent World Models