A Geometric Perspective on Reward Function Updates in Inverse Reinforcement Learning

1 Introduction

• The analysis focuses on entropy-regularized reverse-KL distribution matching between expert and agent occupancies.Source: §1 · p. 1
• Because occupancy updates are not directly realizable, policy optimization on an updated reward must induce movement from current occupancy toward an estimated optimum.The paper centers analysis on occupancy discrepancy, where the optimum is explicitly characterized.Source: §1 · pp. 1–2
• Preconditioning the dual gradient with average Fisher information along the e-geodesic yields a natural-gradient-style update equivalent to interpolating the old reward with an estimated optimal reward.The paper states this can be realized efficiently without explicitly computing the Fisher information matrix.Source: §1 · p. 2
• The geometric account explains why function-space reward updates, previously studied and reported as superior to gradient updates, can be effective [Arenz et al., 2016, Diwan et al., 2026].The paper presents its interpretation as the first explanation of their effectiveness.Source: §1 · p. 2
• Policy search over geometry-aware reward updates is equivalent to entropic Mirror Descent, while the reward update is Mirror Descent in reverse-KL-induced geometry.The paper also describes the occupancy-space update as Mirror Descent with a time-varying proximal term.Source: §1 · pp. 2–3
• The paper presents convergence results for IRL methods built on these geometry-aware reward updates.Source: §1 · p. 3

2 Background

• The setup is a discounted infinite-horizon MDP with stationary stochastic policies and normalized discounted state-action occupancies.Source: §2 · p. 3
• MaxCausalEnt-IRL seeks a reward and policy by maximizing expected reward plus causal entropy while matching expert reward expectations; its optimal policy has a Boltzmann form.Source: §2 · p. 3
• The paper studies entropy-regularized reverse-KL minimization between expert and agent occupancies, a divergence-minimization setting used in prior work.Source: §2 · p. 3
• Classical Lagrangian optimization alternates policy optimization and reward-gradient steps; reverse-KL relaxation instead changes the target from expert occupancy to a reward-dependent distribution.Source: §2 · p. 3

3 Occupancy-Space Reward Function Updates

• Under the exponential-family occupancy assumption, the target occupancy parameter is determined by expert features and the current reward, and reward updates aim to move current occupancy toward it.The analysis assumes a bounded, Lipschitz dual and non-flat feature covariance along the occupancy path; the stronger eigenvalue bound is not always necessary.Source: §3 · p. 4, §3.1 · pp. 4–5

3.2 Reward Update Analysis

• The occupancy-space trajectory is an e-geodesic in natural parameters, while the vanilla dual gradient follows an m-geodesic in expectation parameters.Source: §3.2 · p. 5
• The resulting reward update is mirror descent with a time-varying squared Mahalanobis proximal term induced by the average Fisher metric.This gives a geometric explanation for the reported superiority of function-space reward updates over standard reward-parameter gradient steps.Source: §3.2 · p. 6

3.2.1 Connection to Natural Gradient

• The occupancy-space direction corresponds to a natural-gradient-style reward update preconditioned by the average Fisher information along the e-geodesic.Its preconditioned gradient is a convex interpolation of the old reward and the log-density-ratio target, avoiding explicit Hessian or FIM estimation.Source: §3.2.1 · p. 6, §3.2 · p. 6

3.2.2 Connection to Gradient Descent & Direct Optimisation

• The functional update is first-order accurate, with contraction governed by the condition number of the preconditioned curvature matrix, whereas vanilla descent uses the local curvature condition number.The paper attributes the difference to average-Fisher rescaling along the path to the target occupancy.Source: §3.2.2 · p. 7
• The final result is a functional reward update whose contraction depends on residual curvature after average-Fisher preconditioning, while vanilla gradient contraction depends on local curvature.The authors argue that anisotropic IRL geometry therefore favors functional updates over identity-preconditioned vanilla gradients.Source: §3.2.2 · p. 7

4 Policy Search over Functional Reward Updates

• The policy-side analysis asks which policy update results from optimizing the policies induced by geometry-aware reward iterates.Source: §4 · pp. 7–8

4.1 Assumptions & Remarks

• A negative Shannon entropy mirror map makes the occupancy-space Bregman divergence equal to KL divergence between state-action occupancies.The policy and its state-action occupancy have a unique relationship, and the occupancy mirror map connects occupancies with action-value functions.Source: §4.1 · p. 8

4.2 Policy Update Analysis

