🤖 machine learning

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

The paper proposes SCALE, a lightweight training-time regularizer that aligns the geometry of end-to-end learned latent embeddings with a standardized task-relevant state space, thereby improving JEPA planning performance by ensuring that high-variance directions in the embedding space effectively capture state information without requiring additional planning-time overhead.

Original authors: Jiaming Hu, Yan Zheng, Tian Wang

Published 2026-08-18

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jiaming Hu, Yan Zheng, Tian Wang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the quest to build machines that can think and act like living beings, researchers have long relied on a concept known as the "world model." Imagine a robot that does not just react to what it sees right now, but instead runs a mental simulation of the future. Before it moves a single muscle, it asks itself, "If I reach for that cup, what will happen?" It imagines the sequence of events, predicts the outcome, and only then chooses the action that leads to the best result. For these simulations to work, the machine needs a way to represent the world that is compact enough to calculate quickly but detailed enough to be useful. If the representation is too messy, the robot gets lost in the noise; if it is too simple, it misses the crucial details needed to succeed. The challenge lies in finding a balance where the machine's internal map of the world is not just accurate, but also organized in a way that makes planning easy.

A recent study by researchers at Boston University and Unity Technologies tackles a subtle but critical flaw in how these mental maps are currently built. They discovered that a machine can learn to understand a task perfectly well, yet still fail to use that understanding when it tries to plan. The problem is not that the information is missing, but that it is hidden in the wrong part of the machine's internal geometry. To fix this, the team developed a new method called SCALE, which reshapes the machine's internal map so that the most important details for a task are the ones that stand out the most when the machine is making decisions.

The researchers began by looking at two different ways machines learn to represent the world. One approach, used in a system called DINO-WM, relies on features learned from vast amounts of data before the robot ever sees a specific task. This method naturally creates a map where the differences between important states are easy to see. The other approach, known as LeWorldModel, teaches the machine from scratch to predict the future. While this method is flexible and learns quickly, the researchers found that the resulting map often hides the most important information. Even though the machine could theoretically extract the correct answer if asked directly, the way it measures distance between different future possibilities was dominated by irrelevant details. It was as if the machine had a library where every book was cataloged correctly, but the librarian only looked at the color of the book covers to decide which one to pick, ignoring the actual content.

To understand why this happens, consider how a machine decides which action to take. It imagines many possible futures and measures the distance between where it thinks it will end up and where it wants to go. In the systems that learn from scratch, the machine's internal measurement of distance is heavily influenced by the most common variations in the data, which are often unimportant for the specific task at hand. The crucial information about the task, such as the position of an object or the angle of a joint, might be encoded in the machine's memory, but it sits in a quiet corner of the map that the planning process rarely notices. The machine can decode the information if forced to, but it does not use it when it needs to make a choice.

The team proposed a solution to force the machine to organize its map differently. They introduced a training rule that directly links the machine's internal sense of distance to the actual physical distance between states in the real world. During training, the system is shown pairs of situations and told to ensure that if two situations are far apart in reality, they must also be far apart in the machine's internal representation. This does not replace the machine's learning process; instead, it acts as a gentle guide, nudging the machine to rearrange its internal map so that the most task-relevant differences become the most prominent features. This ensures that when the machine calculates the cost of a future action, the calculation is driven by the details that actually matter for success.

The results of this approach were striking. The researchers tested their new method, SCALE, across five different environments, ranging from pushing a block to navigating a maze, using three different planning algorithms. In every single test, the system trained with SCALE outperformed the standard system that learned from scratch. The improvement was consistent, regardless of how much computing power was available or which planning algorithm was used. To prove that the success was due to the new geometry and not just better memory, they compared it against a different method that simply tried to make the task information easier to read from the machine's memory. That alternative method could indeed read the information accurately, but it failed to improve the planning performance as consistently as SCALE did. This confirmed that the key to success was not just having the information, but having it arranged in a way that the planner could actually use.

The study reveals a fundamental insight about artificial intelligence: the quality of a representation is not just about what information it contains, but how that information is structured. A machine can hold the correct answer in its mind, but if the path to that answer is obscured by the wrong kind of noise, the machine will fail. By calibrating the internal geometry of the world model to match the reality of the task, the researchers showed that planning becomes significantly more reliable. This work suggests that for machines to truly master complex tasks, we must teach them not only to see the world, but to see it in a way that makes sense for the decisions they need to make.

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

1. Problem Statement

Joint-Embedding Predictive Architectures (JEPA) enable planning by predicting future latent representations and scoring candidate action sequences based on the distance between a predicted terminal embedding and a goal embedding. While models like LeWorldModel (LeWM) successfully learn non-collapsed, end-to-end representations using anti-collapse regularization (SIGReg), a critical disconnect exists between decodability and geometric influence.

The paper identifies that while task-relevant state information is often decodable from the full embeddings of LeWM, it does not necessarily shape the latent geometry used by the planner. Under the standard squared Euclidean cost function, directions with high variance dominate the distance metric. If task-relevant variations are encoded in low-variance directions, they remain easy to decode via a regression head but have negligible influence on the planner's candidate selection. This leads to a situation where a model "knows" the state but fails to "use" it effectively for planning.

