arXiv CS open access

Efficient Hypergradient Descent for Inverse Reinforcement Learning

2026 · arxiv_cs

arXiv CS · Papers · 2026

deep-learning machine-learning neural-networks

Conceptio Open Archive

Analyze & explore related research in Conceptio

Connect this record directly into your research workflow or AI agent. Explore citations across open repositories, query full text, and export structured proof bundles.

Explore in Search →Connect via MCP / API ↗

Published as a conference paper at ICOMP 2026

E FFICIENT H YPERGRADIENT D ESCENT VERSE R EINFORCEMENT L EARNING

FOR

IN-

Nikita Sevriukov, Anna Barabanova, Uliana Gagarina, Karina Ivanova & Sofiia Kasaeva HSE University Moscow, Russia {adbarabanova,uigagarina,kaalivanova,sakasaeva,nvsevriukov}@edu.hse.ru Ilya Levin & Marina Sheshukova HSE University Moscow, Russia {ivlevin,msheshukova}@hse.ru

arXiv:2608.11052v1 [cs.LG] 11 Aug 2026

A BSTRACT 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 inverseHessian-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 discreteand 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.

I NTRODUCTION

Inverse reinforcement learning (IRL) seeks to recover a reward function from expert demonstrations. Unlike methods that only imitate observed behavior, learning an explicit reward can support policy re-optimization when the environment dynamics or downstream constraints change Zeng et al. (2022). We study maximum-likelihood IRL (ML-IRL), which admits the bilevel formulation min Louter (θ⋆ (ϕ)) ϕ

s.t.

θ⋆ (ϕ) ∈ arg min Linner (θ, ϕ) , θ

where the inner problem learns an entropy-regularized policy under reward rϕ , while the outer objective measures its fit to expert demonstrations. Bilevel IRL can be approached through nested optimization, simultaneous single-loop updates, or differentiation through the inner solution. Zeng et al. (2022) propose a single-loop ML-IRL method that alternates one policy-improvement step with a stochastic reward update and enjoys finite-time convergence guarantees. This avoids repeatedly solving the inner problem. A complementary approach is implicit differentiation, which explicitly accounts for the sensitivity of the optimized policy to changes in the reward. Its main bottleneck is an inverse-inner-Hessian-vector product: explicit construction requires quadratic memory, while iterative approximations may require many Hessianvector products and can be sensitive to conditioning Ghadimi & Wang (2018). 1

Published as a conference paper at ICOMP 2026

In some problems, the hypergradient can be computed efficiently. For example, Natural Hypergradient Descent (NHGD) Kong et al. (2026) replaces the inner Hessian with a Fisher-information surrogate for likelihood-based inner problems. However, its standard setting assumes a fixed data distribution, whereas the ML-IRL inner objective is a reverse KL divergence over a policy-induced trajectory distribution, whose sampling distribution depends on the policy parameters. The Hessian– Fisher relation is therefore not immediate. We show that, under exact inner optimality and realizability, ∇2θ Linner (θ⋆ (ϕ), ϕ) = αFθ⋆ (ϕ) , where Fθ⋆ (ϕ) is the discounted trajectory Fisher information matrix of the policy. This identity yields a Fisher-based implicit hypergradient for ML-IRL. To avoid the remaining quadratic cost of a dense Fisher matrix, we approximate the required inverse-Fisher-vector product using a streaming SCFD sketch Chen et al. (2020). Our contributions are: • We derive sample-based estimators for the ML-IRL implicit hypergradient and prove that the inner Hessian is proportional to the trajectory Fisher information matrix at a realizable inner optimum. To our knowledge, this is the first use of this identity for implicit hypergradient computation in ML-IRL. • We propose an SCFD-based solver that streams policy-score vectors and avoids constructing either the score matrix or the dense Fisher matrix, reducing storage from O(d2θ ) to O(mdθ ) for sketch size m. • We compare the method with single-loop ML-IRL on CartPole and LQR under matched computational budgets. The Fisher-based methods achieve competitive policy quality and strong reward-ranking performance, while moderate sketches reduce memory and can improve runtime relative to an explicit Fisher solver.

R ELATED W ORK

Inverse reinforcement learning. Inverse reinforcement learning (IRL) infers a reward function from expert demonstrations. Classical approaches include maximum-margin, Bayesian, maximumlikelihood, and maximum-entropy formulations. Maximum-entropy IRL is especially relevant to our setting because it models trajectory probabilities through cumulative rewards and connects reward learning with entropy-regularized control Levine (2018); Bloem & Bambos (2014). The closest related work is the ML-IRL framework of Zeng et al. (2022), which formulates reward learning as a bilevel problem with expert-action likelihood at the outer level and entropy-regularized policy optimization at the inner level. For linear rewards, this formulation is dual to maximumentropy IRL. They propose a single-loop, two-timescale algorithm with finite-time convergence guarantees. We consider a complementary strategy based on implicit differentiation, exploiting inner-level curvature to compute the hypergradient. Related bilevel RL methods model the dependence of an outer reward-design or alignment objective on a learned policy. PARL applies implicit hypergradients to policy alignment Chakraborty et al. (2024), while Shen et al. (2025) replace lower-level optimality with a penalty formulation, and Yang et al. (2024) exploit the soft Bellman fixed point to obtain fully first-order hypergradients. In iterative RLHF, Gauthier et al. (2026) show that ignoring the policy’s influence on subsequent reward-model updates omits an important steering term. These works address broader reward-design and alignment settings, whereas we specialize to ML-IRL and exploit the Fisher geometry of its reverse-KL inner problem. Hypergradient methods. Bilevel hypergradients are commonly estimated either by differentiating through a truncated inner optimization trajectory or by implicit differentiation. The latter requires an inverse-Hessian–vector product, often approximated using Neumann-series estimators or iterative solvers such as conjugate gradients Ghadimi & Wang (2018); Kong et al. (2026). Unrolling may require storing the inner optimization trajectory, while implicit approximations generally require repeated Hessian-vector products. 2

Published as a conference paper at ICOMP 2026

