Do-JEPA: From Masking to Intervention in Latent World Models

Hossein Resani; Javen Qinfeng Shi

cs.CVSep 29, 2026

Do-JEPA: From Masking to Intervention in Latent World Models

Authors: Hossein Resani, Javen Qinfeng Shi

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

Abstract

Latent world models are trained to predict what happens next, so nothing in their objective separates what an action caused from what merely co-occurred with it. Object-masking models such as C-JEPA intervene on what the predictor can see; we intervene on what physically happens. From one saved simulator state we run the dynamics under an action $a$ and under a reference action $a_{\varnothing}$ , and train the model to predict the difference $Δz=z^{a}-z^{a_{\varnothing}}$ between the two latent futures. The resulting objective, Do-JEPA, has an effect loss, a support loss (where the action enters), a propagation loss (where its effect travels) and invariance losses (what must not change). In a synthetic system with object-aligned variables, support supervision finds the directly intervened object in 99.95% of test cases, where a sparse action mask sends the action to a nuisance slot in every case, and response-onset supervision recovers the ring-shaped propagation graph (edge AUROC 0.975 vs. 0.624). From pixels, the effect loss beats a control trained on exactly the same data: it lowers latent effect error by 28.4% on an end-to-end LeWM model and physical effect error by 13.5% when trained and tested on natural action sequences, and on three independently generated CausalWorld benchmarks it lowers responsive effect error by about 20% under physics shifts and the latent context sensitivity of predicted effects by 66%. Trained from scratch it costs factual accuracy; fine-tuning an existing model with it removes this cost. Together, these results show that intervening on the world, rather than on what the model sees, helps latent world models predict what their actions cause.

Figures & tables

Figure 1: Masking versus intervention. Both methods see the same history, slots and action; orange marks where each intervenes. Left: C-JEPA ( Nam et al., 2026 ) hides an object’s slots and predicts them from the others, with targets from the same logged rollout: it changes what the model sees. Right: Do-JEPA runs one saved simulator state forward under $a$ and under $a_{\varnothing}$ and supervises the difference $\Delta z$ of the two latent futures (Eq. 3 ): it changes what happens. Frames: Push-T.

Figure 2: Do-JEPA training. Top: the simulator runs the saved state $s_{t}$ under $do(a)$ and $do(a_{\varnothing})$ ; the encoder $E$ (frozen in the slot models shown here) gives the target $\Delta z$ . Bottom: the world model encodes the observation into slots $z_{t}$ , routes the action into a slot through the mask $m$ and propagates its effect through the gates $G_{ji}$ ; $\widehat{\Delta z}$ is the difference of its two predictions. $\mathcal{L}_{\Delta}$ : what the action caused; $\mathcal{L}_{\mathrm{sup}}$ : where it entered; $\mathcal{L}_{\mathrm{edge}}$ : how it travelled; $\mathcal{L}_{\mathrm{inv}},\mathcal{L}_{\mathrm{ctx}}$ : what must not change.

	Top-1 all [obj.] $\uparrow$	Target $F_{1}$ $\uparrow$	Nuis. mask $\downarrow$	Nuis. effect $\downarrow$	Effect MSE $\downarrow$	Pred. MSE $\downarrow$
Mask (global action)	—	—	—	$\underline{6.3\times 10^{-3}}$	$1.5\times 10^{-3}$	$\underline{3.14\times 10^{-3}}$
Mask $+\,\mathcal{L}_{\Delta}$	—	—	—	$\mathbf{2.9\times 10^{-3}}$	$\mathbf{3.5\times 10^{-4}}$	$\mathbf{2.67\times 10^{-3}}$
Sparse mask	$0.000$ [ $0.243$ ]	$0.000$	$1.000$	$2.5\times 10^{-2}$	$1.1\times 10^{-2}$	$8.70\times 10^{-3}$
Sparse mask $+\,\mathcal{L}_{\Delta}$	$0.000$ [ $0.234$ ]	$0.000$	$1.000$	$1.6\times 10^{-2}$	$1.0\times 10^{-2}$	$8.55\times 10^{-3}$
Sparse mask $+\,\mathcal{L}_{\Delta}+\mathcal{L}_{\mathrm{sup}}$	$\mathbf{0.9995}$ [ $0.9995$ ]	$\mathbf{0.9998}$	$\mathbf{8.2\times 10^{-5}}$	$9.1\times 10^{-3}$	$\underline{9.3\times 10^{-4}}$	$3.43\times 10^{-3}$

Table 1: Where does the action enter? Synthetic SCM, reversed-correlation test split, seed 7. Top-1 over all seven slots (chance

1/7

) and, in brackets, over the four objects (chance

0.25

); nuis. mask: mask mass on nuisance slots; nuis. effect: mean RMS of the predicted effect per nuisance slot. All rows are evaluated with one history slot masked, as in Mask training; mask temperature is

1.0

in the sparse-mask rows and

0.7

in the last (App. E.1 ). Best in bold, second best underlined.

Figure 3: Effect loss on LeWM, pixel Push-T. Aug and Paired see the same paired images; Paired adds

\mathcal{L}_{\Delta}

