GeoWorld: Geometric World Models

Zeyu Zhang; Danning Li; Ian Reid; Richard Hartley

cs.CVFeb 26, 2026

GeoWorld: Geometric World Models

Authors: Zeyu Zhang, Danning Li, Ian Reid, Richard Hartley

Organizations: ANU · MBZUAI

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.

Figures & tables

Panel of Figure 2 — Figure 2: Energy landscape comparison for V-JEPA 2 [ 3 ] and GeoWorld. We visualize the energy by sweeping two orthonormal tangent-space directions $(\Delta x,\Delta y)$ around a reference latent state. GeoWorlds yields a structured, curvature-aware energy landscape that better reflects geometric structure and hierarchical relations among latent states and improves energy-based planning. For more details see Appendix 4 .

Figure 3: Overview of GeoWorld. Our geometric world model integrates Hyperbolic JEPA for geometry-preserving latent dynamics and Geometric Reinforcement Learning for geodesic-consistent multi-step refinement. Together with energy-based planning using CEM, GeoWorld enables stable and geometry-aware long-horizon visual planning.

Method	CrossTask Dataset [ 88 ]						COIN Dataset [ 71 ]
	T=3			T=4			T=3			T=4
	SR	mAcc	mIoU	SR	mAcc	mIoU	SR	mAcc	mIoU	SR	mAcc	mIoU
Random [ 15 ]	0.01	0.94	1.66	0.01	1.83	1.66	0.01	0.01	2.47	0.01	0.01	2.32
Retrieval-Based [ 87 ]	8.05	23.30	32.06	3.95	22.22	36.97	–	–	–	–	–	–
LLM-Based
LFP [ 42 ]	30.55	59.59	76.86	15.97	50.70	75.30	30.64	54.72	76.86	15.97	50.70	75.30

Table 1: Goal-conditioned visual planning with images on CrossTask [ 88 ] and COIN [ 71 ] datasets. We evaluate multi-step planning over a horizon

T

under the procedural planning setup [ 15 ] , where both observations and goals are specified as images.

Method	CrossTask Dataset [ 88 ]						COIN Dataset [ 71 ]
	T=3			T=4			T=3			T=4
	SR	mAcc	mIoU	SR	mAcc	mIoU	SR	mAcc	mIoU	SR	mAcc	mIoU
LLM-Based
InternVL3.5-241B [ 76 ]	44.03	70.01	84.41	27.65	63.54	80.13	36.54	57.22	89.02	25.46	55.30	88.22
Qwen3-VL-Max [ 79 ]	45.47	70.93	86.18	28.76	62.91	81.51	37.56	57.80	90.46	26.17	57.13	87.56
Gemini 2.5 Pro [ 20 ]	48.91	73.82	90.30	31.53	60.58	84.56	42.07	61.02	92.94	30.20	60.13	84.82

Table 2: Goal-conditioned visual planning with videos on CrossTask [ 88 ] and COIN [ 71 ] datasets. We evaluate multi-step planning over a horizon

T

under the visual planning with videos [ 59 ] setup, where both observations and goals are specified as video clips.

Method	Successful Rate (SR, %)
Method	T=3	T=4	T=5	T=6
Procedural Planning (PP)
Random [ 15 ]	0.01	0.01	0.01	0.01
Retrieval-Based [ 87 ]	8.05	3.95	2.40	1.10
DDN [ 15 ]	12.18	5.97	3.10	1.20
P 3 IV [ 86 ]	23.34	13.40	7.21	4.40

Table 3: Long horzion planning on CrossTask [ 88 ] .

C^{(n)}=\left\|P_{\phi}(\hat{a}_{1:T}^{(n)};s_{1}^{e},s_{1}^{x})-s_{1+T}^{x}\right\|_{1}.

Algorithm 1 Energy-Based Planning with Cross-Entropy Method (CEM)

Figure 1: Gromov $\delta$ -hyperbolicity on CrossTask [ 88 ] .

Method	T=3			T=4
Method	SR	mAcc	mIoU	SR	mAcc	mIoU
GeoWorld ViT-L	44.80	70.54	86.30	30.63	65.46	79.73
w/ FFT	45.20	71.17	87.16	31.34	67.16	80.38
GeoWorld ViT-H	47.79	74.42	88.84	34.51	68.89	82.95
w/ FFT	48.46	74.94	89.10	34.95	69.42	83.47
GeoWorld ViT-g	49.23	76.64	90.61	35.49	71.00	84.50

Table 1: Ablation of frozen encoder vs. fully fine-tuned model for visual planning with videos on CrossTask [ 88 ] .

Method	T=3			T=4
Method	SR	mAcc	mIoU	SR	mAcc	mIoU
V-JEPA 2 ViT-g 384 [ 3 ]	50.16	74.86	91.73	35.01	70.24	85.05
GeoWorld ViT-g 384
SFT Only	50.42	75.13	91.94	35.92	70.79	85.88
GRL Only	51.04	76.48	92.42	36.33	71.04	86.31
SFT + GRL	51.71	77.30	92.95	37.04	71.35	87.04