Natural Hypergradient Descent (NHGD) Kong et al. (2026) exploits likelihood structure by replacing the inner Hessian with a Fisher approximation. Its standard formulation assumes a forward-KL or negative-log-likelihood objective under parameter-independent sampling, whereas ML-IRL involves a reverse KL over policy-induced trajectories. We show that, under exact inner optimality and realizability, the inner Hessian equals α times the discounted trajectory Fisher matrix, adapting Fisher-based hypergradients to ML-IRL. A streaming spectral sketch avoids storing this matrix densely. Scalable Fisher approximations. The Fisher information matrix also underlies natural policygradient methods, which precondition policy updates by Fθ−1 ∇θ J(θ) van Heeswijk (2022). Scalable variants include ACKTR, which uses layerwise Kronecker-factored curvature approximations Wu et al. (2017), and rank-one approximations to the inverse Fisher matrix Huo et al. (2026). These methods approximate the natural policy gradient rather than a bilevel hypergradient and impose architectural or fixed-rank structure on the Fisher approximation. Matrix sketching offers a complementary, data-dependent approximation. Frequent Directions Ghashami et al. (2015) maintains a compact spectral approximation of a matrix from streamed rows, while Spectral Compensation Frequent Directions (SCFD) Chen et al. (2020) compensates for spectral information discarded during compression. Related sketches have been used for covariance and inverse-design-matrix approximation in contextual bandits Kuzborskij et al. (2019); Shustova et al. (2025). We apply SCFD to streamed, discounted policy-score vectors and use the resulting sketch to approximate the damped inverse-Fisher-vector product inside the ML-IRL hypergradient, without constructing either the score matrix or the dense Fisher matrix.

P RELIMINARIES

A Markov Decision Process (MDP) is a tuple (S, A, P, r), where S is the state space, A is the action space, P is the transition probability kernel and r : S × A → R is the reward function. For each state s ∈ S and action a ∈ A, P(·|s, a) stands for a distribution over the states in S, that is, the distribution over the next states given that action a is taken in state s. For each action a ∈ A and state s ∈ S, r(s, a) gives the reward received when action a is taken in state s. A policy π is a map from states to probability distributions over actions, that is, π : S → ∆(A), where ∆(A) denotes the set of probability distributions over A, and π(· | s) is the distribution over actions in state s. An MDP describes the interaction of an agent and its environment. At time t, given the current state Xt ∈ S, the agent chooses an action At ∼ π(· | Xt ), and the next state is sampled according to Xt+1 ∼ P(· | Xt , At ). Let πθ be a differentiable policy parameterized by θ ∈ Θ ⊆ Rd . Let pπθ denote the trajectory distribution induced by P and policy πθ . We define the discounted trajectory Fisher information matrix by "∞ # X t−1 ⊤ Fθ := Eτ ∼pπθ γ ∇θ log πθ (at | st )∇θ log πθ (at | st ) . t=1

P ROBLEM STATEMENT

We consider an infinite-horizon MDP M = (S, A, P, ν, r) with initial state distribution ν and transition kernel P(s′ | s, a). To incorporate discounting, we use the standard absorbing-state construction. Introduce an absorbing state sf and a fixed action af , such that P(sf | sf , af ) = 1,

rϕ (sf , af ) = 0,

πθ (af | sf ) = 1.