. Latent metrics; means over three fresh seeds (dots: seeds); red: relative change (absolute for the cosine, taken over all pairs). Panels 1–3: held-out pairs from the training distribution (half responsive); panel 4: outcome-unfiltered holdout (

6.5\%

responsive).

OOD split	Resp. effect RMSE	Effect cosine	Latent ctx. shift	Factual RMSE	Response AUROC	Signature
	Benefit			Cost
Visual	$3.9\%$ lower [2/3]	$+0.018$ [3/3]	$66.2\%$ lower [3/3]	$52.8\%$ higher [3/3]	$-0.054$ [3/3]	$2/3$
Mechanism ∗	$20.3\%$ lower [3/3]	$+0.030$ [2/3]	$65.4\%$ lower [3/3]	$51.5\%$ higher [3/3]	$-0.070$ [3/3]	$2/3$
Composed	$19.5\%$ lower [3/3]	$+0.020$ [2/3]	$66.2\%$ lower [3/3]	$51.1\%$ higher [3/3]	$-0.073$ [3/3]	$2/3$

Table 2: CausalWorld: Paired vs. Aug on three independently generated benchmarks (dataset seeds 17/27/37, model seed 107,

40

k updates), locked OOD splits. Mean relative changes; brackets count instances that move in the stated direction; “Signature” counts instances with the full preregistered benefit and cost. ∗ The mechanism split also decorrelates the visual context (App. C.5 ).

	Factual RMSE		Resp. effect RMSE		Response	Latent ctx.	Factual
Regime	IID	OOD	IID	Physics	AUROC	shift	within $10\%$
From scratch, $\alpha=1$	$+51.9\%$	$+51.8\%$	$-9.0\%$	$-19.9\%$	$-0.066$	$-66.0\%$	$0/3$
From scratch, $\alpha=0.25$	$+10.3\%$	$+10.6\%$	$-9.1\%$	$-12.8\%$	$-0.016$	$-33.6\%$	$1/3$
Fine-tuned from Aug , $\alpha=1$	$-0.6\%$	$-0.9\%$	$-8.2\%$	$-13.3\%$	$+0.008$	$-26.2\%$	$3/3$

Table 3: When does the effect loss cost factual accuracy? CausalWorld, Paired against the Aug model of the same regime, mean relative change over the three benchmark instances. OOD: mean over the three OOD splits; physics: mechanism and composed splits. Last column: instances whose factual error rises by at most

10\%

on every split. The last two rows are post-hoc diagnostics (App. E.6 ).

Appendix figures & tables24 assets

Supplementary material from the paper’s appendix.

Appendix

Finding	Evidence	Where
Support supervision finds where the action enters	direct target found in $99.95\%$ of test cases; a sparse mask picks a nuisance slot in every case	Table 1
Onset supervision recovers where the effect travels	edge AUROC $0.975$ vs. $0.624$ ; the four ring edges recovered exactly on $3/3$ seeds	Table 12
Gates confine effects to the right objects	state Push-T edge AUROC $0.994$ vs. $0.500$ ; predicted effect on nuisances $1.9\times 10^{-3}\to 2.1\times 10^{-7}$	Table 14
Context twins remove a nuisance shortcut	on the two seeds not used to pick the weight: OOD error $44$ – $55\%$ lower, $72$ – $85\%$ of the IID-to-OOD gap closed, IID cost $\leq 4.2\%$	App. E.3
Intervention beats masking on C-JEPA’s own slots	all six metrics, $3/3$ seeds: factual error $30\%$ lower, AUROC $+0.21$ , pixel context shift $31\%$ lower	Table 15
Effect loss helps from pixels (LeWM)	latent effect error $28.4\%$ lower ( $3/3$ ); physical effect error $13.5\%$ lower when trained and tested on natural sequences ( $3/3$ )	Fig. 3 , Table 18

Appendix

Table 4: Main results at a glance. Each row gives the strongest evidence for one finding and where it is reported. Seeds are model seeds; “instances” are independently generated benchmarks.

Symbol	Meaning	Shape	Role	When
$s_{t}$	physical simulator state	env.-specific	observed	build/eval
$a,\ a_{\varnothing}$	action; reference action (a physical no-op unless stated)	$\mathcal{A}$	input	train, test
$c,\ c^{\prime}$	context (nuisance) value and its twin	env.-specific	input	train (twin), test ( $c$ )
$u$	confounder: affects action choice and outcome	scalar	hidden	build
$\varepsilon_{t}$	exogenous simulator noise, shared by both branches	—	hidden	build
$o_{t}$	observation $g(s_{t},c)$ (image or state vector)	env.-specific	observed	train, test

Appendix

Table 5: Notation.

K

: number of slots;

d

: slot dimension;

H

: history length;

k

: rollout horizon;

N

: batch size.

