PaperlayerView exact revision on arXiv

Source-linked AI summary

SCALE: State-Calibrated Latent Embeddings for JEPA Planning in the Right Geometry

Jiaming Hu, Yan Zheng, Tian Wang

arXiv:2608.16287v1cs.LG

TL;DR

Planning with latent representations depends on whether task-relevant information shapes the planner-facing geometry, not merely whether it is decodable. SCALE calibrates LeWM latent distances to task-relevant state differences and improves all 15 task–solver averages over LeWM.

  • Problem

    Task-relevant state can be decodable yet exert little influence on Euclidean planning costs, leaving unclear what makes latent geometry useful for planning.

  • Method

    SCALE trains an end-to-end LeWM representation by correlating pairwise latent distances with distances in task-relevant state space.

  • Results

    15/15 task–solver averages improve over LeWM across five environments, three solvers, and five compute budgets.

  • Takeaways & Limitations

    For Euclidean latent-space planning, representation quality depends on how task-relevant information is organized geometrically, not only on what information is encoded.

Abstract

from arXiv · show

Joint-embedding predictive world models plan by scoring predicted terminal embeddings against a goal embedding using a cost defined on the representation itself. Two prominent strategies for obtaining non-collapsed representations are to inherit a pretrained feature space, as in DINO-WM, and to learn an embedding end to end with anti-collapse regularization, as in LeWorldModel (LeWM) with SIGReg. These strategies show complementary strengths across tasks. Although task-relevant state is decodable from the full embeddings of both models, DINO-WM's leading principal components usually retain substantially more state information than LeWM's. Because Euclidean planning costs are dominated by high-variance directions, this difference affects how strongly state can influence candidate selection. We propose SCALE (State-CAlibrated Latent Embeddings) to give the end-to-end LeWM representation the favorable geometric property observed in DINO-WM. SCALE induces this property by correlating sampled pairwise latent distances with distances in a standardized task-relevant state space, without replacing LeWM's learned encoder. Across five tasks, three planning solvers, and five compute budgets, SCALE improves every task--solver average over LeWM. A latent-to-state regression control matches or exceeds SCALE's full-embedding decodability yet leaves latent--state distance alignment essentially unchanged and yields less consistent planning gains. SCALE adds a single lightweight training-time regularizer and no planning-time overhead. These results show that planning depends not only on whether task-relevant information is present, but also on whether it shapes the geometry consumed by the planner.

1 Introduction

The introduction argues that planning quality depends not only on task-relevant information being decodable, but also on whether it shapes the latent geometry consumed by Euclidean costs. SCALE addresses this gap by correlating pairwise latent distances with task-relevant state distances and consistently improves planning over LeWM.

  • Motivation: LeWM can retain task-relevant state information while SIGReg does not ensure that this information strongly shapes planning distances.SIGReg encourages a globally wellspread, approximately isotropic embedding distribution, but does not determine which variation factors dominate planner-facing distances.
  • Motivation: DINO-WM reflects task-state variation more strongly in latent distances than LeWM, despite both retaining substantial task-relevant information.The comparison separates decodability from the geometric influence of information on planning.
  • Metric leverage: Under the squared Euclidean metric, higher-variance latent directions contribute more to distances, so weakly varying task-relevant directions can have little effect on candidate ranking.The planner acts directly through the planning cost and has no learned decoder to recover and reweight latent information.
  • SCALE: SCALE trains an end-to-end JEPA to correlate pairwise latent distances with distances in task-relevant state space.It is designed to acquire the favorable latent geometry observed in DINO-WM without replacing the end-to-end learned representation.
  • Results: 15 task–solver averages improved with SCALE over LeWM, with gains persisting across compute budgets.An auxiliary latent-to-state regression control matched or exceeded SCALE’s full-embedding decodability on two tasks but produced less consistent planning improvements.

2 Related Work

