Distributional Inverse Reinforcement Learning

Summary

DistIRL proposes an offline inverse-RL framework in which the per-step reward is treated as a random variable and the expert-vs-agent comparison is lifted from expected returns to a first-order-stochastic-dominance (FSD) criterion on full return distributions. The paper wraps the FSD violation integral in an energy-based likelihood, derives an ELBO for a variational reward distribution q_phi(r|s,a), and — because the indicator-based FSD objective is intractable — replaces it in the policy step with a distortion risk measure (CVaR by default), justified by the claim that FSD is equivalent to dominance under every distortion function (Proposition 4.6). A convergence analysis under standard axioms on the DRM (monotone, translation-equivariant, 1-Lipschitz in $\|\cdot\|_\infty$) yields an $O(\varepsilon^{-2})$ iteration bound to an $\varepsilon$-stationary point. Empirical evidence covers a 5x5 gridworld, a neurobehavioural mouse dataset with time-aligned dopamine, and risk-sensitive plus risk-neutral D4RL tasks (HalfCheetah, Hopper, Walker2d), with BIRL, BC, ValueDICE, and Offline ML-IRL as baselines.

The conceptual move — turning reward-distribution learning into a distributional-ordering problem and then using DRMs as a tractable surrogate — is genuinely novel and well motivated. The variance-in-dopamine and contact-uncertainty examples give a clean rationale for going beyond point rewards, and the FSD $\Leftrightarrow$ all-DRM equivalence provides a coherent bridge from the intractable ideal to the practical objective. My main reservations are about how much the reported evidence and derivations actually establish. The intuitive step from "match return distributions" to "recover reward distributions" is presented as a direct consequence, but a return-level objective does not by itself identify the per-step reward law. The one-sided FSD-violation loss is described as "aligning higher-order moments", which overstates what a positive-part integral of quantile gaps can enforce. The main-text proof sketch of Proposition 4.6 proves the easy forward direction only. And several load-bearing empirical numbers — the "0.3" dopamine correlation, the gridworld variance recovery — are used to support broader claims about decoding variability that the reported metrics do not quantitatively substantiate. The theory itself has small but real inconsistencies (stepsize schedule, notation drift, the DRM/FSD approximation gap softened in Limitations), and the reproducibility surface is thin at the very datasets the paper leans on hardest.

Strengths

• Conceptual contribution. Moving from expected-return matching to FSD-based comparison of full return distributions is a clean and non-trivial reformulation of the IRL problem, and the reduction to a tractable DRM surrogate via Proposition 4.6 is the right conceptual bridge to a practical algorithm.
• Unified framework. The paper cleanly combines an EBM-derived ELBO for the reward distribution, a distributional critic, and a DRM-shaped policy update, and Table 4 makes it explicit that the combination of "learn reward distribution + learn return distribution + risk-aware policy" is not covered by any single prior work.
• Broad empirical footprint. Discrete gridworld, a real neuroscience dataset (Markowitz dopamine), and continuous D4RL benchmarks under both risk-sensitive and risk-neutral rewards give a much wider evaluation than most IRL papers, and the low-data ablation (Appendix E.2) is a nice robustness check.
• Convergence analysis. Under the stated DRM axioms, the paper works through the nested-DRM contraction, softmax-Lipschitz, and reward-gradient recursion in enough detail to reach an $O(\varepsilon^{-2})$ rate — a rare complete treatment for an offline distributional IRL algorithm.
• Sensible baseline choice. Directly comparing against BIRL (parameter-uncertainty over deterministic rewards) is the right foil for the "recover full reward distribution" claim and produces a targeted qualitative gap in the mouse and gridworld experiments.

Weaknesses