Loss	Question	Target	Source of the target
$\mathcal{L}_{\mathrm{pred}}$	what happens next?	$z^{a}_{t+1}$ , $z^{a_{\varnothing}}_{t+1}$	observed next frames
$\mathcal{L}_{\Delta}$	what did the action change?	$\Delta z=z^{a}-z^{a_{\varnothing}}$	paired rollouts, Eq. 3
$\mathcal{L}_{\mathrm{sup}}$	where did it enter?	$y^{\mathrm{sup}}_{i}=e_{i}/\sum_{i^{\prime}}e_{i^{\prime}}$	first step of the pair
$\mathcal{L}_{\mathrm{edge}}$	where did it travel?	$y^{\mathrm{edge}}_{ji}\in\{0,1\}$	onset order in the pair
$\mathcal{L}_{\mathrm{inv}}$	what must not change?	$\widehat{\Delta z}^{\,j}=0$ for $j\notin R$	response set of the pair
$\mathcal{L}_{\mathrm{ctx}}$	what must ignore context?	$\widehat{\Delta z}(c)=\widehat{\Delta z}(c^{\prime})$	context twin

Appendix

Table 6: Losses and their targets. Every target is computed from paired rollouts or context twins at training time.

e_{i}=\lVert\Delta z^{\,i}_{t+1}\rVert^{2}

is the immediate response energy of slot

i

;

R

is the response set; “onset order” is defined in Sec. 3.2 .

Information	Training	Test	Build/eval only
Observation history, action $a$	input	input	—
Reference branch $o^{a_{\varnothing}}$	input ( Aug ), target ( $\mathcal{L}_{\Delta}$ )	—	—
Context twin $o^{\prime}$	input ( Aug ), target ( $\mathcal{L}_{\mathrm{ctx}}$ )	—	—
Support, onset and response labels	target (from the pair)	—	—
Object count and slot alignment	architecture (slot models)	architecture	synthetic metrics
Ground-truth graph, object identities	—	—	metrics

Appendix

Table 7: Where information enters. “Input” means the model receives it; “target” means a loss uses it; “build/eval” means it is used only to construct benchmarks or to compute metrics. Test-time model inputs are the observation history and the action only.

Setting	Model input	Intervention vs. reference	Shift at test	Terms used	Train/val/test	Seeds
Synthetic SCM	$7$ slots, $d{=}16$	impulse at a noisy contact point vs. none	reversed action–nuisance corr.	$\mathcal{L}_{\Delta},\mathcal{L}_{\mathrm{sup}},\mathcal{L}_{\mathrm{edge}},\mathcal{L}_{\mathrm{inv}}$	$5$ k/ $1$ k/ $2$ k	$1$ – $3$
Temporal SCM	$9$ slots, $d{=}18$	impulse vs. none	observation stride $\kappa$	$\mathcal{L}_{\Delta},\mathcal{L}_{\mathrm{sup}},\mathcal{L}_{\mathrm{edge}},\mathcal{L}_{\mathrm{inv}}$	$5$ k/ $1.2$ k/ $2$ k	$3$
State Push-T	state + $3$ nuisance slots	action vs. no-op	reversed corr. ( $\rho{=}0.95$ )	$\mathcal{L}_{\Delta},\mathcal{L}_{\mathrm{edge}},\mathcal{L}_{\mathrm{inv}},\mathcal{L}_{\mathrm{ctx}}$	$20$ k/ $2.5$ k/ $2.5$ k	$3$
Vision Push-T	$4$ frozen slots, $d{=}128$	action vs. no-op	reversed pattern corr. ( $\rho{=}0.95$ )	$\mathcal{L}_{\Delta},\mathcal{L}_{\mathrm{inv}},\mathcal{L}_{\mathrm{ctx}}$	$12$ k/ $2$ k/ $2.5$ k	$3$
LeWM Push-T	$224^{2}$ RGB, $K{=}1$ , $d{=}192$	$5$ -step macro vs. zero macro	none; mass, visual (planning)	$\mathcal{L}_{\Delta}$	$9$ k/ $1$ k pairs	$3$
CausalWorld	$64^{2}$ RGB, $K{=}1$ , $d{=}192$	first joint command vs. mirrored one	visual, mechanism, composed	$\mathcal{L}_{\Delta}$ ( $+\mathcal{L}_{\mathrm{ctx}}$ )	$6$ k/ $1$ k/ $1.5$ k per split	$3$ model, $3$ data

Appendix

Table 8: Settings. Terms are the losses of Eq. 10 that are switched on. Seeds are model seeds unless stated.

Figure 4: Nuisances and shifts in each benchmark. In every benchmark the nuisance is correlated with the outcome during training and the correlation is reversed (or, for mechanism shifts, the physics is changed) at test time. The physics of the two members of a context twin is identical by construction.

Corpus	States	$n$	Resp.	Action source
Enriched (E) training	episode replay	$10{,}000$	$50.0\%$	balanced toward/away proposal
Outcome-unfiltered holdout	episode replay	$5{,}000$	$5.9\%$	i.i.d. from expert action marginal
Fresh outcome-unfiltered holdout	episode replay	$5{,}000$	$6.5\%$	i.i.d. from expert action marginal
Expert contiguous	expert prefix	$2{,}500$	$50.8\%$	contiguous expert macro
Expert shuffled	expert prefix	$2{,}500$	$50.8\%$	same macro, order shuffled
Expert i.i.d. marginal	expert prefix	$2{,}500$	$36.2\%$	i.i.d. from expert marginal

