Distributional Inverse Reinforcement Learning

• Distributional IRL is a framework that infers distributions over rewards, policies, or returns, capturing uncertainty from expert demonstrations.
• Techniques range from Bayesian preference elicitation and stochastic IRL to latent reward modeling, each preserving different aspects of reward ambiguity.
• Recent methods integrate risk-aware policy learning with full return distribution modeling using FSD and quantile critics for robust performance.

Searching arXiv for the cited papers and closely related work on distributional inverse reinforcement learning. First, retrieving the 2025 paper explicitly titled “Distributional Inverse Reinforcement Learning”. Searching for foundational Bayesian and reward-distribution IRL papers mentioned in the source material. Searching for related papers on latent/distributional reward modeling, occupancy-distribution IRL, and distributional-critic imitation learning. Distributional inverse reinforcement learning designates a family of inverse reinforcement learning formulations in which the inferred object is not restricted to a single deterministic reward estimate. In the literature, the “distributional” object has taken several forms: a posterior distribution over reward functions, policies, and optionally realized rewards in Bayesian preference-elicitation models (Rothkopf et al., 2011); a probability distribution over reward-function parameters in stochastic IRL (Ju, 2019); a transition-level latent reward distribution learned by amortized variational inference in CWAE-IRL (Kusari, 2019); a target distribution over rewards or discounted future rewards in optimality-profile matching (Haug et al., 2020); expert visitation distributions used as side-information in IRL reductions (Swamy et al., 2023); distributional critics embedded in implicit-reward imitation learning (Karimi et al., 27 Feb 2025); and, in a stronger recent formulation, joint modeling of reward distributions and full return distributions with risk-aware policy learning (Wu et al., 3 Oct 2025). The topic therefore does not denote a single settled formalism, but a set of technically distinct attempts to preserve uncertainty, heterogeneity, or distributional structure in reward inference.

1. Conceptual scope

Classical IRL is typically introduced as the problem of inferring a reward function from expert demonstrations under which the expert policy is optimal or near-optimal. The surveyed literature agrees on the usual underdetermination: many rewards can explain the same behavior, and this motivates distribution-valued rather than point-valued reward inference. What differs sharply across papers is the identity of the distribution that is maintained or matched.

Formulation	Distributional object	Representative paper
Bayesian preference elicitation	Posterior over $\rho$, $\pi$, and optionally $r^T$	(Rothkopf et al., 2011)
Stochastic IRL	Distribution over reward weights $\mathcal W$	(Ju, 2019)
Latent-variable IRL	Conditional latent reward distribution per transition	(Kusari, 2019)
Optimality-profile IRL	Distribution over rewards or discounted future rewards on demonstrations	(Haug et al., 2020)
Expert-reset IRL	Expert visitation distributions $\rho_E^t$	(Swamy et al., 2023)
Implicit-reward imitation with distributional critic	Return distribution $Z^\pi(s,a)$ used through its expectation	(Karimi et al., 27 Feb 2025)
Full distributional IRL	Reward distribution $q_\phi(r \mid s,a)$ and return distribution $Z^\pi$	(Wu et al., 3 Oct 2025)

A central distinction runs through the field. Some methods are distributional over latent explanatory variables of behavior, such as reward functions, reward parameters, or demonstrator policies. Others are distributional over returns or over occupancy-like objects. A further distinction concerns whether the distribution is a posterior with explicit probabilistic semantics, an amortized latent representation, a fitted parametric family, or a side-information distribution used only to change computational complexity. Much of the technical literature consists precisely in clarifying these distinctions rather than collapsing them into one notion.

2. Bayesian and reward-distribution foundations

A clear early formulation appears in the preference-elicitation view of IRL, where the environment is a known controlled Markov process $\nu=(\mathcal S,\mathcal A,\tau)$, demonstrations are fully observed state-action trajectories $D=(a^T,s^T)$, and latent preferences are encoded by a stochastic reward function $\pi$0 together with a latent policy $\pi$1 (Rothkopf et al., 2011). The utility is discounted return,

$\pi$2

with $\pi$3, and the prior is hierarchically factorized as

$\pi$4

