World Modeling through Spectral Alignment

Holger Molin; William Peng; Marco Bagatella; Randall Balestriero

cs.LGOct 11, 2026

World Modeling through Spectral Alignment

Authors: Holger Molin, William Peng, Marco Bagatella, Randall Balestriero

Organizations: Stanford University · ETH Zürich · Brown University

Abstract

How do we tell a world model what information to preserve? We study this question through the lens of spectral representation learning. We introduce SpecWM, a JEPA framework that explicitly specifies desired relationships between observations through a target similarity kernel and trains the encoder to reproduce those relationships in its embedding geometry. Different kernels provide different supervision, ranging from physical state information to a fully self-supervised temporal kernel requiring only the ordering of observations within trajectories. We characterize the global optima of the objective, establish conditions for linear recoverability and invariance to nuisance information, and connect temporal-kernel similarity to the behavioral successor measure under an idealized setting. Empirically, the choice of kernel strongly controls which physical quantities are recoverable from the representation. More importantly, temporal spectral alignment improves CEM planning success over LeWorldModel from 52% to 68% on OGBench Cube, 49% to 77% on OGBench Scene, and 29% to 51% on CALVIN. These results suggest that explicitly specifying representation geometry provides a simple way to steer what latent world models preserve for control.

Figures & tables

Figure 1: Spectral alignment framework. Pairwise distances (target) between samples define a kernel that an encoder is trained to reconstruct through embedding similarities.

Figure 2: Spectral alignment on a single trajectory. Top: sampled observations from OGBench Scene ( Park et al., 2025 ) . Bottom: the temporal target kernel and encoder Gram matrix (before centering), the spectral embedding of $K$ (top three eigenvectors), and the encoder rows $h_{i}$ (top three principal components), colored by frame index. Matching pairwise similarities encourages the encoder to recover the target’s geometry.

Figure 3: Average CEM planning success rate on 100 fixed tasks per environment. Bars report the mean over the 20k, 24k and 28k checkpoints and 3 seeds, with whiskers showing one cross-seed standard deviation. SpecWM consistently outperforms LeWM.

Figure 4: Linear probe

R^{2}

of the manipulated object’s position from each model’s final checkpoint embedding. Bars report the mean over 3 seeds, with whiskers showing one cross-seed standard deviation. Random is an untrained encoder of the same architecture.

	Underspecified					Fully specified
	Temporal	Quasimetric	State	Proprio	Depth	State	Proprio	Depth
Scene (avg)	0.84	0.78	0.85	0.16	0.15	0.93	0.16	0.50
cube	0.79	0.51	0.57	0.10	0.09	0.96	0.08	0.09
arm	0.67	0.66	0.69	0.69	0.60	0.73	0.67	0.59
drawer	0.91	0.90	0.98	0.01	0.01	0.99	0.01	0.87
window	0.87	0.88	0.99	0.00	0.01	0.99	0.00	0.91

Table 1: Linear probe

R^{2}

on held-out OGBench Scene episodes. Bold marks the best result in each row and results within 0.01 of it.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Linear probe

R^{2}

between eigenvectors and embeddings from encoders trained on fully specified state, proprioception, and depth. Underspecified variants and an untrained ViT are included for comparison.

Figure 6: Effect of a stop-gradient on the prediction target. Bars report the mean over 3 seeds, with whiskers showing one cross-seed standard deviation. The stop-gradient generally lowers planning success and raises object recoverability in some cases while diminishing recoverability in others.

Figure 7: Effect of the failure data. CEM success (top) and object position probe

R^{2}

(bottom) for LeWM and SpecWM (temporal), trained with the failure rollouts (solid, 3 seeds) or on the play data alone (hatched, 1 seed), at matched updates and frames per update. Failure data raises planning success for both models; SpecWM outperforms LeWM in both cases.

Figure 8: Effect of the choice of distance measure on CEM planning results. Cosine similarity and L2 distance between current embedding and target embedding are evaluated on both LeWM and SpecWM. The choice of distance measure appears to not be a major contributor to planning ability. Three training seeds are evaluated at their respective final three checkpoints for each setting.

Figure 9: Effect of choice of normalization layer in the MLP converting the CLS token into an embedding on downstream CEM planning. One seed is used per run and CEM is reported as the average of the final three checkpoints of each run.

Parameter	Scene	Cube	CALVIN
Candidate sequences	300	300	300
CEM iterations	30	30	30
Elite sequences	30	30	30
Initial variance scale	1.0	1.0	1.0
Planning horizon (blocks)	5	5	5
Environment steps per block	5	5	10