Appendix

Table 9: Push-T intervention corpora for the LeWM studies. All use exact same-state branching with audited replay. “Resp.” is the fraction of responsive pairs.

Setting	Settings
Hard SCM	Transformer predictor (width $64$ , depth $2$ , $4$ heads), history $3$ , $2$ propagation steps ( $\gamma=0.55$ ); $25$ epochs, batch $128$ , lr $3\times 10^{-4}$ , wd $10^{-4}$ ; $\lambda_{\Delta}=5$ , $\lambda_{\mathrm{sup}}=2$ , mask temperature $1.0$ (sparse-mask rows without $\mathcal{L}_{\mathrm{sup}}$ ) or $0.7$ (with $\mathcal{L}_{\mathrm{sup}}$ , and all gated models), entropy regulariser $0.02$ (a cardinality term, also $0.02$ , is zero under the softmax mask), no-op mask penalty $0.05$ ; every model is trained and evaluated with one random history slot masked (masked-history loss $0.25$ ); edge runs add $\lambda_{\mathrm{inv}}=5$ (threshold $0.05$ ), $\lambda_{\mathrm{edge}}=5$ , gate invariance $5$ , gate $\ell_{1}$ $0.02$ , onset threshold $0.08$ , gate temperature $0.4$ . Sel.: lowest IID prediction MSE. Seeds $7$ (ablations), $7/17/27$ (gated model).
Easy SCM	Transformer predictor (width $64$ , depth $2$ , $4$ heads); $25$ epochs, batch $128$ , lr $3\times 10^{-4}$ , wd $10^{-4}$ ; $\lambda_{\Delta}=5$ , $\lambda_{\mathrm{sup}}=2$ , mask temperature $0.7$ , entropy regulariser $0.02$ ; $5{,}000$ / $1{,}000$ / $2{,}000$ samples. Sel.: lowest IID prediction MSE. Seeds $7/17/27/37/47$ .
Temporal SCM	Width $64$ , depth $2$ , $4$ propagation steps; $25$ epochs, batch $128$ , lr $3\times 10^{-4}$ , wd $10^{-4}$ ; $\lambda_{\Delta}=5$ , $\lambda_{\mathrm{sup}}=2$ , $\lambda_{\mathrm{inv}}=5$ , $\lambda_{\mathrm{edge}}=5$ . Sel.: lowest IID prediction MSE. Seeds $7/17/27$ .
State Push-T	Width $96$ , depth $3$ , $4$ heads; gate temperature $0.7$ ; $30$ epochs, batch $256$ , lr $3\times 10^{-4}$ , wd $10^{-4}$ ; reference-branch weight $0.5$ , $\lambda_{\Delta}=5$ , reward head $1$ , $\lambda_{\mathrm{inv}}=5$ (threshold $0.03$ ), $\lambda_{\mathrm{edge}}=2$ , gate invariance $2$ , gate $\ell_{1}$ $0.01$ , $\lambda_{\mathrm{ctx}}=0.25$ (context twins by shuffling nuisances within a batch; evaluated by sign flip). Sel.: IID prediction and effect error (context not used). Seeds $7/17/27$ .
Vision Push-T	Frozen slots; transformer predictor (width $256$ , depth $6$ , $8$ heads); $40$ epochs, batch $512$ , lr $3\times 10^{-4}$ , wd $10^{-5}$ ; $\lambda_{\Delta}=4$ , $\lambda_{\mathrm{inv}}=1$ (soft weights, temperature $0.35$ ), context-twin prediction weight $0.5$ , history loss $1$ , $\lambda_{\mathrm{ctx}}=0.5$ (sensitivity run $0.25$ ). Sel.: IID validation. Seeds $7/17/27$ .
LeWM Push-T ( Paired vs. Aug )	Released LeWM checkpoint (ViT-tiny/14 encoder, $192$ -D latent, history $3$ , $6$ -layer predictor); $10$ epochs, batch $64$ , lr $10^{-5}$ , wd $10^{-3}$ , SIGReg $0.09$ , reference-branch weight $1.0$ , $\lambda_{\Delta}=0.5$ , gradient clipping $1.0$ , $10\%$ validation split, expert replay in every update. Seeds $7$ (development), $17/27/37$ (confirmation). Sel.: lowest validation paired-prediction loss, the same rule for Aug and Paired (no effect or planning metric).

Appendix

Table 10: Training settings. “Sel.” is the checkpoint-selection rule; all selection uses IID validation data. lr: learning rate; wd: weight decay.

	Pred. MSE $\downarrow$	Effect MSE $\downarrow$	Nuis. effect $\downarrow$	Nuis. mask $\downarrow$
Mask (global action)	$0.00612$	$0.00308$	$0.0875$	—
Mask $+\,\mathcal{L}_{\Delta}$	$0.00280$	$\mathbf{0.000084}$	$0.0110$	—
Sparse mask	$0.00446$	$0.00258$	$0.00606$	$0.0257$
Sparse mask $+\,\mathcal{L}_{\Delta}$	$\underline{0.00279}$	$0.000261$	$\underline{0.00240}$	$\underline{0.0112}$
Sparse mask $+\,\mathcal{L}_{\Delta}+\mathcal{L}_{\mathrm{sup}}$	$\mathbf{0.00276}$	$\underline{0.000232}$	$\mathbf{0.000039}$	$\mathbf{0.000197}$

