PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning Commute-Time-Preserving World Models for Planning

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

arXiv:2610.01373v1cs.LG

TL;DR

The paper addresses whether scalable, reconstruction-free world models can learn planning-suited geometry, given that scaled Laplacian embeddings represent commute times in Euclidean space. CTWM learns that geometry through latent displacement prediction and log-determinant regularization; across goal-reaching tasks it matches or outperforms LeWM with half the parameters.

  • Problem

    Whether reconstruction-free joint world models can learn planning-suited geometry was uninvestigated, despite scaled Laplacian embeddings representing commute-time distances in Euclidean space.

    Source: §1 · p. 2, §3.1 · p. 4

  • Method

    CTWM predicts latent displacements and regularizes representation covariance with a log-determinant term to recover scaled Laplacian geometry.

    Source: §5 · p. 10

  • Results

    Across six goal-reaching tasks, CTWM matches or outperforms LeWM while using half the parameters.

    Source: §5 · p. 10

  • Takeaways & Limitations

    The learned representation can be used by planners that score candidate trajectories by goal distance.

    Source: §5 · p. 10

    The stated recovery result assumes deterministic dynamics and that the predictor is at its fixed point.

    Source: §3.5 · p. 6

Abstract

from arXiv · show

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.

1 INTRODUCTION

The introduction asks whether scalable, reconstruction-free world models can learn representations suited to planning. It presents CTWM as a method for learning commute-time-preserving geometry and reports theoretical and benchmark results.

1 INTRODUCTION

  • World models shaped by fixed rewards can adapt poorly when objectives change, motivating task-agnostic models; the introduction notes that collapse avoidance alone does not ensure planning-useful geometry.The paper frames the gap as steering scalable, joint-embedding learning toward geometry optimal for planning.Source: §1 · pp. 1–2
  • Scaled Laplacian representations make Euclidean distances reflect graph commute times, which account for temporal distances between states and support planning around environmental obstacles.This representation requires eigenvalue-dependent rescaling.Source: §1 · p. 2
  • CTWM jointly predicts action-induced latent displacements and regularizes the representation to recover correctly scaled Laplacian embeddings under deterministic, reversible dynamics.The introduction describes the regularizer as entropy-maximizing and states that the resulting latent distances match commute-time distances.Source: §1 · p. 2
  • Across goal-reaching benchmarks, CTWM outperforms LeWM on half the parameter budget, including two additional visual benchmarks evaluated from pixels without oracle goal representations.The experiments cover LeWM’s benchmarks and OGBench Scene and Pointmaze.Source: §1 · p. 2

2 RELATED WORK

Related work spans reward-conditioned world models, reward-free self-supervised approaches, spectral representations, and planning-oriented latent distances. The paper distinguishes CTWM by combining a latent state predictor with a target geometry tied to environment commute times.

2 RELATED WORK

  • World models progressed from learning representations and dynamics separately to jointly trained reconstruction-based models, while reward dependence can hinder generalization when tasks change.Later work also explored reconstruction-free representation learning and joint-embedding predictive architectures.Source: §2 · p. 2
  • Task-agnostic models learn without reward, using either pretrained self-supervised encoders or jointly trained predictors; some use temporally abstract transitions to limit long-horizon rollout error.Joint-embedding predictive architectures train the encoder and dynamics predictor together.Source: §2 · pp. 2–3
  • Collapse-prevention methods include contrastive losses, asymmetric predictor and stop-gradient designs, and regularizers; many encourage isotropic embeddings, which can challenge scaled representations.SIGReg explicitly matches projected embedding distributions to an isotropic Gaussian.Source: §2 · p. 3
  • Laplacian methods recover spectral geometry from transitions, but do not necessarily learn a latent dynamics model or guarantee eigenvalue-dependent scaling.ALLO is described as explicitly recovering both Laplacian eigenvectors and eigenvalues through constrained optimization.Source: §2 · p. 3
  • Temporal-distance embeddings target planning-relevant distances without end-to-end latent predictors, whereas successor-feature distances depend on policy rather than intrinsic environment structure.The passage contrasts these approaches with CTWM’s recovery of environment commute-time distance.Source: §2 · p. 3

3 CONSTRUCTING COMMUTE-TIME-PRESERVING WORLD MODELS

CTWM combines residual latent-displacement prediction with log-det regularization to learn embeddings whose Euclidean distances preserve commute-time geometry. Its scaled-Laplacian recovery is established under deterministic dynamics at the predictor’s fixed point, and the torus experiment tests the predicted scaling and distances.

3.1 PLANNING IN METRIC LATENT SPACES

  • The scaled Laplacian representation makes commute-time distance Euclidean, providing a latent metric for planning that accounts for temporal distance between states.Its eigenvalue-dependent scaling contrasts with the isotropic embeddings commonly encouraged by existing SSL methods.Source: §3.1 · p. 4

