GeoWorld: Geometric World Models

AI summary

Overview

• Research area: Computer vision, specifically predictive world models and goal-conditioned visual planning, drawing on hyperbolic geometry, energy-based models, and reinforcement learning.
• Technical level: Advanced. The paper assumes familiarity with JEPA-style joint-embedding predictive architectures, Riemannian/Poincaré-ball hyperbolic geometry, Riemannian geodesic distances, energy-based models, value functions, and the Cross-Entropy Method for trajectory optimization.
• Scope: The paper introduces GeoWorld, a predictive world model that represents latent states on a hyperbolic manifold and refines its predictor with a geometry-aware reinforcement learning objective, and evaluates it on multi-step goal-conditioned visual planning over CrossTask and COIN.

What This Paper Is About

Predictive world models such as JEPA-based approaches plan by minimizing an energy landscape in latent space rather than generating pixels, which lets them reason over multiple future steps. However, their latent representations are learned in ordinary Euclidean space, so they do not capture the geometric or hierarchical relations between states, and their predictions degrade quickly as the planning horizon grows. GeoWorld addresses both problems by moving latent representations onto a hyperbolic manifold and adding a geometric reinforcement learning stage that keeps predicted rollouts on geodesic paths.

Key Contributions

1. Geometric World Model (GeoWorld) with Hyperbolic JEPA (H-JEPA): maps encoder outputs from Euclidean space onto a hyperbolic manifold and learns latent dynamics along hyperbolic geodesics, producing a geometry-consistent energy landscape for multi-step prediction.
2. Geometric Reinforcement Learning (GRL): an energy-based optimization framework that directly refines the predictor (without training an additional policy or reward model) through hyperbolic energy minimization plus triangle-inequality regularization, encouraging geodesic-consistent rollouts.
3. Strong long-horizon planning results: consistent improvements over the previous state-of-the-art predictive world model V-JEPA 2 on CrossTask and COIN, reported as around 3% SR improvement in 3-step planning and 2% SR improvement in 4-step planning.
4. A structured, curvature-aware energy landscape: visualization shows GeoWorld's energy landscape reflects geometric structure and hierarchical relations among latent states better than the V-JEPA 2 landscape (Figure 2).

Main Findings

• Stated headline improvement: Around 3% SR improvement in 3-step planning and 2% SR improvement in 4-step planning compared to V-JEPA 2 across CrossTask and COIN.
• Procedural planning (images as observations and goals, Table 1), CrossTask: GeoWorld ViT-g 384 reaches SR 47.47 at T=3 and 31.48 at T=4, versus V-JEPA 2 ViT-g 384 at 45.58 and 31.36. Across model scales, GeoWorld ViT-L scores 43.89/27.64, ViT-H 45.33/29.19, ViT-g 46.25/30.63.
• Procedural planning, COIN: GeoWorld ViT-g 384 reaches SR 34.85 at T=3 and 27.79 at T=4, versus V-JEPA 2 ViT-g 384 at 34.08 and 23.43. The same variant reaches mAcc 61.86 and mIoU 89.88 at T=3.
• Large mIoU gains on CrossTask procedural planning: GeoWorld ViT-g 384 records mIoU 86.55 at T=3 and 82.48 at T=4, against V-JEPA 2 ViT-g 384 at 69.42 and 69.21. Lower-scale GeoWorld variants show the same pattern (ViT-L 82.93/79.43, ViT-H 84.70/80.16).
• Visual planning with videos (Table 2), CrossTask: GeoWorld ViT-g 384 scores SR 51.71 at T=3 and 37.04 at T=4, versus V-JEPA 2 ViT-g 384 at 50.16 and 35.01. It also exceeds the strongest listed LLM-based planner, GPT-5, which scores 50.03 at T=3 and 30.20 at T=4.
• Visual planning with videos, COIN: GeoWorld ViT-g 384 scores SR 45.29 at T=3 and 33.29 at T=4, versus V-JEPA 2 ViT-g 384 at 42.74 and 31.63, and versus GPT-5 at 43.84 and 32.64.
• Long-horizon stability (Table 3, CrossTask): as the horizon grows from T=3 to T=6, GeoWorld ViT-g 384 achieves SR 51.71, 37.04, 24.83, and 18.26, while V-JEPA 2 ViT-g 384 achieves 50.16, 35.01, 23.17, and 16.88. The paper attributes the degradation of existing models to accumulated geometric drift in Euclidean latent space.
• Ablation details are not in the provided text: the paper directs readers to Appendix 5 for the ablation study, and the appendix content is not included here.
• Energy landscape visualization: sweeping two orthonormal tangent-space directions around a reference latent state shows GeoWorld producing a structured, curvature-aware energy landscape relative to V-JEPA 2 (Figure 2, with details deferred to Appendix 4).

Methodology in Plain English

GeoWorld starts from a pretrained, frozen video encoder that turns each observation into a vector in ordinary Euclidean space. Because a Euclidean vector can be read as a direction at the origin of a hyperbolic space, the model applies an exponential map to push that vector onto a Poincaré ball, producing a hyperbolic latent state where distances between states naturally encode hierarchical relations. The curvature of the ball is treated as a learnable parameter, so the exponential map is implemented as a differentiable hyperbolic projection layer.

