PaperlayerView exact revision on arXiv

Source-linked AI summary

IQ-Learn: Inverse soft-Q Learning for Imitation

Divyansh Garg, Shuvam Chakraborty, Chris Cundy, Jiaming Song, Matthieu Geist, Stefano Ermon

arXiv:2106.12142v4cs.LGcs.AI

TL;DR

Imitation learning from limited expert data is difficult because behavioral cloning ignores environment dynamics, while dynamics-aware alternatives can require difficult adversarial training. IQ-Learn instead learns one soft Q-function that implicitly represents reward and policy, achieving strong offline and online results while recovering rewards correlated with ground truth.

  • Problem

    Behavioral cloning ignores sequential environment dynamics, while many dynamics-aware imitation-learning methods require difficult adversarial optimization over reward and policy functions.

  • Method

    IQ-Learn learns a single Q-function that implicitly defines both the reward and policy, replacing min-max optimization over separate functions with minimization over Q.

  • Results

    IQ-Learn achieves state-of-the-art imitation-learning results across sparse-data, offline, online, and complex image-based settings, while learned rewards show high positive correlation with ground-truth rewards.

  • Takeaways & Limitations

    The method supports imitation learning and inverse reinforcement learning with sparse expert data, including complex image-based environments and fewer environment interactions than prior methods.

  • Takeaways & Limitations

    The simplification used to derive the update is not applicable in the online imitation-learning setting.

Abstract

from arXiv · show

In many sequential decision-making problems (e.g., robotics control, game playing, sequential prediction), human or expert data is available containing useful information about the task. However, imitation learning (IL) from a small amount of expert data can be challenging in high-dimensional environments with complex dynamics. Behavioral cloning is a simple method that is widely used due to its simplicity of implementation and stable convergence but doesn't utilize any information involving the environment's dynamics. Many existing methods that exploit dynamics information are difficult to train in practice due to an adversarial optimization process over reward and policy approximators or biased, high variance gradient estimators. We introduce a method for dynamics-aware IL which avoids adversarial training by learning a single Q-function, implicitly representing both reward and policy. On standard benchmarks, the implicitly learned rewards show a high positive correlation with the ground-truth rewards, illustrating our method can also be used for inverse reinforcement learning (IRL). Our method, Inverse soft-Q learning (IQ-Learn) obtains state-of-the-art results in offline and online imitation learning settings, significantly outperforming existing methods both in the number of required environment interactions and scalability in high-dimensional spaces, often by more than 3x.

1 Introduction

Imitation learning seeks policies that reproduce expert behavior, but behavioral cloning ignores dynamics and can compound errors. IQ-Learn addresses this with stable, non-adversarial Q-function learning and reports strong sparse-data, high-dimensional, and reward-recovery results.

  • Behavioral cloning treats imitation learning as supervised action prediction, ignoring dynamics and risking compounding errors after departure from expert states.
  • IQ-Learn uses stable, non-adversarial training for dynamics-aware imitation learning and achieves state-of-the-art benchmark performance.
  • IQ-Learn replaces adversarial reward-policy optimization with a single Q-function that implicitly represents both reward and policy.
  • IQ-Learn surpasses prior methods with one expert demonstration offline and scales to image-based Atari tasks, reaching expert performance.
  • IQ-Learn supports online and offline imitation learning, outperforming prior methods by 3-7x on Atari while requiring 3x fewer environment steps.
  • The method introduces a modified Q-learning update, a framework for IPMs and f-divergences, and rewards highly correlated with ground-truth rewards.

2 Background

The paper formulates sequential decision-making with MDPs, soft Q-functions, maximum-entropy reinforcement learning, and maximum-entropy inverse reinforcement learning. These foundations connect expert occupancy measures and reward-policy optimization to statistical distances between expert and learned behavior.

  • An MDP comprises state and action spaces, initial-state distribution, transition dynamics, reward function, and discount factor.
  • Policies induce discounted occupancy measures over state-action pairs, while expert behavior is represented through demonstrations because the expert policy is unknown.
  • Soft Q-functions satisfy a soft Bellman relation, and the soft Bellman operator defines a unique Q-function for a fixed reward and policy.
  • Maximum-entropy reinforcement learning maximizes expected discounted reward plus policy entropy, with the optimal policy determined by the optimal soft Q-function.
  • Maximum-entropy inverse reinforcement learning searches for rewards favoring expert behavior while optimizing the policy for those rewards in an inner loop.
  • For unrestricted reward functions, inverse reinforcement learning can be expressed as minimizing a statistical distance between expert and policy occupancy measures.