• Reward-distribution identifiability from return-level FSD is asserted, not argued. The motivation is that stochasticity lives in the per-step reward, but the objective compares return distributions. Different per-step reward laws can produce identical returns after discounting, so "compare return distributions $\Rightarrow$ recover reward distribution" is a substantial modelling claim. This gap runs through the abstract, the intro's "thus capture the full landscape of reward distributions", and the "aligning higher-order moments" phrasing around Eq. 3; none of them acknowledges the identifiability step.
• One-sided FSD-violation loss overclaimed as moment alignment. The reward loss integrates $[F^{-1}_{Z^E}-F^{-1}_{Z^\pi}]_+$, which enforces one-sided quantile dominance at optimality — not equality of variance/skewness/tails. The paper repeatedly frames this as "capturing higher-order moments" or "aligning the full distribution", which is stronger than the loss provably delivers. This is a headline-level overclaim.
• Proposition 4.6 main-text proof establishes the easy forward direction. The load-bearing converse (every-DRM dominance implies FSD) is what motivates the DRM surrogate. The main-text proof sketch shows only that FSD implies every-DRM dominance — the direction that is already trivial. The full converse proof is in the appendix, but the main-text sketch as written is misleading.
• Dopamine-decoding conclusion outruns the reported metrics. The Pearson correlation of the mean is $\approx 0.3$ (under 10% of variance), and the "distribution shape" evidence is a handful of qualitative panels plus a relative W-1 comparison against BIRL. Claiming that "the variability in dopamine release can be decoded from behaviour alone" requires a per-mouse or per-state quantitative comparison of inferred vs. observed higher moments, which is not reported. The dopamine skew is also described as both "left-skewed" and "right-skewed and multimodal" within the same subsection, so even the target that the skew-normal is supposed to fit is inconsistently characterised.
• Gridworld variance-recovery claim is asserted rather than measured. The top-right port is $\mathcal{N}(1,1)$ and the risk-averse expert visits it in 1/10 episodes, so effective sample size for its variance is tiny. The paper claims DistIRL "captures the variance" but reports no numerical comparison and no seed-variance for the "9/10" statistic. The policy side of the contribution (does DistIRL reproduce the 9/10 risk-averse choice on this MDP?) is also not evaluated.
• Convergence-theorem stepsize schedule is internally inconsistent. Theorem 5.4 states $\eta_k = \eta = \eta_0 K^{-\sigma}$ (a constant, horizon-tied schedule) while Theorem 5.6 uses $\eta_k = \eta_0 k^{-\sigma}$ (iterate-decaying); the surrounding text mixes both. The two give different constants and different operational readings of the $O(\varepsilon^{-2})$ bound.
• Assumption C.11 is framed as a near-corollary of Theorem 5.5. The bound tying reward-gradient bias to the critic tracking error $E_k$ is presented as "consistent with (5.5)", but it actually requires an additional simulation-lemma-style step (log-policy gap $\to$ occupancy mismatch) and Lipschitzness of $\nabla L_r$ in the occupancy measure. It is an assumption, not a consequence.
• DRM-vs-FSD approximation gap softened in Limitations. Appendix B.6 correctly notes that a single-DRM surrogate is strictly weaker than FSD. The Limitations section frames the difficulty as non-differentiability only, omitting the more fundamental gap.
• Notation drift. Multiple substantive symbol overloads: $r(s,a)$ used as both a deterministic function and a random-variable sample; Assumption 5.1 bounds $|q_\phi(s,a)|$ as if $q_\phi$ were a scalar sample rather than a distribution; the reward-family sentence uses $p_\theta$ where $\theta$ is elsewhere the critic parameter; the indicator function around Eq. 5 is corrupted; DistIRL-qtr vs. DistIRL-qrt appears inconsistently. Individually minor, collectively they make the paper harder to verify.
• EBM partition function is not addressed at introduction. $p(D|r) \propto \exp(-L_{\rm FSD}(\pi,r))$ generally has an $r$-dependent normaliser, and the ELBO drops it without comment. Whether $\log Z(r)$ vanishes or is absorbed elsewhere is not stated.
• "Boltzmann distributions over returns", "strictly generalizes standard IRL", "first offline IRL to learn reward distributions" phrasings. Small but real overclaims: the first is inconsistent with the paper's own Section 3.1; the second confuses a modelling-assumption reduction with a method reduction; the third rests on a dismissive one-liner about adjacent distributional-critic IRL work whose actual objectives are not summarised.
• Related-work table cross-reference and column ordering are broken. The prose calls Table 4 "Table A" and labels "Infer risk-aware policy" the third column when it is the second in the actual table. This is exactly the check-mark table on which the novelty claim rests.

Reproducibility & code

The veritas evidence rates every claim (headline and supporting) as at most "partial" support with medium-to-high reproducibility risk; the released materials cover the DistIRL algorithm and hyperparameters but do not close the loop on baselines, dataset preprocessing, or seed listings.

• Expert-training procedure for Table 1 is thinly specified. The risk-averse D4RL expert is trained with "a variant of DSAC" using a CVaR objective, but the variant, CVaR level, network sizes, and training budget are absent, and no seeds are listed. Since Table 1 is the paper's headline number, reproducing the Expert row within the reported $\pm$ is unlikely without the training script.
• Baseline configurations (Offline ML-IRL, ValueDICE, BC) are undocumented. Table 1's "outperforms baselines" claim is load-bearing, but the baselines' hyperparameters and architectures are not disclosed. The strong claim about Offline ML-IRL's transition-model mismatch cannot be independently checked.
• DistIRL-qrt / DistIRL-qtr parameterisation is only named. The competitive quantile-reward variant is used in the main table and highlighted as a bridge to diffusion families, but its architecture, loss, and prior are not described.
• Reward-range inconsistency. Appendix D prose says the MuJoCo reward range is $[-10,10]$ while Table 6 says $[-5,5]$. The bound directly matters (it is $R_{\max}$ in Assumption 5.1) and it changes the $\tanh$ clip.
• Half-Cheetah stochastic penalty is under-specified. $B_{0.1}$ is used without definition; the LHS is $R_t$ while the "where" clause defines $\bar r_t$. To reproduce the risk-sensitive dataset a reader has to guess Bernoulli(0.1) i.i.d. per step.
• Offline critic tuple $(s_t,a_t,r_t,s_{t+1},a^\pi_{t+1})$ hides the source of $s_{t+1}$. In a strictly offline setup this matters: the next-state must come from the expert data (or a learned model), and if the action was sampled from the current policy rather than the expert, some importance / off-policy correction is needed. None of this is stated.
• Mouse-dataset preprocessing is not released. The dopamine experiment converts raw Markowitz recordings into 10 syllable states, 10 actions, 159 sequences, with dopamine aligned to state-action transitions. No preprocessing script, no specification of which recordings/mice, and no numerical Wasserstein-1 or per-state variance values are provided; the "0.3" correlation cannot be regenerated.
• Footnote result $11239\pm 539$ lacks a protocol. The claim that DistIRL matches Offline ML-IRL when given "the same amount of data" needs a data-volume and data-source specification.
• Ablation Table 3 reports only min-max scaled numbers. Raw pre-scaled returns are absent, so the ablation cannot be re-anchored to Table 1.
• Empirical convergence not shown. Theorem 5.6 gives an $O(K^{-1/2})$ rate on the averaged squared gradient, but no training curves are reported to check that the practical fixed-learning-rate schedule (Table 6) matches the schedule the theorem requires.

