Efficient Hypergradient Descent for Inverse Reinforcement Learning

This paper addresses the computational challenges of inverse reinforcement learning (IRL) formulated as a bilevel optimization problem. The authors propose an efficient hypergradient descent method that exploits the structure of the inner objective to avoid expensive inverse-Hessian-vector products.

Problem Setting: The paper studies maximum-likelihood IRL (ML-IRL), which is formulated as a bilevel optimization problem:

• Inner problem: Learns an entropy-regularized policy under reward rϕ, minimizing Linner(θ, ϕ) = Eτ∼pπθ[Σ γ(t−1)(α log πθ(atst) − rϕ(st, at))]

• Outer problem: Measures the fit of the induced policy to expert demonstrations, minimizing Louter(θ) = −Eτ∼pexpert[Σ γ(t−1) log πθ(atst)]

Key Theoretical Contribution: The main theoretical result is Proposition 4.2, which establishes that at the inner optimum, the Hessian of the inner objective is proportional to the Fisher information matrix of the policy. Specifically, under exact inner optimality and realizability (where the policy class is rich enough that DKL(p̃πθ*(ϕ) ∥ p̃ϕ) = 0), the paper proves that:

∇2θ Linner(θ⋆(ϕ), ϕ) = αFθ⋆(ϕ)

where Fθ⋆(ϕ) is the discounted trajectory Fisher information matrix. The authors note this is the first use of this identity for implicit hypergradient computation in ML-IRL. This identity is significant because the standard Natural Hypergradient Descent (NHGD) framework assumes a fixed data distribution, whereas ML-IRL involves a reverse KL divergence over policy-induced trajectory distributions where the sampling distribution depends on policy parameters.

Methodological Contributions:

1. Fisher-based Hypergradient Descent (Algorithm 1): The paper derives sample-based estimators for all components of the implicit hypergradient:

• Outer gradient estimation using Monte Carlo samples from expert demonstrations (Equation 6)

• Fisher matrix estimation from agent trajectories (Equation 5)

• A damped Fisher system (αF̂θ + λI)v̂ = ĝθ for computing the inverse-Hessian-vector product (Equation 7)

• Mixed-derivative-vector product estimation using prefix score-vector products (Equation 8)

2. Fisher Sketching (Algorithm 2): To address the O(d2θ) memory cost of explicitly constructing the Fisher matrix, the authors apply Spectral Compensation Frequent Directions (SCFD) to streamed, discounted policy-score vectors. This reduces storage from O(d2θ) to O(mdθ) for sketch size m, avoiding explicit construction of either the score matrix or the dense Fisher matrix.

Experimental Results:

The paper evaluates the approach on two environments: CartPole (discrete actions, low-dimensional) and LQR (continuous control, larger policy parameterization), comparing against ML-IRL as the baseline.

Damping analysis: A grid search over the damping parameter λ revealed a stability–fidelity trade-off. Weak damping leaves the Fisher system poorly conditioned and destabilizes outer optimization; strong damping stabilizes the solve but drives (αF̂θ + λI)(−1) ≈ (λI)(−1), so the Fisher information contributes little while remaining expensive to compute.

Sketch size analysis: With λ = 10−3, all sketch sizes achieved lower final outer loss than Explicit Fisher, with m = 64 performing best. The paper notes larger sketches do not improve quality monotonically and increase cost; m = 256 is even slower than Explicit Fisher despite a lower loss.

Computational efficiency: Sketching provided the largest benefit in the higher-dimensional LQR setting, reducing peak memory by up to 1.31×, with small sketches preserving or improving runtime. On CartPole, sketching yielded a 1.29× speedup.

Reward comparison: Under matched 24-hour training budgets, on LQR, ML-IRL achieved the best PolicyNLL and EnvReturn, though its margin over Fisher with Sketching is small, while the latter attains the highest RankCorr. On CartPole, all methods reached expert-level EnvReturn; Fisher with Sketching substantially improved PolicyNLL, while Explicit Fisher achieved the best RankCorr. The paper concludes that the Fisher-based methods remain competitive in policy quality while providing strong reward-ranking performance.