3 Inverse soft Q-learning (IQ-Learn) Framework

IQ-Learn transforms inverse reinforcement learning from reward-policy space into Q-policy space using an inverse soft Bellman operator. Because both policy and reward can be represented through Q, the nested optimization reduces to a concave optimization over one function while retaining the IRL saddle point.

  • The naive inverse reinforcement learning solution alternates reward learning with inner-loop reinforcement learning to find an optimal policy.
  • The optimal policy and reward can both be obtained from Q, reducing the nested min-max problem to optimization over the Q-function.
  • The inverse soft Bellman operator is bijective for a fixed policy, establishing a one-to-one correspondence between rewards and soft Q-functions.
  • IQ-Learn defines the transformed objective J(π,Q) by substituting the inverse soft Bellman mapping into the original inverse reinforcement learning objective.
  • The framework analyzes broad function classes but restricts regularizers to convex-function-induced forms because general Q-policy regularizers are difficult to analyze.
  • Under the stated regularizer restriction, the transformed objective has a unique saddle point corresponding to the inverse reinforcement learning solution.
  • For fixed Q, policy optimization is closed-form and equivalent to maximum-entropy reinforcement learning with reward T^πQ.
  • After substituting the minimizing policy, J*(Q) is concave and has a unique maximum Q* that recovers the required saddle point.

4 Approach

IQ-Learn recovers an optimal soft Q-function from expert data, using it to represent policy and reward while optimizing statistical distances without the usual separate reward-policy min-max formulation.

  • Inverse soft-Q learning: IQ-Learn recovers the optimal soft Q-function from an expert distribution, first using an energy-based policy and later an explicit actor-critic policy.The explicit policy approximates the optimal Q-induced policy in continuous action spaces.
  • General inverse RL objective: The framework minimizes statistical distances between expert and learned distributions, including Integral Probability Metrics and f-divergences.The choice of regularizer determines the induced statistical distance.
  • Choice of statistical distances: Bounded rewards, smooth functions, and norm penalties correspond respectively to Total Variation, Wasserstein-1, and χ2-divergence choices.These reward properties are enforced through the associated regularizers.
  • Inverse soft-Q learning: Each Q-function uniquely determines a reward, so every optimization update can be interpreted as finding a better IRL reward.The recovered reward uses Q-values and expected next-state soft values.
  • Inverse soft actor-critic update: The actor-critic variant alternates Q optimization under a fixed policy with a SAC update that moves the policy toward the Q-induced policy.The policy update also minimizes the inverse soft-Q objective.
  • Inverse soft actor-critic update: Unlike ValueDICE, IQ-Learn does not update the actor adversarially, addressing an objective that may not always converge.This comparison concerns the actor update and convergence behavior described for ValueDICE.

5 Practical Algorithm

The practical IQ-Learn algorithm modifies standard soft-Q or SAC implementations with a compact Q objective and provides separate offline and online estimators for the required value terms.

  • Algorithm: IQ-Learn can be implemented in 15 lines of code by changing the Q-function objective in standard soft DQN or SAC implementations.Default hyperparameters generally work well, with entropy regularization requiring tuning and target networks helping continuous control.
  • Online implementation: Online training estimates the value difference from replay-buffer transitions, using policy and expert samples to stabilize convergence.The best-performing sampling scheme uses half policy-buffer samples and half expert-distribution samples.
  • Algorithm: Algorithm 1 alternates Q-function training with an optional SAC-style policy update for the actor-critic variant.The Q update uses V* for Q-learning and Vπ for actor-critic.
  • Recover policy and reward: The learned transition reward is recovered as Qθ(s,a) − γVπ(s′), requiring next states from the environment or a dynamics model.Marginalizing over next states recovers the corresponding state-action reward.
  • Offline implementation: Offline training estimates the value difference using only expert transitions, avoiding the need for policy samples while producing state-of-the-art offline IL results.This estimator is used because directly estimating the initial-state value can overfit.
  • Statistical-distance implementations: The χ2-divergence objective adds a squared regularization term to the expert transition expression, and its offline simplification matches the CQL objective using zero rewards.The simplification does not apply in the online setting.

