Counterfactual Action Evaluation, Observation Bottlenecks, and Representation Geometry in Joint-Embedding Predictive World Models

Arjun Subramanian

cs.LGOct 2, 2026

Counterfactual Action Evaluation, Observation Bottlenecks, and Representation Geometry in Joint-Embedding Predictive World Models

Authors: Arjun Subramanian

Organizations: Computer Science and Artificial Intelligence Laboratory (CSAIL) Massachusetts Institute of Technology

Abstract

Low latent prediction error does not establish that a world model distinguishes the consequences of its actions. We introduce an evaluation protocol that traces the same intervention through simulator state, raster observations, target embeddings, and predictor outputs. Exact simulator-state forks in a controlled deformable-physics testbed reveal distinct bottlenecks. Changed commands alter particle motion, yet 41.5% of one-step raster pairs are identical. Observation loss is not the whole explanation: among 579 high-visibility counterfactuals, median predictor-to-target response is 0.0051 and 0.0217 across two seeds, falling to 0.0027 and 0.0116 after variance normalization. An isotropic state perturbation matched to the target counterfactual embedding shift produces 190x and 53x larger predictor changes on the same visible pairs, isolating action-path under-use rather than a dead or globally shrunk predictor. MSE-only training gives 8.36x lower 10-step latent error in matched seeds, but in spectrally concentrated spaces; one VICReg target encoder is also strongly concentrated, so neither error nor rank alone certifies physical state. Finally, stiffness remains near chance even from full-resolution rasters and mechanical state while privileged material parameters decode perfectly, indicating weak identifiability under this excitation rather than encoder discard. These results motivate auditing physical effect, observation visibility, representation geometry, and action dependence separately.

Figures & tables

Figure 1: An action intervention must be measured at multiple levels. From the same physical state, fork recorded and alternative commands. Particle differences need not survive rasterization or encoding. Separately compare the predictor outputs at fixed encoded state. Small latent sensitivity alone cannot locate the bottleneck.

Figure 2: Physical response and observable response differ. Original simulator forks across 200 validation episodes. Panels (a,b) show medians with interquartile bands, not confidence intervals. Panel (c) includes exact raster equality. Horizons 1/3/5 use 3,000/2,600/2,200 pairs; they do not have identical temporal support. Position differences grow, while many rasters remain unchanged.

Figure 3: Visible counterfactuals separate observation loss from action-path under-use. Original simulator forks, 579 high-visibility transitions per seed. (a) Median target and predictor response on the top third of nonzero raster-MSE contrasts. (b) Dividing by each branch’s mean latent variance makes the gap larger, not smaller. (c) Perturbing the encoded state in an isotropic direction with

\ell_{2}

norm matched to the target counterfactual shift produces much larger predictor-output changes than replacing the action. Ratios compare magnitudes, not response direction or control quality.

Figure 4: Four-seed action/no-action comparison on the original validation dump. Each point is a final epoch-200 checkpoint; gray segments connect the matched action/no-action models within a seed, not across seeds. (a) MSE@10 ordering changes with seed. (b) Horizon-1 predicted rank also varies. (c) Replacing actions at command-change times produces only small signed changes in recorded-future error.

Figure 5: Matched objective comparison, seeds 42/43, on the original validation dump. (a,b) Individual seed curves are shown directly; no two-seed error bar is interpreted as sampling uncertainty. MSE-only has much lower latent error with and without target-variance normalization. (c) Horizon-1 effective rank is concentrated in the MSE runs, but VICReg is not immune: seed 43 has a target rank near 2.1. Neither low error nor spectral breadth alone certifies useful physical state.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Density-channel examples from the recorded-action replay. Episodes 1800–1805 at observation 12, covering both shapes and all three modulus values. These are simulator observations, not model reconstructions. The small source array is packaged so this figure regenerates without the full tensor.

Component	Configuration
Encoder	ViT; 4 blocks; width 128; 4 heads; MLP ratio 2
Patches / representation	$8\times 8$ patches; CLS token; latent width 128
Predictor	4 residual-MLP blocks; width 128; MLP ratio 2
Action embedding	Width 64; concatenation fusion
EMA	Momentum 0.996; updated each training step
Dropout / stochastic depth	0 / 0.05

Appendix

Table 1: Corrected training configuration shared across the reported checkpoints. The no-action model replaces action fusion with a state-only projection.