The core problem is that existing end-to-end JEPA training objectives do not explicitly ensure that the metric leverage (the weight a latent direction carries in the planning cost) aligns with task-relevant state variations.

2. Methodology: SCALE

The authors propose SCALE (State-CAlibrated Latent Embeddings), a training-time regularizer designed to reshape the latent geometry so that pairwise latent distances correlate with distances in a standardized task-relevant state space.

Key Components:

State Calibration: During training, observations are paired with simulator state variables ( $q$ ). These states are standardized (e.g., using sine-cosine for angles, excluding velocities) to define a target geometry.
Correlation Loss ( $L_{corr}$ ): SCALE samples pairs of frames (both within and across episodes) and computes the squared Euclidean distance in the latent space ( $x_k = \|z_i - z_j\|^2$ $x_{k} = ∥ z_{i} - z_{j} ∥^{2}$ ) and the standardized state space ( $y_k = \|q_i - q_j\|^2$ $y_{k} = ∥ q_{i} - q_{j} ∥^{2}$ ).
- The objective minimizes the correlation loss between the standardized profiles of these distances:
  $L_{corr} = 1 - \frac{1}{K} \sum_{k=1}^K \tilde{x}_k \tilde{y}_k$
- This forces the latent space to organize such that larger state differences result in proportionally larger latent distances.
Integration: The total objective combines the standard LeWM prediction loss ( $L_{pred}$ ), the SIGReg anti-collapse term, and the new correlation term:
$L_{SCALE} = L_{pred} + \lambda_{sig} \text{SIGReg}(Z) + \lambda_{corr} L_{corr}$
Zero Inference Overhead: SCALE operates solely as a training-time regularizer. The encoder and predictor remain image-only at test time; no auxiliary heads or state information are required during planning.

Mechanistic Control (Auxiliary Baseline)

To distinguish between "encoding information" and "shaping geometry," the authors introduce an Aux baseline. This baseline attaches a training-only regression head to predict state $q$ from the latent $z$ . This ensures the state is decodable but leaves the latent metric unconstrained. The results show that high decodability (Aux) does not guarantee improved planning, confirming that geometric alignment is the distinct factor.

3. Key Contributions

Identification of Metric Leverage: The paper formally defines metric leverage as the expected contribution of a latent direction to the squared Euclidean planning cost, proving it is exactly proportional to the direction's variance ( $2\lambda_j$ ). This explains why task-relevant information in low-variance directions fails to influence planning.
Introduction of SCALE: A novel regularizer that directly correlates pairwise latent distances with task-relevant state distances, effectively "calibrating" the latent geometry without replacing the learned encoder.
Empirical Validation: Extensive experiments across five environments (Push-T, Reacher, Cube, Two-Room, PointMaze), three planning solvers (CEM, iCEM, MPPI), and five compute budgets demonstrate that SCALE consistently outperforms the baseline LeWM.

4. Results

Planning Performance: SCALE improves the success rate over LeWM on all 15 task–solver averages. Gains persist across different compute budgets. For example, on the Two-Room task, SCALE achieved a mean success rate of 92.9% compared to LeWM's 82.9% under iCEM.
Comparison with DINO-WM: While DINO-WM (which uses a pretrained feature space) shows strong alignment in some tasks, it underperforms in others (e.g., Push-T). SCALE, however, never falls below LeWM on any task-solver pair, demonstrating a more consistent improvement without inheriting the task-specific biases of a fixed pretrained encoder.
Geometric Analysis:
- Variance Redistribution: SCALE shifts variance toward the leading principal components (high metric leverage) compared to LeWM and the Aux baseline.
- State Predictability: The leading principal components of SCALE are significantly more predictive of task-relevant state than those of LeWM or Aux.
- Rollout Quality: SCALE improves the Kendall's $\tau$ ranking agreement between predicted terminal embeddings and true simulator outcomes, indicating that the reorganized geometry survives the prediction rollout and improves candidate ranking.
Decodability vs. Geometry: The Aux baseline matches or exceeds SCALE in full-embedding state decodability but fails to match SCALE's planning gains. This confirms that simply making state decodable is insufficient; the state must shape the metric used by the planner.

5. Significance and Claims

The paper claims that planning performance depends not only on the presence of task-relevant information in the representation but on how that information is organized geometrically.

Geometric Influence: The authors argue that for Euclidean planners, the "quality" of a representation is defined by whether task-relevant variations occupy high-variance directions (high metric leverage).
End-to-End Flexibility: SCALE achieves the favorable geometric properties observed in pretrained-feature models (like DINO-WM) while retaining the flexibility and adaptability of an end-to-end learned representation.
Minimal Overhead: The method adds only a single lightweight regularizer during training and introduces no computational overhead during the planning phase.

The work concludes that effective latent-space planning requires explicit alignment between the latent metric and the task-relevant state space, a property that standard anti-collapse regularization alone does not guarantee.