Understanding and Improving Hyperbolic Deep Reinforcement Learning

🍞 Top Bread (Hook): You know how family trees keep splitting into branches as you go down the generations? The farther you look, the more people there are. That shape isn’t flat like a sheet of paper; it spreads out fast like a growing web.

🥬 Filling (The Actual Concept):

• What it is: Many decision problems in games and robots grow like trees, where each move splits into many future possibilities.
• How it works: 1) At any moment, you can choose different actions; 2) Each action creates a new branch; 3) After a few steps, there are tons of branches; 4) This branching grows exponentially.
• Why it matters: Flat (Euclidean) space doesn’t have enough “room” to place all these branches without squishing distances, making it hard for AI to learn clean patterns.

🍞 Bottom Bread (Anchor): In chess, each move opens many possible games. Trying to lay those states on flat paper makes connections warped; using a space that expands faster (hyperbolic) better matches the tree.

🍞 Top Bread (Hook): Imagine you’re riding a scooter on a safe path that says, “Don’t turn the handle too much each second.” That rule keeps you from crashing.

🥬 Filling (The Actual Concept):

• What it is: PPO (Proximal Policy Optimization) is a popular RL training rule that gently updates decisions so they don’t change too much at once (a trust region).
• How it works: 1) Collect a batch of experiences; 2) Compare new policy to old using a ratio; 3) “Clip” that ratio if change is too big; 4) Update only within safe bounds.
• Why it matters: Without this, the policy can zig-zag wildly and stop learning.

🍞 Bottom Bread (Anchor): It’s like practicing basketball shots while only adjusting your aim a little each time; big swings can make you miss more.

🍞 Top Bread (Hook): Picture sticking pins into a rubber sheet (flat) versus into a stretchy trampoline (curved). The trampoline can fit more pins without crowding.

🥬 Filling (The Actual Concept):

• What it is: Hyperbolic geometry is a kind of curved space where area grows fast with distance, perfect for tree-shaped data.
• How it works: 1) Points farther from the center have more room around them; 2) Distances reflect hierarchy naturally; 3) Two common models are the Poincaré Ball and the Hyperboloid; 4) They’re different views of the same space (isometric).
• Why it matters: Matching the space to the data’s shape makes learning simpler and more efficient.

🍞 Bottom Bread (Anchor): In BIGFISH (a ProcGen game), growing can’t be undone and creates a natural order; hyperbolic space can arrange these states neatly.

🍞 Top Bread (Hook): Ever try to draw tiny details near the edge of a balloon? Small moves can stretch into big changes.

🥬 Filling (The Actual Concept):

• What it is: Training instability in hyperbolic RL often happens because gradients can explode or vanish, especially near model boundaries.
• How it works: 1) The Poincaré Ball has a “conformal factor” that blows up near its edge; 2) The Hyperboloid avoids that but still has an exponential map whose Jacobian can grow fast when features get big; 3) In PPO, updates are only constrained on sampled states; 4) Big unseen changes can slip through.
• Why it matters: Instability breaks PPO’s “don’t change too much” promise and crashes learning.

🍞 Bottom Bread (Anchor): In experiments, unregularized hyperbolic agents showed spiking KL-divergence, high clip fraction, and entropy collapse—signs the scooter handle turned too hard.

🍞 Top Bread (Hook): Think of a backpack getting heavier as you walk. If it gets too heavy, every step feels unstable.

🥬 Filling (The Actual Concept):

• What it is: Large Euclidean feature norms (the size of the vectors before mapping into hyperbolic space) are the heavy backpack causing gradient trouble.
• How it works: 1) Encoder makes features; 2) Exponential map sends them into hyperbolic space; 3) If features are large, hyperbolic math (like the conformal factor or hyperbolic sines/cosines) magnifies gradients; 4) Instability follows.
• Why it matters: Without controlling feature size, even clever algorithms wobble.

