GeoWorld: Geometric World Models

2026-02-26

AI Summary

This work addresses the limitations of existing energy-based predictive world models, which learn latent representations in Euclidean space and struggle to capture the geometric and hierarchical relationships among states, leading to significant performance degradation in long-horizon prediction. To overcome this, we propose GeoWorld, the first energy-based world model that incorporates hyperbolic geometry. By mapping latent representations onto a hyperbolic manifold through a hyperbolic Joint Embedding Predictive Architecture (JEPA), GeoWorld effectively preserves hierarchical structures inherent in the environment dynamics. Furthermore, we introduce a geometry-aware reinforcement learning method tailored to this space, enabling stable multi-step planning. Empirical results on the CrossTask and COIN datasets demonstrate that GeoWorld improves planning success rates by approximately 3% and 2% over V-JEPA in 3-step and 4-step tasks, respectively.

Technology Category

• Humans and AI: Human-Aware Planning and Behavior Prediction
• Knowledge Representation and Reasoning: Geometric, Spatial, and Temporal Reasoning
• Planning, Routing, and Scheduling: Planning with Language Models

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Energy-based predictive world models provide a powerful approach for multi-step visual planning by reasoning over latent energy landscapes rather than generating pixels. However, existing approaches face two major challenges: (i) their latent representations are typically learned in Euclidean space, neglecting the underlying geometric and hierarchical structure among states, and (ii) they struggle with long-horizon prediction, which leads to rapid degradation across extended rollouts. To address these challenges, we introduce GeoWorld, a geometric world model that preserves geometric structure and hierarchical relations through a Hyperbolic JEPA, which maps latent representations from Euclidean space onto hyperbolic manifolds. We further introduce Geometric Reinforcement Learning for energy-based optimization, enabling stable multi-step planning in hyperbolic latent space. Extensive experiments on CrossTask and COIN demonstrate around 3% SR improvement in 3-step planning and 2% SR improvement in 4-step planning compared to the state-of-the-art V-JEPA 2. Project website: https://steve-zeyu-zhang.github.io/GeoWorld.

Problem

• world models
• geometric structure
• hierarchical relations
• long-horizon prediction
• latent representations

Innovation

• Hyperbolic JEPA
• Geometric World Models
• Energy-based Planning
• Hierarchical Representation
• Geometric Reinforcement Learning