For JEPA World Models, Better Planning Starts with Better Geometry

Introduction

A world model used for model-based reinforcement learning has to do more than predict what happens next. It must also provide a useful basis for choosing actions. In many JEPA-style systems, observations are encoded into a latent representation, an action-conditioned predictor models the transition, and candidate action sequences are ranked by the Euclidean distance between a predicted terminal representation and a goal representation.

That pipeline raises a basic question: does geometric proximity in the latent space correspond to task quality? The paper Anisotropic Representations Improve Planning in JEPA World Models argues that the answer is not guaranteed.

The overlooked role of representation geometry

A representation can support accurate predictions and avoid collapse while still inducing a poor planning cost. Isotropic Gaussian regularization treats latent directions similarly, but the task may not. Some dimensions may encode factors that are crucial for control, while others capture differences with much less impact on the final outcome. If all directions are effectively weighted alike, the planner may rank feasible outcomes in an order that disagrees with the actual task cost.

This reframes a familiar world-model problem. The issue is not only whether the dynamics model predicts correctly, but also how prediction errors are measured when the model is used for decision-making. The latent geometry becomes part of the planner’s objective, even when the planner itself is unchanged.

AnisoWM and ΛReg

The proposed AnisoWM changes the training target rather than the planning procedure. Its ΛReg replaces a fixed isotropic Gaussian target with a learnable diagonal covariance. The covariance is subject to fixed-trace and anisotropy constraints, which limit degenerate solutions while allowing different latent directions to receive different amounts of variance.

The predictor architecture, prediction objective, and Euclidean planner remain unchanged. The learned target is used during training to shape the representation; at planning time, the system can still use the same Euclidean distance. This makes the approach a relatively focused intervention: instead of adding a more complicated planner, it attempts to make the existing distance meaningful by improving the geometry it operates on.

The paper also analyzes how prediction demands and the training distribution influence the allocation of target variance. It identifies conditions under which the resulting metric can reduce planning regret, offering a theoretical explanation for why anisotropy may help rather than simply serving as another regularization choice.

Why it matters

Across four visual control environments, AnisoWM improves planning success over LeWorldModel in all four. Its latent planning cost also agrees more closely with task outcomes. The supplied material does not report environment names or detailed numerical results, so the main takeaway should remain methodological: the gain is attributed to representation geometry, not to a new planner or predictor.

The work suggests that evaluating world models should go beyond prediction accuracy and collapse prevention. A useful diagnostic is whether latent distances rank outcomes in the same direction as task performance. For JEPA world models and other latent-control systems, aligning representation geometry with decision costs may become as important as improving dynamics prediction itself.

Source: Hugging Face Daily Papers