Imitation and Inverse Reinforcement Learning

Imitation learning is the problem of learning a policy $\pi_\theta(a \mid s)$ from a dataset of expert demonstrations $\mathcal{D} = \{\tau_1, \dots, \tau_N\}$, where each trajectory $\tau_i = (s_0, a_0, s_1, a_1, \dots)$ was generated by an expert policy $\pi_E$, typically without access to a reward signal or even to the environment's transition dynamics during training. The simplest instantiation, behavioral cloning (BC), discards the sequential structure of the problem entirely and treats $\mathcal{D}$ as an i.i.d. supervised dataset of $(s,a)$ pairs, minimizing

This is easy to optimize but ignores that the test-time state distribution is generated by rolling out $\pi_\theta$ itself, not by $\pi_E$ — a mismatch that is the central technical obstacle in this entire chapter.

Inverse reinforcement learning (IRL) takes a different stance: instead of directly mimicking actions, it posits that the expert is (approximately) optimal under some unknown reward function $r^*(s,a)$, and tries to recover $r^*$ from the demonstrations. Once a plausible $r^*$ is found, an ordinary RL algorithm is run against it to produce a policy. The appeal is generalization: a reward function is a more compact, more transferable description of what the expert is trying to achieve than a giant lookup table of state-action pairs, and a policy trained to optimize that reward is robust to distribution shift by construction — if it drifts into an unfamiliar state, it can still evaluate what action locally increases reward, whereas a cloned policy has no such fallback. IRL is fundamentally ill-posed: infinitely many reward functions (including the trivial $r \equiv 0$) make any fixed policy optimal. Resolving this ambiguity is the technical heart of the field, addressed principally through the maximum entropy principle, which selects, among all rewards consistent with the observed behavior, the one whose induced trajectory distribution has maximum entropy subject to matching the expert's feature statistics — a well-justified choice because it makes no assumptions beyond what the data supports (any more "peaked" distribution would encode extra, unsupported preferences). Adversarial imitation learning (notably GAIL) reframes the entire IRL-then-RL pipeline as a two-player game between a policy (generator) and a discriminator, sidestepping the need to ever explicitly represent or interpret $r^*$, much as Deep Policy Gradient Methods use a value or advantage estimate without needing a closed-form value function.

1. Behavioral cloning and why it fails

BC minimizes negative log-likelihood over the expert dataset with ordinary supervised learning machinery (cross-entropy for discrete actions, Gaussian NLL for continuous control). The failure mode is compounding error: suppose the learned policy makes a state-prediction-conditional mistake with probability $\epsilon$ at each of $T$ timesteps. Once it makes a mistake, it enters a state slightly outside the expert's data distribution, where its error probability is not bounded by the same $\epsilon$ (no training signal there), so errors compound. This yields expected total cost that scales as $O(\epsilon T^2)$ in the worst case, versus the $O(\epsilon T)$ one would hope for. This quadratic blowup is the reason self-driving prototypes trained by pure BC tend to drift off-road: a tiny steering error puts the car in a state (partially off-lane) never seen in expert data, where the policy's behavior is essentially undefined and the error compounds.

2. DAgger: closing the distribution-shift gap

Dataset Aggregation (DAgger) fixes this by iteratively collecting states from the current policy's own rollouts, querying the expert for the correct action at those states, and retraining — so the training distribution converges to the policy's own state distribution rather than staying frozen at the expert's. ``` Algorithm: DAgger (Dataset Aggregation) Input: expert policy π_E (queryable), environment, number of iterations N Initialize dataset D ← ∅ Initialize policy π_1 arbitrarily (e.g., random or BC-pretrained) for i = 1 to N:

# Roll out the CURRENT learned policy (not the expert) to collect states Sample T-step trajectories using policy π_i (possibly mixed with π_E early on) for each visited state s in these trajectories: query expert action a = π_E(s) D ← D ∪ {(s, a)}

# Retrain policy by supervised learning on ALL data seen so far π_{i+1} ← argmin_θ (1/|D|) Σ_{(s,a) ∈ D} -log π_θ(a | s) return best π_i on held-out validation rollouts ``` Note the key structural difference from BC: the states in $D$ are drawn from mixtures of the learner's own induced state distribution across iterations, not solely from $\pi_E$'s distribution. This requires an interactive expert (a human in the loop, or a scripted oracle, or a privileged-information controller) — a real cost that pure offline BC does not incur, connecting this tradeoff directly to when Offline Reinforcement Learning must be used instead (no further environment or expert interaction is possible).