Recommended Changes

Essential

• Rewrite the "reward-distribution recovery" argument. In the intro/Sec 4 (the "thus capture the full landscape of reward distributions" and "aligning higher-order moments" sentences), explain what part of the per-step reward distribution is identifiable from return-level FSD and what part is not, and reframe the claim as targeted distributional matching + prior/KL identifiability rather than direct recovery.
• Retract or precisify the "aligning higher-order moments" phrasing. Explicitly note that the one-sided FSD-violation loss enforces one-sided dominance, not moment equality; use the DRM/all-distortion equivalence (Prop. 4.6) as the route to full-distribution matching, and reserve moment-alignment claims for that setting.
• Fix the main-text proof sketch of Proposition 4.6. Show the constructive direction (build a distortion $\tilde\xi$ concentrated on the FSD-violation set) rather than the trivial forward direction.
• Rewrite the indicator-function definition around Eq. 5. Present $I(v)$ cleanly and immediately state which quantile ordering it flags, so the reader can sign-check the "reward minimises / policy maximises FSD violations" loop.
• Reconcile the stepsize schedule in Theorems 5.4 and 5.6. State one schedule per theorem, verify that the geometric-sum step in Theorem 5.4 uses the schedule the theorem states, and adjust the practical Table 6 discussion so it matches (or explicitly deviates from) the theorem.
• Add per-mouse / per-state quantitative recovery-of-variance and recovery-of-skewness metrics to the dopamine experiment, and soften the "variability in dopamine release can be decoded from behaviour alone" conclusion until such metrics support it. Also resolve the "left-skewed" vs. "right-skewed and multimodal" contradiction about the target dopamine distribution.
• Release baseline hyperparameters, expert-training scripts, and mouse-data preprocessing. Specifically: DSAC-CVaR expert variant, Offline ML-IRL / ValueDICE / BC configurations, seeds for the 5-seed averages, the syllable-segmentation and dopamine-alignment code (or a cached processed dataset), the DistIRL-qtr architecture and loss.
• Reconcile the reward range used for MuJoCo in Appendix D (prose $[-10,10]$ vs. Table 6 $[-5,5]$) and fix the corresponding $R_{\max}$ used in Assumption 5.1.
• Fully define the risk-sensitive D4RL penalty: state that $B_{0.1}\sim \mathrm{Bernoulli}(0.1)$ i.i.d. per step, resolve the $r_t$/$\bar r_t$ notation inside the same equation, and confirm the sampling procedure.

Suggested

• State how $s_{t+1}$ is obtained in the offline critic update, and, if actions come from the current policy while transitions come from the expert dataset, describe the correction (or the lack thereof) explicitly.
• Discuss the EBM partition function $Z(r)$ at the point $p(D|r)\propto\exp(-L_{\rm FSD})$ is introduced, and state which term of the ELBO derivation absorbs it.
• Report per-seed frequency for the gridworld "9/10" statistic and add a table of true vs. inferred reward variance per state; also report the DistIRL policy's top-right / bottom-right visitation to evaluate the risk-aware policy side of the contribution.
• Report raw pre-scaled returns for the Table 3 ablation alongside the scaled 0–1 numbers.
• Fix small textual inconsistencies: "Table A" $\to$ "Table 4"; column order in the prose surrounding Table 4; "DistIRL-qrt" / "DistIRL-qtr" spelling; the "$p_\theta$" symbol in the reward-family sentence; the "or vice versa depending on the formulation" hedge in the FSD desideratum; "Boltzmann distributions over returns"; "strictly generalizes standard IRL"; "blackprint" $\to$ "blueprint"; the notation collision in Assumption 5.1 ($|q_\phi(s,a)|$ vs. $|r_t|$); quantile fractions $\tau_i = i/N$ vs. midpoints.
• Expand the Limitations paragraph to acknowledge the DRM-vs-FSD approximation gap (single-DRM surrogate is strictly weaker than FSD, per Appendix B.6), alongside the current non-differentiability point.
• Add a one-sentence, evidence-anchored characterisation of the adjacent distributional-critic IRL works (Lee et al. 2022; Karimi & Ebadzadeh) so the "first offline IRL to learn full reward distributions" claim can be verified without deferring to a dismissive one-liner.