Table 2: Ablation of Supervised Fine-Tuning (SFT) vs. Geometric Reinforcement Learning (GRL) for visual planning with videos on CrossTask [ 88 ] .

Method	T=3			T=4
Method	SR	mAcc	mIoU	SR	mAcc	mIoU
V-JEPA 2 ViT-g 384 [ 3 ]	50.16	74.86	91.73	35.01	70.24	85.05
GeoWorld ViT-g 384
$\lambda=1,~1-\lambda=0$	50.16	74.88	91.79	34.65	69.48	84.10
$\lambda=0.9,~1-\lambda=0.1$	50.19	74.91	91.84	34.95	70.05	84.85
$\lambda=0.8,~1-\lambda=0.2$	50.25	74.96	91.88	35.31	70.40	85.27

Table 3: Ablation of weighting hyperparameter

\lambda

in Supervised Fine-Tuning (SFT) Only for visual planning with videos on CrossTask [ 88 ] .

Method	T=3			T=4
Method	SR	mAcc	mIoU	SR	mAcc	mIoU
V-JEPA 2 ViT-g 384 [ 3 ]	50.16	74.86	91.73	35.01	70.24	85.05
GeoWorld ViT-g 384
SFT Only	50.42	75.13	91.94	35.92	70.79	85.88
$\beta=0,~\gamma=0.99$	50.48	75.27	91.99	36.07	70.94	86.07
$\beta=0.05,~\gamma=0.99$	51.04	76.21	92.39	36.58	71.13	86.45

Table 4: Ablation of discount factor

\gamma

and the regularization weight

beta

in Geometric Reinforcement Learning (GRL) for visual planning with videos on CrossTask [ 88 ] .

Method	T=3	T=4	T=5	T=6	T=7	T=8
V-JEPA 2 ViT-g 384	50.16	35.01	23.17	16.88	8.26	4.95
SFT (Hyperbolic)	50.42	35.92	23.64	16.97	14.88	11.51
GRL (Euclidean)	50.26	35.47	23.85	17.03	15.12	12.74
GRL (Hyperbolic)	51.04	36.33	24.05	17.82	15.54	13.10
SFT + GRL	51.71	37.04	24.83	18.26	16.09	13.81

Table 5: SR of long horzion planning on CrossTask [ 88 ] videos.

World Model Learning

World model learning builds predictive models of how environments change under actions to support planning, simulation, and decision-making.

See all 414 papers

Reperesentation Geometry Matters for Planning with JEPA World Models
Aug 17, 2026Jiaming Hu, Yan Zheng, Shi Bo +9Representation Learning Neural Representation Geometry
Boston University · Unity Technologies
ProWorld: Progress-Aware Hyperbolic World Models for Long-Horizon Visual Goal Reaching
Aug 3, 2026Zihan Liu, Yuzhe Zhuang, Yuanzu Li +4World Model-Based Planning Long-Horizon Planning
The Hong Kong University of Science and Technology (Guangzhou) · Independent Researcher · The Chinese University of Hong Kong +1
Fast LeWorldModel
Jun 24, 2026Yuntian Gao, Xiangyu XuLatent Dynamics Modeling World Model-Based Planning
Xi’an Jiaotong University

Aug 17, 2026J/K move · Enter open · S save

Reperesentation Geometry Matters for Planning with JEPA World Models

Jiaming Hu, Yan Zheng, Shi Bo, Tian Wang+8

Boston University · Unity Technologies

Representation Learning Neural Representation Geometry World Models for Robotics World Model-Based Planning Joint-Embedding Predictive Architecture Model-Based Planning

Open paper AskPDF arXiv

Joint-embedding predictive world models support planning through latent predictions, but unconstrained joint training can collapse distinct observations to identical embeddings. Two prominent strategies for avoiding collapse are to inherit pretrained features, as in DINO-WM, or to learn representations end-to-end with anti-collapse regularization, as in LeWorldModel (LeWM). Yet avoiding collapse does not ensure that latent distances distinguish outcomes in ways that matter for the task. In object manipulation, for example, success depends on the object's position and orientation relative to the goal. Such task-relevant state information can remain accurately decodable while barely influencing latent distance. The resulting planning cost may fail to reflect how close a predicted outcome is to the task goal. In this paper, we propose SCALE (State-CAlibrated Latent Embeddings), a method that correlates sampled pairwise latent distances with distances in task-relevant state space. Added to LeWM's existing objective, SCALE preserves its architecture, requires privileged state only during training, and adds no planning-time computation. We show that SCALE improves planning success over LeWM across manipulation and navigation tasks with multiple solvers and provide a comprehensive analysis of how SCALE reshapes representation geometry to support planning.

Abstract

World Model Learning

Explore similar work

Reperesentation Geometry Matters for Planning with JEPA World Models

ProWorld: Progress-Aware Hyperbolic World Models for Long-Horizon Visual Goal Reaching

Fast LeWorldModel