Prior work uses latent-predictive world models for control and adds supervision to make JEPA representations more useful for planning. SCALE aligns latent distances with logged simulator state without requiring rewards, value estimates, or policies.

  • Latent-predictive world models: Latent-predictive world models support control by learning policies from imagined trajectories or optimizing action sequences through model rollouts.Reconstruction-free methods predict future latent representations rather than pixels; DINO-WM predicts pretrained DINOv2 patch features and plans by optimizing action sequences.
  • Supervising latent geometry: Several methods supervise JEPA latent geometry, including value-guided alignment with goal-conditioned value and SCALE alignment with logged simulator state.SCALE requires no reward, value estimate, or policy; concurrent methods PhyLatent and PSG-JEPA use privileged state through auxiliary prediction heads.

3 Preliminaries

The model encodes observations into latents, predicts action-conditioned future embeddings, and ranks candidate sequences by representation-space distance to a goal. Although SIGReg prevents collapse by maintaining global latent spread, neither training term ensures that task-relevant differences govern the learned geometry.

  • Latent prediction: The encoder maps observations to embeddings, while the predictor rolls them forward under candidate action sequences.At test time, the rollout starts from the current observation embedding and applies the action-conditioned predictor for each horizon step.
  • Representation-space planning: The planner ranks candidates using a scalar cost computed entirely from the learned representation, specifically squared Euclidean distance in LeWM and DINO-WM.Because planning decisions access the representation only through this cost, the encoder’s distance geometry directly determines candidate selection.
  • Anti-collapse regularization: The one-step latent-prediction loss admits collapsed solutions in which every observation maps to the same point.The target embedding is produced by encoding the next observation, so prediction training alone does not prevent representational collapse.
  • Anti-collapse regularization: SIGReg projects batch embeddings onto random unit directions and penalizes deviations from a univariate standard normal, driving an approximately isotropic Gaussian distribution with global spread.LeWM combines the one-step prediction objective with SIGReg to maintain non-collapsed embeddings.
  • Geometry limitation: Neither latent prediction nor SIGReg specifies which task-relevant differences should dominate pairwise latent distances.The resulting representation may organize its geometry around variation unrelated to the task while still satisfying both objectives.

4 Where Task-Relevant State Lives

Task-relevant state can be accurately decodable from an embedding yet weakly reflected in the latent distances that rank planning candidates. DINO-WM aligns latent distances with state differences more closely than LeWM on most tasks, motivating attention to metric geometry rather than decodability alone.

  • Planner interface: The planner accesses representations only through latent distance, so state matters for planning insofar as it affects the scalar candidate-ranking cost.The analysis measures Spearman rank correlation between held-out pairwise latent distances and task-relevant state distances.
  • Empirical contrast: DINO-WM’s latent distances track task-state differences more closely than LeWM’s on most tasks, despite both models accurately decoding much of the same state from full embeddings.This exposes a dissociation between full-embedding decodability and the metric consumed by the planner.
  • Metric leverage: Under squared Euclidean cost, direction j contributes expected cost 2λj, so eigenvalue magnitude determines its metric leverage in candidate ordering.High-variance directions therefore receive greater influence in the scalar cost.
  • Decodability versus influence: Task-relevant variables stored in low-variance directions can remain highly decodable while contributing negligibly to planning cost, leaving candidates with different states nearly indistinguishable.The planner has no learned head to reweight these directions.
  • Implication: Planning performance therefore depends on whether task-relevant state registers in latent distance, not merely on whether the representation contains that state.This motivates an objective that gives task-relevant variation direct influence over the metric.

5 Method