🍞 Bottom Bread (Anchor): The paper’s plots show bigger norms matched with unstable gradients and higher PPO clipping—clear signs of “too heavy features.”

🍞 Top Bread (Hook): Imagine we tried band-aids before: smaller steps, smoothing, or guardrails all over the track.

🥬 Filling (The Actual Concept):

• What it is: Past fixes like SpectralNorm-everywhere (S-RYM) add speed bumps to all layers to control smoothness.
• How it works: 1) Normalize layer weights by their largest singular value; 2) Make the whole network Lipschitz-smooth; 3) Reduce worst-case jumps.
• Why it matters: It helps, but it can slow training, reduce flexibility, and still leaves some weak spots (like the Poincaré conformal factor).

🍞 Bottom Bread (Anchor): In ProcGen, SpectralNorm helped compared to nothing, but HYPER++ still achieved better stability and about 30% faster wall-clock.

The Gap and Real Stakes: Before this paper, we knew hyperbolic spaces fit hierarchical RL problems, but training was shaky. The missing piece was a principled, simple way to keep features and gradients in check without choking the whole network. That matters for everyday uses: smarter game agents that generalize across levels, robots that plan reliably, or assistants that make step-by-step choices without flipping moods. If training keeps crashing, none of those work in the real world. This paper provides a clear diagnosis (large feature norms + nasty Jacobians) and a clean treatment (RMSNorm + learned scaling + categorical value loss + Hyperboloid), turning hyperbolic deep RL from fragile to practical.

🍞 Top Bread (Hook): Imagine building a treehouse on a sturdy tree that matches the shape of your design, and using the right tools so your screws don’t strip. The right shape and the right tools make the build safe and fast.

🥬 Filling (The Actual Concept):

• What it is (one sentence): The key insight is that stabilizing hyperbolic RL means taming feature sizes and choosing geometry- and loss-friendly tools, so the math doesn’t explode and PPO’s small-step promise holds.
• How it works: 1) Use the Hyperboloid model to avoid the Poincaré Ball’s conformal-factor blow-ups; 2) Keep Euclidean features well-behaved with RMSNorm right before the hyperbolic mapping; 3) Add a learned scaling gate to safely use more of the space without hitting dangerous edges; 4) Train the critic with a categorical loss that matches hyperbolic logits (signed distances) instead of real-number regression; 5) Combine these in PPO (and DDQN) to get steady, strong learning.
• Why it matters: Without these pieces, gradients misbehave, PPO’s trust region leaks, and learning collapses.

🍞 Bottom Bread (Anchor): With HYPER++, ProcGen scores rise and training time drops about 30% versus prior hyperbolic agents, and Atari-5 gets a solid boost too.

Multiple Analogies (three ways):

1. Road and Car: Hyperbolic space is the right road for tree-shaped trips. RMSNorm are the shock absorbers, learned scaling is the speed governor, and the categorical loss is smoother steering—together they prevent skids.
2. Kitchen and Oven: The Hyperboloid is an oven that heats evenly (fewer hot spots than Poincaré). RMSNorm measures ingredients precisely. Learned scaling adjusts the flame. The categorical loss is a recipe that matches the oven—no more burnt edges.
3. Backpacking Map: The world is a branching canyon (hyperbolic). RMSNorm keeps your backpack light. Learned scaling lets you carry a bit more safely. The categorical loss is a clearer legend for your map. You travel farther without getting lost.

Before vs After:

• Before: Hyperbolic RL promised better structure but stumbled: exploding/vanishing gradients, entropy collapse, high clipping, and KL spikes—especially near Poincaré’s boundary.
• After: With HYPER++, features stay bounded, updates stay within PPO’s comfort zone, and the critic’s target is steadier. Learning becomes smooth and strong across ProcGen and Atari-5.

Why It Works (intuition):