• Policy iterates induced by the reward iterates perform approximate entropic mirror descent, with error bounded by reverse KL divergence to the previous iterate.For KL-constrained max-ent RL, the trust-region constraint already enforces this KL bound.Source: §4.2 · p. 8

5 Related Work

• The paper identifies the natural gradient of the IRL dual with respect to reward parameters as largely unexplored.Source: §5 · p. 8
• Prior policy-geometry methods use natural policy gradients or mirror descent for policies but ordinary reward gradients, including NPG-GAIL and TRPO-GAIL approaches [2020b, 2021].Other cited approaches also keep geometry-aware updates on the policy side [2022].Source: §5 · p. 9
• Existing reward-side mirror-descent methods use proximal terms, whereas this paper derives reward interpolation from the natural gradient of the IRL dual.The paper characterizes its update as mirror descent with a time-varying Mahalanobis proximal term induced by average Fisher information along the occupancy e-geodesic.Source: §5 · p. 9
• Neu and Szepesvári [2007] construct a reward-parameter natural gradient for policy matching, but require differentiating the optimal-policy mapping and pseudo-inverting its metric.This paper instead preconditions the IRL-dual gradient using average Fisher information and gives a closed-form update without explicitly estimating or inverting the FIM.Source: §5 · p. 9

6 Convergence Analysis

• The convergence proof bounds policy-player regret, controls changes in the integral FIM, bounds reward-player regret, and combines these terms through saddle-point theory.The argument treats both updates as regret-minimizing algorithms.Source: §6 · pp. 9–10
• The analysis assumes a bounded reward-parameter domain and a Lipschitz Fisher information matrix over the specified parameter set.Source: §6 · p. 9
• With trust-region tolerance ζ = 1/T^2, the average-iterate saddle point residual converges sublinearly up to a fixed residual level.Source: §6 · p. 10

7 Discussion & Limitations

• The analysis establishes sublinear convergence and shows that preconditioned reward updates induce approximate entropic mirror descent on policies, with errors controlled by policy KL-divergences.The average Fisher information along the occupancy e-geodesic can be incorporated without additional computational overhead.Source: §7 · p. 10
• The analysis does not explicitly model function-approximation, data-availability, or density-ratio-estimation errors, and assumes finite features with linear occupancies and rewards.The feature construction can itself introduce approximation error, and the reward-player regret bound relies on Assumption 6.1.Source: §7 · p. 10

A.1 Feature Learning

• Features are learned with a neural encoder and energy network by minimizing denoising score matching loss.The procedure repurposes a method previously used to learn energy-based reward functions.Source: §A.1 · p. 14

A.2 Error Propagation under Misspecified Features

• The appendix uses maximum-likelihood fitting of exponential-family occupancies, for which the estimated natural parameters satisfy moment matching.The text states that this M-projection analysis applies because maximum-likelihood estimation and the feature-learning loss share a population minimizer.Source: §A.2 · p. 14
• Expert occupancy misspecification propagates into the algebraically constructed target occupancy, making it an erroneous target for the update.The estimated exponential-family natural parameter describes the fitted distribution, not the true target occupancy, so their feature expectations can differ.Source: §A.2 · pp. 14–15
• The target error enters the functional update through an additive residual term.The analysis frames violations of the finite-feature assumption as misspecification in the occupancy exponential family.Source: §A.2 · pp. 14–15
• A controlled violation makes expert misspecification uncorrelated with every feature and varies its magnitude relative to the in-span component.The appendix reports measuring mean-gradient error and angular deviation between the ideal and misspecified update directions, but gives no outcome values in these passages.Source: §A.2 · p. 15

A.3 Additional Results and Experiment Details

• The Figure 2 experiment uses a 25-state gridworld with two-dimensional features and numerically estimates occupancy parameters, Fisher information, and the dual Hessian.Finite differences and least-squares fitting introduce minor approximation errors, while the actual preconditioned reward update remains a convex interpolation.Source: §A.3 · p. 15
• Functional and natural gradients are often comparable in alignment with direct optimization, and both significantly outperform vanilla gradients.Table 2 reports angle deviation to the direct optimization direction in degrees, with lower values better.Source: §A.3 · p. 15, Table 2 · p. 16
• Functional updates usually have a much better contraction coefficient than vanilla gradient descent, even after accounting for changing norms.Table 3 compares contraction coefficients and condition numbers for the two update types, with lower values better.Source: §A.3 · p. 16, Table 3 · p. 16

B Section 2 Details