SCALE trains image-only latent representations whose pairwise distances align with standardized task-relevant state distances, using detached state targets and stratified frame-pair sampling. Its correlation regularizer complements prediction and SIGReg, while an auxiliary decodability control leaves metric geometry unconstrained.

  • State target: SCALE constructs standardized task-relevant states from selected dimensions, sine–cosine angle pairs, and fixed training-set statistics, excluding velocities.The state target is detached during training; the encoder and planner remain image-only at test time, and equal dimensionality or coordinate-wise correspondence is unnecessary.
  • Pair sampling: SCALE samples half of its distinct frame pairs within sub-trajectories and half across episodes to obtain local and global geometric comparisons.Each minibatch contains B sub-trajectories of T frames, flattened into N = BT frames before pair sampling.
  • Distance alignment: The correlation loss standardizes latent and state distance profiles so it depends on relative distance structure rather than global scale.In the ideal case, larger state differences produce proportionally larger latent distances and Lcorr = 0.
  • Training objective: The complete objective combines action-conditioned latent prediction, SIGReg-based anti-collapse regularization, and Lcorr-based redistribution of metric leverage toward task-relevant variation.LSCALE updates the encoder with all three roles, whereas the predictor is updated only by Lpred.
  • Decodability control: Aux replaces Lcorr with a training-only nonlinear state-decoding head, allowing state information to occupy representation directions with negligible metric leverage.The control uses the same selected and standardized state, architecture, and training protocol, testing whether improved decodability alone can match SCALE.

6 Experiments

Across five tasks, three planners, and five compute budgets, SCALE consistently improves LeWM planning by reshaping latent geometry so task-relevant state occupies high-leverage directions. Its gains persist through rollout and candidate ranking, whereas auxiliary state regression can improve decodability without producing comparable planning benefits.

  • Planning performance: SCALE improves LeWM on every task–solver average (15/15), with iCEM gains of 4.3, 3.1, 5.2, 10.0, and 3.1 percentage points across the five tasks.The same improvement holds under CEM and MPPI across a 300× range of rollout budgets.
  • Planning performance: Aux encodes supervised state successfully, yet improves only 13 of 15 task–solver averages and trails SCALE on four of five tasks.Aux exceeds SCALE’s full-embedding decodability on Push-T and Cube, showing that decodability alone does not explain planning gains.
  • Planning performance: SCALE never falls below LeWM on any task–solver pair, unlike DINO-WM, whose relative performance depends on the task.DINO-WM is significantly better on Two-Room and PointMaze, significantly worse on Push-T (p = .031), and not significantly different on Reacher or Cube.
  • Representation geometry: SCALE raises variance in leading principal directions relative to LeWM and Aux on most tasks, increasing their expected leverage under squared Euclidean planning costs.A principal direction with eigenvalue λ_j contributes 2λ_j in expectation to the squared distance between independently sampled embeddings.
  • Representation geometry: SCALE makes the dominant latent subspace more predictive of task-relevant state than LeWM on every task and substantially more predictive than Aux on most listed tasks.Aux remains accurate when predicting state from the full embedding, but organizes less state information in the high-variance subspace.
  • Rollout and candidate ranking: SCALE increases Kendall ranking agreement over LeWM on every task, with especially larger improvements on Cube, Two-Room, and PointMaze.Aux has lower rollout error on Push-T and Cube but smaller Kendall gains, while DINO-WM’s rollout and ranking behavior varies by task.

7 Conclusion · A Implementation details

The paper shows that task-relevant information can be decodable yet geometrically weak for latent-space planning. SCALE calibrates latent distances to task-relevant state differences while preserving LeWM’s end-to-end learning framework.

  • 7 Conclusion: Task-relevant information may be readily decodable from an embedding while remaining geometrically weak.
  • 7 Conclusion: SCALE directly calibrates latent distances with task-relevant state differences during training.
  • 7 Conclusion: The resulting representations preserve LeWM’s end-to-end learning framework while making task-relevant variation geometrically stronger.

A.1 SCALE training pseudocode

