Source-linked AI summary
Imitation Learning as $f$-Divergence Minimization
Liyiming Ke, Sanjiban Choudhury, Matt Barnes, Wen Sun, Gilwoo Lee, Siddhartha Srinivasa
TL;DR
The paper studies imitation learning from multi-modal demonstrations, where learning every mode may be unnecessary and existing objectives can interpolate between modes. It proposes estimating and minimizing f-divergences, emphasizing reverse KL as an I-projection. Empirically, reverse KL collapses to a demonstrator mode more safely and efficiently than KL and JS in multi-modal settings.
Problem
Imitation learning must determine how to imitate multi-modal demonstrations when it may suffice to learn a single expert mode.
Method
The paper introduces a general f-divergence minimization framework, with iterative estimator updates and learner policy-gradient updates, and examines reverse KL.
Results
Reverse KL safely and efficiently collapses to a subset of modes, whereas KL and JS often produce unsafe behavior in multi-modal settings.
Takeaways & Limitations
Choosing the divergence changes learned behavior, and reverse KL is supported as a practical choice for multi-modal inputs.
Takeaways & Limitations
The framework minimizes an approximate divergence estimate, notably a lower bound, and its optimization may be slow to converge.
Abstract
from arXiv · showhide
We address the problem of imitation learning with multi-modal demonstrations. Instead of attempting to learn all modes, we argue that in many tasks it is sufficient to imitate any one of them. We show that the state-of-the-art methods such as GAIL and behavior cloning, due to their choice of loss function, often incorrectly interpolate between such modes. Our key insight is to minimize the right divergence between the learner and the expert state-action distributions, namely the reverse KL divergence or I-projection. We propose a general imitation learning framework for estimating and minimizing any f-Divergence. By plugging in different divergences, we are able to recover existing algorithms such as Behavior Cloning (Kullback-Leibler), GAIL (Jensen Shannon) and Dagger (Total Variation). Empirical results show that our approximate I-projection technique is able to imitate multi-modal behaviors more reliably than GAIL and behavior cloning.
1 Introduction
Multi-modal demonstrations can make behavior cloning and GAIL interpolate between expert modes, whereas reverse KL is proposed to select a single mode. The paper introduces an f-divergence framework and empirically shows reverse KL collapses to one demonstrator mode while KL and JS interpolate unsafely.
- Multi-modal demonstrations arise from diverse near-optimal solutions, variable inputs, and differences in expert skill or preferences.
- Behavior cloning can interpolate between left- and right-steering modes, causing the learner to drive into an obstacle.
- GAIL exhibits a similar interpolation trend, while recovering every latent mode can be prohibitively slow as the number of modes grows.
- The framework views imitation learning as minimizing an f-divergence between expert and learner trajectory distributions.
- Reverse KL has mode-seeking behavior and can collapse onto a subset of modes, unlike mode-covering KL and JS divergences.
2 Related Work
The paper frames imitation learning as choosing an objective for similarity to expert behavior. It situates behavior cloning, DAgger, inverse reinforcement learning, and adversarial methods as different responses to this problem, each with limitations or trade-offs.
- Imitation learning asks what objective best captures similarity between learner behavior and expert demonstrations.
- Behavior cloning treats imitation as supervised action prediction, but small errors can create large distribution mismatch.
- DAgger can alleviate distribution mismatch through interactive expert labeling, but on-policy labels may be impractical and its objective has counter-examples.
- Inverse reinforcement learning and value matching formalize imitation through recovered rewards, Q-values, or feature expectations, including maximum-entropy formulations.
- Adversarial game-theoretic formulations connect inverse reinforcement learning to min-max optimization, enabling methods such as GAIL, SAM, AIRL, and EAIRL.
3 Problem Formulation
The paper formulates imitation as matching expert and learner trajectory distributions with an f-divergence. Because trajectory distributions are difficult to estimate, it instead uses average state-action distributions, whose divergence lower bounds the trajectory-level objective.
- f-Divergences provide a broad family of distributional discrepancy measures, including KL, Jensen-Shannon, and Total Variation divergences.
- The problem is to match the learner trajectory distribution ρπ(τ) to the expert trajectory distribution ρπ∗(τ).
- Different choices of f-divergence produce different learned policies.
- Accurately estimating trajectory divergence requires many samples because trajectory-space size grows exponentially with the time horizon.
- The f-divergence between average state-action distributions lower bounds the corresponding divergence between trajectory distributions.
4 Framework for Divergence Minimization
The framework estimates divergences from expert demonstrations and learner rollouts, then jointly optimizes an estimator and learner policy. Its variational formulation supports any selected f-divergence and recovers algorithmic instances through the same meta-algorithm.
- The framework estimates a divergence lower bound from samples because the expert policy is unknown and only demonstrations are available.
- A variational representation converts f-divergence estimation into an optimization using samples from the compared distributions.
- The estimator is trained as a discriminative classifier between expert demonstrations and learner rollouts, with a lower bound induced by Jensen’s inequality and estimator-class restriction.
- Variational Imitation (VIM) jointly finds a learner policy and discriminator as a saddle point for the chosen divergence.
- The f-VIM meta-algorithm updates the estimator through supervised learning and the learner through policy gradients, yielding different algorithms for different divergences.
1. KL–VIM: Minimizing forward KL divergence
This section presents RKL–VIM as a method for minimizing reverse KL divergence, alongside a table organizing the divergences and estimator components used in the framework.
- RKL–VIM minimizes reverse KL divergence.
- Table 1 lists the f-divergences used in the framework.The table also lists their conjugates, optimal estimators, and activation functions.
- The framework organizes each divergence with its corresponding estimator and activation function.
3. JS–VIM: Minimizing Jensen-Shannon divergence
The framework recovers existing imitation-learning methods through different divergence choices, including GAIL through Jensen–Shannon divergence and DAgger through Total Variation distance.
- JS-VIM is exactly the GAIL optimization without the entropic regularizer.
- DAgger corresponds to Total Variation distance within the divergence framework.Pinsker’s inequality and the metric property of Total Variation provide an associated upper bound.
- DAgger addresses the non-i.i.d. learning problem through iterative supervised learning with an interactive expert.The paper explains counter-examples to DAgger as artifacts of this divergence.
5 Multi-modal Trajectory Demonstrations
For multi-modal demonstrations, fixed-variance Gaussian policies may be unable to match both modes safely, so divergence choice determines whether the learner collapses to one mode or covers both.
- The expert avoids a tree by turning left or right with equal probability, while the Gaussian policy class may be unable to achieve zero divergence.The paper frames objective design as choosing how a globally optimal but imperfect policy should fail safely.
- RKL produces mode collapsing, whereas KL/JS produce mode covering in the illustrated policy behavior.
- For safe behavior, RKL receives only a small penalty while KL receives an infinite penalty.For unsafe behavior that causes a learner crash, the penalty relationship reverses.
6 Experiments
Experiments compare divergences and optimization methods in bandit, gridworld, and Mujoco settings. Reverse KL generally selects single demonstrator modes, while KL and JS tend to interpolate; divergence estimates also underestimate KL and JS values.
- Low dimensional tasks: Policy enumeration found that RKL selects a single mode in bandit and gridworld, whereas KL and JS interpolate between modes.The selected RKL policies were A and Right; KL and JS selected M and Up.
- Low dimensional tasks: RKL strongly prefers mode collapse, KL strongly prefers mode covering, and JS lies between them as control noise changes.The preferred policy can change with the stochasticity parameter ϵ0, especially for RKL at high control noise.
- Divergence estimation: For preferred policies, sample-based KL and JS estimates are often much lower than their true divergences.The paper attributes this pattern possibly to the sampling issue discussed in Appendix F.
- Low dimensional tasks: RKL-VIM produces policies that collapse to one mode, whereas JS-VIM and KL-VIM do not.This result is reported for a continuously parameterized policy class using policy gradient.
- High dimensional continuous control task: On Mujoco, both algorithms perform similarly on Humanoid, Reacher, and Walker2d, while RKL-VIM converges to higher values on Ant, Hopper, and HalfCheetah.RKL-VIM can converge more slowly because it upweights expert-visited states, but terminates with higher reward on those tasks.
7 Discussion
The framework unifies imitation learning algorithms through f-divergence minimization and supports reverse KL for multimodal demonstrations. Its approximation has important estimation and support-related limitations, motivating smoothing and alternative metrics.
- Framework: The framework recovers behavior cloning, GAIL, and DAgger as KL, JS, and TV divergence methods, respectively.It presents imitation learning as f-divergence minimization between learner and trajectory distributions.
- Multimodal demonstrations: In multimodal settings, RKL collapses to a subset of demonstrator modes, whereas KL and JS often produce unsafe interpolation.The discussion characterizes RKL as a practical I-projection choice for these demonstrations.
- Limitations: The framework minimizes an approximate divergence, notably a lower bound that is tight only when the function approximator has sufficient capacity.KL is the only divergence the authors can actually measure in their stated analysis.
- Limitations: For reverse KL, unbounded derivatives may require exponentially many samples, while finite samples cannot generally identify distribution support accurately.Without assumptions on expert or learner distributions, the resulting upper bound is vacuous.
- Practical remedies: Adding Gaussian noise can make learner and expert distributions absolutely continuous and upper-bound divergence magnitude, smoothing the minimized cost.This is presented as a practical remedy because minimizing divergence may matter more than measuring it exactly.
A Lower Bounding f-Divergence of Trajectory Distribution with State Action Distribution
The paper proves that trajectory-level f-divergence is lower bounded by divergence between average state-action distributions, while action-level bounds require additional conditions. The gap reflects information discarded when trajectories are aggregated.
- Lower bound: A generalized log-sum inequality and information-loss lemma establish the lower-bound relationship between divergences over richer and reduced spaces.The proof assumes one random variable contains all information about the other and applies convexity-based inequalities.
- Lower bound: Trajectory divergence is lower bounded by state-action divergence, and the same theorem extends to state distributions.A trajectory contains the state-action information, so marginalizing temporal structure cannot increase the relevant divergence under the theorem’s conditions.
- Gap analysis: The gap equals a conditional f-divergence over trajectories containing a given state-action pair, scaled by the state-action probability.The difference arises because trajectories are treated as separate events while their state-action representation merges them.
- Upper bound: The upper bound requires querying expert labels on states visited by the learner, making an interactive expert necessary.This contrasts with settings where only fixed expert demonstrations are available.
- Upper bound: For divergences satisfying the triangle inequality, trajectory divergence is upper bounded by expected action-distribution divergence on learner-induced states.Total Variation and squared Hellinger distance are examples of divergences with this property.
C Existing algorithms as different f-divergence minimization
The framework interprets existing imitation-learning algorithms as minimizing different f-divergences and develops reverse-KL methods aimed at selecting rather than interpolating among modes. Interactive and density-ratio variants provide alternative optimization routes with distinct assumptions and guarantees.
- Existing algorithms: Behavior cloning minimizes KL divergence, GAIL minimizes JS divergence, and DAgger minimizes Total Variation within the unified f-divergence framework.The corresponding objectives arise by plugging each divergence into the framework.
- Existing algorithms: DAgger’s Total Variation objective can admit policies that leave the expert’s track yet imitate the expert perfectly on the resulting bad states.The paper presents this as an undesirable behavior and a counterexample to relying on DAgger’s objective alone.
- Reverse KL methods: RKL-VIM approximates the I-projection by minimizing a lower bound on reverse KL, but uses two lower bounds and may require complex estimators and slow minimax optimization.The approximation first moves from trajectory to state-action divergence and then applies a variational lower bound.
- Reverse KL methods: RKL-iVIM queries the expert on learner-induced states and collects fresh expert and learner data at every iteration.This interactive design directly targets action-distribution divergence on states visited by the learner.
- Reverse KL methods: Algorithm 3 estimates π_n(a|s)/π*(a|s) with density-ratio estimation, aggregates state-cost pairs, and updates the policy by cost-sensitive classification.Its performance depends on density-ratio estimation, classification, and no-regret learning quality.
E.2 An Example of Density Ratio Estimation and Its Finite Sample Analysis
The density-ratio analysis assumes bounded support and realizability, then uses finite-class uniform convergence to control estimation error. The paper also highlights divide-by-zero failures when sampled and target distributions have mismatched support.
- Setup and assumptions: The paper adds a finite-sample analysis for the Nguyen et al. density-ratio estimation algorithm.This analysis was identified as missing from the earlier work.
- Setup and assumptions: Assumption E1 requires p and q to be lower bounded by c, ensuring full support and ratios in [c, 1/c].The bound makes p(z)/q(z) well defined for every element z.
- Setup and assumptions: Assumption E2 requires the true ratio r(z)=p(z)/q(z) to belong to a finite realizable function class G.The discrete-class assumption supports the stated uniform-convergence analysis.
- Finite-sample analysis: For continuous function classes, log(|G|) can be replaced by complexity measures such as Rademacher complexity.The paper presents the finite class as an analytical simplification rather than an essential restriction.
- Finite-sample analysis: The analysis applies concentration and union bounds over G to control empirical risk and obtain a high-probability ratio-estimation guarantee.The proof uses Hoeffding’s inequality and the finite size of G.
- Support mismatch: If p(x)>0 while q(x)=0, sample-based estimators may miss infinite f-divergence contributions, especially for KL and JS.This is the central divide-by-zero issue in divergence estimation.