Metric	Seed	Original	Replay
Mean $D_{\mathrm{pred}}\times 10^{6}$	42	3.1832	3.1832
Mean $D_{\mathrm{pred}}\times 10^{6}$	43	4.1197	4.1197
Mean recorded-future error	42	0.07686	0.07686
Mean recorded-future error	43	0.07025	0.07025
Predicted effective rank	42	12.268	12.268
Predicted effective rank	43	5.776	5.776

Appendix

Table 2: Cross-platform consistency checks between original cached metrics and the earlier Linux recorded-action replay. The local original-RNG replay is separately bitwise identical to the original dump.

Model	Seed	MSE@1	MSE@5	MSE@10	VN@10	PredR@1	TgtR@1
VICReg + action	42	0.07551	0.61990	1.21103	4.447	5.97	4.64
VICReg + action	43	0.06968	0.61641	1.10064	7.152	2.44	2.11
VICReg, no action	42	0.07474	0.59499	1.14593	4.282	6.16	4.98
VICReg, no action	43	0.09871	0.71149	1.27055	4.466	6.54	5.48
VICReg + action	44	0.07080	0.50848	1.06394	4.820	9.24	7.94
VICReg, no action	44	0.07525	0.54499	1.06228	4.317	6.93	5.81

Appendix

Table 3: Every corrected epoch-200 checkpoint on the original validation dump. VN@10 denotes target-variance-normalized error. PredR@1 and TgtR@1 are horizon-1 effective ranks for prediction and EMA-target branches.

Model	Seed	P1	T1	P3	T3	P5	T5	P10	T10
VICReg+a	42	5.97	4.64	8.26	4.66	9.38	4.61	9.91	4.64
VICReg+a	43	2.44	2.11	3.19	2.14	3.70	2.13	4.49	2.10
VICReg-noa	42	6.16	4.98	8.00	4.99	8.97	5.01	8.11	4.94
VICReg-noa	43	6.54	5.48	8.91	5.43	10.94	5.36	11.17	5.37
VICReg+a	44	9.24	7.94	10.44	7.96	10.87	7.95	10.08	7.93
VICReg-noa	44	6.93	5.81	8.86	5.81	9.95	5.75	8.98	5.73

Appendix

Table 4: Per-horizon effective rank for prediction (P) and EMA target (T). Pooled horizon rank is omitted here because it can inflate apparent diversity.

Model	Seed	Geometry	Stiffness	Contact	Shuffle change (%)
VICReg + a	42	99.90	30.05	39.70	+0.017
VICReg + a	43	100.00	30.30	39.60	+0.011
VICReg, no a	42	100.00	30.05	38.60	+0.000
VICReg, no a	43	100.00	35.35	40.70	+0.000
VICReg + a	44	100.00	31.10	39.70	+0.039
VICReg, no a	44	100.00	33.10	43.40	+0.000

Appendix

Table 5: All corrected forecast probes and action interventions on the original validation dump. Accuracies are percentages averaged over horizons 1–10. The last column is the signed percentage change in one-step recorded-future error at command changes. No-action zeros are structural; exact machine values remain in the packaged CSV.

Seed	$n_{vis}$	$D_{tgt}$	$D_{pred}$	Raw	VN	State/action	State/action vis
42	579	4.13e-04	2.14e-06	0.0051	0.0027	59.5	189.5
43	579	1.35e-04	3.04e-06	0.0217	0.0116	16.9	53.2

Appendix

Table 6: Phase-2 action-sensitivity controls. Raw and VN are median

D_{\mathrm{pred}}/D_{\mathrm{tgt}}

on high-visibility pairs. State/action ratios are medians for

D_{\mathrm{tgt}}>10^{-8}

and high-visibility subsets.

Intervention	Pairs	RMS pos. $\times 10^{5}$	RMS vel. $\times 10^{3}$	Identical raster (%)
Shuffled	3000	1.191	2.386	41.53
Mean action	3000	0.788	1.577	57.93
Force reversed	3000	1.576	3.156	52.03
Force rotated	3000	1.111	2.232	54.57
Force scaled	3000	0.394	0.789	40.40
Contact shifted	3000	0.464	0.903	47.33

Appendix

Table 7: All original one-step intervention families. Position and velocity columns are medians in simulator units. The first two rows are primary contrasts. Secondary force/contact perturbations are not pooled into the primary conclusions.

Figure 7: Fixed readouts of observed EMA embeddings on the replay. Individual conditions for corrected VICReg seeds 42/43. Dashed horizontal lines show