SCALE trains LeWM with a correlation regularizer that aligns standardized pairwise latent distances with standardized task-state distances, while retaining the predictive and SIGReg objectives. Training samples balanced within-subtrajectory and cross-episode pairs, using fixed architectural and optimization settings.

  • Correlation regularizer: SCALE computes latent and detached state pairwise Euclidean distances, standardizes them, and minimizes Lcorr = 1 − mean(ẋ ⊙ ẏ).If either distance distribution has standard deviation below δ, Lcorr is set to a graph-connected zero instead.
  • Architecture: The encoder is a ViT-Tiny with patch size 14, 12 layers, 3 attention heads, and width 192, followed by a linear–BatchNorm projector.The predictor is a six-layer causal transformer with 16 attention heads, dropout 0.1, history length 3, and zero-initialized AdaLN action injection.
  • Pair sampling: 4096 pairs per training step are sampled with equal within- and cross-episode proportions for the correlation objective.The algorithm samples K/2 pairs within sub-trajectories and K/2 across episodes.
  • Optimization: Training uses four-frame sub-trajectories, batch size 128, 10 epochs, Adam learning rate 5 × 10−5, and weight decay 10−3.SIGReg uses M=1024 random unit projections with λsig = 0.09; λcorr = 0.1 except for Reacher, where it is 0.15.

B Additional representation and state-selection results

Additional analyses show that probing full frozen embeddings distinguishes supervised from broader state recovery, while expanding the supervised state selection does not improve planning success. On Reacher and Cube, larger selections slightly reduce average success, and reducing Push-T’s selected state also fails to help.

  • Representation results: Held-out R2 probes use full frozen embeddings to evaluate recovery of task-relevant state variables, with broader unsupervised targets on Reacher and Cube.Linear probes use closed-form ridge regression; supervised variables are identified in each caption.
  • State selection: 66.17 to 65.75 average success on Reacher when both joint angles are supervised instead of only the first.The selected state originally supervises a strict subset of the available simulator state.
  • State selection: 62.65 to 62.55 average success on Cube when gripper and block coordinates are added to the five selected variables.Enlarging the supervised selection therefore does not help planning success.
  • State selection: Reducing the six-dimensional Push-T state does not improve performance either.The passage presents this outcome as expected under the account of Sec. 4.2.

B.1 Held-out latent–state distance alignment

Held-out evaluation directly tests the geometric quantity optimized by SCALE: rank alignment between pairwise latent and task-state distances. SCALE increases this alignment over LeWM on every task, whereas Aux leaves alignment essentially unchanged despite high full-embedding state recoverability.

  • Held-out latent–state distance alignment: SCALE directly targets held-out Spearman rank correlation between pairwise latent distances and task-state distances as a consistency check of its geometric objective.The evaluation uses the protocol of Sec. 4.1.
  • Held-out latent–state distance alignment: SCALE substantially increases latent–state rank alignment over LeWM on every task.Table 9 reports held-out Spearman ρ over 1.12M frame pairs per task, with LeWM and DINO-WM repeated for comparison.
  • Held-out latent–state distance alignment: Aux leaves latent–state alignment essentially unchanged despite making task state highly recoverable from the full embedding.This contrasts with SCALE’s alignment improvement over LeWM.
  • Held-out latent–state distance alignment: The Aux result shows that state decodability alone does not ensure that state differences are expressed in the latent metric.It provides an additional check on the distinction motivating SCALE.

C Full planning results

Tables 10–14 report success rates across five tasks, three planning solvers, and five compute tiers, with DINO-WM included as the pretrained-feature reference. Results are summarized as per-tier mean ± SD and an average across T1–T5.

  • Evaluation protocol: CEM and iCEM average six paired evaluation sets, whereas MPPI averages five held-out evaluation sets under its per-task temperature protocol.The tables include DINO-WM as the pretrained-feature reference.
  • Task coverage: Tables 10–14 cover Push-T, Reacher, Cube, Two-Room, and PointMaze success rates for all solvers.Each table reports results by task and solver, with the five tasks represented individually.
  • Reporting format: Each table reports mean ± SD per compute tier and averages results across T1–T5 in the final column.Bold identifies the best method per column within each solver.
Loading 2608.16287v1…