• Big Euclidean features are the spark that lights gradient explosions through exponential maps and curvature-sensitive terms. RMSNorm snuffs that spark by normalizing just where it matters (final encoder output), keeping expressivity elsewhere.
• Learned scaling is a safe “expansion valve”: it lets the model use more hyperbolic volume without touching the dangerous outer rim.
• The Hyperboloid avoids the Poincaré conformal-factor minefield, removing a major source of instability.
• The critic’s categorical loss turns wobbly regression into sturdy classification over value bins, matching hyperbolic MLR’s distance-based outputs.

Building Blocks (each as a sandwich):

• 🍞 Hook: You know how holding a ruler at the middle is steadier than at the tip? 🥬 Concept: RMSNorm at the final encoder layer. What: normalize the root-mean-square of features. How: rescale the vector by its RMS so magnitudes stay reasonable. Why: keeps feature norms bounded, preventing gradient blow-ups. 🍞 Anchor: The last feature vector stops swinging wildly, so PPO updates don’t yank the policy around.

• 🍞 Hook: Like a faucet handle that sets max flow even if you turn it too far. 🥬 Concept: Learned feature scaling. What: a single trainable gate multiplies features to expand usable space but caps the max safely. How: apply a sigmoid gate and a preset maximum so mapped points never reach risky zones. Why: avoids the “edge of the world” where math explodes. 🍞 Anchor: The agent can explore more of the hyperbolic ball without running into the cliff at the boundary.

• 🍞 Hook: Multiple-choice tests are sometimes easier to grade than open-ended essays. 🥬 Concept: Categorical value loss (HL-Gauss/C51-style idea). What: predict a distribution over value bins, not a single number. How: turn the critic into a classifier using hyperbolic distances as logits. Why: reduces instability from nonstationary targets and fits the hyperbolic MLR output better than MSE. 🍞 Anchor: The critic stops oscillating, and the actor gets steadier advantages.

• 🍞 Hook: Choose hiking boots made for rocky trails, not sandals. 🥬 Concept: Hyperboloid model. What: a hyperbolic model with no conformal factor, more numerically stable. How: map Euclidean features into its tangent space and use the exponential map to land on the manifold. Why: avoids Poincaré’s edge blow-ups and makes training calmer. 🍞 Anchor: Training curves smooth out and beat prior hyperbolic and Euclidean baselines.

At a high level: Observations → Euclidean Encoder → RMSNorm → Learned Scaling → Exponential Map to Hyperboloid → Hyperbolic Actor & Critic (MLR layers) → PPO (or DDQN) updates with categorical value loss for the critic.

Step 1: Input and Euclidean Encoder

• What happens: The agent receives images (like 64×64 RGB frames in ProcGen or 84×84 gray frames in Atari). A standard CNN (Impala-ResNet for ProcGen, NatureCNN for Atari) extracts features. These are still in flat, Euclidean space.
• Why this exists: CNNs are great at turning pixels into meaningful vectors (edges, shapes, objects), giving the agent a compact summary of “what it sees.”
• Example: In BIGFISH, the CNN might produce a 512-dim feature capturing fish positions and size.

Step 2: RMSNorm at the Final Encoder Layer

• What happens: Right before going hyperbolic, apply RMSNorm to the final Euclidean feature vector. This scales the vector by its root-mean-square, keeping its overall magnitude controlled.
• Why this exists: Large Euclidean norms lead to unstable hyperbolic math (via exponential maps and curvature-sensitive terms). RMSNorm tames feature sizes right where they matter most without limiting the whole encoder.
• Example: A feature vector that sometimes spikes to norm 20 now sits calmly near 1–2, preventing later gradients from exploding.

Step 3: Learned Feature Scaling (Safety Valve)

• What happens: After RMSNorm, pass the features through a single learned gate (a scalar between 0 and a safe cap). This lets the network expand usable hyperbolic space but never approach dangerous boundaries.
• Why this exists: Hyperbolic volume grows fast. We want to use more of it for richer representations, but safely. The gate ensures points stay comfortably away from the rim where math blows up.
• Example: With a cap tuned to allow up to ~95% of the safe radius, the agent gets vastly more volume in higher dimensions without crossing the red line.