The resulting posterior is explicitly distributional: $\pi$5 That formulation also supports a reward-augmented model with posterior target $\pi$6, extending uncertainty from reward laws to realized reward sequences (Rothkopf et al., 2011).

This Bayesian line is distributional in the strongest posterior sense among the early papers. It preserves uncertainty over reward functions and policies rather than collapsing immediately to a MAP reward. It also explicitly handles demonstrator suboptimality by adopting a softmax policy model

$\pi$7

with a prior over the inverse temperature $\pi$8, rather than assuming exact optimality (Rothkopf et al., 2011). At the same time, it assumes known dynamics and requires repeated MDP solution inside MCMC, so scalability is limited.

A different but closely related response to IRL non-identifiability is to model a distribution over reward parameters directly. Stochastic inverse reinforcement learning reformulates IRL as estimation of a conditional probability distribution over reward-function weights $\pi$9, written as $r^T$0, rather than a single reward (Ju, 2019). Its central optimization problem is

$r^T$1

where $r^T$2 is a representative trajectory-subset class and the learned reward distribution is instantiated as a Gaussian mixture model. This makes multimodality structurally available through

$r^T$3

with parameters fitted by a two-stage MCEM-style procedure (Ju, 2019). The method is directly distributional over rewards, but it is not exact Bayesian posterior inference; the paper presents it as a latent-variable maximum-likelihood-style procedure rather than a fully specified posterior model.

3. Latent reward distributions and amortized inference

A separate line treats reward as a latent variable inside a conditional generative model. CWAE-IRL starts from an MDP $r^T$4, uses expert transition tuples $r^T$5, and reinterprets the reward at each transition as a continuous latent variable (Kusari, 2019). Under the faithful latent-variable interpretation described in the paper, the model factorizes as

$r^T$6

with amortized posterior

$r^T$7

The encoder performs backward inference from a demonstrated transition to a latent reward, while the decoder performs forward inference by reconstructing $r^T$8 from $r^T$9 (Kusari, 2019).

The variational reading is explicit. CWAE-IRL adopts the VAE logic

$\mathcal W$0

and uses reparameterization

$\mathcal W$1

Because KL regularization was reported as unstable, especially with small samples, the paper replaces it with a Wasserstein-style discrepancy through Maximum Mean Discrepancy and optimizes a CWAE-IRL objective that combines next-state reconstruction with latent distribution matching toward a Gaussian prior (Kusari, 2019).

The method is distributional only in a limited but technically significant sense. It does not maintain a posterior over global reward functions $\mathcal W$2 or $\mathcal W$3. Instead, it learns a latent Gaussian reward distribution per transition and uses the encoder outputs $\mathcal W$4 as an amortized approximation to reward uncertainty (Kusari, 2019). The paper’s reported strengths are avoidance of inner-loop planning, no requirement for explicit dynamics knowledge, and applicability in objectworld and pendulum. Its limitations are equally central to the concept of distributional IRL: the reward uncertainty is local and amortized, the link to expert optimality is indirect because training is based on next-state reconstruction, and the formulation does not establish a coherent posterior over reward functions.

4. Return-distribution and profile-matching formulations

Another use of “distributional” in IRL replaces expert optimality assumptions with supervision in the form of a target distribution over reward-related quantities. In optimality-profile matching, the demonstrations may be suboptimal and heterogeneous, and the learner receives an “optimality profile” $\mathcal W$5: either a distribution over immediate rewards seen during demonstrations,

$\mathcal W$6

or, more generally, a distribution over cumulative discounted future rewards over suffixes,

$\mathcal W$7

Reward learning is then posed as

$\mathcal W$8

with an entropic Wasserstein loss in experiments, plus pairwise-ranking and fixed-return supervision losses to break non-identifiability (Haug et al., 2020).

This formulation is distributional in that it matches a scalar distribution over rewards or discounted future rewards, not because it learns state-conditional return distributions in the sense of distributional RL. The paper is explicit that matching only the pushforward distribution is highly non-identifiable: if $\mathcal W$9 preserves the state occupancy measure, then $\rho_E^t$0 equals $\rho_E^t$1, so additional supervision is required (Haug et al., 2020). In experiments, the target profile is synthesized from the ground-truth reward function, which makes the method a form of weakly supervised distribution-matching IRL rather than a pure demonstration-only procedure.

