Distributional Inverse Reinforcement Learning

We propose a distributional framework for offline Inverse Reinforcement Learning (IRL) that jointly models uncertainty over reward functions and full distributions of returns. Unlike conventional IRL approaches that recover a deterministic reward estimate or match only expected returns, our method captures richer structure in expert behavior, particularly in learning the reward distribution, by minimizing first-order stochastic dominance (FSD) violations and thus integrating distortion risk measures (DRMs) into policy learning, enabling the recovery of both reward distributions and distribution-aware policies. This formulation is well-suited for behavior analysis and risk-aware imitation learning. Theoretical analysis show that the algorithm converge with $\mathcal{O}(\varepsilon^{-2})$ iteration complexity. Empirical results on synthetic benchmarks, real-world neurobehavioral data, and MuJoCo control tasks demonstrate that our method recovers expressive reward representations and achieves state-of-the-art imitation performance.

Many AI systems learn from examples of what an expert did, but they often assume there is one fixed score, or reward, explaining those actions. In real settings, that score can be uncertain: a robot may face rare failures, and an animal’s brain reward signal can vary from trial to trial. This paper introduces Distributional Inverse Reinforcement Learning, a method that learns a full range of possible rewards from previously collected demonstrations instead of only an average reward. It compares the expert’s and learner’s whole outcome distributions, so it can notice differences in variability and rare outcomes that ordinary methods miss. The method also trains policies that account for risk, such as preferring reliable actions when the demonstrations suggest caution. In tests on toy navigation tasks, mouse behavior with dopamine recordings, and robot control benchmarks, the method recovered more realistic reward patterns and matched expert behavior better than several offline imitation-learning baselines. This matters because it gives researchers a more expressive tool for studying behavior and building safer imitation-learning systems when uncertainty and risk are part of the task.