Conclusion: The paper develops a Fisher-based implicit hypergradient method for bilevel IRL, establishes the Hessian-Fisher identity under inner optimality, and demonstrates that moderate sketch sizes improve the stability–efficiency trade-off relative to the dense Fisher solver while remaining competitive with ML-IRL under matched computational budgets.

Improvements to AI Systems:

1. Sample-Efficient Inverse Reinforcement Learning (IRL): The Fisher-based hypergradient method enables AI systems to infer reward functions from expert demonstrations using far fewer environment interactions, because it avoids expensive inverse-Hessian-vector products and uses a damped Fisher system for stable gradient estimation. This makes IRL practical for real-world robotics and autonomous driving where expert data is scarce and environment rollouts are costly.

2. Scalable Reward Learning for High-Dimensional Policies: The Fisher sketching algorithm (SCFD) reduces memory from O(d2θ) to O(mdθ), allowing AI systems to learn rewards for large neural network policies (e.g., deep RL agents with millions of parameters) on a single GPU, which was previously infeasible due to memory constraints. This enables reward inference for complex tasks like manipulation or navigation with continuous action spaces.

3. Stable Bilevel Optimization for Meta-Learning: The Hessian-Fisher identity (∇2θ Linner = αFθ) provides a principled way to compute implicit gradients in any bilevel optimization where the inner problem is entropy-regularized policy optimization. This can be generalized to meta-reinforcement learning, where an AI system learns to adapt quickly to new tasks by optimizing outer objectives (e.g., fast adaptation performance) through inner policy updates, with improved numerical stability and faster convergence.

4. Robust Reward Ranking for Interpretable AI: The method’s strong RankCorr performance (correlating learned rewards with true rewards) allows AI systems to produce reward functions that correctly order states/actions by desirability, even when policy performance is similar. This improves interpretability in safety-critical applications (e.g., healthcare or finance), where understanding why an action is preferred is as important as the action itself.

5. Memory-Aware Adaptive Optimization: The damping analysis reveals a stability–fidelity trade-off, which can be used to build an AI optimizer that automatically tunes λ based on the conditioning of the Fisher system. This leads to an adaptive IRL system that self-adjusts its hyperparameters to maintain stable training across different environments, without manual tuning.

6. Efficient Continuous Control from Demonstrations: On LQR and CartPole, the method achieves expert-level performance with lower PolicyNLL (better imitation fidelity) under matched computational budgets. An improved AI system can learn control policies for drones, robotic arms, or autonomous vehicles directly from human demonstrations, with faster wall-clock training and lower memory footprint, making it deployable on embedded hardware.

7. Offline Reward Learning from Logged Data: Since the outer gradient uses Monte Carlo samples from expert demonstrations (not online rollouts), the method can be adapted to offline IRL—learning rewards purely from a static dataset of expert trajectories. This enables AI systems to improve from historical human interaction logs (e.g., clickstreams, surgical recordings) without requiring live environment access, which is crucial for privacy-sensitive or high-risk domains.

Inverse reinforcement learning (IRL) aims to recover a reward function under which the resulting policy reproduces the behavior observed in expert demonstrations. A natural approach is to formulate IRL as a bilevel optimization problem, in which the inner level corresponds to policy optimization under the learned reward and the outer level measures the discrepancy between the induced policy and expert data. However, this formulation is computationally challenging in practice because the outer update requires a hypergradient involving an inverse-Hessian-vector product for the inner objective. We address this challenge by showing that, at the inner optimum, the Hessian of the inner objective is proportional to the Fisher information matrix of the policy, yielding a structured Fisher-based hypergradient closely related to Natural Hypergradient Descent. To address the resulting scalability bottleneck associated with large Fisher matrices, we approximate the required inverse-Fisher-vector product using a streaming spectral sketch, avoiding explicit construction of the Fisher matrix. We evaluate our approach against a first-order stochastic bilevel baseline across discrete- and continuous-control environments. The results demonstrate competitive policy performance and strong reward-ranking quality, while Fisher sketching reduces curvature-storage complexity and can improve computational efficiency relative to an explicit Fisher solver.