Understanding and Improving Hyperbolic Deep Reinforcement Learning

AI summary

Overview

Research area: Deep reinforcement learning (RL) combined with hyperbolic representation learning — specifically, why proximal policy optimization (PPO) agents built on hybrid Euclidean–hyperbolic encoders are unstable to train, and how to fix them.

Technical level: Advanced. The paper's core argument rests on closed-form Riemannian gradient derivations in the Poincaré Ball and Hyperboloid models, though the proposed fixes (normalization, a scaling layer, and a categorical value loss) are straightforward to state.

Scope in one sentence: The paper derives where gradients explode or vanish in hyperbolic PPO agents, proposes the agent Hyper++ based on those derivations, and evaluates it on ProcGen with PPO and PPG and on Atari-5 with Double DQN.

What This Paper Is About

Sequential decision-making produces hierarchical, tree-like data: each state branches into multiple possible next states, and this structure grows exponentially with depth. Euclidean space only grows polynomially with radius, so it distorts these hierarchies, while hyperbolic space grows exponentially and can embed them with far less distortion. Hyperbolic deep RL, however, suffers severe optimization failures, and the authors state that formal analysis of why optimization fails has been lacking. This paper supplies that analysis and builds an agent, Hyper++, that trains stably.

Key Contributions

1. Characterization of training issues. For both the Poincaré Ball and the Hyperboloid, the authors formally analyze key operations (hyperbolic multinomial logistic regression layers and the exponential map) and link their gradients to training instability in deep RL, extending earlier analyses by Guo et al. (2022) and Mishne et al. (2023) with new expressions for PPO.
2. Principled regularization. They study the weaknesses of existing stabilization techniques (notably SpectralNorm and the S-RYM recipe of Cetin et al., 2023) and propose RMSNorm combined with a novel learned scaling layer, which bounds embedding norms while avoiding SpectralNorm's stability–capacity trade-off.
3. Hyper++, a strong and general hyperbolic agent. The agent combines RMSNorm with the learned scaling layer, the Hyperboloid model instead of the Poincaré Ball, and a categorical value loss. It is evaluated under PPO and PPG on ProcGen and under Double DQN on Atari-5.

Main Findings

