Behavioral Cloning and Imitation Learning

Table of Contents

• 1. The Imitation Learning Setting
• 2. Behavioral Cloning
  • 2.1 The Algorithm
  • 2.2 BC as KL Divergence Minimization
• 3. The Covariate Shift Problem
  • 3.1 Distribution Shift Under Execution
  • 3.2 Error Compounding: The O(T²) Bound
• 4. DAgger: Online Imitation Learning
  • 4.1 The Algorithm
  • 4.2 No-Regret Guarantee
• 5. Uses of BC in Reinforcement Learning
  • 5.1 Policy Warm-Starting
  • 5.2 Offline RL: BC as a Regularizer
  • 5.3 Generative Adversarial Imitation Learning
  • 5.4 LLM Post-Training: The SFT Stage
• References

1. The Imitation Learning Setting 🎯

In standard RL the reward function \(r(s, a)\) is given and the agent must discover good behavior through interaction. Imitation learning (IL) inverts this: instead of a reward, we have access to a set of expert demonstrations, and the goal is to recover a policy that behaves like the expert.

Formal setup. Let \(\mathcal{S}\) be the state space, \(\mathcal{A}\) the action space, and \(T\) the episode horizon. An expert policy \(\pi^*\) generates a dataset of trajectories:

\[\mathcal{D} = \{\tau_i\}_{i=1}^N, \quad \tau_i = (s_0, a_0, s_1, a_1, \ldots, s_T)\]

where \(a_t \sim \pi^*(a | s_t)\) and \(s_{t+1} \sim P(s | s_t, a_t)\) from the environment dynamics. The state-action pairs \((s_t, a_t)\) are the demonstrations.

The goal is to learn a policy \(\pi_\theta\) that matches \(\pi^*\) without access to the reward function — only to \(\mathcal{D}\).

IL methods fall into three families, ordered by increasing complexity and data requirements:

2. Behavioral Cloning 🤖

Behavioral cloning (BC) is the most direct approach: treat the demonstrations as a supervised learning dataset and train \(\pi_\theta\) to predict expert actions from states.

2.1 The Algorithm

Given: a dataset \(\mathcal{D} = \{(s_t, a_t)\}\) of expert state-action pairs.

Minimize the negative log-likelihood of expert actions under the learned policy:

\[\mathcal{L}_\text{BC}(\theta) = -\frac{1}{|\mathcal{D}|}\sum_{(s,a) \in \mathcal{D}} \log \pi_\theta(a \,|\, s)\]

This is exactly maximum likelihood estimation of the conditional distribution \(\pi^*(a|s)\) from i.i.d. samples. No environment interaction, no reward — just supervised learning.

2.2 BC as KL Divergence Minimization

The BC objective has a clean probabilistic interpretation. Treating each \((s, a)\) pair as drawn from the joint expert distribution \(\rho^*(s, a) = d^*(s)\,\pi^*(a|s)\) (where \(d^*(s)\) is the state marginal induced by the expert), minimizing \(\mathcal{L}_\text{BC}\) is equivalent to minimizing:

\[D_\text{KL}(\rho^*(s,a) \,\|\, d^*(s)\,\pi_\theta(a|s)) = \mathbb{E}_{s \sim d^*, a \sim \pi^*}\!\left[\log\frac{\pi^*(a|s)}{\pi_\theta(a|s)}\right]\]

The first term (entropy of \(\pi^*\)) is constant in \(\theta\), so this reduces to maximizing \(\mathbb{E}[\log \pi_\theta(a|s)]\) — exactly \(-\mathcal{L}_\text{BC}\).

📐 Key observation: BC minimizes KL divergence between \(\pi^*\) and \(\pi_\theta\) under the expert’s state distribution \(d^*\). It makes no guarantees about \(\pi_\theta\)’s behavior under its own state distribution \(d^{\pi_\theta}\) — the distribution it induces when actually executed. This asymmetry is the root of the covariate shift problem.

3. The Covariate Shift Problem ⚠️

3.1 Distribution Shift Under Execution

BC minimizes loss under the expert’s state distribution \(d^*(s)\). But at test time, \(\pi_\theta\) induces its own state distribution \(d^{\pi_\theta}(s)\), which differs from \(d^*\) whenever \(\pi_\theta \neq \pi^*\).

The problem compounds over time. Suppose \(\pi_\theta\) makes a small error at time \(t=1\), landing in a state \(s_1' \neq s_1^\text{expert}\). From \(s_1'\), the policy was never trained (it saw only expert states), so it may make a larger error at \(t=2\), landing even further from the expert trajectory. Each mistake enlarges the gap between \(d^{\pi_\theta}\) and \(d^*\), causing subsequent mistakes to be larger still.

