Geometric Reinforcement Learning, or GRL, is the GeoWorld stage that refines the action-conditioned predictor for multi-step planning in a Hyperbolic Latent Space.
Main idea
GeoWorld defines the energy cost of a transition as the hyperbolic distance between the predicted next state and the target next state.
- lower energy means the prediction is closer to the target
- reward is defined as negative energy cost
- the path value combines these rewards across the planning horizon
This aligns lower energy with higher expected reward.
What GRL trains
GRL trains the predictor. It does not add a separate policy or a separate reward model. The visual encoder stays frozen.
The predictor learns to produce multi-step paths with low cumulative hyperbolic distance to the target trajectory.
Triangle-inequality regularization
The regularizer checks consecutive predicted states and encourages their distances to remain consistent with the triangle inequality. This discourages invalid shortcuts and supports geodesic-consistent rollouts.
The regularizer does not reduce the number of predicted steps. It shapes the geometry of the fixed-length predicted path.
Role in GeoWorld
GeoWorld combines supervised teacher-forcing and Latent Rollout training with GRL. Its ablations show that hyperbolic supervised training and GRL each help, while their combination gives the strongest long-horizon results.
GRL affects training. The Cross-Entropy Method performs the action search during planning.
My way of understanding GRL
The key phrase for me is "aligning lower energy with higher expected reward." This confirms my initial guess that reward has a part to play in GeoWorld.
Energy cost is defined as the cost of moving from the current state to the next state. Reward is then defined as the negative energy cost. My initial thought was that negative energy means the model is gaining energy. The clearer interpretation is that maximizing negative cost is the same as minimizing a positive cost.
The path value function extends this idea over a longer horizon. It measures the total discounted reward, or equivalently the total energy cost, from the current state toward the goal.
I initially thought triangle-inequality regularization reduces the total number of predicted steps. The correction is that the planning horizon stays fixed. The regularizer encourages the predicted states to form a consistent path that stays close to a hyperbolic geodesic.
This helps me separate the two roles. GRL trains the predictor to produce better multi-step paths, while CEM searches for the action sequence during planning.