Imitation Learning as Return Distribution Matching

2025-09-15

AI Summary

This paper addresses risk-sensitive imitation learning by precisely matching the expert’s return distribution—not merely its expectation—via the Wasserstein distance, thereby jointly replicating both average performance and risk preferences (e.g., variance, skewness). To enable tractable distributional matching, we introduce a rich yet efficiently implementable class of non-Markovian policies. We propose two theoretically grounded algorithms: RS-KT, designed for settings with known environment dynamics and offering substantially improved sample complexity; and RS-RL, applicable when the reward function is unknown. Experiments demonstrate that our approach outperforms existing baselines in both sample efficiency and fidelity of risk-behavior modeling. Notably, it achieves, for the first time, end-to-end, distribution-level imitation of expert risk preferences—capturing higher-order statistical properties of the return distribution beyond mean and variance.

Technology Category

• Machine Learning: Imitation Learning & Inverse Reinforcement Learning
• Intelligent Robots: Learning & Optimization for ROB
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• Responsible Web: Human-perceived consequences of algorithmic deployment on the web
• Systems and Infrastructure for Web, Mobile and WoT: Experiences and lessons learnt from Web-based algorithms and system deployments
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

We study the problem of training a risk-sensitive reinforcement learning (RL) agent through imitation learning (IL). Unlike standard IL, our goal is not only to train an agent that matches the expert's expected return (i.e., its average performance) but also its risk attitude (i.e., other features of the return distribution, such as variance). We propose a general formulation of the risk-sensitive IL problem in which the objective is to match the expert's return distribution in Wasserstein distance. We focus on the tabular setting and assume the expert's reward is known. After demonstrating the limited expressivity of Markovian policies for this task, we introduce an efficient and sufficiently expressive subclass of non-Markovian policies tailored to it. Building on this subclass, we develop two provably efficient algorithms, RS-BC and RS-KT, for solving the problem when the transition model is unknown and known, respectively. We show that RS-KT achieves substantially lower sample complexity than RS-BC by exploiting dynamics information. We further demonstrate the sample efficiency of return distribution matching in the setting where the expert's reward is unknown by designing an oracle-based variant of RS-KT. Finally, we complement our theoretical analysis of RS-KT and RS-BC with numerical simulations, highlighting both their sample efficiency and the advantages of non-Markovian policies over standard sample-efficient IL algorithms.

Problem

• Matching expert return distribution beyond expected value
• Developing risk-sensitive imitation learning algorithms
• Addressing expressivity limitations with non-Markovian policies

Innovation

• Non-Markovian policies for distribution matching
• Wasserstein distance matching expert return distribution
• Two algorithms RS-BC and RS-KT