6 Related Work

Related imitation-learning methods span classical supervised cloning, dynamics-aware online adversarial approaches, and offline objectives that differ in how they use dynamics and optimize policies and Q-functions.

  • Classical IL: Classical imitation learning matches expert actions through supervised learning, while RL/IRL methods infer rewards before finding policies that maximize them.Early RL/IRL approaches required hand-designed MDP featurizations.
  • Online IL: Online IL methods use environment interaction for dynamics information; GAIL formulates learning as an adversarial max-min problem over reward and policy.Several later methods build on this adversarial approach, while others fix a reward or policy component.
  • Offline IL: Offline IL removes environment access: behavioral cloning remains dynamics-agnostic, whereas ValueDICE and related methods incorporate dynamics through objectives involving policy and Q-functions.ValueDICE uses adversarial optimization and a biased gradient estimator according to the passage.

7 Experiments

Experiments evaluate IQ-Learn across offline and online imitation learning, low-dimensional control, MuJoCo, and Atari. IQ-Learn achieves strong performance with sparse demonstrations, fast convergence, reward recovery, and generalization under distribution shifts.

  • Experimental settings: IQ-Learn is evaluated across low-dimensional control, MuJoCo, Atari, and both offline and online imitation-learning settings.The comparisons include CartPole, Acrobot, LunarLander, HalfCheetah, Hopper, Walker, Ant, and visually challenging Atari tasks.
  • Offline IL: IQ-Learn strongly outperforms prior methods on three offline control tasks and achieves expert performance on Acrobot using one expert trajectory.It reaches near-expert performance on CartPole with the same amount of expert data.
  • MuJoCo control: IQ-Learn achieves expert-level performance across MuJoCo control tasks using a single expert trajectory, outperforming ValueDICE and GAIL.SQIL was not competitive in this setting and is omitted from the comparison.
  • Atari: On Atari, IQ-Learn reaches expert performance on Space Invaders and near-expert performance on Pong and Breakout using 20 expert demonstrations.It obtains 3-7x normalized score compared with SQIL and converges in approximately 300k environment steps, versus more than 1M for Q-learning-based methods.
  • Recovered rewards: IQ-Learn recovers rewards that correlate strongly with ground truth, including a Pearson correlation of 0.99 on Hopper.Recovered rewards are also visualized in a discrete GridWorld with 30 expert demonstrations.
  • Generalization: IQ-Learn remains robust to shifts in initial-state and goal distributions, generalizing to unseen goals and supporting meta-learning for imitation.These distribution-shift experiments are detailed in Appendix F.

8 Discussion and Outlook

IQ-Learn provides a principled framework for learning soft-Q functions in imitation learning while recovering policies and rewards. Its rewards can be highly correlated with actual rewards, but depend on environment dynamics, limiting direct reward transfer.

  • IQ-Learn learns Q-functions for imitation learning while recovering the optimal policy and reward.
  • IQ-Learn recovers rewards highly correlated with actual rewards, but those rewards depend on environment dynamics and cannot be trivially transferred across reward settings.
  • State-only rewards: The framework’s theoretical development uses a modified objective for learning Q while recovering state-only rewards.
  • State-only rewards: The state-only reward objective removes action dependence by replacing the expert policy with the learner policy.
  • State-only rewards: Recovered state-only rewards are high around all sides of the target, indicating they are not directional or action-dependent.

A.2 Proofs for Section 3 and Section 4