1/3

uniform stiffness chance. Lines connect descriptive accuracies, not independent samples or temporal learning curves. No validation-based choice of a best time, probe, or history is used for a primary claim.

Input	Readout	Conditions	Stiffness (%)	Geometry (%)
EMA, single frame	Linear	14	32.0–35.0	100.0–100.0
EMA, single frame	MLP	14	30.0–39.0	100.0–100.0
EMA, history	Linear	8	31.5–36.5	100.0–100.0
EMA, history	MLP	8	30.5–37.5	100.0–100.0
Raster, single frame	Linear	7	33.5–33.5	100.0–100.0
Raster, single frame	MLP	7	33.5–33.5	100.0–100.0

Appendix

Table 8: All probe-condition ranges. “Conditions” counts settings per task, including both model seeds where applicable. These ranges summarize the full packaged 120-row table, not error bars or confidence intervals.

Input	Stiffness (%)	Geometry (%)
Mechanical $x,v,C,F$	28.5–39.5	100.0–100.0
Mechanical history	33.0–33.5	100.0–100.0
Float $64\times 64$ raster	33.0–33.5	100.0–100.0
Float raster history	33.0–33.5	50.5–100.0
uint8 $64\times 64$ raster	32.5–34.0	100.0–100.0
uint8 raster history	33.5–33.5	50.5–100.0

Appendix

Table 9: Upstream positive controls for material identifiability. Ranges span fixed linear/MLP readouts and tested times/histories. The validation-majority stiffness baseline is 33.5%. Privileged

(\mu,\lambda)

contains the defining material parameters and is an oracle control.

World Model Evaluation

World Model Evaluation measures whether learned simulators predict consistent, physically plausible, action-responsive futures that help agents plan and act.

See all 214 papers

UWM-JEPA: Predictive World Models That Imagine in Belief Space
May 25, 2026Santosh Kumar Radha, Oktay GoktasWorld Model Learning Latent Dynamics Modeling
AgentField AI, Toronto, Canada
Learning Counterfactual World Models for Embodied Reasoning under Partial Observability
Sep 5, 2026Todd Y. Zhou, Daniel ZhangEmbodied World Models World Model-Based Planning
Harvard University
Do-JEPA: From Masking to Intervention in Latent World Models
Sep 29, 2026Hossein Resani, Javen Qinfeng ShiCausal Effect Estimation Causal World Models
Australian Institute for Machine Learning, Adelaide University · Responsible AI Research Centre

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

UWM-JEPA: Predictive World Models That Imagine in Belief Space

Santosh Kumar Radha, Oktay Goktas

AgentField AI, Toronto, Canada

World Model Learning Latent Dynamics Modeling Joint-Embedding Predictive Architecture Action-Conditioned World Models

Open paper AskPDF arXiv

World models for partially observed environments must imagine multiple compatible hidden futures and steer between them under counterfactual actions. Joint Embedding Predictive Architectures (JEPAs) do this in latent space, but a vector-valued latent has no internal structure for carrying the belief over hidden continuations through blind rollout. We introduce the Unitary World Model JEPA (UWM-JEPA), a JEPA world model with a density-matrix latent on a joint system-environment space and a learned unitary predictor. The construction preserves the joint-state spectrum exactly during rollout, so the predictor itself cannot dissipate the represented uncertainty. On a hidden-velocity indicator task requiring five-step forward simulation under a given action sequence with the target observation masked, UWM-JEPA reaches 0.77 accuracy and degrades monotonically as actions are perturbed; a parameter-matched LSTM-JEPA trained under the same counterfactual-target objective and action head collapses to majority-class accuracy (0.53) under every action condition. Under blind rollout, UWM-JEPA loses fewer than ten points of probe R^2 at short horizons while vector-latent baselines lose forty-one and sixty-eight; both nevertheless tie on a held-out context probe, locating the separation in the predictor rather than the encoder. Action sensitivity itself requires training against counterfactual rather than teacher-forced targets, a finding that applies beyond the unitary parameterisation. For JEPA world models to imagine under partial observability, latent geometry and predictor dynamics matter, not frozen context-encoding capacity alone.

Abstract

World Model Evaluation

Explore similar work

UWM-JEPA: Predictive World Models That Imagine in Belief Space

Learning Counterfactual World Models for Embodied Reasoning under Partial Observability

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