GeoWorld

GeoWorld is a predictive World Model for goal-conditioned visual planning. It does not generate pixels. It encodes the current observation and the goal, predicts how candidate actions change the latent state, and searches for the action sequence with the lowest energy.

Core flow

current observation
        ↓
frozen visual encoder
        ↓
Euclidean latent state
        ↓ exponential map
hyperbolic latent state
        ↓ candidate actions
action-conditioned predictor
        ↓
predicted future latent states
        ↓
energy against the goal latent state
        ↓
lowest-energy action sequence

The model has three main additions:

• Hyperbolic JEPA maps Euclidean visual representations into a Hyperbolic Latent Space and predicts transitions along that space
• Geometric Reinforcement Learning refines the predictor so that lower hyperbolic energy corresponds to higher reward over a multi-step path
• Energy-based planning uses the Cross-Entropy Method to search candidate action sequences

The encoder and predictor form an Energy-Based Model. The predictor performs a Latent Rollout to imagine the result of each candidate action sequence.

Why GeoWorld exists

V-JEPA 2 can predict future visual representations from the current representation, robot state, and candidate actions. Its main long-horizon problems are:

• autoregressive prediction errors accumulate when a prediction becomes the input for the next prediction
• the number of possible action sequences grows rapidly as the planning horizon increases
• its ordinary Euclidean latent space does not explicitly represent the branching and hierarchical relations between possible future states

GeoWorld addresses these problems with hyperbolic geometry and rollout-based refinement. The representation geometry and the rollout-based training are complementary. The strongest result comes from using supervised fine-tuning and GRL together.

GeoWorld compared with V-JEPA 2

Part	V-JEPA 2 and V-JEPA 2-AC	GeoWorld
Visual representation	Learns video representations in Euclidean latent space	Maps encoder outputs into a hyperbolic latent space
Action-conditioned prediction	V-JEPA 2-AC predicts future latent states from the current state and candidate actions	Uses an action-conditioned predictor over hyperbolic latent states
Multi-step training	Uses teacher forcing and a two-step rollout loss	Uses teacher forcing, rollout loss, hyperbolic distance, and GRL
Planning	Uses CEM and executes the first action before replanning	Also uses CEM, but scores paths with hyperbolic energy
Demonstrated setting	V-JEPA 2-AC controls a physical Franka arm with visual goals	Evaluated on instructional-video planning with CrossTask and COIN
Main weakness addressed	Prediction drift and a rapidly growing action search space	Tries to make long rollouts more stable and geometry-consistent

GeoWorld reports higher success rates than the matched V-JEPA 2 baseline. The gain is modest at short horizons and becomes more useful as the horizon grows. The ablations also show that the gain does not come from hyperbolic geometry alone.

Meaning of hierarchy

GeoWorld does not implement a high-level planner followed by a low-level controller. Its hierarchy comes from multi-step future expansion. One state can lead to many possible next states, so the future forms a branching tree.

High-level task labels, mid-level actions, and low-level end-effector control remain possible future extensions.

Link to SmolVLA

SmolVLA is a policy that directly maps its observations and instructions to an action chunk. GeoWorld is a world model that takes a candidate action sequence and predicts the future representation that would result from those actions.

A possible combined system is:

SmolVLA proposes action chunks
            ↓
GeoWorld-style model predicts each consequence
            ↓
reject high-energy consequences
            ↓
execute the selected chunk and observe again

The visual state is mapped into hyperbolic space. The action chunk conditions the predictor, so the action chunk does not need the same hyperbolic embedding as the visual state.

Careful interpretations

• A geodesic is the shortest path under the latent-space distance rule, not a guarantee of the shortest physical robot motion
• The exponential map is a geometric mapping, not the ordinary scalar function $e^x$
• Energy is a compatibility score used for search, while a training loss still teaches the model how that score should behave
• GeoWorld tests planning on instructional-video datasets, so camera noise, control error, contact failure, and changing physical environments still need embodied evaluation
• Joint-Embedding Predictive Architecture (JEPA) is one family of predictive models, not another name for every world model

My way of understanding GeoWorld

What I understand from this is that GeoWorld is like an energy-type planner that produces actions over a hyperbolic energy landscape. It is not like a diffusion model that generates pixels.

My way of visualizing the architecture is that the initial image goes through the encoder to get its representation in latent space. The goal image is also encoded and represented in the same latent space. There is then an extra step that converts the ordinary latent representation into hyperbolic space. The action-conditioned predictor uses candidate actions to imagine how this representation will change.

At this point, I understand the motivation for GeoWorld. V-JEPA 2 already predicts future states in latent space, but its errors accumulate over longer horizons. GeoWorld tries to improve this with Hyperbolic JEPA and Geometric Reinforcement Learning. The results also suggest that the representation geometry and rollout-based training are complementary in making the model better.

I also noticed that GeoWorld does not use hierarchical planning in the sense of a high-level planner and a low-level controller. Its hierarchy comes from multi-step future expansion, where one state can branch into many possible future states.

Since my FYP involves manipulation tasks, one idea I find interesting is to combine SmolVLA with a GeoWorld-style model. SmolVLA can produce action chunks, while GeoWorld can predict the consequences before execution. It can then reject action chunks that produce a high energy error. The correction to my initial idea is that the action chunks do not need to be represented in hyperbolic space in the same way as the visual states. They condition the predictor instead.