• Large embedding norms destabilize PPO, not just bad luck. Growing embedding norms destabilize gradient-based training in both the Poincaré Ball and the Hyperboloid models and cause trust-region violations despite PPO's clipping mechanism. PPO only constrains importance-sampling ratios on the sampled states in a batch, so the heuristic trust region can fail across unseen states.
• Unregularized hyperbolic PPO collapses early. In the ProcGen BigFish environment, unregularized hyperbolic PPO shows early entropy collapse, a rapid rise in entropy variance across batch states, larger update KL-divergence, and more trust-region violations. The Poincaré Ball's conformal factor λ = 2/(1 − c‖x‖²) explodes, and the paper identifies it as a key driver of instability.
• The Poincaré Ball's conformal factor sits in the gradient denominator. In the HNN++ MLR formulation, the term (1 − c‖x_H‖²)² in the gradient causes explosion near the Poincaré Ball boundary as ‖x_H‖ → 1/√c. Clipping the conformal factor is undesirable because HNN++ MLR logits depend on it, which shifts decision boundaries and leads to performance plateaus. HNN++ removes over-parameterization but does not by itself resolve PPO training instabilities.
• The Hyperboloid avoids conformal-factor problems but not all of them. The Hyperboloid MLR score contains no conformal factor (it is not conformal to Euclidean space) and neither multiplies nor divides by the Euclidean feature norm. However, sinh and cosh in the Jacobian of the Hyperboloid exponential map grow exponentially, faster than √c‖x_E‖, so the Hyperboloid exponential map can still destabilize gradients when Euclidean feature norms grow.
• Regularizing Euclidean embeddings before the hyperbolic layers is necessary. The exponential map Jacobian of the Poincaré Ball decays like O(‖x_E‖⁻¹), but its directional term is highly sensitive to growing ‖x_E‖. The paper's S-RYM discussion concludes that regularizing Euclidean embeddings before the hyperbolic layers is a necessity for stable hyperbolic PPO agents.
• SpectralNorm cannot fix this on its own. Lemma 4.1 shows that SpectralNorm applied only to the last linear layer does not prevent embedding norms from growing through preceding layers; applying it to every layer constrains the Lipschitz constant of all layers and reduces expressivity globally, and it adds computational overhead from power-iteration steps at each forward pass.
• RMSNorm plus learned scaling bounds norms without hurting capacity. Proposition 4.2 shows that applying RMSNorm before the activation of the encoder's last linear layer bounds both the Euclidean embedding norm and the conformal factor, while leaving the other layers' expressivity intact. For common 1-Lipschitz activations such as TanH and ReLU the bounds reduce to ‖x̂‖₂ < 1 and ‖exp₀(x̂)‖ < (1/√c) tanh(√c). RMSNorm is chosen over LayerNorm because mean-centering would distort the hierarchical structure of the hyperbolic embeddings.
• The learned scaling layer recovers lost volume. With ReLU as the last encoder activation and curvature c = 1, the Proposition 4.2 bound restricts the Poincaré Ball radius to ‖x_H‖₂ ≤ 0.76. Since d-ball volume scales as r^d, this costs exponentially in d. Introducing a learnable scalar ξ_θ with ρ_max = atanh(α)/√c and setting α = 0.95 (c = 1) expands the usable radius from 0.76 to 0.95, a volume gain of (0.95/0.76)^d — approximately 1.2 × 10³ more volume at d = 32.
• Removing normalization causes complete failure. In the ProcGen ablations, removing RMSNorm and the 1/√d feature scaling (-RMSNorm) causes complete learning failure, manifesting as large embedding norms and near-zero gradients in the encoder's final layer. The next most important architectural choice is learned scaling (-Scaling), attributed to its synergy with RMSNorm.
• SpectralNorm variants both fail entirely. Replacing RMSNorm with SpectralNorm applied to the complete Euclidean encoder (Hyper++ (SN Full)) or only to the penultimate layer (Hyper++ (SN Penultimate)) results in agents that fail to learn entirely.
• Loss-function variants underperform HL-Gauss. Replacing the categorical HL-Gauss loss with MSE (+MSE) degrades performance, though not uniformly across all games — consistent with Farebrother et al. (2024), who observed that HL-Gauss does not consistently improve performance on all environments. Substituting the C51 distributional loss performs even worse than MSE. Using the Poincaré ball instead of the Hyperboloid (+Poincaré) leads to a modest drop, which the authors call expected given the two models are isometric.
• Head-to-head results on ProcGen. Hyper++ outperforms Poincaré agents with and without S-RYM and the Euclidean baseline with PPO, and wins head-to-head against Hyper+S-RYM in 8/16 games on the train set and 11/16 games on the test set. With PPG, Hyper++ outperforms a strong Euclidean baseline in all metrics, while Hyper+S-RYM is substantially worse than Euclidean; notably, Hyper++ with PPO achieves a higher test IQM than Hyper+S-RYM with PPG.
• Reported headline improvements. The paper reports that Hyper++ improves test return by 52% (PPO + ProcGen), reduces forward pass time by 30% (with the abstract reporting approximately 30% reduction in wall-clock time), and that gains transfer to Double DQN (Atari-5) and PPG (ProcGen). On Atari-5 with Double DQN, Hyper++ strongly outperforms Euclidean and hyperbolic baselines.
• Hyperbolic geometry's advantage over well-regularized Euclidean agents is small in this ablation. Equipping Euclidean agents with HL-Gauss, RMSNorm, and the full regularization combination yields a test IQM of 0.35, slightly better than Hyper++ with the Poincaré ball (IQM = 0.34). Without RMSNorm-type regularization, the Euclidean HL-Gauss variant (Euclidean+Categorical) performs worse than MSE, whereas adding RMSNorm to Euclidean agents improves performance.

Methodology in Plain English

The authors start from a known recipe: an RL agent whose shared encoder is Euclidean, but whose final actor and critic layers operate in hyperbolic space (the hybrid architecture used by Cetin et al., 2023). They then take that architecture apart mathematically.

