Joint-embedding predictive architectures (JEPAs) learn latent dynamics for planning and avoid representation collapse by matching features to maximum-entropy distributions such as isotropic Gaussians, yielding dense representations. However, it is unclear whether dense representations are the most favorable geometry for modeling dynamics. In this work, we ask whether a different geometry, sparse representations, can make action-conditioned latent dynamics easier to model, and what dynamical structure emerges from such representations. We first show that nonlinear Lipschitz dynamics can be approximated arbitrarily well by action-conditioned linear dynamics in a sufficiently high-dimensional one-hot latent space, with rollout error vanishing as the dimension grows. This motivates distributed sparse representations as a practical relaxation of one-hot sparsity. We introduce LpWorldModel (LpWM), a JEPA model regularized with Rectified Distribution Matching Regularization (RDMReg) to match encoder features to a Rectified Generalized Gaussian distribution, yielding non-negative sparse codes. Empirically, sparsity lowers the predictor complexity required for successful planning: on PushT, sparse LpWM outperforms dense LeWM by up to 57% in planning success at intermediate predictor capacities. This advantage also extends beyond Gaussian distribution matching, with LpWM outperforming dense VICReg representations across multiple predictor families. We further find that the learned sparse representations are mode-factored, with support encoding discrete dynamical regimes and feature magnitudes capturing continuous within-regime state. Together, these results suggest that sparse representations can reduce the predictor complexity required for control while revealing interpretable structure.
Figures & tables
Figure 1: LpWorldModel: a JEPA model with sparse representations. (a) An encoder f maps observation ot to non-negative, sparse latents zt via a terminal (Rep)ReLU function. The predictor g takes inputs zt and actions at to predict z^t+1 , compared against the encoded target zt+1=f(ot+1) . RDMReg ( 24 ) matches per-timestep latent marginals to a Rectified Generalized Gaussian, preventing collapse and inducing sparsity. (b). Arrows indicate the success-rate difference between LpWM and LeWM: ↑ denotes higher performance for LpWM, while ↓ denotes lower performance. LpWM outperforms LeWM with more linear predictor variants (MLP ∘ LTV (k) , MLP ∘ LTI (k) , and LTI (k) . See full specifications in Table 1 ) by up to 57% in success rate, while performing on par with LeWM using the high-capacity DiT predictors. This suggests that sparse representations induce a latent space with simpler dynamics to learn compared to their dense counterparts.
Figure 2: The sparse support recovers the piecewise structure of the dynamics. (a) Support Jaccard heatmap over a 20×20 grid: the Jaccard index of each cell’s support with a fixed reference cell (white + ) is high within the reference zone and drops across zone boundaries (dashed), for a model trained with rendered zones (top) and one with no visual cues (bottom). (b) The binary support decodes the ground-truth zone identity as accurately as the full embedding, while agent position is decoded better from the magnitudes than from the support. (c) Planning success vs. horizon (closed-loop, 3-seed mean ± std). Goal locations of the dot agent are sampled randomly. Both sparse variants beat dense LeWM at every horizon.
Predictor
Layers
History ( k )
Functional form
Deep-AdaLN (k)
6
3
z^t+1=σ(AdaLN(6)(zt−k+1:t,at))
Shallow-AdaLN (k)
1
3
z^t+1=σ(AdaLN(1)(zt−k+1:t,at))
MLP ∘ LTV (k)
1
3
z^t+1=σ(WReLU(∑i=0k−1Ai(zt)zt−i+B(zt)at))
MLP ∘ LTI (k)
1
3
z^t+1=σ(WReLU(∑i=0k−1Aizt−i+Bat))
LTI (k)
1
3
z^t+1=σ(∑i=0k−1Aizt−i+Bat)
LTI (1)
1
1
z^t+1=σ(Azt+Bat)
Table 1: Predictors with Different Levels of Complexity. z^t+1 is the predicted latent, zt−i past latents, and at the action. AdaLN predictors are Transformer-based models with Adaptive Layer Normalization for action conditioning, i.e. DiT architecture ( 30 ) . The Deep-AdaLN (k) predictor is the one used in LeWM ( 28 ) . LTI denotes a linear time-invariant predictor with fixed operators Ai,B , while LTV denotes a linear time-varying predictor with state-dependent operators Ai(zt),B(zt) . We use MLP ∘ to denote one extra nonlinear transform. The history k denotes context length, with 1 the single-frame case. The output link is σ=identity for dense LeWM and σ=RepReLU for sparse LpWM. See Appendix B for details.
Predictor
D=384
D=768
D=1536
D=2048
D=4096
Deep-AdaLN (k)
25.8 M
62.2 M
166.9 M
260.2 M
822.4 M
Shallow-AdaLN (k)
5.6 M
13.0 M
33.1 M
50.4 M
151.1 M
MLP ∘ LTV (k)
0.81 M
3.10 M
12.1 M
21.4 M
84.7 M
MLP ∘ LTI (k)
0.74 M
2.95 M
11.8 M
21.0 M
83.9 M
LTI (k)
0.59 M
2.36 M
9.44 M
16.8 M
67.1 M
LTI (1)
0.30 M
1.18 M
4.72 M
8.39 M
33.6 M
Table 2: Predictor parameter counts across latent dimensions D . Counts are for the predictor alone for each D used in our experiments. LTI, MLP ∘ LTI, and MLP ∘ LTV are attention-free and far smaller than the AdaLN Transformers.
Figure 3: Linear predictor is sufficient for the wall environment regardless of dense or sparse embeddings, and LpWM outperforms VICReg-based JEPA across predictor settings in PushT. (a) We plot the planning success rate of the linear LTI (1) predictor against the nonlinear Deep-AdaLN (k) , where each dot denotes one setting of method and feature dimension. For simplicity, we use μ=∗ to denote LpWM with the target distribution parameter μ . We show the open-loop evaluation on top and closed-loop results at bottom. Both are ran using CEM planner. We observe that the closed-loop success rate for LTI (1) predictor using SIGReg is already approaching 100% , indicating that the wall environment can already be described well by a linear dynamics model in the latent space. (b) We have already shown LeWM vs. LpWM in Figure 1(b) . Here we compare LpWM with a variant of LeWM where we replaced SIGReg with VICReg in the PushT environment across predictor settings. As oppposed to SIGReg, we observe consistent improvement from LpWM over VICReg even for the Deep-AdaLN (k) predictor, and the gap between sparse LpWM and dense methods (i.e. VICReg here) continues to hold for MLP ∘ LTV (k) and MLP ∘ LTI (k) predictors. See full hyperparameter sweeps in Appendix H.2 .
Figure 4: Temporal-Jaccard (TJ) makes the cube support track contact, not effector motion. (a) As TJ strength increases, support-instability correlation with effector motion drops ( 0.87→0.40 ) while cube-motion ( 0.26→0.80 ) and gripper-contact ( 0.05→0.61 ) rise—at no change in planning success. (b) One example episode: without TJ (top) support instability tracks effector displacement ( reff=0.86 ); with TJ (bottom) it tracks cube displacement ( rcube=0.80 ), spiking inside the shaded contact windows. (c) Support raster: without TJ the support rapidly transitions across frame; with TJ it is smoothly varying in time, reorganizing almost during the shaded contact window.
Appendix figures & tables13 assets
Supplementary material from the paper’s appendix.
Appendix
D
Deep-AdaLN (k)
Shallow-AdaLN (k)
MLP ∘ LTV (k)
MLP ∘ LTI (k)
LTI (k)
LTI (1)
384
0.34
0.33
0.37
0.28
0.41
0.63
768
0.44
0.37
0.54
0.53
0.41
0.63
1536
0.49
0.39
0.49
0.54
0.41
0.60
2048
0.42
0.46
0.51
0.46
0.43
0.61
4096
0.48
0.49
0.47
0.43
0.44
0.56
Appendix
Table 3: Sparsity of the LpWM encoder representations across predictor types and latent dimensions. Each entry is the fraction of non-zero coordinates E[∥z∥0]/D in the sparse code for the target distribution with {μ=0,σ=1/2,p=1} , measured on the validation set at the tuned-best cell for the LpWM run paired with each predictor at latent dimension D . Dense LeWM has active fraction 1.0 by construction. The learned code stays ∼30 – 65% active regardless of predictor or dimension, confirming the codes underlying the planning comparison are genuinely sparse.
Goals from Evaluation Set
Random Goals
Env
Method
H=R=5
H=5,R=1
H=R=5
Piecewise ( 2×2 )
LpWM (Ours)
99.33±1.15
84.67±4.16
59.33±2.31
LpWM+TJ (Ours)
99.33±1.15
82.67±5.03
68.67±4.62
LeWM
99.33±1.15
65.33±4.16
36.00±3.46
Piecewise ( 3×3 )
LpWM (Ours)
100.00±0.00
58.67±1.15
56.00±4.00
LpWM+TJ (Ours)
100.00±0.00
60.00±2.00
58.00±5.29
Appendix
Table 4: Success Rate (%) over 50 Test Samples in the Piecewise Environments with Varying Numbers of Zones . Piecewise ( n×n ) denotes the environment with varying numbers of zones. Goal images are either sampled from the evaluation dataset or randomly sampled from all possible locations of the dot agent. H denotes the horizon and R represents the receding-horizon. Results are shown as mean ± std over 3 seeds. Bold denotes the best performance.
Figure 5: Environments. (a) Wall: 2D navigation through a doorway (red agent). (b) PushT: push a T-shaped block to a target pose (green). (c) Piecewise: an open room partitioned into force-field zones (colored) with a known ground-truth segmentation. (d) OGBench-Cube: a robot arm manipulating a cube.
Section 3
Section 4
Wall
PushT
Piecewise
OGBench-Cube
Encoder
from-scratch ViT, 12 layers, width 384
from-scratch ViT-Tiny, width 192
Predictor
See Table 1
Deep-AdaLN (k)
Dimension ( D )
{384,768,1536,2048,4096}
192
History k
1
3
3
3
Frameskip
5
5
5
5
Appendix
Table 5: Experimental configurations for Section 3 and Section 4 . We list out the experimental pipelines in the following table.
Figure 6: Deep-AdaLN (k) Predictor. Hyperparameter Sweeps for the PushT Environment.
Figure 7: Shallow-AdaLN (k) Predictor. Hyperparameter Sweeps for the PushT Environment.
Figure 8: MLP ∘ LTV (k) Predictor. Hyperparameter Sweeps for the PushT Environment.
Figure 9: MLP ∘ LTI (k) Predictor. Hyperparameter Sweeps for the PushT Environment.
Figure 10: LTI (1) Predictor. Hyperparameter Sweeps for the PushT Environment.
Figure 11: LTI (k) Predictor. Hyperparameter Sweeps for the PushT Environment.
Figure 12: Deep-AdaLN (k) Predictor. Hyperparameter Sweeps for VICReg-based JEPA models in the PushT Environment.
Figure 13: MLP ∘ LTV (k) Predictor. Hyperparameter Sweeps for VICReg-based JEPA models in the PushT Environment.
Figure 14: MLP ∘ LTI (k) Predictor. Hyperparameter Sweeps for VICReg-based JEPA models in the PushT Environment.
Latent world models learn action-conditioned dynamics in representation space and often score candidate actions by Euclidean distance to a goal representation. Joint training typically regularizes the representation to prevent collapse, but the resulting representation geometry also determines how terminal errors are weighted during planning. We show that accurate prediction and noncollapsed representations do not guarantee a task-aligned latent planning cost: isotropic Gaussian regularization can induce a geometry that ranks feasible outcomes differently from the task cost. To address this mismatch, we introduce AnisoWM with ΛReg, which replaces the fixed isotropic Gaussian target with a learnable diagonal covariance under fixed-trace and anisotropy constraints. The prediction objective, predictor architecture, and Euclidean planner remain unchanged; the target is used only during training. Our analysis characterizes the prediction-driven allocation of target variance, its dependence on the training distribution, and the conditions under which the induced metric reduces planning regret. Across four visual control environments, AnisoWM improves planning success over LeWorldModel in all four. Its latent planning cost also shows better agreement with task outcomes. Project website: https://rkdrn79.github.io/AnisoWM-page/
Mingu Kang, Yoori Oh, Sookyung Kim +1
Seoul National University · Ewha Womans University
Compact JEPA world models enable efficient latent-space planning, but low-dimensional representation trained under reward-free self-supervision must encode both action-conditioned dynamics and predictable visual context. This competition can entangle controllable state with high-rank nuisance appearance and degrade planning as scenes become more complex. We introduce LRC-JEPA, a lightweight end-to-end world model that routes information into a compact predictive latent z and learned-query residual-context embeddings u. Only z is propagated by the dynamics model and used for planning, while u captures temporally persistent information for cross-attention reconstruction; a differentiable residual connection encourages the latent to retain complementary dynamic content. Under explicit assumptions, we show that the resulting representation is sufficient, minimal, nuisance-invariant, and disentangled. Across four simulated control environments, LRC-JEPA improves average planning success over a parameter-matched JEPA baseline by 9 percentage points and matches or exceeds substantially larger pretrained models. On the real-world Bridge-v2 set, its 5.5M-parameter active encoder outperforms DINO-WM (22.1M) and V-JEPA2 (303.9M) encoders while also enabling faster planning. Physical-state probes, reconstruction interventions, and ablations confirm the effectiveness of LRC-JEPA's representation disentanglement.
We introduce Adaptive LeWorldModel (ALeWM), a world model based on a joint-embedding predictive architecture (JEPA) that learns to concentrate predictive information in compact prefixes of a wide latent representation. To encourage this ordering, ALeWM learns a sequence-conditioned distribution over prefix lengths and trains the predictor to estimate the full next embedding from a sampled input prefix. As standard anti-collapse objectives encourage variation across latent coordinates and do not organize them by predictive importance, we also introduce MixSIGReg. MixSIGReg regularizes the masked embeddings against a prior-weighted mixture with Gaussian active prefixes and zeros in the remaining coordinates. As a result, the ALeWM objective encourages early coordinates to retain information useful for prediction and recursive planning. Our analysis shows that the mixture distribution used by MixSIGReg assigns higher variance to earlier coordinate blocks and lower variance to later ones. In addition, we show that, under specified assumptions, prediction error is minimized by placing the information most useful for prediction in earlier blocks. Empirically, we study the behavior of ALeWM in a controlled dynamical system with known state variables and in goal-conditioned visual control. We show that ALeWM consistently achieves higher mean success rates than tuned fixed-width LeWM, with lower planning capacity on average.