• The entropy-regularized occupancy objective has strong duality because its entropy and negative-KL terms are concave and its Bellman flow constraints are linear.The section also states that the expected policy entropy is concave in the occupancy measure.Source: B Section 2 Details · p. 16

C Section 3 Proofs

• For exponential-family occupancies, the log-partition function and Fisher information characterize the geometry used in the update analysis.The proofs invoke the fundamental theorem of calculus to relate gradients along the parameter path and establish self-adjointness under the Fisher-information inner product.Source: C Section 3 Proofs · pp. 17–18
• The functional update has a best one-step contraction result, and a separate corollary accounts for the worst-case scaling caused by changing norms between iterations.The appendix also extends a contraction result to the Euclidean norm using bounds relating it to the iteration-dependent norm.Source: C Section 3 Proofs · p. 19

The worst-case scaling factor for the unsquared norm is hence

• Under identity preconditioning, the proof derives the best one-step contraction by analyzing the eigenvalues of the average Hessian.The derivation uses the symmetric positive-definite Hessian and its eigendecomposition.Source: The worst-case scaling factor for the… · pp. 20–21

C.1 An Alternate Proof of First-Order Contraction

• The functional update is first-order accurate relative to direct optimization under the stated Lipschitz assumptions.These include local Lipschitz assumptions on the dual Hessian and occupancy parameter.Source: §C.1 · p. 21

Triangle inequality

• Under equal curvature at optimality and step size ϵ = 1, the geometry-aware update is second-order accurate relative to direct optimization.The equal-curvature condition makes the optimal step size tend to 1, causing the first term in the accuracy bound to vanish.Source: Triangle inequality · pp. 22–23

C.2 Connection to Pseudo-Gradient Optimization in Euclidean Geometry

• In Euclidean geometry, δ(i) is a pseudo-gradient of the vanilla dual gradient, with alignment lower-bounded by the vanilla-gradient norm under Assumption 3.3B.The update direction is also stated to be an ascent direction for the dual.Source: §C.2 · pp. 23–24

• The gradient of the occupancy objective at the current occupancy equals −δ(i) up to a constant and potential-shaping term, inducing the same optimal policy.The proof uses the Boltzmann form of the optimal policy under Shannon entropy regularization and policy invariance under potential shaping.Source: D Section 4 Proofs · pp. 24–25
• Optimizing the reward is equivalent to optimizing the log-policy, while optimizing the occupancy objective uses the same update direction δ(i).The reward-policy relationship and gradient correspondence establish these links.Source: D Section 4 Proofs · pp. 24–25
• With mirror-descent step size ω = 1/ϵ, the trust-region update is analogous to occupancy-space mirror descent, with deviation K bounded by the enforced trust-region bound ζ.Source: D Section 4 Proofs · p. 27

E Section 6 Proofs

• Theorem 6.6’s proof uses bounded regret for both players and standard Nash-equilibrium theory to establish its result.The supplied proof passage does not include the complete resulting bound.Source: E Section 6 Proofs · p. 29

NeurIPS Paper Checklist

• The authors say the abstract’s claims about e-geodesic Fisher information, approximate entropic mirror descent, and convergence and regret bounds are supported by theorems and lemmas.Source: NeurIPS Paper Checklist · p. 30
• The stated limitations are omitted function-approximation and limited-data errors, linear finite-feature assumptions, and a reward-player regret bound dependent on Assumption 5.1.Source: NeurIPS Paper Checklist · p. 30
• The authors state that all paper claims have corresponding proofs, with assumptions and remarks given in the main text and some in the appendix.Source: NeurIPS Paper Checklist · p. 31
• Experimental setup details are in Appendix A.3, IRL hyperparameters in supplementary material, and the IRL and Figure 2 code is also supplied there.The paper says experiments illustrate the theoretical mechanism and that their compute requirements are negligible.Source: NeurIPS Paper Checklist · pp. 32–33
• The checklist notes that IRL and imitation learning may support safer alignment but may also reproduce or amplify biases, errors, or unsafe demonstrated behavior.It emphasizes demonstration quality and appropriate oversight.Source: NeurIPS Paper Checklist · p. 34
• The paper reports no crowdsourcing or human-subject research, so it marks participant-risk disclosure and IRB review as not applicable.The checklist notes that the study does not involve human subjects.Source: NeurIPS Paper Checklist · p. 36
• The paper marks LLM-use declaration as not applicable because LLMs were not an important, original, or non-standard component of its core method.Source: NeurIPS Paper Checklist · p. 36