Step 4: Exponential Map to Hyperboloid

• What happens: The (scaled) Euclidean vector becomes a tangent vector at the Hyperboloid’s origin. The exponential map then places it onto the curved manifold. The result is a hyperbolic embedding that respects hierarchical distances.
• Why this exists: To benefit from hyperbolic geometry, we must live on the manifold. The exponential map is the standard, smooth way to do that.
• Example: Two similar game states end up close on the Hyperboloid; branching futures can be represented farther out where there’s more room.

Step 5: Hyperbolic Actor and Critic via MLR

• What happens: Both policy (actor) and value head (critic) are Hyperboloid multinomial logistic regression (MLR) layers. They compute logits from signed distances to learned hyperbolic hyperplanes (one per action or per value bin).
• Why this exists: In hyperbolic space, representing decisions as margins to curved hyperplanes aligns with the geometry and avoids extra factors (like Poincaré’s conformal term).
• Example: For actions left/right/up/down, each action’s hyperplane yields a distance logit; softmax turns those into action probabilities.

Step 6: Critic with Categorical Value Loss

• What happens: Instead of predicting one number with MSE, the critic predicts a distribution over value bins (HL-Gauss/C51-style). Targets are constructed from returns and projected onto bins.
• Why this exists: RL targets move over time (nonstationary), making regression unstable. Classification over bins is steadier and matches the critic’s distance-based logits in hyperbolic space.
• Example: If the true return is about 7.2, the target softly lights up bins near 7; the critic learns to place mass there.

Step 7: PPO (or DDQN) Optimization

• What happens (PPO): Use the clipped objective to update the actor within a safe change window and train the critic with the categorical loss. Entropy bonus encourages exploration. Gradients flow back through hyperbolic layers, the exponential map, the learned scaling, and RMSNorm into the encoder.

• Why this exists: PPO’s small, safe steps prevent wild policy swings. The categorical critic provides steadier advantages, reducing actor-critic tug-of-war.

• Example: On BIGFISH, update KL and clip fraction stay low and stable, meaning fewer trust-region violations.

• What happens (DDQN variant): For Atari-5, replace PPO with Double DQN (value-based). Keep the same hyperbolic representation, RMSNorm, and learned scaling. Use DDQN targets to reduce overestimation bias.

• Why this exists: To show the method isn’t PPO-only. Stability from feature control and Hyperboloid geometry helps value-based learning too.

• Example: On Q*BERT and NameThisGame, HYPER++ learns faster and higher than Euclidean and prior hyperbolic baselines.

The Secret Sauce (why these steps together):

• RMSNorm + learned scaling = bounded, robust features → stable hyperbolic mappings and gradients.
• Hyperboloid MLR = avoids Poincaré conformal-factor spikes → calmer, geometry-aligned logits.
• Categorical critic loss = steady targets matching hyperbolic distances → smoother actor updates.
• Combined inside PPO or DDQN, they reduce entropy collapse, clipping spikes, and KL jumps, leading to faster, more reliable learning.

Sandwich mini-explanations for new terms used above:

• 🍞 Hook: Like promising to only take baby steps when learning to skateboard. 🥬 Concept: Trust region in PPO. What: a rule that keeps updates small. How: clip the policy ratio to limit change. Why: avoids wiping out what you’ve learned. 🍞 Anchor: Lower “clip fraction” means the agent is respecting the baby-steps rule.

• 🍞 Hook: Think of a magnifying glass that gets too strong near the edge. 🥬 Concept: Conformal factor (Poincaré). What: a scaling of lengths that explodes at the boundary. How: as points get close to the rim, gradients blow up. Why: makes training unstable. 🍞 Anchor: HYPER++ switches to the Hyperboloid to remove this problem.