A stronger recent formulation moves from matching scalar profiles to joint modeling of stochastic rewards and full return distributions. Distributional Inverse Reinforcement Learning assumes a stochastic immediate reward

$\rho_E^t$2

and a return random variable

$\rho_E^t$3

then replaces expectation-based IRL criteria with minimization of first-order stochastic dominance violations (Wu et al., 3 Oct 2025). The central objective is

$\rho_E^t$4

where the integrated positive CDF gap vanishes when the expert return distribution first-order stochastically dominates the learner’s (Wu et al., 3 Oct 2025). The key identity

$\rho_E^t$5

converts the loss into a quantile form suitable for empirical estimation from sampled returns (Wu et al., 3 Oct 2025).

A plausible synthesis is that optimality-profile matching and FSD-based DistIRL occupy different points on the same spectrum. The former matches a one-dimensional pushforward distribution on demonstrations and relies on auxiliary supervision; the latter explicitly models stochastic rewards, learns return distributions through a distributional critic, and couples reward inference to risk-aware policy optimization through distortion risk measures (Haug et al., 2020, Wu et al., 3 Oct 2025).

5. Occupancy and visitation distributions as computational side-information

In another strand of the literature, “distributional IRL” refers neither to reward uncertainty nor to return distributions, but to expert visitation distributions used as side-information inside the IRL reduction. Inverse Reinforcement Learning without Reinforcement Learning studies a finite-horizon MDP $\rho_E^t$6 and the standard reward-moment matching game

$\rho_E^t$7

but assumes access to time-indexed expert visitation distributions $\rho_E^t$8 and the ability to reset the learner to states sampled from them (Swamy et al., 2023).

That side-information changes the computational structure of IRL. Rather than repeatedly solving full RL problems under adversarial rewards, the learner performs local cost-sensitive policy improvement on expert-relevant states. The paper proposes MMDP, a backward sequence of local moment-matching games, and NRMM, a no-regret policy-iteration variant with expert resets. The central theoretical result is an exponential speedup in horizon dependence of the interaction complexity of the inner loop: for classical primal and dual IRL reductions there exist instances requiring $\rho_E^t$9 interactions, whereas with expert resets the per-iteration sample complexity becomes polynomial,

$Z^\pi(s,a)$0

with imitation guarantees of order $Z^\pi(s,a)$1 or $Z^\pi(s,a)$2 depending on the variant (Swamy et al., 2023).

The relevance of this paper to distributional IRL is therefore classificatory rather than semantic. The distributional object is $Z^\pi(s,a)$3, an expert visitation distribution or occupancy marginal, not a reward posterior and not a return distribution. The paper is explicit that this is not distributional RL in the Bellemare sense and not a method for inferring a distribution over reward functions (Swamy et al., 2023). Its importance lies in showing that occupancy-like distributions can fundamentally alter the computational profile of IRL.

6. Distributional critics, implicit rewards, and full distributional IRL

A more recent connection between IRL and distributional RL appears in implicit-reward imitation methods. RIZE is formulated in the discounted MDP setting $Z^\pi(s,a)$4, explicitly grounded in the Maximum Entropy IRL objective

$Z^\pi(s,a)$5

but replaces explicit reward learning with an implicit reward derived from a critic through the inverse soft Bellman relation (Karimi et al., 27 Feb 2025). The key quantity is

$Z^\pi(s,a)$6

RIZE then introduces an adaptive squared TD-style regularizer

$Z^\pi(s,a)$7

and replaces the scalar critic with a distributional critic $Z^\pi(s,a)$8 represented with IQN-style quantiles (Karimi et al., 27 Feb 2025).

The distributional Bellman operator used in RIZE is

$Z^\pi(s,a)$9

with

$q_\phi(r \mid s,a)$0

This is technically distributional RL, but the policy update and the implicit reward both depend on the expectation of the value distribution rather than on a genuinely distributional control criterion (Karimi et al., 27 Feb 2025). For that reason, RIZE is best read as a MaxEnt-IRL-derived imitation learner with a distributional critic, not as a strong-form distributional IRL method in which reward inference itself is distribution-valued.