A roughly 300M-parameter transformer predictor then takes a sequence of these hyperbolic latent states plus a sequence of actions and predicts the next hyperbolic latent states. Training is supervised first: a teacher-forcing term pulls each one-step prediction toward the encoded ground-truth next state, measured with the hyperbolic distance, and a two-step rollout term feeds the model's own predictions back in to enforce consistency farther ahead. The two are combined with a weighting hyperparameter, and both V-JEPA 2 and GeoWorld are trained with this same two-stage recipe for a fair comparison.

The second stage is Geometric Reinforcement Learning. Here the cost of moving between two states is defined as the hyperbolic distance between the predicted next state and the true next state, and reward is simply the negative of that cost. Maximizing expected cumulative reward therefore becomes equivalent to minimizing the total hyperbolic distance along the predicted trajectory. Because hyperbolic distance obeys the triangle inequality, the authors add a regularization term that penalizes violations across consecutive triplets of predicted states, pushing rollouts to approach geodesic paths. This objective directly optimizes the world model's predictor rather than learning a separate policy or reward network.

At test time, the frozen encoder and trained predictor serve as the world model, and planning searches for the action sequence that minimizes the hyperbolic energy between the predicted final state and the encoded goal state. The search is carried out with the Cross-Entropy Method, using 800 samples, 80 elites, and 10 refinement iterations. Training ran on 4 nodes each with 8 NVIDIA H100 GPUs, 48-core Intel Xeon Platinum 8469C CPUs, and 230 GB of RAM, with inference on a single H100. Both the baseline and GeoWorld use frozen encoders pretrained on VideoMix22M. The supervised stage uses AdamW with weight decay 0.04, a warmup from 7.5×10⁻⁵ to 4.25×10⁻⁴ over 4500 iterations, 85,500 constant iterations, a decay to zero over the final 4500 iterations, and batch size 256. The GRL stage uses warmup from 5.0×10⁻⁵ to 2.0×10⁻⁴ over 2,000 iterations, 18,000 constant iterations, decay over the final 5,000 iterations, batch size 128, discount factor γ=0.99, and triangle-inequality regularization weight β=0.1.

Why This Matters

The paper argues that geometry, not just scale, is a missing ingredient in predictive world models. If latent representations are learned in Euclidean space, the induced energy landscape cannot express geodesic distances or hierarchy, and errors compound over long rollouts. GeoWorld shows that placing latents on a hyperbolic manifold and optimizing the predictor with a geodesic-consistency objective improves multi-step planning, which matters for any system that must plan ahead from video without generating pixels.

Real-world applications (as implied by the paper's task setups):

• Instructional and procedural video understanding, since CrossTask and COIN consist of human daily activities with fine-grained action labels and timestamps.
• Assistive and step-by-step guidance systems, where observations are videos and goals can be specified in natural language (the VPA setup discussed in the related work).
• Robotics and embodied agents that need multi-step action sequences from visual observations without rendering intermediate frames.
• Video-based planning assistants that receive a video clip of the current state and a video clip of the goal and must output the intervening actions.

Industry relevance: the strongest LLM and VLM baselines in the video setup (InternVL3.5-241B, Qwen3-VL-Max, Gemini 2.5 Pro, GPT-5) are large general-purpose models evaluated zero-shot, and GeoWorld's ViT-g 384 variant surpasses them on the reported CrossTask and COIN success rates. That positions geometry-aware predictive world models as a competitive, more targeted alternative for planning tasks where a language-centric model may not capture physical and geometric structure. The efficiency angle also matters: training used 32 H100 GPUs across 4 nodes and inference a single H100, so the approach fits within typical industrial compute for model post-training rather than pretraining from scratch.

Future Directions

• Scaling the horizon further: the paper evaluates up to T=6 on CrossTask; whether the geodesic-consistency property continues to hold at substantially longer horizons is not reported in the provided content.
• Extending beyond goal-conditioned planning: the GRL formulation is tied to a goal state and a hyperbolic energy cost; whether it transfers to settings without an explicit goal is an open question the paper does not address here.
• Cross-modal goals: all evaluation here uses images or video clips as observations and goals. The related work section discusses language-specified goals (VPA), but the provided content does not report GeoWorld experiments in that setting.
• Understanding what the hyperbolic manifold buys: the paper reports an ablation in Appendix 5 that is not included in the provided content, and the energy-landscape visualization details are deferred to Appendix 4, so the precise contribution of each geometric component remains to be examined in the full text.

Target Audience

Researchers and graduate students working on world models, video representation learning, and goal-conditioned visual planning, particularly those already familiar with JEPA-style architectures or energy-based planning. It is also relevant to practitioners in robotics and embodied AI who need multi-step action prediction from video, and to anyone interested in applying non-Euclidean geometry as an inductive bias in deep learning. Because the method section leans on Poincaré-ball geometry and Riemannian distances, readers without that background will need the preliminaries in the appendix to follow the derivations.

Authors’ 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.