Efficient Hypergradient Descent for Inverse Reinforcement Learning

Technical Summary: Efficient Hypergradient Descent for Inverse Reinforcement Learning

Problem Statement
Inverse Reinforcement Learning (IRL) seeks to recover a reward function from expert demonstrations such that the optimal policy under this reward reproduces the observed behavior. The authors focus on Maximum-Likelihood IRL (ML-IRL), which can be formulated as a bilevel optimization problem. The inner level optimizes an entropy-regularized policy under a learned reward $r_\phi$, while the outer level minimizes the discrepancy (specifically, the KL divergence) between the induced policy and expert data.

The primary computational bottleneck in solving this bilevel problem via implicit differentiation is the requirement to compute the hypergradient. This involves an inverse-Hessian-vector product of the inner objective. Explicitly constructing the Hessian requires quadratic memory ($O(d_\theta^2)$), while iterative approximations (e.g., conjugate gradients) can be computationally expensive and sensitive to conditioning. While Natural Hypergradient Descent (NHGD) has been proposed to replace the Hessian with a Fisher information matrix surrogate, standard NHGD assumes a fixed data distribution. In ML-IRL, the inner objective involves a reverse KL divergence over a policy-induced trajectory distribution, making the direct application of NHGD non-trivial.

Methodology
The paper proposes a method to efficiently compute the hypergradient for ML-IRL by leveraging the geometric structure of the inner problem and employing matrix sketching.

1. Fisher-Hessian Equivalence: The authors prove that under exact inner optimality and realizability (where the policy class is rich enough to perfectly model the expert distribution), the Hessian of the inner objective is proportional to the discounted trajectory Fisher information matrix ($F_\theta$). Specifically, $\nabla^2_\theta L_{inner} = \alpha F_\theta$. This allows the inverse-Hessian-vector product required for the hypergradient to be replaced by an inverse-Fisher-vector product.
2. Hypergradient Estimation: The method decomposes the hypergradient calculation into three components:
   • The outer gradient with respect to policy parameters.
   • The solution to a damped linear system involving the Fisher matrix and the outer gradient.
   • A mixed-derivative term involving the reward gradient and the solution to the linear system.
3. Streaming Spectral Sketching (SCFD): To avoid the $O(d_\theta^2)$ memory cost of storing the dense Fisher matrix, the authors approximate the required inverse-Fisher-vector product using a Streaming Spectral Compensation Frequent Directions (SCFD) sketch. Instead of materializing the Fisher matrix, the algorithm streams weighted policy-score vectors from agent trajectories into a compact sketch of size $m$. This reduces the storage complexity to $O(md_\theta)$. The sketch is then used to solve the damped linear system approximately.

Key Contributions

• Theoretical Derivation: The paper derives sample-based estimators for the ML-IRL implicit hypergradient and proves that the inner Hessian is proportional to the trajectory Fisher information matrix at a realizable inner optimum. This is presented as the first application of this identity for implicit hypergradient computation in ML-IRL.
• Scalable Solver: The authors propose an SCFD-based solver that streams policy-score vectors. This approach avoids constructing the score matrix or the dense Fisher matrix, significantly reducing memory requirements.
• Empirical Validation: The method is evaluated against a first-order stochastic bilevel baseline (single-loop ML-IRL) and an explicit Fisher solver on discrete (CartPole) and continuous (LQR) control environments.

Results
The experiments were conducted under matched computational budgets (24-hour training limit) with identical architectures and hyperparameter tuning.

• Policy Performance: The Fisher-based methods achieved competitive policy quality. On CartPole, all methods reached expert-level returns. On LQR, the single-loop ML-IRL achieved the best Policy Negative Log-Likelihood (NLL) and Environment Return, though the margin over the "Fisher with Sketching" method was small.
• Reward Ranking: The Fisher-based methods demonstrated strong reward-ranking quality. On LQR, the "Fisher with Sketching" method achieved the highest RankCorrelation (0.972) between learned and ground-truth trajectory returns, outperforming both the single-loop baseline and the explicit Fisher solver.
• Efficiency:
  • Memory: Sketching provided the most significant benefit in the higher-dimensional LQR setting, reducing peak memory usage by up to 1.31$\times$ compared to the explicit Fisher solver.
  • Speed: On CartPole, sketching yielded a 1.29$\times$ speedup. On LQR, moderate sketch sizes (e.g., $m=32$) offered a slight speedup (1.04$\times$), while larger sketches ($m=256$) became slower than the explicit solver due to the overhead of the sketching process, despite lower final loss values.
  • Stability: The authors found that the sketched solver remained stable under substantially weaker damping parameters ($\lambda$) than the explicit formulation, suggesting the sketching process may act as an implicit regularizer.

Significance and Claims
The paper claims to address the scalability bottleneck of bilevel IRL by establishing a direct link between the inner Hessian and the Fisher information matrix in the context of reverse-KL optimization. By combining this theoretical insight with streaming matrix sketching, the authors demonstrate that it is possible to compute implicit hypergradients without the quadratic memory cost of explicit Hessian or Fisher matrix construction.

The authors modestly conclude that while the Fisher-based methods are competitive with single-loop ML-IRL in terms of final policy quality, they offer a distinct advantage in reward-ranking performance and computational efficiency (specifically memory reduction) for higher-dimensional problems. The work suggests that moderate sketch sizes offer the optimal trade-off between optimization quality, memory usage, and runtime.