3. Maximum Entropy Inverse Reinforcement Learning

MaxEnt IRL assumes the expert's trajectory distribution follows a Boltzmann form over cumulative reward:

where $r_\theta$ is typically linear in features, $r_\theta(s,a) = \theta^\top \phi(s,a)$ (or a neural network in deep variants). Learning $\theta$ maximizes the log-likelihood of the observed expert trajectories under this model:

The gradient of this objective has a clean and interpretable form:

i.e., the difference between the expert's feature expectations and the feature expectations under the current reward's induced (soft-optimal) policy. At the optimum this difference is zero — the learned reward makes the model's expected feature counts match the expert's exactly, which is precisely the maximum-entropy solution consistent with the observed statistics (by a Lagrangian/duality argument: matching feature expectations under a maximum-entropy trajectory distribution is the dual of maximizing the exponential-family log-likelihood above). ``` Algorithm: MaxEnt IRL (tabular / feature-linear version) Input: expert trajectories D, feature map φ(s,a), MDP transition dynamics Initialize θ arbitrarily repeat until convergence:

# 1. Solve for the soft-optimal policy under current reward r_θ = θᵀφ Compute soft value function via soft Bellman backup: Q_soft(s,a) = r_θ(s,a) + γ E_{s'}[ V_soft(s') ] V_soft(s) = log Σ_a exp(Q_soft(s,a)) # softmax, not max Derive soft policy: π_θ(a|s) = exp(Q_soft(s,a) - V_soft(s))

# 2. Compute expected state-visitation frequencies under π_θ

# (forward pass through the MDP, propagating visitation probabilities) μ_θ(s,a) ← state-action visitation frequencies induced by π_θ

# 3. Gradient step: match feature expectations grad ← (1/|D|) Σ_{τ∈D} Σ_t φ(s_t,a_t) - Σ_{s,a} μ_θ(s,a) φ(s,a) θ ← θ + α · grad return θ, π_θ ``` This tabular version requires solving a full soft-value-iteration inside every gradient step — expensive, and infeasible in continuous or high-dimensional state spaces (the inner loop is itself a planning problem). This computational bottleneck is exactly what motivates adversarial imitation learning.

4. Generative Adversarial Imitation Learning (GAIL)

GAIL reframes IRL as: find a policy $\pi_\theta$ whose state-action occupancy measure $\rho_{\pi_\theta}(s,a)$ matches the expert's occupancy measure $\rho_{\pi_E}(s,a)$, measured via Jensen–Shannon divergence, mirroring exactly the generator/discriminator structure of a GAN. A discriminator $D_\phi(s,a) \in (0,1)$ is trained to distinguish expert transitions from policy-generated ones; the policy is trained, via any policy-gradient method, to maximize a reward signal derived from fooling the discriminator. ``` Algorithm: GAIL Input: expert trajectories D_E, environment, policy π_θ, discriminator D_φ Initialize θ, φ randomly for iteration = 1, 2, ...:

# 1. Roll out current policy to collect trajectories Sample trajectories τ_π using π_θ → dataset D_π of (s,a) pairs

# 2. Update discriminator to separate expert from policy data

# D_φ(s,a) → 1 means "looks like expert", → 0 means "looks like policy" φ ← φ + α_D · ∇_φ [ E_{(s,a)~D_E}[ log D_φ(s,a) ] + E_{(s,a)~D_π}[ log(1 - D_φ(s,a)) ] ]

# 3. Define an imitation reward from the discriminator, e.g.

# r(s,a) = -log(1 - D_φ(s,a)) (or log D_φ(s,a) )

# then update the policy with any policy-gradient method (e.g. TRPO/PPO)

# to maximize E[ Σ_t γ^t r(s_t,a_t) ] (plus optional causal-entropy bonus) θ ← θ + α_π · PolicyGradientUpdate(π_θ, r, D_π) return π_θ ``` The discriminator implicitly defines a reward: as $\pi_\theta$ becomes indistinguishable from $\pi_E$, $D_\phi(s,a) \to 1/2$ everywhere, and the imitation reward saturates — a signal, analogous to GAN convergence, that occupancy measures have matched. Crucially, GAIL never explicitly represents $r^*$ as an interpretable object; the discriminator is a reward classifier, not a portable reward function, which is a genuine limitation relative to explicit IRL when transfer to a new environment or task is the goal.