graph LR
    s0["s₀ (in-distribution)"] --> s1e["s₁* (expert)"]
    s0 --> s1l["s₁' (BC error)"]
    s1l --> s2e["s₂* (expert)"]
    s1l --> s2l["s₂'' (further drift)"]
    s2l --> sdots["…out-of-distribution"]

3.2 Error Compounding: The O(T²) Bound

Ross & Bagnell (2010) formalize this intuition with a sharp bound.

Setup. Let \(\epsilon\) be the BC policy’s expected per-step mistake rate under the expert state distribution:

\[\epsilon = \max_t\, \mathbb{E}_{s \sim d_t^*}\!\left[\mathbf{1}[\pi_\theta(s) \neq \pi^*(s)]\right]\]

Theorem (Ross & Bagnell, 2010). Under mild Lipschitz assumptions on the dynamics, the total cost of the BC policy over a horizon \(T\) satisfies:

\[J(\pi_\theta) \leq J(\pi^*) + O(\epsilon T^2)\]

The error grows quadratically in \(T\), not linearly. Even a tiny per-step error rate \(\epsilon\) is catastrophic for long horizons.

Proof sketch. At time \(t\), the TV distance between \(d_t^{\pi_\theta}\) and \(d_t^*\) accumulates at rate \(\sim \epsilon\) per step: \(\|d_t^{\pi_\theta} - d_t^*\|_\text{TV} \leq \epsilon t\). The cost difference at step \(t\) is bounded by \(\epsilon t \cdot C_\text{max}\). Summing over \(t = 0, \ldots, T-1\):

\[J(\pi_\theta) - J(\pi^*) \leq C_\text{max} \sum_{t=0}^{T-1} \epsilon t = C_\text{max}\,\epsilon\,\frac{T(T-1)}{2} = O(\epsilon T^2)\]

Contrast with on-policy RL. A policy that can query the reward at each step suffers only \(O(\epsilon T)\) degradation, because it corrects at every step rather than compounding. BC’s \(O(T^2)\) penalty is the price of being purely offline.

4. DAgger: Online Imitation Learning 🔄

DAgger (Dataset Aggregation, Ross et al. 2011) fixes the covariate shift problem by querying the expert at the states the learned policy actually visits, rather than only at expert states.

4.1 The Algorithm

Initialize: \(\mathcal{D}_0 = \mathcal{D}\) (initial expert demonstrations), \(\pi_1 = \hat\pi_1(\mathcal{D}_0)\) (BC policy).

For iteration \(n = 1, 2, \ldots, N\):

1. Execute \(\pi_n\) in the environment to collect a trajectory \(\tau_n = (s_0, s_1, \ldots, s_T)\)
2. Query the expert for labels at the visited states: \(a_t^* = \pi^*(s_t)\) for all \(t\)
3. Aggregate: \(\mathcal{D}_n = \mathcal{D}_{n-1} \cup \{(s_t, a_t^*)\}\)
4. Train a new policy \(\pi_{n+1} = \hat\pi(\mathcal{D}_n)\) by BC on the full aggregated dataset

The key move: by labeling states from \(d^{\pi_n}\) (not just \(d^*\)), the training distribution gradually covers the states the policy will encounter. The policy learns to recover from its own mistakes because it has seen those states during training.

4.2 No-Regret Guarantee

DAgger is analyzed as an online learning algorithm. Let \(\hat\ell_n(\pi) = \mathbb{E}_{s \sim d^{\pi_n}}[\mathbf{1}[\pi(s) \neq \pi^*(s)]]\) be the per-step loss of policy \(\pi\) under the state distribution of the \(n\)-th iterate.

Theorem (Ross et al., 2011). After \(N\) iterations, the average policy \(\bar\pi_N = \frac{1}{N}\sum_{n=1}^N \pi_n\) satisfies:

\[J(\bar\pi_N) \leq J(\pi^*) + O\!\left(\epsilon^* T + \sqrt{\frac{\log N}{N}} \cdot T^2\right)\]

where \(\epsilon^* = \min_\pi \mathbb{E}_{s \sim d^\pi}[\ell(\pi, s)]\) is the best achievable loss in the policy class.

Key improvement over BC: The dominant term is \(O(\epsilon^* T)\) — linear in \(T\), matching the on-policy RL bound — rather than BC’s \(O(\epsilon T^2)\). The \(\sqrt{\log N / N}\) term vanishes as \(N \to \infty\).

The proof maps DAgger to the Follow-the-Leader online learning algorithm: at each round the learner picks the policy that minimizes loss on all data seen so far. Follow-the-Leader achieves \(O(\sqrt{N})\) regret on the per-step losses, which translates (after multiplying by the horizon \(T\)) to the bound above.

