Anisotropic Representations Improve Planning in JEPA World Models
Mingu Kang, Yoori Oh, Sookyung Kim, Joonseok Lee
Seoul National University · Ewha Womans University
Published September 29, 2026 · on the feed September 30, 2026 · ▲ 32 · arXiv 2609.37441
The abstract
Latent world models learn action-conditioned dynamics in representation space and often score candidate actions by Euclidean distance to a goal representation. Joint training typically regularizes the representation to prevent collapse, but the resulting representation geometry also determines how terminal errors are weighted during planning. We show that accurate prediction and noncollapsed representations do not guarantee a task-aligned latent planning cost: isotropic Gaussian regularization can induce a geometry that ranks feasible outcomes differently from the task cost. To address this mismatch, we introduce AnisoWM with ΛReg, which replaces the fixed isotropic Gaussian target with a learnable diagonal covariance under fixed-trace and anisotropy constraints. The prediction objective, predictor architecture, and Euclidean planner remain unchanged; the target is used only during training. Our analysis characterizes the prediction-driven allocation of target variance, its dependence on the training distribution, and the conditions under which the induced metric reduces planning regret. Across four visual control environments, AnisoWM improves planning success over LeWorldModel in all four. Its latent planning cost also shows better agreement with task outcomes. Project website: https://rkdrn79.github.io/AnisoWM-page/
The Brief
2:13A short podcast summary of the paper. Written and voiced by AI.
Transcript
Anisotropic Representations Improve Planning in JEPA World Models. By Mingu Kang and three others, from Seoul National University and Ewha Womans University. Published on the twenty-ninth of September. No code repository is linked to it.
Latent world models learn action-conditioned dynamics in representation space, and they often score candidate actions by Euclidean distance to a goal representation. Joint training usually regularizes the representation to prevent collapse. But the geometry that results also determines how terminal errors are weighted during planning. The authors show that accurate prediction, and a representation that has not collapsed, still do not guarantee a planning cost aligned with the task. Isotropic Gaussian regularization can induce a geometry that ranks feasible outcomes differently from the way the task cost ranks them. A planner can then prefer an outcome that is close in latent distance but worse for the task, and do so with a world model that predicts perfectly well. The fault is not in the predictions. It is in the yardstick used to compare them.
AnisoWM addresses that mismatch with a new regularizer. It replaces the fixed isotropic Gaussian target with a learnable diagonal covariance, under fixed-trace and anisotropy constraints. Directions in the representation no longer have to carry equal variance. The total is held fixed, so the representation cannot simply grow or shrink overall, but how that total is shared out across directions is now learned.
Everything else stays as it was. The prediction objective, the predictor architecture and the Euclidean planner are unchanged, and the new target is used only during training. The analysis characterizes how prediction drives the allocation of target variance, how that allocation depends on the training distribution, and the conditions under which the induced metric reduces planning regret.
Across four visual control environments, AnisoWM improves planning success over LeWorldModel in all four. Its latent planning cost also shows better agreement with task outcomes, which is the mismatch the method set out to reduce.
The planner measures distance, but the regularizer decides what distance means. By letting that geometry be anisotropic, AnisoWM brings the latent cost closer to the task's own ranking, and it does so with a change confined to training, leaving the planner exactly as it was.
Paper Quest · 7 chapters · about 40 min
Play this paper
Learn it by doing: make the calls the authors made, guess the result before you see it, and defend it in a viva graded on your own words. Chapter 1 is right here.
Chapter 1 🧩 · The problem
Your predictions are right, and your plans are still wrong. Find out why.
You are building a planner that works inside a learned latent space. An encoder maps images to vectors, a predictor rolls out candidate actions, and the planner picks the outcome closest to the goal by Euclidean distance.
The road ahead
- 1 · The problem ~7 min A world model can predict the future accurately and still steer its planner wrong, because the distance it plans with weights directions unevenly.
- 2 · What came before ~3 min
- 3 · The key idea ~8 min Let the regularizer's target give each latent direction its own learned variance, within limits, while leaving everything else unchanged.
- 4 · Choose your own research ~3 min
- 5 · Build the method ~8 min Train the encoder and predictor with the learned-variance regularizer, drop the target, then plan by rolling out actions and scoring Euclidean distance.
- 6 · The experiment ~6 min You’ll guess: Across four visual control environments, in how many did AnisoWM improve planning success over LeWorldModel?
- 7 · The viva ~5 min Three questions from the examiners, answered in your own words. Win it for the paper’s trophy.
Free, with Google or an email address. Your stars and progress are kept.