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
Figure 2: Energy landscape comparison for V-JEPA 2 [ 3 ] and GeoWorld. We visualize the energy by sweeping two orthonormal tangent-space directions (Δx,Δ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, %)
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)=Pϕ(a^1:T(n);s1e,s1x)−s1+Tx1.
Algorithm 1 Energy-Based Planning with Cross-Entropy Method (CEM)
Figure 2: Geometric effects and curvature dynamics : (a) Poincaré disk geodesics connecting x and y under different curvatures K . As the curvature K becomes less negative (i.e., closer to 0 ), the hyperbolic distance between x and y increases, and the geodesic paths bend less and shift closer toward the origin. (b) Geodesic patterns induced by different boundary anchor points. Varying the anchor location produces a characteristic geodesic fan in the Poincaré disk. (c) As the curvature becomes less negative, the space flattens and the distance between x and y decreases. (d) Learnable curvature c during supervised training, showing a gradual decrease from its initialization and convergence to a stable value 0.3.
Method
T=3
T=4
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
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
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
λ=1,1−λ=0
50.16
74.88
91.79
34.65
69.48
84.10
λ=0.9,1−λ=0.1
50.19
74.91
91.84
34.95
70.05
84.85
λ=0.8,1−λ=0.2
50.25
74.96
91.88
35.31
70.40
85.27
Table 3: Ablation of weighting hyperparameter λ in Supervised Fine-Tuning (SFT) Only for visual planning with videos on CrossTask [ 88 ] .
Method
T=3
T=4
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
β=0,γ=0.99
50.48
75.27
91.99
36.07
70.94
86.07
β=0.05,γ=0.99
51.04
76.21
92.39
36.58
71.13
86.45
Table 4: Ablation of discount factor γ 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.
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.
JEPA-style visual world models offer an effective paradigm for visual goal planning by predicting future latent representations. Existing methods typically learn local transition consistency through next-step representation prediction. However, in long-horizon tasks, accurate local prediction alone need not ensure sustained progress toward the goal. First, multi-step rollouts can remain locally plausible while drifting away from goal-relevant trajectories. Second, locally similar future states can correspond to substantially different long-term progress, making them difficult to distinguish in a latent space optimized mainly for local consistency. To address these challenges, we introduce goal-conditioned progress order, a relative ordering of states according to how they advance toward a given goal. This order exhibits an asymmetric, coarse-to-fine structure: early states retain broader future possibilities, while later states concentrate on more specific goal-relevant regions. Such a structure is well suited to hyperbolic geometry. Motivated by this observation, we propose ProWorld, a progress-aware hyperbolic visual world model. ProWorld leverages goal-conditioned progress order to organize visual latent-space dynamics, maintains directional progress within trajectories via hyperbolic entailment learning, and mitigates progress ambiguity among locally similar future states via hyperbolic future discrimination. Furthermore, we design a progress-aware planning objective that scores candidate rollouts by jointly considering proximity to the goal and sustained progress across intermediate states. Experiments on four visual goal-reaching tasks demonstrate that ProWorld achieves an average absolute success-rate gain of 9.67 over LeWM. The code will be released after the paper is accepted.
Zihan Liu, Yuzhe Zhuang, Yuanzu Li +4
The Hong Kong University of Science and Technology (Guangzhou) · Independent Researcher · The Chinese University of Hong Kong +1
Joint-Embedding Predictive Architectures (JEPAs), including recent LeWorldModel (LeWM), have become a promising foundation for reconstruction-free visual world models. For visual planning, however, LeWM evaluates candidate action sequences by repeatedly applying a local one-step latent transition model. This autoregressive rollout makes planning computationally expensive and exposes the predicted trajectory to accumulated latent errors as the horizon grows. We propose Fast LeWorldModel (Fast-LeWM), a fast latent world model that replaces repeated local rollout with action-prefix prediction. Given the current latent and a candidate action sequence, Fast-LeWM encodes its prefixes and predicts the future latents reached after executing those prefixes in parallel. By making action prefixes the basic prediction unit, Fast-LeWM directly models action effects accumulated to different extents over multiple horizons. This prefix-level supervision forces the model to learn how states continuously evolve under different action prefixes, rather than only fitting one-step state transitions. During planning, the predictor can use the prefix token from the encoded action sequence to evaluate the corresponding future latent without explicitly rolling through each intermediate imagined state. Across multiple tasks, Fast-LeWM improves average success over LeWM while substantially reducing planning time, achieving lower open-loop latent loss whose growth becomes significantly slower as the rollout horizon increases.