Appendix

Table 11: Easy SCM , reversed-correlation test split, mean over five seeds. Every mask-based variant reaches object-only target top-1 of

1.0

. “Nuis. effect” is the RMS of the predicted effect on the nuisance slot. Mask with

\mathcal{L}_{\Delta}

has the lowest effect MSE but

282\times

more predicted nuisance effect than the full model. Best in bold, second best underlined.

	Edge AUROC $\uparrow$	True-edge gate $\uparrow$	Off-path gate $\downarrow$	Nuis. in-gate $\downarrow$	Nuis. effect $\downarrow$	Direct-effect MSE $\downarrow$
Gates without $\mathcal{L}_{\mathrm{edge}}$	$0.624$	$0.974$	$0.465$	$1.0\times 10^{-3}$	$5.6\times 10^{-5}$	$0.0129$
$+\,\mathcal{L}_{\mathrm{edge}}$ (onset labels)	$0.975$	$0.955$	$0.053$	$2.7\times 10^{-4}$	$3.9\times 10^{-5}$	$0.0134$
$+$ gate term of $\mathcal{L}_{\mathrm{inv}}$	$0.975$	$0.947$	$0.052$	$1.0\times 10^{-4}$	$2.0\times 10^{-5}$	$0.0135$
same, 3 seeds	$0.9751{\pm}0.0002$	$0.950$	$0.051$	$8.9\times 10^{-5}$	$1.9\times 10^{-5}$	$0.0137$

Appendix

Table 12: Where does the effect travel? Synthetic SCM, reversed-correlation test split; seed 7 except the last row (mean over seeds 7/17/27). Every row also uses

\mathcal{L}_{\Delta}+\mathcal{L}_{\mathrm{sup}}

and the effect part of

\mathcal{L}_{\mathrm{inv}}

. Edge AUROC and the true-edge gate are scored against the test pairs’ onset labels. “Off-path gate”: mean gate on slot pairs that are not on the rollout’s propagation path (this includes correctly learned but unused ring edges); “nuis. in-gate”: mean gate into nuisance slots; “nuis. effect”: as in Table 1 .

$\kappa$	$1$	$2$	$3$	$4$
Direct-target $F_{1}$	$0.997$	$0.997$	$0.997$	$0.997$
Descendant-effect $F_{1}$	$0.947$	$0.889$	$0.889$	$0.889$
One-hop structural $F_{1}$	$\mathbf{0.999}$	$0.549$	$0.502$	$0.337$
One-hop structural AUROC	$1.000$	$0.910$	$0.884$	$0.843$
Onset-label $F_{1}$	$1.00$	$0.40$	$0.40$	$0.25$
Schedule-aware $F_{1}$	$0.999$	$0.935$	$0.992$	$0.894$

Appendix

Table 13: Temporal SCM , reversed-correlation test split, mean over seeds 7/17/27.

\kappa

causal stages are hidden in one observed step. “Onset-label

F_{1}

” scores the training labels themselves as one-hop edges. The schedule-aware relation is reachability within the observation interval.

	Pred. MSE $\downarrow$	Agent pos. $\downarrow$	Nuis. effect $\downarrow$	Edge AUROC $\uparrow$	Nuis. in-gate $\downarrow$
Obs (global action)	$2.77\times 10^{-4}$	$5.49$	$1.9\times 10^{-3}$	—	—
Global action, $+\mathcal{L}_{\Delta}+\mathcal{L}_{\mathrm{inv}}$	$2.78\times 10^{-4}$	$3.16$	$4.8\times 10^{-4}$	—	—
Action routed to agent, no gates	$\underline{2.75\times 10^{-4}}$	$\mathbf{0.93}$	$\underline{2.2\times 10^{-4}}$	$0.500$	$1.000$
Routed, $+$ gates, $\mathcal{L}_{\mathrm{edge}}$ , gate invariance	$\mathbf{2.72\times 10^{-4}}$	$\underline{0.95}$	$\mathbf{2.1\times 10^{-7}}$	$\mathbf{0.994}$	$\mathbf{3.1\times 10^{-5}}$

Appendix

Table 14: State-based Push-T , reversed-correlation split, mean of seeds 7/17/27. “Agent pos.” is the agent-position MSE in px 2 ; “nuis. effect” is the mean per-slot RMS of the predicted effect on nuisance slots; edge AUROC is scored against the test pairs’ onset labels. The agent is the known direct target in all rows. Arrows give the better direction. Best in bold, second best underlined.

Figure 5: Context weight in state Push-T , strong confound. Left: IID and OOD prediction error against

\lambda_{\mathrm{ctx}}

(seed 7); the circled point (both panels) is

\lambda_{\mathrm{ctx}}=0.25