First, they write down the gradient of the agent's loss with respect to the last Euclidean layer's weights as a product of four terms: the loss derivative, the hyperbolic layer's Jacobian, the exponential map's Jacobian, and the Euclidean layer's derivative. Each term is then analyzed in closed form for both the Poincaré Ball and the Hyperboloid. This decomposition is what lets them say that a specific denominator in the Poincaré MLR gradient explodes near the ball's boundary, and that the Hyperboloid's exponential map has sinh and cosh terms that grow exponentially with the Euclidean feature norm.

Second, they turn each diagnosed failure into a specific fix. Because large Euclidean feature norms are the trigger, they place RMSNorm before the last encoder layer's activation and prove a bound on the resulting embedding norm and conformal factor. Because bounding the radius this tightly would shrink the usable hyperbolic volume exponentially with dimension, they add a learnable scalar that rescales the tangent embeddings up to a controlled radius (α = 0.95). Because the Poincaré Ball's conformal factor is the source of the numerical trouble, they switch to the Hyperboloid and show via the isometry between the two models that the same regularization bounds the Hyperboloid's time component, which in turn bounds its space component. Because value learning in hyperbolic MLR layers is geometrically a hyperplane-distance computation rather than a continuous regression, they replace MSE with a categorical value loss (HL-Gauss).

Third, they evaluate. On ProcGen they train PPO agents for 25M time steps across all 16 environments, normalize returns between random performance and a theoretical or empirically determined maximum, and aggregate with the rliable library (median, interquartile mean, mean, and optimality gap defined as 1 − IQM), reporting means over six seeds with one standard deviation. They repeat with PPG, run an ablation grid on both hyperbolic and Euclidean variants, and test off-policy transfer with Double DQN on Atari-5.

Why This Matters

Impact on research. The paper supplies the missing formal account of why hyperbolic deep RL fails to train, and shows that the answer lies in the interaction between growing Euclidean feature norms and model-specific gradient terms. It also challenges a common assumption: existing stabilization recipes such as SpectralNorm can be insufficient (Lemma 4.1, plus the finding that both SpectralNorm variants fail to learn at all), and a well-regularized Euclidean agent can come within a hair of the hyperbolic one (IQM 0.35 vs 0.34 for the Poincaré variant) — which reframes the debate from "hyperbolic versus Euclidean" to "which geometry, with which regularization."

Real-world applications (the paper itself reports benchmark experiments, not deployments):

• Agents that must reason over hierarchically branching futures, such as the tree of moves in chess used as the paper's opening example.
• Game-playing agents on procedurally generated, long-horizon benchmarks, exemplified by ProcGen BigFish, where an agent grows by eating smaller fish along an irreversible hierarchy of states.
• Any sequential decision-making system where state spaces are naturally tree-structured and Euclidean embeddings distort those relationships — the structural setting the paper argues contributes to deep RL's data inefficiency.

Industry relevance. The reported approximately 30% reduction in wall-clock time (30% reduction in forward pass time) and the avoidance of SpectralNorm's power-iteration overhead matter directly for training cost, since hyperbolic agents have historically been slower to train than Euclidean ones. The claimed transfer across PPO, PPG, and Double DQN indicates the fixes are not tied to one algorithm, which lowers the risk of adopting them in existing RL stacks. Code is released at https://github.com

Authors’ abstract

The exponential volume growth of hyperbolic geometry can embed the hierarchical relationships between states in reinforcement learning (RL) with far less distortion than Euclidean space. However, hyperbolic deep RL faces severe optimization challenges, and formal analysis of why optimization fails is lacking. We identify key factors that determine the success and failure of training hyperbolic deep RL agents. By analyzing the gradients of core operations in the Poincaré Ball and Hyperboloid models of hyperbolic geometry, we show that large-norm embeddings destabilize gradient-based training, leading to trust-region violations in proximal policy optimization (PPO). Based on these insights, we introduce Hyper++, a new hyperbolic deep RL agent that consists of three components: (1) feature regularization guaranteeing bounded norms while avoiding the curse of dimensionality from clipping; (2) a categorical value loss for stable critic training; and (3) a more optimization-friendly formulation of hyperbolic network layers. On ProcGen, we show that Hyper++ guarantees stable learning, outperforms prior hyperbolic agents, and reduces wall-clock time by approximately 30%. On Atari-5 with Double DQN, Hyper++ strongly outperforms Euclidean and hyperbolic baselines. We release our code at https://github.com/Probabilistic-and-Interactive-ML/hyper-rl.