Recent empirical and theoretical advances suggest that joint-embedding predictive architectures (JEPAs) may learn meaningful representations for action-conditioned prediction of future outcomes, thus becoming one of the foundational structures for world models. However, accurate prediction does not, in general, necessarily imply recovery of underlying causal states that give rise to the observed dynamics. This work investigates when and how JEPAs can recover the underlying causal states from observations. We first introduce a latent variable model, in which high-dimensional observations are generated from latent causal states whose dynamics are governed by action-conditioned transition mechanisms. Based on this formulation, we develop a general information-theoretic objective that combines conditional likelihood maximization for learning transition dynamics with entropy maximization for preserving latent state information. We then establish identifiability conditions under which representations learned by this general objective recover the underlying latent causal states up to component-wise invertible transformations and permutation. One key condition for such identifiability is sufficient action-induced variation in the transition mechanisms. Guided by this finding, we instantiate the general objective with an action-modulated Gaussian additive-noise model, yielding action-modulated JEPA (A-JEPA). Experiments on synthetic environments verify the theoretical findings under the identifiability conditions and robustness to moderate violations, while visual benchmarks demonstrate improved state recovery and transfer to unseen transition mechanisms.
Figures & tables
Figure 1: Analysis of latent-state recovery. Left: effect of the trade-off between transition likelihood and information preservation. Middle: effect of increasing action-induced mechanism variation. Right: component-wise recovery across latent dimensions compared with JEPA-style baselines.
Figure 2: Analysis of latent structure recovery. Left: effect of parent-gate sparsity on graph recovery. Middle: F1, precision, and recall across latent dimensions. Right: SHD across latent dimensions.
Figure 3: Assumption and model-specification analysis. Left: latent-state recovery degrades as the observation mapping becomes increasingly non-invertible. Middle: reducing action-induced variation in either the conditional mean or variance degrades recovery. Right: the proposed method remains robust to Laplace and Student- t noise, while uniform noise leads to a larger degradation.
Figure 4: Representation and transition recovery across four visual control environments.Columns correspond to TwoRoom , OGB , Reacher , and PushT , while rows report component-wise state recovery (MCC), joint-state information preservation (nonlinear state R2 ), and multi-step action-conditioned prediction (rollout R2 ), respectively. A-JEPA consistently improves component-wise recovery while retaining strong state information and predictive dynamics.
Figure 5: Held-out robot transfer on RoboNet. A-JEPA maintains substantially more stable latent rollout predictions on unseen Franka and Sawyer robots. See Sec. J for details.
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 6: Latent-dimension selection on TwoRoom . We sweep dz∈{6,7,8,9,10} and evaluate both component-wise recovery (MCC) and rollout R2 at horizon 5 . We use dz=10 for the main experiments because it provides a strong balance between component recovery and multi-step transition prediction.
Figure 7: Latent-dimension selection on Reacher . We sweep dz∈{6,8,10,12,14,16} and evaluate grouped circular MCC and rollout R2 at horizon 5 . We use dz=6 for the main experiments because it achieves strong component-wise recovery and multi-step prediction with the most compact representation.
Figure 8: Latent-dimension selection on PushT . We evaluate both component-wise recovery and rollout R2 at horizon 5 . We use dz=10 for the main experiments because it provides the strongest joint component-recovery and multi-step prediction performance in the dimension-selection run.
Figure 9: Latent-dimension selection on OGBench Cube . We sweep dz∈{10,12,14,16,18} and evaluate component-wise recovery and rollout R2 at horizon 5 . We use dz=14 for the main experiments because it achieves the strongest joint recovery and multi-step prediction performance.
Figure 10: Learned latent transition graph on Reacher . MCC assigns z0 to the proximal joint factor and z3 to the distal joint factor. The learned graph contains the directed edge z0→z3 , consistent with the proximal-to-distal dependency of the two-link arm.
Figure 11: Latent intervention diagnostics on Reacher . Top: intervention on the proximal-aligned coordinate z0 changes the proximal joint and induces corresponding changes in the downstream distal arm segment. Bottom: intervention on the distal-aligned coordinate z3 produces a more localized distal-joint change while leaving the proximal configuration relatively unchanged. This asymmetric response is consistent with the learned directed dependency z0→z3 .
Horizon
A-JEPA
LeWM
PLDM
Held-out Franka
h=1
0.9691±0.0088
0.6154±0.2418
0.4486±0.4487
h=2
0.9605±0.0112
0.6033±0.1973
−0.0679±0.4816
h=3
0.9534±0.0128
0.5728±0.1905
−0.8018±1.1256
h=4
0.9468±0.0141
0.5428±0.1928
−1.7901±2.0465
h=6
0.9387±0.0158
0.4878±0.1995
−3.9602±3.8622
Appendix
Table 1: Held-out robot transfer on RoboNet. Models are trained on Baxter, Kuka, and WidowX and evaluated without adaptation on Franka and Sawyer. We report latent rollout R2 at selected horizons and the average over h=1,…,20 . Values are mean ± standard deviation over three seeds.
School of Vehicle and Mobility, Tsinghua University · 2Didi Voyager Labs, Didi Autonomous Driving · School of Computing and Data Science, The University of Hong Kong +2