, chosen as the point on the IID–OOD trade-off curve closest to the ideal, so this seed’s OOD split informed the choice. Right: the context-induced block displacement falls monotonically with

\lambda_{\mathrm{ctx}}

	Factual (px) $\downarrow$	Resp. effect (px) $\downarrow$	Cosine $\uparrow$	AUROC $\uparrow$	Latent ctx. $\downarrow$	Phys. ctx. (px) $\downarrow$
Mask (C-JEPA-style)	$39.02\pm 1.24$	$20.25\pm 0.96$	$0.600\pm 0.108$	$0.485\pm 0.020$	$\underline{1.003\pm 0.054}$	$19.91\pm 1.50$
Aug	$\mathbf{24.88\pm 0.58}$	$\mathbf{16.79\pm 0.11}$	$\mathbf{0.763\pm 0.010}$	$\underline{0.666\pm 0.008}$	$1.024\pm 0.014$	$\mathbf{12.59\pm 0.51}$
Do-JEPA	$\underline{27.19\pm 0.16}$	$\underline{17.89\pm 0.20}$	$\underline{0.744\pm 0.013}$	$\mathbf{0.691\pm 0.008}$	$\mathbf{0.819\pm 0.019}$	$\underline{13.74\pm 0.68}$

Appendix

Table 15: Vision Push-T , frozen VideoSAUR slots, reversed-pattern OOD split, mean

\pm

sd over seeds 7/17/27. Block effects are decoded to pixels by a context-balanced probe trained on IID data. “Latent ctx.” is the context shift of predicted effects in slot space; “Phys. ctx.” is the same shift decoded to pixels. Arrows give the better direction. Best in bold, second best underlined.

	Factual (px) $\downarrow$	Resp. eff. (px) $\downarrow$	Cos. $\uparrow$	AUROC $\uparrow$	Lat. ctx. $\downarrow$	Phys. ctx. (px) $\downarrow$
Object-centric JEPA, no mask	$36.68$	$21.13$	$0.560$	$0.481$	$1.021$	$20.45$
Mask	$40.41$	$21.30$	$0.479$	$0.465$	$\underline{1.008}$	$21.22$
Paired pred.	$33.23$	$20.51$	$0.620$	$0.529$	$1.027$	$17.07$
Paired pred. $+$ ctx. twins ( Aug )	$\mathbf{24.31}$	$\mathbf{16.83}$	$\mathbf{0.756}$	$\underline{0.672}$	$1.009$	$\mathbf{12.30}$
Paired	$37.15$	$20.62$	$0.597$	$0.537$	$1.066$	$17.43$
Do-JEPA ( $\lambda_{\mathrm{ctx}}=2$ )	$\underline{26.78}$	$\underline{18.68}$	$\underline{0.727}$	$\mathbf{0.703}$	$\mathbf{0.515}$	$\underline{13.71}$

Appendix

Table 16: Vision Push-T screen (seed 7, OOD). “Paired pred.” predicts both branches as separate targets (an Aug without context twins); “+ ctx. twins” adds context-swapped copies as data ( Aug ). Do-JEPA here uses

\lambda_{\mathrm{ctx}}=2

. The true latent effect is itself context-sensitive (target context shift

2.13

). Best in bold, second best underlined.

Figure 6: Context weight in vision Push-T (seed 7, OOD). Do-JEPA across

\lambda_{\mathrm{ctx}}

(orange) against C-JEPA-style masking (dashed) and Aug (dotted) on the same split. Do-JEPA beats masking on every metric at every weight, and a larger weight halves the latent context shift. The circled weight was selected on IID data as the smallest one with at least

20\%

lower latent context shift and at most

5\%

higher effect RMSE,

0.02

lower AUROC and

0.03

lower cosine; OOD values are shown.

	Held-out training-distribution pairs			Outcome-unfiltered holdout
Seed	Effect RMSE $\downarrow$	Cosine $\uparrow$	Factual MSE	Effect RMSE $\downarrow$	Cosine $\uparrow$	Factual MSE
$17$	$0.2123\to 0.1507$	$0.669\to 0.738$	$-40.9\%$	$0.1031\to 0.0914$	$0.177\to 0.149$	$-17.7\%$
$27$	$0.2118\to 0.1512$	$0.683\to 0.744$	$-38.4\%$	$0.1020\to 0.0906$	$0.169\to 0.145$	$-17.3\%$
$37$	$0.2160\to 0.1565$	$0.663\to 0.730$	$-39.9\%$	$0.1020\to 0.0907$	$0.173\to 0.145$	$-17.3\%$
Mean change	$-28.4\%$	$+0.066$	$-39.7\%$	$-11.2\%$	$-0.027$	$-17.5\%$

Appendix

Table 17: Effect loss on LeWM, pixel Push-T, per seed ( Aug

\to

Paired ). Left: held-out pairs from the training distribution (half responsive). Right: fresh outcome-unfiltered holdout (

5{,}000

pairs,

323

responsive). All metrics are latent, in each model’s own latent space; cosines are over all pairs. Factual columns give the relative change in factual MSE.