3.2 LEARNING TASK-AGNOSTIC WORLD MODELS WITH SSL

  • CTWM learns an encoder and latent dynamics model from task-agnostic transitions, using residual prediction and entropy regularization to prevent collapse.The transition data are collected without rewards or task labels; collapse prevention is implemented with a log-determinant covariance regularizer.Source: §3.2 · p. 4, §3.4 · p. 5

3.3 LEARNING WORLD MODELS VIA RESIDUAL PREDICTIONS

  • Residual prediction trains the model to predict the displacement between successive embeddings, while a penalty on large predictor outputs encourages consecutive latent states to remain close.The loss also includes a regularizer that prevents representational collapse.Source: §3.3 · pp. 4–5

3.4 REGULARIZATION TO PREVENT COLLAPSE

  • Unlike regularizers that encourage isotropic covariance, log-det regularization constrains the product of covariance eigenvalues and supports the scaling needed to capture commute-time distances.The passage contrasts it with hard orthonormality constraints and VCReg, which are not consistent with the correctly scaled representation.Source: §3.4 · pp. 5–6, §3.5 · p. 6

3.5 CONVERGENCE TO THE SCALED LAPLACIAN REPRESENTATION

  • Under deterministic dynamics and at the predictor’s fixed point, the objective is minimized by the scaled Laplacian representation, up to rotation and translation.This scaling makes squared L2 distances correspond to commute-time distances, up to a global multiplicative constant.Source: §3.5 · p. 6
  • Figure 2 reports that only CTWM tracks the ground-truth eigenvalue scaling, and that CTWM achieves the highest R2 for latent versus commute-time distances.The variance ratio is compared with λ1/λk; the distance comparison uses squared L2 latent distances against ground-truth commute times.Source: Figure 2 · p. 7

4 EXPERIMENTS

Across discrete and continuous experiments, CTWM learns commute-time geometry and matches or outperforms LeWM on six goal-reaching tasks with half the parameters. Its advantages become more pronounced at longer goal distances.

4.1.1 CTWM RESULTS IN SUPERIOR PLANNING PERFORMANCE

  • Across six benchmark tasks, CTWM matched or outperformed LeWM while using half the parameters.The comparison covers Reacher, Cube, Two-Room, PushT, Scene, and Pointmaze.Source: Figure 3 · p. 8, §4.1.1 · pp. 8–9

4 EXPERIMENTS

  • On the discrete toroidal graph, CTWM’s commute-time fit reached R2 = 0.95, versus 0.84 for VICReg, 0.73 for WM-GDO, and 0.67 for ALLO.CTWM’s subspace similarity was higher than VICReg’s and on par with ALLO’s; degeneracy of the graph’s eigenvalues caps the similarity near 0.91.Source: §4 · pp. 7–8

4.1 LEARNING CTWM IN CONTINUOUS RL ENVIRONMENTS

  • The six pixel-observation benchmarks span navigation and manipulation, using datasets or trajectories from heuristic policies, trained policies, and demonstrations.CTWM used 9M parameters versus 18M for LeWM, with a 64-dimensional latent space.Source: §4.1 · p. 8

4.1.1 CTWM RESULTS IN SUPERIOR PLANNING PERFORMANCE

  • At 100-step goal distance, CTWM retained 98.3% success on Two-Room and 98.8% on Pointmaze, while LeWM fell to 17.3% and 28.2%, respectively.The extended setting used a maximum planning budget of 125 steps.Source: §4.1.1 · p. 9, Table 1 · p. 9

4.1.2 COMMUTE-TIME GEOMETRY IN CONTINUOUS SPACES

  • On Pointmaze, CTWM reproduced ground-truth commute-time distances more faithfully and exhibited the predicted λ^-1/2 variance decay, unlike LeWM’s flatter spectrum.The analysis discretized the Pointmaze state space into a graph and compared embedding distances with computed ground truth.Source: Figure 4 · p. 9, §4.1.2 · p. 10

5 CONCLUSION

CTWM uses latent-displacement prediction with covariance log-determinant regularization to recover a commute-time-preserving representation under deterministic dynamics at the predictor’s fixed point. Across six goal-reaching tasks, it matches or outperforms LeWM with half the parameters.

5 CONCLUSION

  • CTWM predicts latent displacements and regularizes their representation with the covariance log-determinant, recovering correctly scaled Laplacian geometry under deterministic dynamics at the predictor’s fixed point.In that setting, latent L2 distances equal commute-time distances on the transition graph.Source: §5 · p. 10
  • The authors report using generative AI assistance for proof writing, parts of the experimental code, readability edits, related-literature identification, and reference formatting, with author review.They report that hypotheses, methodology, experiments, and results interpretation were not AI-generated.Source: AI USE STATEMENT · p. 10