5. Uses of BC in Reinforcement Learning 🔑

BC is not just an alternative to RL — it is a component inside many modern RL systems.

5.1 Policy Warm-Starting

RL from scratch is slow: the policy must discover good behavior through random exploration before the reward signal provides useful gradients. BC provides a shortcut: pre-train \(\pi_\theta\) on demonstrations via BC, then fine-tune with RL.

This warm-starting pattern appears in: - AlphaGo / AlphaStar / OpenAI Five: BC on human games initializes the policy before self-play RL - Robotics: BC on teleoperation data initializes locomotion and manipulation policies - RLHF for LLMs: the supervised fine-tuning (SFT) stage (§5.4)

The benefit: RL starts from a policy that already produces reasonable behavior, so the reward signal can immediately do useful work rather than searching from scratch. The risk: if the BC pre-trained policy is poor, RL may take longer to escape the bad initialization than starting from scratch.

5.2 Offline RL: BC as a Regularizer

In offline RL, the agent must learn from a fixed dataset \(\mathcal{D}\) of transitions without environment interaction. Naively applying Q-learning or policy gradient to offline data leads to overestimation: the policy learns to exploit out-of-distribution actions that the Q-function incorrectly values highly.

TD3+BC (Fujimoto & Gu, 2021) adds a BC regularization term to the standard actor loss:

\[\mathcal{L}_\text{actor}(\theta) = -\lambda \mathbb{E}_{(s,a) \in \mathcal{D}}\!\left[Q_\phi(s, \pi_\theta(s))\right] + \mathbb{E}_{(s,a) \in \mathcal{D}}\!\left[\|\pi_\theta(s) - a\|^2\right]\]

The second term penalizes deviation from the behavioral policy (the policy that collected \(\mathcal{D}\)) — it is a quadratic BC loss. The combined objective maximizes Q-value while staying close to demonstrated behavior, preventing out-of-distribution extrapolation.

This strategy — maximize reward, regularize toward BC — is the offline RL analog of the KL-penalized RL objective in rejection sampling as RL:

\[\mathcal{J}(\pi) = \mathbb{E}[r] - \beta\,D_\text{KL}(\pi \| \pi_\text{BC})\]

where \(\pi_\text{BC}\) plays the role of \(\pi_\text{ref}\).

5.3 Generative Adversarial Imitation Learning

GAIL (Ho & Ermon, 2016) reframes imitation learning as a GAN problem. Instead of directly matching \(\pi_\theta\) to \(\pi^*\) via BC, GAIL trains:

• A discriminator \(D_\phi(s, a) \in [0,1]\) that tries to distinguish expert \((s,a)\) pairs from \(\pi_\theta\)’s rollouts
• The policy \(\pi_\theta\) that tries to fool \(D_\phi\) — i.e., produce \((s,a)\) pairs that look like the expert

The policy objective uses the discriminator output as a surrogate reward and optimizes it with any RL algorithm (e.g., TRPO):

\[\mathcal{L}_\text{GAIL} = \mathbb{E}_{(s,a) \sim \pi^*}[\log D_\phi(s,a)] + \mathbb{E}_{(s,a) \sim \pi_\theta}[\log(1 - D_\phi(s,a))]\]

GAIL avoids specifying a reward function entirely — the discriminator learns the reward implicitly. At the GAN saddle point, \(D_\phi(s,a) = 1/2\) everywhere, meaning \(\pi_\theta\) and \(\pi^*\) are indistinguishable.

5.4 LLM Post-Training: The SFT Stage

The supervised fine-tuning (SFT) stage of RLHF is behavioral cloning on human-written demonstrations.

Setup. Given a dataset \(\mathcal{D}_\text{SFT} = \{(x_i, y_i)\}\) of (prompt, response) pairs written by humans:

\[\mathcal{L}_\text{SFT}(\theta) = -\sum_{(x,y) \in \mathcal{D}_\text{SFT}} \log \pi_\theta(y | x)\]

This is exactly BC: the “state” is the prompt \(x\), the “action” is the response \(y\), and the “expert” is the human annotator. The resulting \(\pi_\text{SFT}\) is then used as \(\pi_\text{ref}\) for the subsequent RL stage (PPO or GRPO).

The SFT → RL pipeline is the warm-starting pattern from §5.1 applied to language models: 1. BC (SFT): learn to imitate good responses 2. RL: fine-tune \(\pi_\text{SFT}\) to maximize reward while staying close (KL penalty)

RAFT itself is a fusion of both stages: it uses RL (reward filtering) to improve the demonstrations, then BC (SFT) to imitate the improved set. Each RAFT iteration is a micro SFT→RL→SFT cycle, as derived in rejection sampling as RL.

References