PaperlayerView exact revision on arXiv

Source-linked AI summary

GeoWorld: Geometric World Models

Zeyu Zhang, Danning Li, Ian Reid, Richard Hartley

arXiv:2602.23058v2cs.CVcs.RO

TL;DR

Existing energy-based predictive world models often neglect geometric and hierarchical structure in latent space and degrade over long planning horizons. GeoWorld combines Hyperbolic JEPA with Geometric Reinforcement Learning to support geometry-aware, stable multi-step planning, reporting around 3% SR improvement for 3-step and 2% SR improvement for 4-step planning over V-JEPA 2 across CrossTask and COIN.

  • Problem

    Existing energy-based predictive world models typically use Euclidean latent representations that neglect geometric and hierarchical relations, while their performance degrades as planning horizons increase.

  • Method

    GeoWorld uses Hyperbolic JEPA to map Euclidean representations onto a hyperbolic manifold and Geometric Reinforcement Learning to optimize energy-based predictions with geodesic constraints.

  • Results

    Around 3% SR improvement in 3-step planning and 2% SR improvement in 4-step planning are reported over V-JEPA 2 across CrossTask and COIN.

  • Takeaways & Limitations

    GeoWorld consistently improves long-horizon goal-conditioned visual planning across the CrossTask and COIN benchmarks.

  • Takeaways & Limitations

    GeoWorld’s hierarchical structure arises from multi-step future expansion rather than explicit high-level and low-level planning, and embodied extensions remain future work.

Abstract

from arXiv · show

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.

1. Introduction

Existing energy-based predictive world models neglect latent geometric and hierarchical structure and degrade as planning horizons increase. GeoWorld addresses these challenges with Hyperbolic JEPA and Geometric Reinforcement Learning, improving long-horizon visual planning on CrossTask and COIN.

  • Existing predictive world models typically learn Euclidean representations that fail to preserve meaningful geodesic distances and hierarchical relations among latent states.
  • One-step training data limits expose a multi-step shortcoming: performance degrades rapidly as the planning horizon increases.The passage attributes this weakness to limited and expensive multi-step video data and inadequate long-term temporal dependency modeling.
  • GeoWorld maps Euclidean latent representations onto a hyperbolic manifold, where geodesic distances encode hierarchical relations and dynamics follow hyperbolic geodesics.This produces a geometry-consistent energy landscape for multi-step prediction and planning.
  • Geometric Reinforcement Learning optimizes an energy-based value function through hyperbolic geodesic minimization and triangle-inequality regularization without an additional policy or reward model.The method is designed to enforce geodesic-consistent rollouts and improve long-horizon stability.
  • Around 3% SR improvement in 3-step planning and 2% SR improvement in 4-step planning are reported over V-JEPA 2 across CrossTask and COIN.

2. Related Works

Prior work approaches video world modeling through generative or predictive paradigms, while goal-conditioned visual planning varies by the modalities of observations and goals.

  • Video world models primarily use generative or predictive approaches, differing in whether they generate future visual content or predict latent representations.
  • Goal-conditioned visual planning has developed into distinct setups based on whether observations and goals are images, videos, or language.

3. Method

GeoWorld combines Hyperbolic JEPA with Geometric Reinforcement Learning to model and refine geometry-consistent latent trajectories. Energy-based planning with CEM then searches for low-cost paths between current and goal states in hyperbolic latent space.

  • Hyperbolic JEPA: H-JEPA maps encoder outputs from Euclidean space onto a hyperbolic manifold, preserving hierarchical relations and learning latent dynamics along hyperbolic geodesics.The encoder output is interpreted as a tangent vector and projected onto the Poincaré-ball manifold using the exponential map.
  • Hyperbolic JEPA: The predictor receives action-conditioned sequences of hyperbolic latent states and predicts next-state representations over a planning horizon.Its supervised objective minimizes hyperbolic distance between predicted and true latent representations, targeting minimum-energy geodesic transitions.
  • Training Objective: The supervised objective combines teacher-forcing and rollout losses to improve one-step accuracy and temporal consistency across multiple future steps.The rollout loss recursively feeds predictions back as inputs, while the total loss weights both terms with λ.
  • Geometric Reinforcement Learning: GRL adjusts the predictor’s energy-based value representation so lower hyperbolic energy corresponds to higher expected cumulative reward.Hyperbolic geodesic minimization and triangle-inequality regularization encourage rollout trajectories to follow geodesically consistent paths.
  • Energy-Based Planning: After training, CEM searches for action sequences that minimize goal-conditioned hyperbolic energy between predicted and target latent states.The frozen encoder and predictor define the planning model, while CEM optimizes the action distribution over the selected horizon.

