Anisotropic Representations Improve Planning in JEPA World Models
cs.RO
Submitted: 2026-09-29
Updated: 2026-09-29
Project page: https://rkdrn79.github.io/AnisoWM-page
License: http://creativecommons.org/licenses/by/4.0/
The gist: Latent world models learn action-conditioned dynamics in representation space and often score candidate actions by Euclidean distance to a goal representation.
Terminology
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/
Sources
- Beyond Isotropy in JEPAs: Hamiltonian Geometry and Symplectic Prediction
- V-JEPA 2: Self-Supervised Video Models Enable Understanding, Prediction and Planning
- LeJEPA: Provable and Scalable Self-Supervised Learning Without the Heuristics
- Fast LeWorldModel
- Reperesentation Geometry Matters for Planning with JEPA World Models
- When Does LeJEPA Learn a World Model?
- Radial-VCReg: More Informative Representation Learning Through Radial Gaussianization
- Rectified LpJEPA: Joint-Embedding Predictive Architectures with Sparse and Maximum-Entropy Representations
- World Model Control by Trajectory Reachability Metrics
- Predictive but Not Plannable: RC-aux for Latent World Models
- Temporally Centered SIGReg Improves LeWorldModel Representations for Robot Policy Learning
- LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels
- KerJEPA: Kernel Discrepancies for Euclidean Self-Supervised Learning
Related papers
- FMT x: An Efficient and Asymptotically Optimal Extension of the Fast Marching Tree for Dynamic Replanning
- MPCFormer: A physics-informed data-driven approach for explainable socially-aware autonomous driving
- RoboLab: A High-Fidelity Simulation Benchmark for Analysis of Task Generalist Policies
- ARC: A Reasoning Recipe for Robot Foundation Models
- HRDexDB: A 4D Dexterous Grasping Dataset Across Human and Multiple Robot Embodiments
- APT: Action Expert Pretraining Improves Instruction Generalization of Vision-Language-Action Policies