A Divergence Minimization Perspective on Imitation Learning Methods

TL;DR

The paper recasts imitation learning through a divergence-minimization lens, introducing f-MAX as a unifying Max-Ent IRL framework that generalizes AIRL and links BC, GAIL, and AIRL via f-divergences. It identifies state-marginal matching as the primary driver of IRL’s superior performance in low-data settings and introduces FAIRL for forward KL optimization to further explore divergence effects. Beyond standard IL, the authors apply state-marginal matching to synthesize diverse behaviors using hand-designed state distributions, without rewards or expert demonstrations. Empirical results on high-dimensional continuous control tasks validate the framework and highlight the practical relevance of state marginal alignment for robust imitation and exploration.

Abstract

In many settings, it is desirable to learn decision-making and control policies through learning or bootstrapping from expert demonstrations. The most common approaches under this Imitation Learning (IL) framework are Behavioural Cloning (BC), and Inverse Reinforcement Learning (IRL). Recent methods for IRL have demonstrated the capacity to learn effective policies with access to a very limited set of demonstrations, a scenario in which BC methods often fail. Unfortunately, due to multiple factors of variation, directly comparing these methods does not provide adequate intuition for understanding this difference in performance. In this work, we present a unified probabilistic perspective on IL algorithms based on divergence minimization. We present $f$-MAX, an $f$-divergence generalization of AIRL [Fu et al., 2018], a state-of-the-art IRL method. $f$-MAX enables us to relate prior IRL methods such as GAIL [Ho & Ermon, 2016] and AIRL [Fu et al., 2018], and understand their algorithmic properties. Through the lens of divergence minimization we tease apart the differences between BC and successful IRL approaches, and empirically evaluate these nuances on simulated high-dimensional continuous control domains. Our findings conclusively identify that IRL's state-marginal matching objective contributes most to its superior performance. Lastly, we apply our new understanding of IL methods to the problem of state-marginal matching, where we demonstrate that in simulated arm pushing environments we can teach agents a diverse range of behaviours using simply hand-specified state distributions and no reward functions or expert demonstrations. For datasets and reproducing results please refer to https://github.com/KamyarGh/rl_swiss/blob/master/reproducing/fmax_paper.md .

Figures (3)

• Figure 1: $r(s,a)$ as the function of the logits of the optimal discriminator, $\ell^\textnormal{opt}(s,a) = \textnormal{log } \frac{\rho^{\textnormal{exp}}(s,a)}{\rho^\pi(s,a)}$. As a reminder, AIRL, GAIL, and FAIRL respectively correspond to the reverse KL, JS, and forward KL divergences.
• Figure 2: (a) Using the Fetch robot we demonstrate that we can train exploration policies through our approach to state-marginal matching. Figures in order are: Fetch environment, target, and two policies' state marginals. Full image region depicts the extent of the table. (b) In the point-mass domain we train policies that exhibit complex and multi-modal trajectories. (c) Using state-marginal matching we train policies for solving the Pusher Push task. Left image is Pusher environment. The next two columns correspond to the target and policy distribution of the arm tip position. Top images are bird's eye view (x-y) and bottom images are side view (y-z) of these distributions. (d) Pusher Draw target and policy distributions. Top images are top-down view of arm tip distribution and bottom images visualize the $z$ coordinate as a function of angle of rotation around the circle.
• Figure 3: Hyperparameter grid search for AIRL and FAIRL in the Ant environment.