4. Experiments

Experiments evaluate GeoWorld on CrossTask and COIN using image- and video-based goal-conditioned planning, multiple baselines, and horizons up to T = 6. GeoWorld consistently improves multi-step planning, with stronger stability at longer horizons.

  • Benchmarks: CrossTask and COIN provide diverse goal-conditioned visual-planning benchmarks with fine-grained action labels and timestamps.CrossTask contains 4.7K videos across 83 tasks, while COIN contains 11,287 videos across 180 tasks.
  • Evaluation Setup: The evaluation covers procedural image-based planning and video-based planning, with models predicting T actions from observations and goals.Both setups evaluate planning over a specified horizon T; image-based planning uses images, while video-based planning uses video clips.
  • Baselines and Configuration: The experiments compare LLM-based, generative, predictive, random, and retrieval-based planning baselines under specified encoder and training configurations.V-JEPA 2 and GeoWorld use frozen VideoMix22M-pretrained encoders, while VideoWorld uses full finetuning.
  • Main Results: GeoWorld consistently improves multi-step goal-conditioned visual planning across CrossTask and COIN, outperforming prior predictive world models.Under procedural planning, gains appear in SR, mAcc, and mIoU for T=3 and T=4; in video-based planning, GeoWorld outperforms V-JEPA 2 across model scales.
  • Long-Horizon Planning: As the horizon increases from T = 3 to T = 6, GeoWorld maintains the best Success Rate while existing predictive and generative models degrade.The passage attributes competing-model degradation to accumulated geometric drift in Euclidean latent space.

5. Conclusion

GeoWorld targets long-horizon visual planning by preserving geometric and hierarchical structure in latent space. Hyperbolic JEPA and Geometric Reinforcement Learning produce stable improvements over strong predictive world models across multiple planning horizons.

  • Conclusion: GeoWorld preserves geometric structure and hierarchical relations in latent space for long-horizon visual planning.It maps Euclidean representations onto a hyperbolic manifold and uses geodesic-aware latent dynamics.
  • Conclusion: Hyperbolic JEPA and Geometric Reinforcement Learning yield geodesic-consistent rollouts and reduce error accumulation across extended horizons.The method combines hyperbolic energy optimization with triangle-inequality regularization.
  • Conclusion: GeoWorld consistently improves success rates over V-JEPA 2 on CrossTask and COIN for planning horizons from T = 3 to T = 6.The conclusion presents these experiments as evidence for geometry-aware predictive world models and reinforcement learning.

1. Preliminaries

Energy-based world models represent plausible futures as low-energy regions in latent space, while JEPA predicts future embeddings without generating pixels. Hierarchical JEPA extends this approach by assigning abstract dynamics to higher levels and fine-grained predictions to lower levels, with hyperbolic geometry suited to hierarchical structure.

  • Energy-based world models assign low energy to plausible future states and high energy to implausible ones, representing multiple futures as low-energy valleys.
  • Hierarchical planning: Hierarchical planning finds minimum-energy trajectories between abstract latent states, then refines them into fine-grained predictions while enforcing consistency across levels.
  • JEPA: JEPA encodes observations into compact latent representations and predicts future embeddings by minimizing an energy or similarity objective rather than reconstructing pixels.
  • Hierarchical-JEPA: Hierarchical JEPA uses higher-level modules for long-term abstract dynamics and lower-level modules for short-term fine details.
  • Hyperbolic learning: Hyperbolic space is negatively curved and exhibits exponential volume growth, making it suitable for hierarchical or tree-like representations.

2. Motivation

Video-derived world-state transitions form exponentially branching trajectory trees as the prediction horizon increases. This hierarchical structure motivates representing states and planning trajectories in hyperbolic space.

  • Each action choice creates a distinct future, producing N_d = B^d possible states after d prediction steps.
  • The resulting futures form an exponentially branching tree whose depth is the prediction horizon and whose branching factor is the action-space size.

3. Baseline Details