A HYPERPARAMETERS

  • The code and dataset access, hyperparameters, assumptions, and complete proofs are reported in the paper’s reproducibility materials and appendices.The code and dataset access are available from the linked project repository.Source: REPRODUCIBILITY STATEMENT · p. 10

C.1 CONVERGENCE OF EQUATION 5

The convergence argument simplifies the CTWM objective under deterministic dynamics and relates its regularized displacement objective to Laplacian representations. With an orthonormality constraint, the latent states converge to the Laplacian representation.

C.1 CONVERGENCE OF EQUATION 5

  • The loss predicts the displacement between successive latent states from the current latent state and action, adds a predictor-norm penalty, and regularizes the encoder.The displayed objective is an expected squared displacement-prediction error plus β times the squared predictor norm and Reg(ϕ).Source: §C.1 · p. 16
  • With deterministic dynamics and an orthonormality constraint, the learned latent states converge to the Laplacian representation.The cited prior results describe convergence up to rotation under the graph-drawing objective.Source: §C.1 · p. 16
  • For deterministic dynamics, the transition covariance is zero and the conditional mean displacement equals the latent displacement, simplifying the objective.The resulting expression contains the expected squared latent displacement and the regularizer.Source: §C.1 · p. 16

C.2 SCALING OF THE REPRESENTATION UNDER DIFFERENT REGULARIZERS (HEURISTICS)

Under an eigenvector-alignment heuristic, the regularizer determines coordinate scaling, while the selected eigendirections are the d slowest modes.

C.2 SCALING OF THE REPRESENTATION UNDER DIFFERENT REGULARIZERS (HEURISTICS)

  • The log-det analysis selects the d slowest eigendirections, and the regularizer determines their eigenvalue-dependent scaling.The analysis notes that its stronger log-det result removes the eigenbasis-alignment assumption.Source: §C.2 · p. 17
  • With equal variance in every direction, orthonormality yields the unscaled Laplacian representation.Source: §C.2 · p. 17
  • VCReg sets the scale to τ above its threshold; below it, the stationary scale is α/(4cλk), while slow modes remain constant.Source: §C.2 · p. 17

C.3 PROOF OF THEOREM 1

The proof separates the represented subspace from its covariance, showing how the optimal scaled Laplacian representation follows under deterministic dynamics and an eigengap.

C.3 PROOF OF THEOREM 1

  • Under deterministic dynamics, every minimizer has the stated scaled-Laplacian form, up to an orthogonal transformation and translation.Source: §C.3 · pp. 17–18
  • The proof reduces the objective to matrix optimization, separates subspace from covariance, optimizes covariance, then selects the subspace using Poincaré separation.Source: §C.3 · p. 18
  • For fixed spectral subspace, the optimal covariance is proportional to the inverse restricted Laplacian; with an eigengap, the unique optimal subspace spans the first d eigendirections.Source: §C.3 · pp. 17, 19

C.4 PROOF OF COROLLARY 1.2

The proof uses centered, orthonormal Laplacian eigenfunctions; the next-state prediction objective can also be written as a latent-displacement prediction loss with a displacement penalty.

C.4 PROOF OF COROLLARY 1.2

  • The nonconstant Laplacian eigenfunctions are centered and have identity covariance under the stationary state distribution.Source: §C.4 · p. 20

D.1 SMOOTHNESS WITH NEXT-STATE PREDICTION

  • Next-state prediction can be reparameterized as latent-displacement prediction, with a squared displacement penalty weighted by β.Source: §D.1 · p. 20

D.2 PLANNING WITH BETA

The section estimates predictor scaling across β values and reports that the PushT estimate is close to the scaling predicted by equation 19.

D.2 PLANNING WITH BETA

  • The fixed-point condition is stated as the basis for the planning method's reliance on a quantity specified in the following equation.Source: §D.2 · p. 21
  • For each trained β value, ordinary least squares estimates βest to minimize the predictor's scaling error.Source: §D.2 · p. 21
  • On PushT, estimated predictor scaling is close to the scaling computed in equation 19.Source: §D.2 · p. 21

E PLANNING WITH FIRST-ORDER METHODS

The section reports planning success-rate evaluations and defines a scale-invariant subspace-similarity measure, with eigenvalue-dependent scaling assessed separately.

E PLANNING WITH FIRST-ORDER METHODS

  • Table 4 reports planning success rates for each model, with a random-policy floor for each environment.Source: Table 4 · p. 21

F SUBSPACE SIMILARITY

  • Mean principal-angle cosines measure alignment between centered learned embeddings and Laplacian eigenvectors, excluding the constant eigenvector.Source: F SUBSPACE SIMILARITY · p. 21
  • The similarity measure is invariant to rotations within degenerate eigenspaces and embedding rescaling; eigenvalue-dependent scaling is assessed separately.Source: F SUBSPACE SIMILARITY · p. 21
Loading 2610.01373v1…