• 🍞 Hook: Folding a map neatly so streets line up. 🥬 Concept: Exponential map. What: sends a tangent vector into the curved space along the shortest path. How: start at origin, walk in that direction with curved geometry. Why: it places features correctly on the manifold. 🍞 Anchor: Two similar frames map to neighboring points in hyperbolic space.

The Test: The authors evaluated whether HYPER++ truly stabilizes learning and improves scores.

• What they measured: episode rewards, trust-region signals (KL divergence, clip fraction), entropy and its variance, gradient norms in the encoder, and wall-clock time per step.
• Why: If the method works, we should see higher rewards, smoother updates (lower KL and clipping), healthier exploration (entropy), smaller/steadier gradients, and faster training.

The Competition: HYPER++ was compared to

• Euclidean agents (standard flat geometry),
• A prior hyperbolic PPO agent with SpectralNorm-based regularization (Hyper+S-RYM),
• An unregularized hyperbolic PPO agent (Hyper),
• Variants of HYPER++ where parts were removed or swapped (ablations),
• On Atari-5 with DDQN, the same family of baselines.

The Scoreboard (with context):

• ProcGen (PPO): Using normalized test rewards and robust aggregates (IQM, median, optimality gap), HYPER++ consistently outperformed all baselines. Think of it as getting an A while others get B or C. The wall-clock time improved by about 30% over the previous hyperbolic approach, like finishing the same homework faster and better.
• Per-game results: Hyperbolic agents shine on BIGFISH and DODGEBALL (more hierarchical structure), while on STARPILOT and FRUITBOT all methods bunched near the top. HYPER++ won head-to-head against Hyper+S-RYM in most train and test games.
• PPO stability metrics: Unregularized agents saw entropy collapse (losing exploration), high clip fractions (bumping against trust-region walls), and rising update KL (big policy jumps). HYPER++ kept these in check, indicating smoother, safer learning.
• Atari-5 (DDQN): On NAMETHISGAME, Q*BERT, BATTLEZONE, PHOENIX, and DOUBLE DUNK, HYPER++ strongly beat Euclidean and hyperbolic baselines in human-normalized return. This shows the approach isn’t tied to PPO; it helps value-based learning, too.

Surprising/Notable Findings:

• RMSNorm is crucial: Removing it (or replacing with SpectralNorm in limited spots) caused learning to fail—feature norms grew, and encoder gradients vanished or exploded. This confirms the core diagnosis: control feature size right before the manifold.
• Learned scaling matters: Without it, even with RMSNorm, performance dropped—showing the need to safely expand usable hyperbolic volume.
• Categorical value loss is geometry-friendly: In hyperbolic agents, swapping it for MSE usually hurt. But in Euclidean agents, MSE remained competitive or better, highlighting that the categorical target fits the hyperbolic MLR’s distance-based logits especially well.
• Hyperboloid vs Poincaré: Switching to the Hyperboloid gave a stability edge by removing conformal-factor issues. The two are isometric, but numerically the Hyperboloid behaved better under training.
• Off-batch PPO metrics: Measuring KL and clipping on fresh, on-policy states looked similar to the in-batch metrics, suggesting the improvements generalize beyond just the sampled batch.

Concrete Illustrations:

• Trust-region health: Lower clip fraction and lower update KL in HYPER++ means fewer “too big” changes per update—akin to smoother steering.
• Entropy behavior: HYPER++ maintained healthy exploration longer, avoiding early collapse that traps agents in bad habits.
• Wall-clock: By avoiding SpectralNorm’s repeated power iterations and training more stably, HYPER++ sped up training. Faster and better is a practical win.

Ablations (what breaks when we remove parts):

• No RMSNorm: training collapses (heavy features, unstable gradients).
• No learned scaling: worse performance; RMSNorm alone isn’t enough to use the space efficiently.
• MSE instead of categorical: critic gets shakier in hyperbolic geometry, hurting scores.
• Poincaré instead of Hyperboloid: modest drop, consistent with the theory and stability diagnosis.
• SpectralNorm variants: applying SpectralNorm only to some layers didn’t stabilize; full-encoder SpectralNorm reduced expressivity and still underperformed RMSNorm + scaling.