DistIRL makes that final step more explicitly. It parameterizes a reward distribution directly, for example with a skew-normal family

$q_\phi(r \mid s,a)$1

trains a quantile critic with the standard quantile Huber objective

$q_\phi(r \mid s,a)$2

and updates the reward model through the variational FSD objective

$q_\phi(r \mid s,a)$3

(Wu et al., 3 Oct 2025). Policy learning then becomes risk-aware through distortion risk measures,

$q_\phi(r \mid s,a)$4

leading to an entropy-regularized objective $q_\phi(r \mid s,a)$5 (Wu et al., 3 Oct 2025).

The conceptual difference between RIZE and DistIRL is sharp. RIZE imports a distributional critic into implicit-reward imitation learning while still optimizing expected $q_\phi(r \mid s,a)$6 (Karimi et al., 27 Feb 2025). DistIRL jointly models a reward distribution and a return distribution, uses FSD violations rather than mean matching, and interprets risk-sensitive policy learning as part of the reward-inference mechanism itself (Wu et al., 3 Oct 2025).

7. Identifiability, limitations, and boundaries of the concept

The principal open issue across all variants is identifiability. Bayesian preference elicitation preserves posterior uncertainty, but still depends on a specified policy class and known dynamics, and its computational bottleneck is repeated dynamic programming within MCMC (Rothkopf et al., 2011). SIRL explicitly targets reward ambiguity, yet the learned distribution is a fitted GMM over reward weights rather than a calibrated posterior derived from a fully explicit generative model (Ju, 2019). CWAE-IRL is distributional only at the transition level; because the reward latent variable primarily helps reconstruct $q_\phi(r \mid s,a)$7, it may absorb transition-correlated factors rather than a semantically global reward unless the setup is strongly constrained (Kusari, 2019).

Return- and profile-matching approaches expose another limitation. Optimality-profile matching uses a target distribution over rewards or discounted future rewards, but distribution matching alone is non-identifiable and therefore requires pairwise ranking or fixed-return supervision; moreover, in the reported experiments the target profile is computed from the ground-truth reward (Haug et al., 2020). DistIRL is the strongest reward-and-return distribution framework among the surveyed papers, yet it does not provide a full identifiability theorem, no full convergence analysis for the alternating optimization, and no guarantee that a single chosen distortion risk measure recovers FSD; the appendix proposition establishes only that dominance under every distortion function implies FSD (Wu et al., 3 Oct 2025).

The literature also draws a boundary around what should not be called distributional IRL. The continuous-time output-feedback IRL method for deterministic linear systems estimates cost parameters online from input-output measurements and inverse HJB regression, but it is a point-estimation method, not a distributional one; it models neither posterior distributions over rewards nor return distributions (Kamalapurkar, 2018). Similarly, occupancy-distribution-based reductions use expert visitation distributions operationally to avoid global exploration, but do not infer reward uncertainty or return laws (Swamy et al., 2023). These boundary cases matter because they show that “distributional” in IRL has been attached to several mathematically distinct objects.

Taken together, the literature suggests a layered taxonomy. At one end are Bayesian models that maintain explicit posterior uncertainty over preferences, policies, and rewards (Rothkopf et al., 2011). Next are reward-distribution models, either as parametric families over reward weights (Ju, 2019) or as amortized latent reward distributions per transition (Kusari, 2019). A separate branch uses distribution matching over scalar rewards or returns as weak supervision (Haug et al., 2020). Another branch uses visitation distributions to alter the computational reduction while leaving reward inference non-distributional (Swamy et al., 2023). Distributional critics in imitation learning sit between IRL and distributional RL but remain expectation-based in their control objective (Karimi et al., 27 Feb 2025). The strongest current formulation is the offline risk-aware framework that jointly learns a reward distribution, a return distribution, and a distribution-aware policy through FSD and distortion risk measures (Wu et al., 3 Oct 2025). This suggests that the term “distributional inverse reinforcement learning” is best understood not as a single algorithmic family, but as an umbrella for IRL methods that preserve and exploit distributional structure in the latent reward, the induced return, the expert occupancy, or some combination of these objects.