Behavioral cloning. Time complexity per training step is that of ordinary supervised learning: $O(|\mathcal{D}| \cdot C_\text{model})$ per epoch, where $C_\text{model}$ is the cost of a forward/backward pass. This is cheap and requires no environment interaction. The correctness caveat is the crucial one: even if the per-step action-prediction error is bounded by $\epsilon$ everywhere the expert data covers, the expected total cost incurred by rolling the learned policy out for $T$ steps is $O(\epsilon T^2)$ in the worst case (Ross & Bagnell's analysis). Sketch: at each timestep there is probability roughly $\epsilon$ of a mistake; once a mistake occurs, the agent may enter a state distribution never seen at training time, after which errors are essentially unconstrained (worst-case cost per remaining step), and a mistake at time $t$ can cost up to $T - t$ additional steps of accumulated regret — summing $\epsilon(T-t)$ over $t = 1, \dots, T$ gives $O(\epsilon T^2)$.

DAgger. Each iteration costs one policy rollout of length $T$ plus one expert query per visited state plus a retraining pass, so total cost after $N$ iterations is $O(N \cdot T \cdot C_\text{query} + N\cdot |\mathcal{D}| \cdot C_\text{model})$, dominated in practice by the number of expert queries, which is the scarce resource (a human demonstrator, or an expensive privileged oracle). The correctness guarantee is the point of the algorithm: because training data accumulates from the learner's own induced state distribution across iterations (an online-learning / no-regret argument over the sequence of retrained policies), DAgger achieves expected total cost $O(\epsilon T)$ — linear, not quadratic, in horizon — assuming the underlying no-regret online learner converges (this follows from a reduction of imitation learning to no-regret online classification, where the sequence of policies $\pi_1, \dots, \pi_N$ constitutes an online learning process over changing state distributions).

MaxEnt IRL (tabular). Each gradient step requires solving soft value iteration to convergence, costing $O(|S|^2|A|)$ per Bellman backup (or $O(|S||A|)$ for sparse transitions) times the number of backups to converge, repeated over $G$ gradient steps — so total cost is roughly $O(G \cdot |S|^2 |A| \cdot K)$ for $K$ backups per inner solve. This is why tabular MaxEnt IRL is restricted to small, enumerable state spaces; scaling to continuous or high-dimensional states requires either sampling-based approximations (Guided Cost Learning) or abandoning explicit reward recovery altogether (GAIL). Correctness: the MaxEnt formulation is a strictly concave maximum-likelihood problem in $\theta$ when the reward is linear in features and the induced trajectory distribution is in the exponential family, so gradient ascent converges to the unique global optimum where model feature expectations equal empirical feature expectations — this is exactly the maximum-entropy distribution consistent with the observed statistics, by the standard duality between maximum entropy under moment constraints and maximum likelihood in an exponential family.

GAIL. Each iteration costs one policy rollout ($O(T)$ environment steps), one discriminator gradient step ($O(|D_E| + |D_\pi|)$ per minibatch), and one policy-gradient update using an algorithm like TRPO or PPO (inheriting whatever per-step complexity that method has — typically $O(T \cdot C_\text{policy})$ for computing gradients plus the cost of a trust-region or clipped-surrogate optimization). The correctness argument mirrors GAN theory: at the Nash equilibrium of the min-max game, $\rho_{\pi_\theta}(s,a) = \rho_{\pi_E}(s,a)$, i.e., the policy's state-action occupancy exactly matches the expert's, provided the discriminator and policy classes are expressive enough and the alternating optimization converges (a caveat as real as it is for GANs — in practice GAIL training can oscillate or mode-collapse, and convergence guarantees are local, not global, once neural function approximation replaces the tabular/linear setting).

Honest tradeoffs. BC is cheapest but least robust; DAgger is more robust but requires an interactive, queryable expert, which is often unavailable (offline logs cannot be re-queried, directly motivating Offline Reinforcement Learning approaches when only static datasets exist); IRL recovers an interpretable, transferable reward but is computationally heavy and only well-posed under strong regularization assumptions; GAIL is sample-efficient in expert demonstrations but sample-inefficient in environment interactions (it needs many on-policy rollouts, exactly as Deep Policy Gradient Methods do), and it sacrifices reward interpretability and transfer.

Consider a 5-state deterministic gridworld corridor, states $s_0, s_1, s_2, s_3, s_4$ (goal), where the expert always moves right ($a = \rightarrow$). We trace how each method handles a single injected error.

Setup. Feature map $\phi(s) =$ one-hot encoding of position. The expert trajectory is $s_0 \to s_1 \to s_2 \to s_3 \to s_4$, all actions "$\rightarrow$".

BC trace. Suppose the cloned policy has learned $\pi_\theta(\rightarrow \mid s_2) = 0.9$, $\pi_\theta(\leftarrow \mid s_2) = 0.1$ — a small 10% error at $s_2$, purely because $s_2$ was underrepresented in training data relative to $s_0, s_1$. On a rollout, with probability 0.1 the agent steps back to $s_1$. But the training data at $s_1$ was collected from expert trajectories that only ever visited $s_1$ moving forward from $s_0$ — the policy has no signal on "what to do at $s_1$ having just failed a step at $s_2$," so let's say its behavior there degrades further, e.g., $\pi_\theta(\leftarrow \mid s_1) = 0.3$ now (never trained on this branch). The point is quantitative: the error at $s_2$ creates a novel state-distribution problem at $s_1$ that compounds, exactly the $O(\epsilon T^2)$ mechanism, not merely a one-off 10% chance of one wasted step.

DAgger trace. After BC pretraining, we roll out $\pi_1$ and it makes exactly this mistake, visiting $s_1$ via the erroneous backward step. DAgger queries the expert at this specific, off-expert-distribution state: $\pi_E(s_1) = \rightarrow$. This $(s_1, \rightarrow)$ pair — critically, reached via the policy's own error, not the expert's nominal path — is added to $D$. After retraining, $\pi_2(\rightarrow \mid s_1)$ is now correctly reinforced even in this "recovery" context, so a future mistake at $s_2$ no longer compounds: the agent self-corrects at $s_1$ and proceeds to the goal. This is the mechanism, concretely, by which DAgger converts $O(\epsilon T^2)$ into $O(\epsilon T)$: every mistake teaches the policy how to recover from that exact mistake.

MaxEnt IRL trace. Suppose instead we are given only the expert trajectory (no queryable expert) and must recover a reward. Empirical feature expectations from the demonstration: state $s_3$ (one step from goal) is visited once, as is every other state, giving uniform empirical counts along the corridor. We initialize $\theta = 0$ (reward zero everywhere), compute the soft-optimal policy (uniform random, since all $Q$-values are equal), and get feature expectations that overweight early states $s_0, s_1$ (a random walker lingers near the start and rarely reaches $s_4$). The gradient $\phi_\text{expert} - \phi_\text{model}$ is strongly positive at $s_3, s_4$ and negative at $s_0, s_1$, so $\theta$ increases reward near the goal. Iterating this, $r_\theta$ converges to a monotonically increasing reward along the corridor — precisely what makes "always move right" the optimal policy without ever having been told "reach $s_4$" as an explicit goal predicate. This reward, unlike a cloned policy, would also correctly guide behavior if the corridor were extended to $s_5, s_6, \dots$ at test time, because it encodes directional progress, not memorized states — the central generalization advantage IRL buys over BC.

Imitation learning traces back to control-theoretic ideas of learning from demonstration in robotics, but its modern theoretical foundation and its connection to inverse optimal control were shaped by a lineage of researchers bridging control theory, statistical learning theory, and deep RL.

Autonomous driving stacks use variants of imitation learning (often DAgger-like online correction and large-scale behavioral cloning from human driving logs) as a core component alongside planning and perception modules, precisely because reward-shaping "drive safely and comfortably" by hand is intractable. Robotic manipulation research relies heavily on learning-from-demonstration pipelines, teleoperating a robot arm and then using GAIL-style or BC-with-DAgger training to generalize beyond the exact teleoperated trajectories. Game-playing agents and NPC behavior modeling use imitation learning to bootstrap policies from human replays before fine-tuning with RL, since pure RL exploration from scratch in enormous action spaces (e.g., real-time strategy games) is prohibitively expensive without a good initialization. Most consequentially, the reward-modeling stage of RLHF pipelines used to align large language models is structurally an inverse-reinforcement-learning problem: human preference comparisons play the role of "expert demonstrations," a learned reward model plays the role of the recovered $r_\theta$, and a policy-gradient method (often PPO) optimizes against that reward exactly as in the IRL-then-RL pipeline described here — a connection made explicit under RLHF Connections. When demonstrations naturally decompose into reusable sub-behaviors (e.g., "grasp," then "lift," then "place"), imitation and IRL methods interact closely with Hierarchical Reinforcement Learning, which structures policies around such recovered or cloned sub-skills.

Ask the tutor: How does Adversarial Inverse Reinforcement Learning (AIRL) modify GAIL's objective to recover a reward function that transfers across changes in environment dynamics?