Big Picture: The experiments show that the three-part recipe—RMSNorm, learned scaling, and categorical critic—paired with the Hyperboloid model converts hyperbolic PPO (and DDQN) from fragile to robust, with strong scores and meaningful speedups.

Limitations:

• Scope of analysis: The paper zeroes in on optimization stability—why training breaks and how to fix it—rather than probing what structures the hyperbolic embeddings learn or when hyperbolic geometry helps most.
• Environment fit: While BIGFISH-like tasks seem to benefit (hierarchical/tree-like), a full map of which RL problems gain the most isn’t provided.
• Algorithm coverage: PPO and DDQN are tested; other families (e.g., actor-critic in continuous control, model-based RL, large-scale distributed systems) remain to be explored.
• Numerical corners: Even the Hyperboloid can become ill-conditioned far from the origin; the method mitigates this but doesn’t eliminate all geometric pitfalls.

Required Resources:

• Hardware: Training used modern GPUs (e.g., A100s). While not extreme by today’s standards, hyperbolic operations add some overhead versus plain Euclidean nets.
• Software: Implementations of hyperbolic layers, exponential maps, and MLR heads. The provided code helps, but users should expect some learning curve.
• Data/Compute: ProcGen (25M steps) and Atari-5 (10M steps) are standard but still nontrivial runs.

When NOT to Use:

• Flat/simple tasks: If the environment doesn’t exhibit hierarchical or branching structure, Euclidean agents may be simpler and equally strong.
• Ultra-tight compute budgets: If you must run on very small devices with minimal math overhead, hyperbolic ops might be a stretch.
• Nonstationary regimes without careful tuning: While categorical loss helps, highly volatile reward scales or sparse-reward setups might still require additional stabilization tricks.

Open Questions:

• Representation probing: What patterns do hyperbolic embeddings capture during RL? Can we visualize or measure hierarchy alignment over time?
• Task taxonomy: Which RL problem families benefit most from hyperbolic geometry, and can we predict that upfront?
• Algorithmic breadth: How do these ideas translate to model-based RL, offline RL, or large-scale distributed actors-learners (e.g., IMPALA-style)?
• Adaptive curvature and gating: Can curvature and scaling caps be learned per task or per layer for even better stability and performance?
• Safety and robustness: Do hyperbolic agents resist distribution shifts or adversarial perturbations better than Euclidean ones due to their geometric bias?

Three-sentence summary:

1. Many RL problems look like branching trees, which fit hyperbolic geometry better than flat space, but training in hyperbolic space used to be unstable. 2) By analyzing where gradients go wrong, this paper shows large feature norms and specific hyperbolic operations (like the Poincaré conformal factor and exponential maps) cause PPO’s trust region to fail. 3) HYPER++ combines RMSNorm, a learned scaling gate, a categorical value loss, and the Hyperboloid model to stabilize training, improve scores, and cut wall-clock time.

Main achievement: Turning hyperbolic deep RL from a promising-but-fragile idea into a practical, strong performer across ProcGen (with PPO) and Atari-5 (with DDQN), through a clear diagnosis and a simple, effective remedy.

Future directions:

• Probe what structures the embeddings learn and when hyperbolic geometry pays off most.
• Extend to more RL families (continuous control, offline RL, model-based planning) and larger scales.
• Explore adaptive curvature, smarter scaling schedules, and hybrid loss designs.

Why remember this: It’s a blueprint for matching geometry to problem shape and aligning training tools to the math. With three focused components—RMSNorm, learned scaling, and categorical critic—hyperbolic RL stops wobbling and starts winning, opening the door to agents that learn faster, generalize better, and spend their compute wisely.