Trained on	Tested on	Latent effect RMSE (%)	Physical effect RMSE (%)	Physical cosine
Natural	natural test ( $2{,}500$ )	$-18.1/-18.4/-18.2$	$-13.6/-13.6/-13.3$	$0.754\to 0.798$
Natural	unfiltered holdout ( $323$ )	$-8.6/-8.1/-8.2$	$-5.7/-5.6/-5.3$	$0.441\to 0.547$
Enriched	natural test ( $2{,}500$ )	$-3.6/-3.1/-3.2$	$+5.8/+5.8/+6.1$	$0.354\to 0.241$
Enriched	unfiltered holdout ( $323$ )	$-1.0/-0.2/-0.3$	$+3.7/+5.3/+5.2$	$0.168\to 0.101$

Appendix

Table 18: The LeWM Push-T gain follows the training intervention distribution. Paired vs. Aug trained with the same recipe on the enriched or the natural corpus; per-seed relative change (seeds 17/27/37) and mean responsive physical cosine ( Aug

\to

Paired ). Latent effect RMSE is over all pairs, physical error over responsive pairs.

	Held-out training-distribution pairs				Outcome-unfiltered holdout
Seed	Effect RMSE $\downarrow$	Cosine $\uparrow$	Zero	Decoded true	Effect RMSE $\downarrow$	Cosine $\uparrow$
$17$	$11.75\to 11.03$	$-0.188\to 0.240$	$10.96$	$7.49$ / $7.24$	$13.21\to 13.56$	$0.206\to 0.124$
$27$	$11.73\to 10.67$	$-0.198\to 0.237$	$10.67$	$6.23$ / $6.02$	$13.29\to 13.71$	$0.177\to 0.112$
$37$	$13.78\to 13.00$	$-0.170\to 0.212$	$12.63$	$9.54$ / $9.33$	$13.38\to 13.77$	$0.160\to 0.084$

Appendix

Table 19: Physical read-out of the LeWM Push-T effect gain ( Aug

\to

Paired , per seed; responsive pairs; px). “Zero”: error of predicting no displacement; “decoded true”: the read-out applied to the true latent effects ( Aug / Paired ).

Hypothesis	Split	Resp. effect RMSE	Cosine	AUROC	Factual RMSE	Ctx. shift	Seeds	Result
H1: Paired vs. Aug	visual	$11.3\%$ lower	$+0.056$	$-0.019$	$17.2\%$ higher	—	$1/3$	not met
	mechanism	$22.0\%$ lower	$+0.067$	$-0.022$	$15.1\%$ higher	—	$1/3$	not met
	composed	$24.3\%$ lower	$+0.059$	$-0.012$	$17.8\%$ higher	—	$1/3$	not met
H2: Do-JEPA vs. Paired	visual	$15.2\%$ higher	$-0.130$	$-0.019$	$33.3\%$ higher	$47.2\%$ lower	$0/3$	not met
	mechanism	$7.1\%$ higher	$-0.113$	$-0.021$	$31.7\%$ higher	$46.1\%$ lower	$0/3$	not met
	composed	$8.8\%$ higher	$-0.125$	$-0.031$	$32.3\%$ higher	$46.6\%$ lower	$0/3$	not met

Appendix

Table 20: CausalWorld fresh-seed confirmation (model seeds 47/57/67,

40

k updates,

\beta=0.25

, frozen selector, OOD opened once). Relative changes of the candidate against the reference, mean over seeds. The rule for each split: at least one benefit (AUROC

+0.02

, cosine

+0.05

10\%

lower responsive effect error) with at most

10\%

higher factual error, on at least

2/3

seeds.

Split	Factual RMSE (%)	Resp. effect RMSE (%)	Effect cosine	Response AUROC	Latent ctx. shift (%)
IID	$-0.5/+1.7/-3.0$	$-6.1/-8.9/-9.5$	$+0.044/+0.051/+0.091$	$-0.006/-0.004/-0.001$	$-24.2/-24.7/-28.3$
Visual	$-0.2/+2.0/-4.4$	$-7.5/-8.7/-10.6$	$+0.060/+0.050/+0.080$	$+0.009/+0.006/+0.010$	$-24.8/-24.4/-28.1$
Mechanism ∗	$+0.1/+1.6/-4.0$	$-11.5/-12.5/-16.2$	$+0.042/+0.041/+0.085$	$+0.003/+0.006/+0.004$	$-25.7/-25.4/-28.4$
Composed	$-0.8/+1.9/-4.2$	$-11.4/-12.7/-15.7$	$+0.042/+0.064/+0.079$	$+0.011/+0.015/+0.011$	$-25.5/-24.9/-28.5$

Appendix

Table 21: CausalWorld, fine-tuning a trained Aug model with the effect loss (

2{,}000

updates; Paired vs. Aug fine-tuned identically). Values are per benchmark instance (dataset seeds 17/27/37): relative changes in % for factual error, responsive effect error and latent context shift; differences for cosine and AUROC. ∗ Also decorrelates the visual context.