Appendix

Table 2: CEM planning configuration.

Component	Configuration
Visual encoder	Randomly initialized ViT-T/14 with output dimension 192
Projection head	MLP $192\rightarrow 2048\rightarrow 192$
Action encoder	$1\times 1$ Conv1D followed by an MLP with hidden dimension 768 and output dimension 192
Dynamics predictor	6-layer action-conditioned transformer
Attention	16 heads with head dimension 64
Transformer MLP dimension	2048

Appendix

Table 3: Shared visual-encoder and action-conditioned dynamics architecture.

Environment	Frame skip	Raw action dim.	Action-block dim.
OGBench Scene	5	5	25
OGBench Cube	5	5	25
CALVIN	10	7	70

Appendix

Table 4: Environment-dependent action representation. Each action block contains all low-level actions between two consecutive sampled observations.

Setting	Configuration
Peak learning rates	$1.461\times 10^{-3}$ for matrices, $2.922\times 10^{-4}$ for convolutions, and $7.30\times 10^{-5}$ for biases and normalization parameters
MuonAdamW	$(\beta_{1},\beta_{2})=(0.9,0.999)$ , $\epsilon=10^{-8}$ , Muon momentum $0.95$ , Nesterov updates, and RMS matching
Regularization	Weight decay $10^{-3}$ on weights; no decay on biases or normalization parameters; gradient clipping at $1.0$
SpecWM	33,375 updates; 552 base clip frames per update ( $24\times 23$ ); $\lambda_{\mathrm{spec}}=1$ for Cube and Scene and $2$ for CALVIN; cosine prediction-anchor weight $2$ .
LeWorldModel	33,375 updates; 552 frames per update ( $138\times 4$ ); prediction-MSE weight $1$ and SIGReg weight $0.09$

Appendix

Table 5: Optimization and student-training settings.

Explore similar work

Learning Commute-Time-Preserving World Models for Planning
Oct 1, 2026Michael Hauri, Peter Buttaroni, Fabian A. Mikulasch +1Representation Learning World Model Learning
Friedrich Miescher Institute for Biomedical Research, Basel, Switzerland · Faculty of Science, University of Basel, Switzerland
Reperesentation Geometry Matters for Planning with JEPA World Models
Aug 17, 2026Jiaming Hu, Yan Zheng, Shi Bo +9Representation Learning Neural Representation Geometry
Boston University · Unity Technologies
Spectral-Target Physical Latent Structuring for JEPA-Style World Models
Sep 2, 2026Penghao Zhu, Salvatore Penachio, Kaustav Mukherjee +1World Model Learning World Model-Based Planning
Kaliber Labs San Francisco, CA 94103, USA

Oct 1, 2026J/K move · Enter open · S save

Learning Commute-Time-Preserving World Models for Planning

Michael Hauri, Peter Buttaroni, Fabian A. Mikulasch, Friedemann Zenke

Friedrich Miescher Institute for Biomedical Research, Basel, Switzerland · Faculty of Science, University of Basel, Switzerland

Representation Learning World Model Learning Latent Dynamics Modeling World Model-Based Planning Model-Based Planning Graph Embedding

Open paper AskPDF arXiv

World models allow agents to plan in latent space by choosing a sequence of actions that most reduces the distance to a given goal state. Thus, planning can benefit from latent representations whose distances mirror commute-times in the environment. The spectral embedding space of the graph Laplacian provides such a representation, if it obeys a specific eigenvalue-dependent scaling. Unfortunately, instantiating the graph Laplacian is intractable in large, continuous environments. Self-supervised learning offers a natural route to such commute-time-preserving embeddings at scale. However, here we show that existing methods, which commonly encourage isotropic representations to prevent representational collapse, tend to degrade the "correct" eigenvalue-dependent scaling, leading to an inaccurate representation of commute times. To address this problem, we introduce Commute-Time-Preserving World Models (CTWMs), combining a latent displacement predictor and a log-determinant regularizer that prevents collapse, which provably recover the correctly scaled Laplacian representation under reversible deterministic dynamics and at the predictor's fixed point. In numerical simulations, CTWM matches or outperforms LeWM, a task-agnostic baseline, on several complex, continuous goal-reaching benchmarks, while using half the parameters.

Abstract

Explore similar work

Learning Commute-Time-Preserving World Models for Planning

Reperesentation Geometry Matters for Planning with JEPA World Models

Spectral-Target Physical Latent Structuring for JEPA-Style World Models