The proofs establish a one-to-one correspondence between soft-Q functions and rewards and characterize the resulting optimization landscape. Under the stated regularity conditions, the objective has a unique saddle point and SAC updates converge toward it.

  • For a fixed policy, the inverse soft Bellman operator bijectively maps Q-functions to rewards.
  • The initial-state identity follows from telescoping occupancy-measure sums and relates expected soft Bellman terms to the initial-state value.
  • The Q-policy objective is concave in Q because the inverse soft Bellman transformation is affine and the inverse-RL objective is concave in rewards.
  • The unique saddle-point result extends from non-decreasing to arbitrary concave φ within the feasibility region.
  • The policy minimizing the objective for fixed Q is the energy-based policy induced by Q.
  • SAC policy updates minimize KL divergence toward the Q-induced policy and are guaranteed to converge with sufficient steps, unlike adversarial updates.

C.3 Effect of various divergences

The paper analyzes how statistical divergences determine the optimization geometry and saddle point, and introduces a modified Reverse-KL representation to avoid biased gradients. Experiments compare divergence choices and variational forms in offline LunarLander imitation.

  • Different regularizers define different Q-maxima manifolds and therefore different saddle points, while the resulting policy remains close to the expert policy.
  • For Integral Probability Metrics, the objective is maximized on the boundary of the reward-function constraint without a closed-form equation.
  • ValueDICE uses the Donsker-Varadhan Reverse-KL form but suffers from biased gradients and adversarial policy updates.
  • Reverse-KL variants: The proposed Modified KL representation preserves the Q-maxima manifold of ValueDICE while avoiding its biasing issue.
  • Reverse-KL variants: In LunarLander with one expert trajectory, KL-fix and DV perform similarly, while the original KL form performs noticeably worse.

D.2 Additional Results

IQ-Learn achieves strong imitation and inverse-reinforcement-learning results across MuJoCo, Atari, reward recovery, tabular settings, and dynamics-aware toy environments. It remains stable under extended training and outperforms behavioral cloning when dynamics matter.

  • IQ-Learn reaches state-of-the-art performance and expert-level rewards on MuJoCo tasks using 10 expert trajectories.
  • IQ-Learn avoids overfitting with one Half-Cheetah expert trajectory, while ValueDICE begins overfitting around 100k update steps.
  • A Pearson correlation of 0.99 between recovered and original Hopper episode rewards indicates that IQ-Learn can support inverse reinforcement learning.
  • In tabular Gridworld, IQ-Learn recovers rewards similarly to Max Entropy IRL while avoiding its expensive backward pass.
  • With one expert trajectory in the Loop MDP, dynamics-aware methods find the optimal policy, whereas behavioral cloning achieves around 50 reward.

E.2 Ablation on Gamma

The ablations show that dynamics-awareness is important for stable imitation and that performance depends on the training setting. IQ-Learn handles some distribution shifts online but remains limited offline when initial states differ.

  • E.2 Ablation on Gamma: Setting γ to 0 removes dynamics-awareness from IQ-Learn, while γ = 0.99 retains it.
  • E.2 Ablation on Gamma: With γ = 0, training is unstable and rewards strongly decay; with dynamics enabled, training is stable and converges properly.
  • E.2 Ablation on Gamma: IQ-Learn is evaluated under initial- and goal-distribution shifts to test robustness beyond the expert distribution.
  • E.2 Ablation on Gamma: In the modified LunarLander initial-shift setting, online learning reaches an average episode reward of 250 with 10 expert demonstrations.
  • E.2 Ablation on Gamma: Applying reward regularization to both expert and policy states can help the implicit reward generalize beyond the expert distribution.

F.1.2 Goal Distribution Shift

Goal distribution shifts make inverse reinforcement learning difficult because demonstrations may not cover the changing targets. IQ-Learn improves with more demonstrations offline and can use online exploration to generalize across goals.

  • F.1.2 Goal Distribution Shift: Reacher_easy changes its target position randomly each episode, making broad goal coverage difficult with few demonstrations.
  • F.1.2 Goal Distribution Shift: Offline IQ-Learn typically overfits to one target with a single demonstration, while imitation performance improves as the number of expert demonstrations increases.
  • F.1.2 Goal Distribution Shift: Online exploration lets IQ-Learn correct behavior across episodes and learn a reward function that generalizes over multiple goals.
  • F.1.2 Goal Distribution Shift: With five expert demonstrations, IQ-Learn matches the mean rewards obtained by BC and GAIL with 50 experts: 325.2 and 440.1, respectively.
Loading 2106.12142v4…