Anisotropic Representations Improve Planning in JEPA World Models

2026-09-29

AI Summary

This study addresses the limitation of existing world models that employ isotropic regularization, which causes a mismatch between latent space geometry and task costs, thereby constraining planning performance. To overcome this, we propose AnisoWM, a framework introducing a novel anisotropic regularization mechanism. Built upon the JEPA architecture, it replaces fixed Gaussian targets with learnable diagonal covariance matrices subject to trace constraints, dynamically adjusting the latent space metric without altering the predictive architecture. Coupled with a Euclidean distance planner, this approach minimizes planning regret. Evaluated across four visual control tasks, AnisoWM consistently outperforms LeWorldModel, significantly improving planning success rates and enhancing the alignment between latent space costs and actual task outcomes.

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/

Problem

• latent world models
• representation geometry
• planning cost
• isotropic regularization
• JEPA

Innovation

• JEPA World Models
• Anisotropic Representations
• Latent Planning
• Diagonal Covariance Regularization
• Visual Control