For a fixed discount factor γ ∈ (0, 1), define the modified transition kernel e ′ | s, a) = γP(s′ | s, a) + (1 − γ)I{s′ = sf }. P(s f = (S ∪ {sf }, A ∪ {af }, P, e ν, r). Under this construction, We denote the corresponding MDP by M the process continues with probability γ at each time step and transitions to the absorbing state with 3

Published as a conference paper at ICOMP 2026

probability 1 − γ. Since both the reward and the log-policy term are zero in the absorbing state, this construction is equivalent to weighting the contribution of time step t by γ t−1 . Let π be a policy and let a trajectory be denoted by τ = (s1 , a1 , . . . , sT (τ ) , aT (τ ) , sf , . . . ),

s1 ∼ ν ,

e it can be easily where we have set T (τ ) = inf{t ≥ 1 : st+1 = sf }. Due to the structure of P, verified that P[T (τ ) < ∞] = 1. The trajectory distribution induced by π in the modified MDP is given by ∞ ∞ Y Y e t+1 | st , at ) , peπ (τ ) = ρ̃(τ ) π(at | st ) , ρ̃(τ ) = ν(s1 ) P(s t=1

t=1

where ρ̃(τ ) denotes the environment-dependent part of the trajectory probability. We parameterize the state-action reward function by ϕ. The learned reward assigned to a trajectory is defined as the undiscounted sum of per-step rewards: ∞ X Rϕ (τ ) = rϕ (st , at ) . t=1

Note that since T is finite almost surely, Rϕ (τ ) is also finite almost surely. Following the work Levine (2018), the learned reward induces a trajectory distribution that assigns higher probability to trajectories with larger cumulative reward. Since the numerical scale of the learned reward is not fixed a priori, we introduce a fixed temperature parameter α > 0: Z Rϕ (τ ) Rϕ (τ ) 1 , Zϕ = ρ̃(τ ) exp dτ < ∞ . ρ̃(τ ) exp p̃ϕ (τ ) = Zϕ α α Here, α is treated as a fixed hyperparameter rather than a learnable reward parameter. Smaller values of α make the induced distribution more concentrated around high-reward trajectories, while larger values of α produce a smoother distribution. This temperature parameter controls the relative scale of the reward term in the maximum-entropy trajectory model. We now assume that the policy is parameterized by θ, i.e. π = πθ . Let p̃πθ (τ ) denote the trajectory f. We also assume access to expert demonstrations distribution induced by this policy in MDP M generated by an expert policy πexpert , which represents the behavior of the real object or agent. We denote the corresponding trajectory distribution by p̃expert (τ ). The next lemma guarantees that the absorbing-state construction allows expectations under the modified process to be evaluated using trajectories from the original MDP. Lemma 4.1 (Absorbing-state representation of discounting). Let f (s, a) be any measurable function such that f (sf , af ) = 0 and the expectations below are finite. Then "∞ # "∞ # X X t−1 Eτ ∼epπθ f (st , at ) = Eτ ∼pπθ γ f (st , at ) . t=1

t=1

Therefore, we formulate the following bilevel optimization problem: DKL p̃expert (τ ) ∥ p̃πθ∗ (ϕ) (τ ) → min ,

(1)

s.t.

θ∗ (ϕ) ∈ arg min DKL (p̃πθ (τ ) ∥ p̃ϕ (τ )) . θ

Now, using Lemma 4.1, we provide an equivalent formulation of this problem in Proposition 4.1. Proposition 4.1. Set "∞ # X t−1 Louter (θ) := −Eτ ∼pexpert γ log πθ (at | st ) , (2) t=1

Linner (θ, ϕ) := Eτ ∼pπθ

"∞ X

# γ

t−1

(α log πθ (at | st ) − rϕ (st , at ))

t=1

The optimization problem defined in 1 is equivalent to Louter (θ∗ (ϕ)) → min s.t. θ∗ (ϕ) ∈ arg min Linner (θ, ϕ) . ϕ

(3)

Published as a conference paper at ICOMP 2026

The proof of Proposition 4.1 is provided in Section B. This type of bilevel IRL problem formulation is not novel and has appeared in, e.g., Zeng et al. (2022). Here, the inner objective is equivalent to entropy-regularized RL, and the outer objective evaluates a policy on expert trajectories. The KL formulation of this problem is primarily interesting in the context of hypergradient descent because it helps us derive explicit formulas for the hypergradient, as described further in this section. We consider a gradient-based approach to solve 3. Computing the gradient ∇ϕ Louter (θ∗ (ϕ)) is not straightforward, since Louter depends on ϕ implicitly through the inner solution θ∗ (ϕ). Changing ϕ changes the reward-induced distribution p̃ϕ , which changes the solution of the inner optimization problem and, in turn, the value of the outer objective. However, it is possible to derive −1 ⊤ ∇ϕ Louter (θ⋆ (ϕ)) = − [∇θ,ϕ Linner (θ⋆ (ϕ), ϕ)] ∇2θ Linner (θ⋆ (ϕ), ϕ) ∇θ Louter (θ⋆ (ϕ)) , (4) which, for completeness, we prove in Section A. In general, this gradient is impractical to compute since it requires estimating the inverse Hessian of the inner loss. However, in some problems it can be computed efficiently by exploiting the structure of the loss, for instance, when the Hessian can be represented by a Fisher information matrix Kong et al. (2026). In our setting, however, the inner problem involves reverse-KL minimization. Nevertheless, it can be shown that at the optimum θ∗ (ϕ), the Hessian of the inner loss is proportional to the Fisher information matrix. We establish this result in Proposition 4.2. Proposition 4.2. Assume that the inner problem is solved exactly and that the policy class is rich enough so that DKL peπθ∗ (ϕ) (τ ) ∥ peϕ (τ ) = 0 . Then the Hessian of the inner objective at the inner optimum is proportional to the Fisher information matrix of the policy-induced trajectory distribution: ∇2θ Linner (θ⋆ (ϕ), ϕ) = αFθ⋆ (ϕ) . The proof of Proposition 4.2 is provided in Section B. Thus, when the inner solution is sufficiently close to the optimum, this result motivates approximating the Hessian using the Fisher information matrix. Given a batch of trajectories Dagent from pπθ of size Nagent , we can estimate the Fisher information matrix as Fbθ =

1 Nagent

|τ | X

γ t−1 ∇θ log πθ (at | st )∇θ log πθ (at | st )⊤ .

(5)

τ ∈Dagent t=1

This representation allows us to estimate the inverse matrix efficiently using rank-one updates, as described in the next section.

P ROPOSED M ETHOD

In this section, we describe how the implicit hypergradient in Equation 4 can be evaluated efficiently from sampled trajectories. Its computation requires three quantities: the outer gradient, the inverse-inner-Hessian-vector product, and the mixed-derivative-vector product. We derive practical estimators for each of them and then introduce a sketching-based approximation that avoids explicit construction of the Fisher matrix. 5.1

F ISHER -BASED H YPERGRADIENT D ESCENT

We first consider the outer-gradient term in Equation 4. Differentiating the discounted outer objective in Equation 2 with respect to the policy parameters gives "∞ # X t−1 ∇θ Louter (θ) = −Eτ ∼pexpert γ ∇θ log πθ (at | st ) . t=1

Therefore, given expert demonstrations Dexpert from pexpert of size Nexpert , we use the Monte Carlo estimate |τ | X X 1 gbθ = − γ t−1 ∇θ log πθ (at | st ). (6) Nexpert t=1 τ ∈Dexpert

Published as a conference paper at ICOMP 2026

The corresponding discounted identity is proved in Proposition B.1. The second term in Equation 4 involves the inverse Hessian of the inner objective. By Proposition 4.2, at the exact inner optimum this Hessian is proportional to the Fisher information matrix. Furthermore, its discounted trajectory form admits the stepwise representation "∞ # X 2 ⋆ t−1 ⊤ ∇θ Linner (θ (ϕ), ϕ) = αEτ ∼pπθ⋆ (ϕ) γ ∇θ log πθ (at | st )∇θ log πθ (at | st ) . t=1

Its derivation is provided in Proposition B.2. The hypergradient does not require the inverse Fisher matrix itself, but only its product with the b as the solution of the damped Fisher system outer gradient. We therefore introduce v bθ = gbθ , αFbθ + λI v (7) where λ > 0 is a regularization parameter. It remains to evaluate the mixed-derivative term in Equation 4. Direct construction of the mixed derivative matrix is unnecessary. For each trajectory and time step, define the prefix score-vector product t X bθ . bθ,t (τ ) = ∇θ log πθ (ak | sk )⊤ v k=1

The mixed-derivative-vector product is then estimated as   |τ | X X 1 b=− u ∇ϕ  γ t−1 rϕ (st , at )bt (τ ) . Nagent t=1

(8)

τ ∈Dagent

The derivation is given in Proposition B.3 and Corollary B.1. Computationally, all products b are obtained without explicitly forming the policy-score Jacobian, and the ∇θ log πθ (at | st )⊤ v prefix values bt (τ ) are then computed by a single cumulative-sum pass. b = −b b The complete procedure Combining these components gives h u and the update ϕ ← ϕ − ηϕ h. is summarized in Algorithm 1. Algorithm 1 Hypergradient Descent with Explicit Fisher Input: expert dataset Dexpert , iterations K, damping λ 1: Initialize θ, ϕ 2: for k = 1, . . . , K do 3: θ⋆ (ϕ) ≈ arg min Linner (ϑ, ϕ) ⋆

4: θ ← θ (ϕ) 5: Collect trajectories Dagent ∼ pπθ b 6: Compute g 7: Compute Fb b 8: Solve for v b 9: Compute u b ← −b 10: h u b and update ϕ ← ϕ − ηϕ h b 11: Clip h 12: end for

5.2

▷ e.g., PPO, SAC, or REINFORCE

▷ Equation 6 ▷ Equation 5 ▷ Equation 7 ▷ Equation 8

H YPERGRADIENT D ESCENT WITH F ISHER S KETCHING

The explicit formulation requires storing Fb ∈ Rdθ ×dθ , resulting in O(d2θ ) memory complexity. To avoid constructing Fbθ explicitly, we instead represent the scaled Fisher matrix as αFbθ = X ⊤ X, where the rows of X are constructed from the weighted policy score vectors. For every transition from an agent trajectory, define r α xt (τ ) = γ t−1 ∇θ log πθ (at | st ). (9) Nagent 6

Published as a conference paper at ICOMP 2026

Stacking these vectors as rows of a matrix X gives αFbθ = X ⊤ X, so Equation 7 becomes b = gb. X ⊤ X + λI v

(10)

This representation allows the rows of X to be generated sequentially and immediately passed to a matrix sketch, without materializing either X or Fb . We employ Spectral Compensation Frequent Directions (SCFD) Chen et al. (2020) to maintain a compact approximation and approximately solve Equation 10. In our setting, SCFD is applied to the sequentially generated weighted policyscore vectors defined in Equation 9. For completeness, the SCFD update and solve procedures are provided in Algorithm 3 in the appendix. After all agent trajectories have been processed, the solver b, while the remaining hypergradient computation is unchanged. The is queried once to obtain v resulting sketched hypergradient descent method is summarized in Algorithm 2. Algorithm 2 Hypergradient Descent with Fisher Sketching Input: expert dataset Dexpert , iterations K, sketch size m, damping λ 1: Initialize θ, ϕ 2: for k = 1, . . . , K do 3: θ⋆ (ϕ) ≈ arg min Linner (ϑ, ϕ) ⋆

θ ← θ (ϕ) Collect trajectories Dagent ∼ pπθ b Compute g Initialize a SCFD sketch for all τ ∈ Dagent do for t = 1,r . . . , |τ | do α 10: x← γ t−1 ∇θ log πθ (at | st ) Nagent 11: A PPEND(x) 12: end for 13: end for b ← S OLVE(b 14: v g) b 15: Compute u b ← −b 16: h u b and update ϕ ← ϕ − ηϕ h b 17: Clip h 18: end for 4: 5: 6: 7: 8: 9:

▷ e.g., PPO, SAC, or REINFORCE

▷ Equation 6 ▷ Algorithm 3

▷ Equation 8

N UMERICAL R ESULTS

We evaluate the proposed approach on two control environments: CartPole Barto et al. (1983) and LQR, a continuous-control environment based on the classical linear-quadratic regulator problem Anderson & Moore (2007). CartPole is a low-dimensional benchmark with discrete actions; LQR is a continuous-control setting with a larger policy parameterization. Unless stated otherwise, all methods within each environment share the same policy and reward architectures, expert data, optimization settings, and computational budget. We compare against ML-IRL Zeng et al. (2022) as the bilevel IRL baseline and consider two variants of our Fisher-based hypergradient: Explicit Fisher, which explicitly constructs the empirical Fisher matrix and solves the damped linear system of Section 5.1, and Fisher with Sketching, which replaces this with the sketched approximation of Section 5.2. We study both optimization behavior and computational efficiency (wall-clock time and peak memory). 6.1

S TABILITY AND E FFICIENCY OF THE F ISHER S OLVER

Effect of Fisher Damping. We first study the effect of the damping parameter λ in Equation 7, using the explicit Fisher formulation to isolate the effect of regularization from sketching error. We perform a grid search over λ and monitor the outer objective. The results reveal a stability–fidelity trade-off. Weak damping leaves the Fisher system poorly conditioned, destabilizing the outer optimization; strong damping stabilizes the solve but drives 7

Published as a conference paper at ICOMP 2026

λ −3

10 10−2 10−1 100 101 102

Louter 532.75 ± 145.02 205.84 ± 34.91 427.76 ± 166.27 185.27 ± 4.38 501.83 ± 410.30 698.50 ± 227.01

Figure 1: Effect of the Fisher damping parameter λ on optimization. Left: validation outer loss throughout training. Right: Louter averaged over the final five outer iterations. (αFbθ + λI)−1 ≈ (λI)−1 , so the Fisher information contributes little to the update while remaining expensive to compute. The explicit solver must therefore trade off numerical stability against preserving the Fisher geometry. Effect of the SCFD Sketch Size. We fix λ = 10−3 , since the sketched solver remains stable under substantially weaker damping than the explicit formulation, and vary the sketch size m.

Method

Louter

Speedup

Explicit 219.15 ± 13.82 1.00× m = 32 187.90 ± 5.22 1.15× m = 64 181.70 ± 2.60 1.17× m = 128 204.96 ± 14.38 1.13× m = 256 194.73 ± 4.64 0.78×

Figure 2: Effect of the SCFD sketch size on optimization for λ = 10−3 . Left: validation outer loss for Explicit Fisher and SCFD at different sketch sizes. Right: Louter averaged over the final five outer iterations, with speedup relative to Explicit Fisher. A small sketch already suffices to match or exceed the optimization quality of Explicit Fisher: all sketch sizes achieve a lower final outer loss, with m = 64 performing best. Larger sketches do not improve quality monotonically and increase cost; m = 256 is even slower than Explicit Fisher despite a lower loss. Moderate sketch sizes therefore offer the best trade-off, with m = 64 attaining the lowest outer objective while remaining faster than the explicit solver. 6.2

C OMPUTATIONAL E FFICIENCY

As shown in Table 1, sketching provides the largest benefit in the higher-dimensional LQR setting, reducing peak memory by up to 1.31×; small sketches also preserve or improve runtime, while larger sketches erode this advantage. On CartPole, memory usage is essentially unchanged, but sketching yields a 1.29× speedup. The benefit of sketching thus grows with the cost of the Fisher representation. 6.3

R EWARD C OMPARISON

We evaluate the recovered rewards under a matched 24-hour training budget, using identical architectures within each environment and separately tuned hyperparameters; a fresh policy is then trained from scratch on each recovered reward under a fixed RL budget. We report PolicyNLL (↓, negative log-likelihood of expert actions under the induced policy), EnvReturn (↑, return of this policy under the ground-truth reward), and RankCorr (↑, rank correlation between learned and ground-truth trajectory returns). ExpertRet and RandomRet give the ground-truth returns of the expert and random policies as reference. 8

Published as a conference paper at ICOMP 2026

Table 1: Computational efficiency of Fisher with Sketching relative to Explicit Fisher. Memory gain is the ratio of peak memory usage, time speedup the ratio of wall-clock time per outer iteration. Values above 1× indicate an improvement over Explicit Fisher. Environment

Method

Memory gain

Time speedup

CartPole

Explicit Fisher Sketching (m = 8)

1.00× 0.97×

1.00× 1.29×

LQR

Explicit Fisher Sketching (m = 32) Sketching (m = 64) Sketching (m = 128) Sketching (m = 256)

1.00× 1.28× 1.31× 1.21× 1.10×

1.00× 1.04× 0.89× 0.99× 0.79×

Table 2: Final IRL performance under a matched computational budget. Arrows indicate the preferred direction. PolicyNLL ↓

EnvReturn ↑

ML-IRL Explicit Fisher Fisher with Sketching

331.20 ± 0.23 281.53 ± 0.21 183.35 ± 1.66

500.00 ± 0.00 500.00 ± 0.00 500.00 ± 0.00

ML-IRL Explicit Fisher Fisher with Sketching

248.25 ± 2.33 284.57 ± 3.24 252.37 ± 1.33

−25.38 ± 1.32 −36.82 ± 1.82 −28.58 ± 1.29

Method

RankCorr ↑

ExpertRet

RandomRet

500.00 ± 0.00

22.47 ± 0.99

−23.23 ± 1.16

−10828.62 ± 721.99

CartPole 0.892 ± 0.009 0.901 ± 0.007 0.870 ± 0.011 LQR 0.908 ± 0.012 0.953 ± 0.007 0.972 ± 0.005

On LQR, ML-IRL achieves 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 reach expert-level EnvReturn; Fisher with Sketching substantially improves PolicyNLL, while Explicit Fisher achieves the best RankCorr. Overall, the Fisher-based methods remain competitive in policy quality while providing strong reward-ranking performance. Figure 3 shows that Fisher with Sketching reaches competitive validation performance early in training. Despite using far fewer outer updates, it closely matches ML-IRL in PolicyNLL, EnvReturn, and RankCorr within the same wall-clock window, whereas Explicit Fisher converges substantially more slowly.

Figure 3: PolicyNLL, EnvReturn, and RankCorr over wall-clock time on the LQR validation set, truncated to the first 20,000 seconds. Fisher with Sketching matches ML-IRL’s performance using an order of magnitude fewer iterations.

C ONCLUSION

We developed a Fisher-based implicit hypergradient method for bilevel IRL and established that, under exact inner optimality and realizability, the inner Hessian equals α times the discounted trajectory Fisher matrix. An SCFD-based streaming approximation reduces curvature storage from O(d2θ ) to O(mdθ ). Experiments on two control tasks indicate that moderate sketches improve the stability–efficiency trade-off relative to the dense Fisher solver while remaining competitive with ML-IRL under matched computational budgets. 9

Published as a conference paper at ICOMP 2026

R EFERENCES Brian D. O. Anderson and John B. Moore. Optimal Control: Linear Quadratic Methods. Dover Publications, 2007. Andrew G. Barto, Richard S. Sutton, and Charles W. Anderson. Neuronlike adaptive elements that can solve difficult learning control problems. IEEE Transactions on Systems, Man, and Cybernetics, 13(5):834–846, 1983. Michael Bloem and Nicholas Bambos. Infinite time horizon maximum causal entropy inverse reinforcement learning. In 53rd IEEE Conference on Decision and Control, pp. 4911–4916, 2014. doi: 10.1109/CDC.2014.7040156. Souradip Chakraborty, Amrit Bedi, Alec Koppel, Huazheng Wang, Dinesh Manocha, Mengdi Wang, and Furong Huang. Parl: A unified framework for policy alignment in reinforcement learning from human feedback. In International Conference on Learning Representations, volume 2024, pp. 24410–24449, 2024. Cheng Chen, Luo Luo, Weinan Zhang, Yong Yu, and Yijiang Lian. Efficient and robust highdimensional linear contextual bandits. In Proceedings of the Twenty-Ninth International Joint Conference on Artificial Intelligence, IJCAI-20, pp. 4259–4265, Yokohama, Japan, 2020. Etienne Gauthier, Francis Bach, and Michael I Jordan. Explaining and preventing alignment collapse in iterative rlhf. arXiv preprint arXiv:2605.04266, 2026. Saeed Ghadimi and Mengdi Wang. Approximation methods for bilevel programming. arXiv preprint arXiv:1802.02246, 2018. Mina Ghashami, Edo Liberty, Jeff M Phillips, and David P Woodruff. Frequent directions: Simple and deterministic matrix sketching. arXiv preprint arXiv:1501.01711, 2015. Yingxiao Huo, Satya Prakash Dash, Radu Stoican, Samuel Kaski, and Mingfei Sun. Rank-1 approximation of inverse fisher for natural policy gradients in deep reinforcement learning. arXiv preprint arXiv:2601.18626, 2026. Deyi Kong, Zaiwei Chen, Shuzhong Zhang, and Shancong Mou. Natural hypergradient descent: Algorithm design, convergence analysis, and parallel implementation. arXiv preprint arXiv:2602.10905, 2026. Ilja Kuzborskij, Leonardo Cella, and Nicolò Cesa-Bianchi. Efficient linear bandits through matrix sketching. In Proceedings of the 22nd International Conference on Artificial Intelligence and Statistics, volume 89 of Proceedings of Machine Learning Research, pp. 177–185, Naha, Okinawa, Japan, 2019. Sergey Levine. Reinforcement learning and control as probabilistic inference: Tutorial and review. arXiv preprint arXiv:1805.00909, 2018. Han Shen, Zhuoran Yang, and Tianyi Chen. Principled penalty-based methods for bilevel reinforcement learning and rlhf. Journal of Machine Learning Research, 26(114):1–49, 2025. Ekaterina Shustova, Marina Sheshukova, Sergey Samsonov, and Evgeny Frolov. Scalable linucb: Low-rank design matrix updates for recommenders with large action spaces. arXiv preprint arXiv:2510.19349, 2025. W. J. A. van Heeswijk. Natural policy gradients in reinforcement learning explained. arXiv preprint arXiv:2209.01820, 2022. Yuhuai Wu, Elman Mansimov, Roger B Grosse, Shun Liao, and Jimmy Ba. Scalable trust-region method for deep reinforcement learning using kronecker-factored approximation. Advances in neural information processing systems, 30, 2017. Yan Yang, Bin Gao, and Ya-xiang Yuan. Bilevel reinforcement learning via the development of hyper-gradient without lower-level convexity. arXiv preprint arXiv:2405.19697, 2024. Siliang Zeng, Chenliang Li, Alfredo Garcia, and Mingyi Hong. Maximum-likelihood inverse reinforcement learning with finite-time guarantees. Advances in Neural Information Processing Systems, 35:10122–10135, 2022. 10

Published as a conference paper at ICOMP 2026

H YPERGRADIENT DERIVATIONS

Proposition A.1 (Implicit hypergradient for bilevel IRL). Assume that θ⋆ (ϕ) is a differentiable local solution of the inner problem, θ⋆ (ϕ) ∈ arg min Linner (θ, ϕ), θ

and that the Hessian

∂ 2 Linner ∂θ2 θ ⋆ (ϕ),ϕ

is invertible. Then the gradient of the induced outer objective Leouter (ϕ) := Louter (θ⋆ (ϕ)) is given by ∇ϕ Leouter

∂ 2 Linner =− ∂ϕ∂θ θ⋆ (ϕ),ϕ ϕ

∂ 2 Linner ∂θ2 θ ⋆ (ϕ),ϕ

#−1 ∇θ Louter |θ⋆ (ϕ) .

Proof. By definition, the outer objective optimized with respect to the reward parameters is Leouter (ϕ) := Louter (θ⋆ (ϕ)),

θ⋆ (ϕ) ∈ arg min Linner (θ, ϕ). θ

By the chain rule: dLeouter dϕ

= ϕ

∂Louter dθ⋆ . ∂θ θ ⋆ (ϕ) dϕ ϕ

Define the inner first-order optimality condition as: ∂Linner (θ, ϕ). ∂θ Since θ⋆ (ϕ) is a local solution of the inner problem, it satisfies: g(θ, ϕ) :=

g(θ⋆ (ϕ), ϕ) = 0. Differentiating this identity with respect to ϕ, we obtain: dg ⋆ (θ (ϕ), ϕ) = 0. dϕ Applying the chain rule gives ∂g dθ⋆ ∂g + = 0. ∂θ θ⋆ (ϕ),ϕ dϕ ϕ ∂ϕ θ⋆ (ϕ),ϕ Therefore,

" #−1 dθ⋆ ∂g ∂g =− . dϕ ϕ ∂θ θ⋆ (ϕ),ϕ ∂ϕ θ⋆ (ϕ),ϕ

Since

∂g ∂ 2 Linner = , ∂θ ∂θ2

∂g ∂ 2 Linner = , ∂ϕ ∂θ∂ϕ

we get " #−1 dθ⋆ ∂ 2 Linner ∂ 2 Linner =− . 2 dϕ ϕ ∂θ ∂θ∂ϕ θ⋆ (ϕ),ϕ θ ⋆ (ϕ),ϕ Substituting this expression into the derivative of the outer objective gives " #−1 dLeouter ∂Louter ∂ 2 Linner ∂ 2 Linner =− . dϕ ∂θ ∂θ2 ∂θ∂ϕ θ⋆ (ϕ),ϕ θ ⋆ (ϕ) θ ⋆ (ϕ),ϕ ϕ

Published as a conference paper at ICOMP 2026

Finally, switching from derivatives to gradients gives  dLeouter ∇ϕ Leouter =  dϕ ϕ

⊤  . ϕ

Moreover, transposing the mixed derivative changes the order of differentiation !⊤ ∂ 2 Linner ∂ 2 Linner = . ∂θ∂ϕ θ⋆ (ϕ),ϕ ∂ϕ∂θ θ⋆ (ϕ),ϕ Using also the symmetry of the Hessian, ∂ 2 Linner ∂θ2 θ ⋆ (ϕ),ϕ

!⊤ =

∂ 2 Linner , ∂θ2 θ ⋆ (ϕ),ϕ

we obtain ∂ 2 Linner ∇ϕ Leouter = − ∂ϕ∂θ θ⋆ (ϕ),ϕ ϕ

∂ 2 Linner ∂θ2 θ ⋆ (ϕ),ϕ

#−1 ∇θ Louter |θ⋆ (ϕ) .

D ISCOUNTED INFINITE - HORIZON DERIVATIONS

Throughout this appendix, pπθ denotes the infinite-horizon trajectory distribution induced by the original transition kernel P, whereas peπθ denotes the distribution induced by the absorbing-state e defined in the main text. We assume throughout that the relevant random variables are kernel P integrable and that differentiation, expectation, and the discounted infinite sums can be interchanged. In particular, the series below are assumed to be absolutely convergent. e the process follows the original dynamics at each step with probabilProof of Lemma 4.1. Under P, ity γ and otherwise moves to the absorbing state. Consequently, the probability of remaining in the original dynamics up to time t is γ t−1 . Conditional on this event, (st , at ) has the same distribution as under pπθ . Since f (sf , af ) = 0, it follows that Eτ ∼epπθ [f (st , at )] = γ t−1 Eτ ∼pπθ [f (st , at )] . Summing over t ≥ 1 and using the assumed integrability gives "∞ # "∞ # X X t−1 Eτ ∼epπθ f (st , at ) = Eτ ∼pπθ γ f (st , at ) , t=1

t=1

which proves the claim. Proof of Proposition 4.1. Under the absorbing-state construction, log peπθ (τ ) = log ρe(τ ) +

∞ X

log πθ (at | st ),

t=1

where the terms after absorption vanish because πθ (af | sf ) = 1. Moreover, ∞

log peϕ (τ ) = log ρe(τ ) +

1X rϕ (st , at ) − log Zϕ . α t=1

Therefore, αDKL (e pπθ ∥ peϕ ) = Eτ ∼epπθ

"∞ X

# (α log πθ (at | st ) − rϕ (st , at )) + α log Zϕ .

t=1

Published as a conference paper at ICOMP 2026

The normalization term is independent of θ. Applying Lemma 4.1 to the expectation shows that minimizing the inner KL divergence is equivalent to minimizing # "∞ X t−1 Eτ ∼pπθ γ (α log πθ (at | st ) − rϕ (st , at )) , t=1

which is Linner (θ, ϕ) in Equation 2. For the outer problem, terms involving only the expert distribution and the environment dynamics are independent of θ, so "∞ # X DKL (e pexpert ∥ peπθ ) = Cexpert − Eτ ∼epexpert log πθ (at | st ) , t=1

where Cexpert does not depend on θ. A second application of Lemma 4.1 gives Louter (θ) in Equation 2. Hence the bilevel KL problem in Equation 1 and the discounted bilevel problem in Equation 3 have the same minimizers. Proposition B.1 (Discounted outer-objective gradient). The discounted outer objective satisfies "∞ # X ∇θ Louter (θ) = −Eτ ∼pexpert γ t−1 ∇θ log πθ (at | st ) . t=1

Proof. The expert trajectory distribution does not depend on θ. Therefore, using the regularity assumptions stated at the beginning of this appendix, "∞ # X t−1 ∇θ Louter (θ) = −∇θ Eτ ∼pexpert γ log πθ (at | st ) t=1

= −Eτ ∼pexpert

"∞ X

# γ

t−1

∇θ log πθ (at | st ) .

t=1

Proposition B.2 (Discounted stepwise form of the trajectory Fisher). Under the assumptions of Proposition 4.2, "∞ # X 2 ⋆ t−1 ⊤ , ∇θ Linner (θ (ϕ), ϕ) = α Eτ ∼pπθ γ gt (τ )gt (τ ) t=1

θ=θ ⋆ (ϕ)

where gt (τ ) := ∇θ log πθ (at | st ). In particular, ∇2θ Linner (θ⋆ (ϕ), ϕ) = αFθ⋆ (ϕ) . Proof. By Lemma 4.1, the discounted inner objective can equivalently be written under the absorbing-state trajectory distribution as Linner (θ, ϕ) = Eτ ∼epπθ [ℓθ,ϕ (τ )] , where ℓθ,ϕ (τ ) :=

∞ X

(α log πθ (at | st ) − rϕ (st , at )) .

t=1

Define the trajectory score and its derivative by Hθ (τ ) := ∇2θ log peπθ (τ ).

Sθ (τ ) := ∇θ log peπθ (τ ), Since the transition kernel is independent of θ, Sθ (τ ) =

∞ X

∇θ log πθ (at | st ) =

t=1

∞ X t=1

gt (τ ).

Published as a conference paper at ICOMP 2026

After absorption, all terms vanish because πθ (af | sf ) = 1, so these sums contain only finitely many nonzero terms almost surely. Using the score-function identity and ∇θ ℓθ,ϕ (τ ) = αSθ (τ ), we obtain ∇θ Linner (θ, ϕ) = Eτ ∼epπθ ℓθ,ϕ (τ ) + α Sθ (τ ) . Differentiating once more gives ∇2θ Linner (θ, ϕ) = Eτ ∼epπθ

i ℓθ,ϕ (τ ) + 2α Sθ (τ )Sθ (τ )⊤ + ℓθ,ϕ (τ ) + α Hθ (τ ) .

Equivalently, h i ∇2θ Linner (θ, ϕ) = Eτ ∼epπθ αSθ (τ )Sθ (τ )⊤ + ℓθ,ϕ (τ ) + α Sθ (τ )Sθ (τ )⊤ + Hθ (τ ) . The trajectory score satisfies Eτ ∼epπθ Sθ (τ )Sθ (τ )⊤ + Hθ (τ ) = 0, R which follows by differentiating peπθ (τ ) dτ = 1 twice with respect to θ. We now evaluate the Hessian at θ = θ⋆ (ϕ). The two trajectory distributions satisfy 1 log peπθ (τ ) − log peϕ (τ ) = ℓθ,ϕ (τ ) + log Zϕ . α Under exact realization, peπθ⋆ (ϕ) = peϕ , and hence ℓθ⋆ (ϕ),ϕ (τ ) = −α log Zϕ a.s. Thus, ℓθ⋆ (ϕ),ϕ (τ ) + α is constant with respect to τ . Using the trajectory score identity above, h i Eτ ∼epπθ⋆ (ϕ) ℓθ⋆ (ϕ),ϕ (τ ) + α Sθ⋆ (ϕ) Sθ⊤⋆ (ϕ) + Hθ⋆ (ϕ) = 0. Consequently, ∇2θ Linner (θ⋆ (ϕ), ϕ) = α Eτ ∼epπθ Sθ (τ )Sθ (τ )⊤

θ=θ ⋆ (ϕ)

It remains to obtain the stepwise form. Expanding the trajectory score gives ∞ X X ⊤ gt gt⊤ + gt gk⊤ . Sθ (τ )Sθ (τ ) = t=1

t̸=k

For t < k, the conditional score identity yields E[gk (τ ) | s1 , a1 , . . . , sk−1 , ak−1 , sk ] = 0. Since gt is measurable with respect to the conditioning variables, E[gt gk⊤ ] = 0. The case k < t follows by transposition. Therefore, "∞ # X ⊤ ⊤ Eτ ∼epπθ Sθ (τ )Sθ (τ ) = Eτ ∼epπθ gt (τ )gt (τ ) . t=1

Finally, applying Lemma 4.1 componentwise gives "∞ # "∞ # X X ⊤ t−1 ⊤ Eτ ∼epπθ gt gt = Eτ ∼pπθ γ gt gt . t=1

t=1 ⋆

Combining the preceding identities and evaluating at θ = θ (ϕ) yields "∞ # X 2 ⋆ t−1 ⊤ ∇θ Linner (θ (ϕ), ϕ) = α Eτ ∼pπθ γ gt (τ )gt (τ ) t=1

θ=θ ⋆ (ϕ)

= αFθ⋆ (ϕ) .

Published as a conference paper at ICOMP 2026

Proposition B.3 (Discounted prefix-score form of the inner cross derivative). The inner cross derivative satisfies !# "∞ ⊤ X t X ∂ 2 Linner ∂rϕ (st , at ) ∂ t−1 = −Eτ ∼pπθ log πθ (ak | sk ) . γ ∂ϕ ∂θ ∂ϕ ∂θ t=1 k=1

Proof. Differentiating the discounted inner objective first with respect to ϕ gives # "∞ X ∂Linner t−1 ∂rϕ (st , at ) . = −Eτ ∼pπθ γ ∂ϕ ∂ϕ t=1 For a quantity determined by the trajectory up to time t, differentiation of its expectation with respect to θ introduces only the policy scores up to that time. Hence, by the score-function identity and causality, !# " ⊤ X t ∂ ∂rϕ (st , at ) ∂rϕ (st , at ) ∂ Eτ ∼pπθ = Eτ ∼pπθ log πθ (ak | sk ) . ∂θ ∂ϕ ∂ϕ ∂θ k=1

Interchanging differentiation with the absolutely convergent discounted series then gives the claimed expression. Corollary B.1 (JVP form of the inner cross-derivative-vector product). Let v ∈ Rdθ and define the scalar sequence t X ∂ log πθ (at | st ) v, bt (τ ) := qk (τ ). qt (τ ) := ∂θ k=1

Then

" ∂ 2 Linner v = −Eτ ∼pπθ ∇ϕ ∂ϕ ∂θ

∞ X

!# γ

t−1

rϕ (st , at )bt (τ )

t=1

Proof. Multiplying the identity in Proposition B.3 by v gives "∞ # ⊤ X ∂ 2 Linner ∂rϕ (st , at ) t−1 v = −Eτ ∼pπθ γ bt (τ ) . ∂ϕ ∂θ ∂ϕ t=1 The assumed absolute convergence justifies the exchange of the matrix-vector product and the series. Since bt (τ ) is independent of ϕ, it can be treated as constant when differentiating with respect to ϕ, which gives the stated JVP representation.

Published as a conference paper at ICOMP 2026

A LGORITHMIC D ETAILS

Algorithm 3 Spectral Compensation Frequent Directions Solver Input: sketch size m, dimension d, regularization λ > 0 1: S ← 0m×d 2: R ← ∅ 3: α ← λ 4: D ← λ−1 Im

▷ compressed sketch ▷ pending rows

5: procedure A PPEND (x) 6: 7:

R x⊤ if |R| = m then R←

S ∈ R2m×d R 9: Compute SVD: [U , Σ, V ] ← SVD(Y ) 2 10: δ ← σm 11: α←α+δ 12: for i = 1,q . . . , m do 13: σ̄i ← max{σi2 − δ, 0} 14: end for ⊤ 15: Update the compressed sketch S ← diag(σ̄1 , . . . , σ̄m )V1:m 1 1 , . . . , 16: Update the diagonal inverse statistics D ← diag 2 +α σ̄12 + α σ̄m 17: R←∅ 18: end if 19: end procedure 8:

Form the matrix Y ←

20: procedure S OLVE(g) 21: if R = ∅ then 22: Z←S 23: M ←D 24: else 25: C ← SR⊤ 26: P ← DC 27: Compute the Schur complement K ← RR⊤ − C ⊤ P + αI|R| 28: Compute K −1 using a Cholesky factorization

D + P K −1 P ⊤ 29: Form the inverse row-space matrix M ← −K −1 P ⊤ S 30: Form the current sketch matrix Z ← R 31: end if 32: Compute the inverse-vector product v ← α−1 g − Z ⊤ M Zg 33: return v 34: end procedure

−P K −1 K −1

Published as a conference paper at ICOMP 2026

A DDITIONAL C OMPUTATIONAL R ESULTS

Table 3: Wall-clock time per outer iteration across the considered environments. For runs executed under the fixed 24-hour computational budget, the iteration time is computed using the number of completed outer iterations. Method

Time / outer iteration (s)

Runs

CartPole ML-IRL Explicit Fisher Fisher with Sketching (m = 8)

510.20 ± 120.77 1464.41 1132.08 ± 130.33

3 1 3

53.61 ± 6.57 1037.76 ± 120.69 995.66 ± 91.16 1171.03 ± 188.30 1042.90 ± 160.41 1316.88

3 8 3 4 2 1

LQR ML-IRL Explicit Fisher Fisher with Sketching (m = 32) Fisher with Sketching (m = 64) Fisher with Sketching (m = 128) Fisher with Sketching (m = 256)

Table 4: Peak memory usage across the considered environments. Peak memory is measured as the maximum resident set size (RSS) observed during training. Method

Peak RSS (MB)

Runs

645.57 ± 1.56 644.57 662.14 ± 30.00

3 1 4

779.83 ± 120.03 929.61 ± 198.11 725.48 ± 99.84 710.89 ± 102.31 767.64 ± 122.08 844.58

3 10 3 4 2 1

CartPole ML-IRL Explicit Fisher Fisher with Sketching (m = 8) LQR ML-IRL Explicit Fisher Fisher with Sketching (m = 32) Fisher with Sketching (m = 64) Fisher with Sketching (m = 128) Fisher with Sketching (m = 256)

Efficient Hypergradient Descent for Inverse Reinforcement Learning

Analyze & explore related research in Conceptio

Related documents