Configuration	Model	Adapted	Effect RMSE $\downarrow$	Cosine $\uparrow$	Planning $\uparrow$
Base LeWM	—	—	—	—	$94\%$
Full fine-tuning	Aug	yes	$0.2145$	$0.659$	$58\%$
	Paired	yes	$\mathbf{0.1522}$	$\mathbf{0.736}$	$54\%$
Expert data only	—	yes	$0.6588$	$0.428$	$\mathbf{94\%}$
Frozen representation	Aug	yes	$0.2998$	$0.810$	$36\%$
	Paired	yes	$0.2889$	$0.863$	$36\%$

Appendix

Table 22: Ways of adding paired supervision to LeWM (seed 7,

50

-episode planning; effect metrics on each configuration’s own validation pairs, so compare within a row group only). “Adapted” says whether the selected checkpoint differs from the base model.

Property	Measured	Check
Prop. 1 : no-op correction	$0.0$ (bitwise)	pass
Prop. 2 : linearity	max. error $\leq 2\times 10^{-7}$	pass
Prop. 2 : operator bound	holds for all pairs; mean tightness $0.037$	pass
Prop. 3 : sampled $\hat{L}$	$0.714$ (lower bound, not certified)	—
Spectral / Frobenius norm of $W_{\phi}$	$0.857$ / $1.203$	—
Rank of $W_{\phi}$ ( $192\times 2048$ )	$\leq 16$ by construction	—

Appendix

Table 23: Numerical checks of the adapter properties (enriched-trained Paired adapter, seed 7,

256

examples).

	Seed 17	Seed 27	Seed 37
Effect RMSE, Aug	$0.5835$	$0.5956$	$0.6043$
Effect RMSE, Paired	$\mathbf{0.5496}$	$\mathbf{0.5640}$	$\mathbf{0.5726}$
Relative reduction	$5.80\%$	$5.31\%$	$5.25\%$
Rollout ratio, Paired	$1.032$	$1.036$	$1.033$
No-op and representation unchanged	yes	yes	yes
Planning, Aug	$81.5\%$	$84.5\%$	$84.5\%$

Appendix

Table 24: Centered adapter, three fresh seeds. Effect RMSE on held-out enriched pairs; planning over

200

episodes (base LeWM:

83.0\%

Explore similar work

PhyLatent: Learning Dynamics-Relevant Representations for JEPA World Models
Aug 6, 2026Xi Zeng, Haojie Ren, Ziying Song +2Latent World Models Joint-Embedding Predictive Architectures
School of Mechanical and Aerospace Engineering, Nanyang Technological University, Singapore · The University of Sheﬃeld, Sheﬃeld, United Kingdom · School of Artiﬁcial Intelligence (School of Software), Yanshan University, Qinhuangdao, China
D-JEPA: A Decision-Aligned Latent World Model
Sep 21, 2026Shuaijun Liu, Chengyu Wu, Qifu Wen +5Latent World Models World Models
The Hong Kong University of Science and Technology (Guangzhou) · Boston University · Shanghai Jiao Tong University
Is Forward Prediction Enough? Physical State Grounding for JEPA World Models
Aug 7, 2026Haodong Yan, Jiaguan Zhu, Mingyuan Jia +12Latent World Models World Models
1The Hong Kong University of Science and Technology (Guangzhou), Guangzhou, China · 2COCO Matrix

Aug 6, 2026, cs.CVJ/K move · Enter open · S save

PhyLatent: Learning Dynamics-Relevant Representations for JEPA World Models

Xi Zeng, Haojie Ren, Ziying Song, Yuanbo Nie+1

School of Mechanical and Aerospace Engineering, Nanyang Technological University, Singapore · The University of Sheﬃeld, Sheﬃeld, United Kingdom · School of Artiﬁcial Intelligence (School of Software), Yanshan University, Qinhuangdao, China

Latent World Models Joint-Embedding Predictive Architectures World Models Inference Dynamics Latent Variable

Open paper AskPDF arXiv

We propose PhyLatent, a dynamics-relevant training objective for Joint-Embedding Predictive Architecture (JEPA) world models. Our key observation is that preventing global latent collapse does not necessarily ensure that the learned representation preserves physically meaningful state and action relationships. We identify three failure modes: sensitivity to appearance changes that leave the physical state unchanged, insufficient separation of distinct physical states, and insufficient separation of different action-conditioned futures. We refer to these as Physical Invariance Collapse, Physical Distinguishability Collapse, and Counterfactual Dynamics Collapse, respectively. PhyLatent targets these failures through three coordinated training pathways, implemented with static visual invariance, physical state grounding, future representation alignment, counterfactual branch separation, and latent denoising. On OGBench-Cube, PhyLatent reduces the three collapse rates by 43.9%, 27.9%, and 47.1%, respectively, while improving model predictive control (MPC) success by 12.0 percentage points (17.2% relative). Across four visual-control tasks, average planning success increases by 6.62 percentage points (8.3% relative). These results show that global non-collapse alone is insufficient for learning a reliable JEPA world-model state space, and that explicitly preserving dynamics-relevant structure can improve closed-loop planning.

Abstract

Explore similar work

PhyLatent: Learning Dynamics-Relevant Representations for JEPA World Models

D-JEPA: A Decision-Aligned Latent World Model

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