The evaluation compares GeoWorld with LLM-based, generative, and predictive planning baselines, alongside random-selection and retrieval-based references. These baselines span language reasoning, visual generation, direct action prediction, and trajectory retrieval.

  • The evaluation groups baselines into LLM-based, generative world-model, and predictive world-model categories.
  • Experiment baselines: Random Selection samples actions uniformly without considering task context.
  • Experiment baselines: Retrieval-Based planning selects the training trajectory with the smallest visual feature distance between start and goal observations.
  • LLM-based baselines: LLM-based methods infer intermediate actions from language or vision-language reasoning over observed task cues.
  • Generative baselines: Generative baselines produce multi-step plans through pixel, latent-token, or action-sequence generation, often using autoregressive or diffusion-based refinement.
  • Predictive baselines: Predictive baselines model procedural transitions or trajectories directly in feature space, including recurrent, autoregressive, and differentiable latent-planning approaches.

4. Energy Landscape

The energy-landscape visualization probes local latent geometry around a reference state using two tangent-space directions and compares Euclidean with hyperbolic formulations. GeoWorld yields a sharper, curvature-aware landscape associated with more stable long-horizon planning.

  • In visual planning, the displacement coordinates represent latent perturbations rather than physical end-effector offsets.
  • The visualization sweeps two orthonormal tangent-space directions around a reference latent state to probe local energy geometry.
  • Visualization: Figure 2 compares a smooth, nearly symmetric Euclidean paraboloid with a sharper, curvature-aware hyperbolic basin.
  • Geometric comparison: Gromov δ-hyperbolicity evaluation on CrossTask shows GeoWorld has a more concentrated distribution of near-zero δ values than V-JEPA 2.
  • Planning implication: Curvature-aware energy landscapes promote more stable long-horizon planning by guiding CEM along GeoWorld’s hyperbolic geodesics.

5. Ablation Study

The ablations examine curvature learning, encoder fine-tuning, supervised rollout weighting, GRL regularization, and long-horizon stability. Results indicate that geometric constraints and rollout-oriented optimization improve planning, while full encoder fine-tuning offers only incremental gains.

  • Curvature: Learned curvature c decreases from near 1 to around 0.3, producing a flatter yet still hyperbolic latent geometry.The paper reports that smaller curvature reduces exponential-map distortion and supports more stable multi-step planning.
  • Frozen Encoder vs. Fully Fine-Tuned: Full encoder fine-tuning improves SR by approximately 0.3−0.8% and mAcc and mIoU by 0.5−1.2% across T=3 and T=4.These modest gains require substantially more trainable parameters and slower optimization, with diminishing returns for larger backbones.
  • Effectiveness of GRL: GRL consistently improves SR, mAcc, and mIoU over SFT alone, while combining SFT and GRL achieves the strongest performance.The results characterize SFT as initialization and GRL as refinement toward energy-minimizing trajectories.
  • SFT Hyperparameters: Rollout supervision improves planning over the pure one-step objective, with λ = 0.5 providing the strongest overall trade-off.The gains become more pronounced at the longer T=4 horizon as λ decreases from 1 to 0.5.
  • GRL Hyperparameters: β > 0 consistently improves performance over β = 0, while β = 0.1 with γ = 0.99 achieves the strongest results.The paper links triangle-inequality regularization to stabilized prediction and long-horizon consistency.

6. Error Accumulation in Long-Horizon Planning

Autoregressive long-horizon planning accumulates prediction errors over time. GeoWorld’s geometric perspective addresses how those errors accumulate by constraining drift through hierarchical latent structure.

  • Error Accumulation in Long-Horizon Planning: In Euclidean latent spaces, small prediction errors cause unconstrained drift that compounds uniformly over time.The paper contrasts this with hyperbolic geometry, which imposes hierarchical structure on the latent space.

7. Limitation and Future Work

GeoWorld’s hierarchical interpretation comes from branching multi-step futures rather than explicit high-level and low-level task hierarchies. Future work will extend the framework toward explicit sub-task hierarchies and embodied planning.

  • Limitation: The proposed hierarchy arises from exponentially branching d-step futures, not explicit high-level planning and low-level execution.The paper notes that annotated CrossTask and COIN action sequences may appear linear while future expansion induces an implicit tree.
  • Future Work: Future work may add high-level task labels, mid-level actions, and low-level end-effectors, while extending the framework to embodied settings.The authors state that the framework is compatible with embodied planning but defer that extension.
Loading 2602.23058v2…