Ahmed Haroon

Learning from Demonstrations

In reinforcement learning, an agent discovers a policy from rewards. But sometimes it is much easier to demonstrate the desired behavior than to write a reward function for it.

Imitation learning starts with expert trajectories

\tau^{(i)} = \left(o_1^{(i)}, a_1^{(i)}, \ldots, o_H^{(i)}, a_H^{(i)}\right),

and learns a policy that predicts the expert's action from the information available at each time step:

a_t \sim \pi_\theta(a_t \mid o_t).

This looks like supervised learning: $o_t$ is the input, $a_t$ is the label, and the policy is a conditional predictive model. The difference is that a policy changes its future inputs:

a_t \sim \pi_\theta(a_t \mid o_t), \qquad s_{t+1} \sim p(s_{t+1} \mid s_t, a_t).

A prediction error can move the agent to an unfamiliar state, where another error becomes more likely. This feedback loop is the central difficulty in imitation learning.

When the policy observes the complete state, $o_t=s_t$ , we can use a Markov policy

\pi_\theta(a_t \mid s_t).

The state is Markov when it contains everything needed to predict the next state:

p(s_{t+1} \mid s_1, a_1, \ldots, s_t, a_t) = p(s_{t+1} \mid s_t, a_t).

In a partially observed environment, the policy only receives

o_t \sim p(o_t \mid s_t).

The same observation can correspond to different hidden states. We may then condition the policy on its history

h_t = (o_1, a_1, \ldots, o_{t-1}, a_{t-1}, o_t),

so that

\pi_\theta(a_t \mid h_t) = \pi_\theta(a_t \mid o_{1:t}, a_{1:t-1}).

Imitation learning does not require us to explicitly learn the observation or transition model. However, those dynamics still determine which inputs the policy encounters after deployment.

Behavioral Cloning

The simplest imitation-learning algorithm is behavioral cloning. We collect a dataset of expert trajectories

\mathcal D = \left\{\tau^{(i)}\right\}_{i=1}^{N}

and maximize the likelihood of the demonstrated actions:

\boxed{ \theta^* = \arg\max_\theta \sum_{i=1}^{N}\sum_{t=1}^{H} \log \pi_\theta\left(a_t^{(i)} \mid o_t^{(i)}\right) }.

Equivalently, we minimize

\boxed{ \mathcal L_{\mathrm{BC}}(\theta) = -\frac{1}{NH} \sum_{i=1}^{N}\sum_{t=1}^{H} \log \pi_\theta\left(a_t^{(i)} \mid o_t^{(i)}\right) }.

See derivation

In the partially observed setting, the trajectory probability factorizes as

p_\theta(\tau) = p(s_1) \prod_{t=1}^{H} p(o_t \mid s_t) \pi_\theta(a_t \mid o_t) p(s_{t+1} \mid s_t, a_t),

where the final transition is omitted if the trajectory ends at $H$ . Taking logs gives

\begin{aligned} \log p_\theta(\tau) &= \log p(s_1) + \sum_{t=1}^{H}\log p(o_t \mid s_t) \\ &\quad + \sum_{t=1}^{H}\log \pi_\theta(a_t \mid o_t) + \sum_{t=1}^{H}\log p(s_{t+1} \mid s_t, a_t). \end{aligned}

\arg\max_\theta \log p_\theta(\tau) = \arg\max_\theta \sum_{t=1}^{H}\log \pi_\theta(a_t \mid o_t).

Thus maximum-likelihood trajectory fitting reduces to maximum likelihood over the expert's observation-action pairs.

\textbf{Behavioral Cloning}

\begin{aligned} &\text{Collect expert demonstrations } \mathcal D.\\[0.25em] &\text{Initialize policy parameters } \theta.\\[0.25em] &\text{Repeat:}\\[0.25em] &\qquad \text{sample a minibatch } \mathcal B \subset \mathcal D,\\[0.25em] &\qquad \mathcal L(\theta) \leftarrow -\frac{1}{|\mathcal B|}\sum_{(o,a)\in\mathcal B} \log\pi_\theta(a\mid o),\\[0.25em] &\qquad \theta \leftarrow \theta-\alpha\nabla_\theta\mathcal L(\theta).\\[0.25em] &\text{Return } \pi_\theta. \end{aligned}

For $K$ discrete actions, a neural network produces logits $f_\theta(o)$ and a categorical policy

\pi_\theta(a=k \mid o) = \frac{\exp(f_k(o))}{\sum_{j=1}^{K}\exp(f_j(o))}.

For an expert action $a^*$ ,

-\log \pi_\theta(a^* \mid o) = -f_{a^*}(o)+\log\sum_{j=1}^{K}\exp(f_j(o)),

which is the usual multiclass cross-entropy loss.

For a continuous action $a\in\mathbb R^d$ , a common policy is

\pi_\theta(a \mid o) = \mathcal N\left(a \mid \mu_\theta(o), \Sigma_\theta(o)\right).

Its negative log-likelihood is

\begin{aligned} -\log \pi_\theta(a \mid o) &= \frac{1}{2}(a-\mu_\theta(o))^\top \Sigma_\theta(o)^{-1}(a-\mu_\theta(o))\\ &\quad +\frac{1}{2}\log\det\Sigma_\theta(o) +\frac{d}{2}\log(2\pi). \end{aligned}

If $\Sigma_\theta(o)=I$ is fixed, maximum likelihood is equivalent to mean-squared-error regression:

\arg\max_\theta \sum_t \log \pi_\theta(a_t \mid o_t) = \arg\min_\theta \sum_t \left\|a_t-\mu_\theta(o_t)\right\|_2^2.

Distributional Shift

Behavioral cloning fits actions on observations generated by the expert:

o_t \sim p_{\pi^*}(o_t).

At deployment, observations are generated by the learned policy:

o_t \sim p_{\pi_\theta}(o_t).

These distributions need not agree:

\boxed{p_{\pi^*}(o_t) \neq p_{\pi_\theta}(o_t)}.

Ordinary supervised learning assumes that predictions do not determine which example appears next. Sequential decision making breaks this assumption. A steering error moves a vehicle away from the lane center, creating a camera image absent from the clean demonstrations. The policy is less reliable there, so another error becomes more likely.

\text{mistake} \longrightarrow \text{new state} \longrightarrow \text{distribution shift} \longrightarrow \text{higher error probability}.

The objective controls error under the expert distribution,

\mathbb E_{s\sim p_{\pi^*}} \left[\ell(\pi_\theta(s),\pi^*(s))\right],

but deployment performance depends on

\mathbb E_{s\sim p_{\pi_\theta}} \left[\ell(\pi_\theta(s),\pi^*(s))\right].

A low behavioral-cloning loss does not guarantee a low on-policy loss.

Compounding errors

Let $c(s,a)=\mathbb 1[a\neq\pi^*(s)]$ . Suppose the learner has error at most $\epsilon$ under expert states:

\mathbb E_{s_t\sim p_{\pi^*}} \left[ \Pr_{a_t\sim\pi_\theta(\cdot\mid s_t)} \left(a_t\neq\pi^*(s_t)\right) \right] \leq \epsilon.

If the $H$ decisions were independent test examples, the expected number of errors would be $O(\epsilon H)$ . But an early policy error changes every later input. In the worst case,

\boxed{ J_H(\pi_\theta) = \sum_{t=1}^{H}\mathbb E[c(s_t,a_t)] = O(\epsilon H^2) }.

See derivation

The probability of making at least one mistake before time $t$ is bounded by

\Pr(\text{an earlier mistake}) \leq (t-1)\epsilon.

If no earlier mistake occurs, the learner remains on an expert-like trajectory and makes an error with probability at most $\epsilon$ . After a mistake, we pessimistically allow the cost to be one. Hence

\mathbb E[c(s_t,a_t)] \leq \epsilon + (t-1)\epsilon =t\epsilon.

Summing over the horizon,

\begin{aligned} J_H(\pi_\theta) &\leq \sum_{t=1}^{H}t\epsilon\\ &= \epsilon\frac{H(H+1)}{2}\\ &= O(\epsilon H^2). \end{aligned}

\Pr\left(\bigcup_{k=1}^{t-1}E_k\right) \leq \sum_{k=1}^{t-1}\Pr(E_k).

The argument is deliberately pessimistic: it permits one mistake to place the learner in states where it fails forever.

The bound explains why imperfect demonstrations can sometimes be more useful than pristine ones: a mistake followed by a good correction teaches the policy how to recover.

Dataset Aggregation

Behavioral cloning trains on expert-induced states. DAgger, or Dataset Aggregation, instead asks the expert to label states induced by the learner:

s_t \sim p_{\pi_\theta}(s_t), \qquad a_t^* = \pi^*(s_t).

Starting with an expert dataset $\mathcal D_0$ , iteration $k$ trains

\theta_k = \arg\max_\theta \sum_{(s,a^*)\in\mathcal D_k} \log\pi_\theta(a^*\mid s).

We then execute $\pi_{\theta_k}$ , ask the expert to label the visited states, and aggregate the new examples:

\boxed{ \mathcal D_{k+1} = \mathcal D_k \cup \left\{(s,\pi^*(s)):s\sim p_{\pi_{\theta_k}}\right\} }.

\textbf{DAgger}

\begin{aligned} &\mathcal D \leftarrow \text{initial expert demonstrations}.\\[0.25em] &\text{Repeat:}\\[0.25em] &\qquad \text{train } \pi_\theta \text{ on } \mathcal D,\\[0.25em] &\qquad \text{roll out } \pi_\theta \text{ to collect states } \mathcal S_\pi,\\[0.25em] &\qquad \text{query } a^*(s)=\pi^*(s) \text{ for every } s\in\mathcal S_\pi,\\[0.25em] &\qquad \mathcal D \leftarrow \mathcal D\cup\{(s,a^*(s)):s\in\mathcal S_\pi\}.\\[0.25em] &\text{Return the final policy.} \end{aligned}

Aggregation retains expert states, states visited by weak policies, recovery states, and states visited by stronger policies. For safer early rollouts, we can execute a mixture

\pi_k^{\mathrm{mix}} = \beta_k\pi^*+(1-\beta_k)\pi_{\theta_k},

where $\beta_k$ decreases over time.

DAgger directly attacks distributional shift, but it requires interactive expert access. Labeling arbitrary off-trajectory states can be expensive or unsafe. A practical variant lets the expert intervene only when needed and stores the corrective segments. This reduces annotation, but no longer guarantees coverage of every learner-induced state.

Policies with Memory

Distributional shift is not the only reason behavioral cloning fails. The policy class may be unable to represent the expert.

A reactive policy assumes

a_t \perp (o_1,a_1,\ldots,o_{t-1},a_{t-1}) \mid o_t.

This is unsuitable when the correct action depends on an earlier instruction, hidden velocity, a completed subtask, or an object that is now occluded.

A fixed window conditions on recent context:

\pi_\theta \left(a_t\mid o_{t-K+1:t},a_{t-K+1:t-1}\right).

A recurrent policy instead maintains a learned summary

h_t=f_\theta(h_{t-1},o_t,a_{t-1}), \qquad \pi_\theta(a_t\mid h_t).

A causal transformer directly processes $(o_1,a_1,\ldots,o_t)$ . In either case, the sequence objective becomes

\mathcal L_{\mathrm{seq}}(\theta) = -\sum_{i=1}^{N}\sum_{t=1}^{H} \log\pi_\theta \left(a_t^{(i)}\mid o_{1:t}^{(i)},a_{1:t-1}^{(i)}\right).

During training, previous expert actions are usually available through teacher forcing. During deployment, the history contains the learner's actions instead. Memory solves partial observability but can therefore introduce another train-test mismatch.

More context does not guarantee better decisions. A policy can exploit a variable that predicts the expert's action without representing its cause.

Suppose a brake indicator is visible whenever a human driver brakes. A model may learn to respond to the indicator rather than the pedestrian that caused the braking. The indicator is correlated with the label, but arrives too late to determine the correct action:

\text{predictive correlation} \neq \text{causal reason}.

History helps when it exposes the true cause, but hurts when it supplies easier spurious shortcuts.

Multimodal Behavior

The same state can admit several valid actions. A driver can pass an obstacle on either side, and a robot can grasp an object from several directions.

Suppose the expert distribution has two modes:

p^*(a\mid o) = \frac12\delta(a-a^{(1)}) +\frac12\delta(a-a^{(2)}).

A deterministic policy trained with squared error predicts

\mu^*(o) = \mathbb E[a\mid o] = \frac{a^{(1)}+a^{(2)}}{2}.

But the average of two good actions can be a bad action. Averaging a left turn and a right turn may point directly at the obstacle. We need a policy distribution that preserves distinct modes.

Autoregressive discretization

Discretizing each of $d$ action dimensions into $K$ bins creates $K^d$ joint categories. The probability chain rule avoids this exponential output layer:

\boxed{ \pi_\theta(a_t\mid s_t) = \prod_{j=1}^{d} \pi_\theta(a_{t,j}\mid s_t,a_{t,<j}) }.

The negative log-likelihood decomposes into $d$ categorical losses:

-\log\pi_\theta(a_t\mid s_t) = -\sum_{j=1}^{d} \log\pi_\theta(a_{t,j}\mid s_t,a_{t,<j}).

This reduces the output size from $K^d$ to roughly $dK$ , while retaining dependencies between action dimensions. The trade-off is sequential decoding and sensitivity to the dimension order.

Conditional latent-variable policies

A generative policy introduces a latent variable $z$ that can select different valid behaviors:

p_\theta(a\mid o) = \int p_\theta(a\mid o,z)p(z\mid o)\,dz.

A conditional VAE uses an approximate posterior $q_\phi(z\mid o,a)$ during training and maximizes

\boxed{ \mathcal L_{\mathrm{CVAE}} = \mathbb E_{q_\phi(z\mid o,a)} \left[\log p_\theta(a\mid o,z)\right] -D_{\mathrm{KL}} \left(q_\phi(z\mid o,a)\,\|\,p(z\mid o)\right) }.

The reconstruction term teaches the decoder to produce the expert action. The KL term aligns the posterior used for training with the prior available at test time. The general ELBO derivation is developed in the variational autoencoder notes.

At deployment,

z\sim p(z\mid o), \qquad a\sim p_\theta(a\mid o,z).

Different samples of $z$ can represent different plans. This only works if the decoder actually uses $z$ ; otherwise the latent variable collapses and the policy becomes unimodal again.

Flow-matching policies

Flow matching learns a vector field that transports noise into the action distribution. Let

a_0\sim\mathcal N(0,I), \qquad a_1\sim p_{\mathrm{data}}(a\mid o).

Along the straight interpolation

a_\tau=\tau a_1+(1-\tau)a_0, \qquad \tau\in[0,1],

the target velocity is

\frac{d a_\tau}{d\tau}=a_1-a_0.

We train a conditional vector field with

\boxed{ \mathcal L_{\mathrm{FM}}(\theta) = \mathbb E \left[ \left\| v_\theta(o,a_\tau,\tau)-(a_1-a_0) \right\|_2^2 \right] }.

At deployment, we sample $a_0\sim\mathcal N(0,I)$ and integrate

\frac{d a_\tau}{d\tau}=v_\theta(o,a_\tau,\tau).

Repeated updates transform noise into a sample from a multimodal action distribution. Diffusion policies use a related sequence of noisy actions, but learn a reverse denoising process instead of deterministic ODE transport.

Action chunking

Rather than predict one action, a policy can predict a coherent chunk:

\boxed{ a_{t:t+K}\sim\pi_\theta(a_{t:t+K}\mid o_t) }.

Chunks capture correlations between adjacent controls and reduce the number of high-level policy queries. The cost is an open-loop interval. Long chunks are smooth but slow to react to disturbances. Receding-horizon control predicts a long chunk, executes only its first few actions, and then replans.

Coverage and Recovery

Perfect demonstrations show what success looks like, but may never show how to recover. A robust dataset should include both ideal states

(s,a^*(s))

and plausible deviation states

(\tilde s,a^*(\tilde s)).

Dataset	Strength	Weakness
Narrow, high-quality demonstrations	Consistent successful behavior	Little state and recovery coverage
Broad, imperfect behavior	Diverse situations and corrections	Contains suboptimal actions

Several techniques try to obtain both quality and coverage.

Data augmentation

An observation transformation $T$ may synthetically create a deviation state. But control labels must usually change with the observation:

\tilde o=T(o), \qquad \tilde a=T_a(a,o).

The augmented objective is

\mathcal L_{\mathrm{aug}}(\theta) =-\mathbb E_{(o,a)\sim\mathcal D} \mathbb E_{T\sim\mathcal T} \left[ \log\pi_\theta(T_a(a,o)\mid T(o)) \right].

A side-facing driving camera can approximate lateral displacement, provided its steering label is corrected accordingly. A visually plausible augmentation is harmful when it changes the correct physical action without changing the label.

Pre-training and post-training

Broad data can first teach representations, environment regularities, and recovery behaviors:

\theta_{\mathrm{pre}} =\arg\min_\theta\mathcal L_{\mathrm{broad}}(\theta).

We then fine-tune from $\theta_{\mathrm{pre}}$ on high-quality demonstrations:

\theta_{\mathrm{post}} =\arg\min_\theta \mathcal L_{\mathrm{expert}}(\theta;\theta_{\mathrm{pre}}).

Pre-training supplies coverage; post-training specifies the behavior we want. Domain mismatch and undesirable behavior in the broad dataset remain important risks.

Goal-Conditioned Imitation

When the same state supports different tasks, the goal must be part of the policy input:

\boxed{\pi_\theta(a\mid s,g)}.

Otherwise, actions for different goals appear as conflicting labels.

Given a successful trajectory

\tau^{(i)}=(s_1^{(i)},a_1^{(i)},\ldots,s_{T_i}^{(i)}),

we can treat its terminal state as the achieved goal $g^{(i)}=s_{T_i}^{(i)}$ . Every earlier transition becomes a goal-conditioned example:

\boxed{ \theta^* =\arg\max_\theta \sum_i\sum_{t=1}^{T_i-1} \log\pi_\theta \left(a_t^{(i)}\mid s_t^{(i)},g=s_{T_i}^{(i)}\right) }.

Multi-task learning can share perception, dynamics, and reusable subskills across goals. It also introduces two kinds of shift:

p_{\mathrm{train}}(s_t\mid g) \neq p_{\pi_\theta}(s_t\mid g), \qquad p_{\mathrm{train}}(g) \neq p_{\mathrm{test}}(g).

The first is policy-induced state shift. The second occurs when deployment requests goals unlike the achieved terminal states used during training.

Latent plans

There may be several ways to reach the same goal. A latent-plan policy samples

z\sim p_\psi(z\mid s_1,g), \qquad a_t\sim\pi_\theta(a_t\mid s_t,g,z).

During training, a recognition model can inspect the demonstrated trajectory and infer $q_\phi(z\mid\tau,g)$ . The objective becomes

\begin{aligned} \mathcal L &=\mathbb E_{z\sim q_\phi(z\mid\tau,g)} \left[ \sum_t\log\pi_\theta(a_t\mid s_t,g,z) \right]\\ &\quad-D_{\mathrm{KL}} \left(q_\phi(z\mid\tau,g)\,\|\,p_\psi(z\mid s_1,g)\right). \end{aligned}

The recognition model explains how the demonstration reached the goal, while the proposal model learns to choose a compatible plan without seeing the future.

Hindsight relabeling

A rollout that fails at its requested goal may still succeed at another. If

\tau=(s_1,a_1,\ldots,s_T)

reaches $s_T$ , relabel it with $g'=s_T$ and train on

\log\pi_\theta(a_t\mid s_t,g').

The trajectory becomes a demonstration for what it actually achieved. Repeating rollout, relabeling, and supervised learning turns the agent's own experience into an expanding goal-conditioned dataset.

Choosing the Right Method

Imitation-learning failures are easier to fix once we separate their causes:

Symptom	Likely cause	Useful response
High demonstration loss	Insufficient capacity or optimization	Improve the model, training, memory, or action distribution
Low training loss but deployment drift	Distributional shift	Add recovery data, augmentation, interventions, or DAgger
Averaged or indecisive actions	Multimodal expert behavior	Use autoregressive, latent-variable, flow, or diffusion policies
Correct action depends on earlier context	Partial observability	Use a recurrent or transformer policy
Jittery one-step control	Missing temporal structure	Predict action chunks or add sequence structure
Poor transfer to new tasks	Narrow state or goal coverage	Use broader pre-training and multi-task data

Plain behavioral cloning is a strong baseline when demonstrations densely cover deployment states, the horizon is short, the system naturally recovers, the action distribution is nearly unimodal, and the current observation contains enough information.

Interactive imitation becomes valuable when learner errors create qualitatively new states and an expert can safely label or correct them. Expressive sequence and generative models solve different problems: memory addresses hidden state, while multimodal distributions address multiple valid actions. Neither one, by itself, removes policy-induced distribution shift.

The core logic is therefore

\underbrace{ \min_\theta \mathbb E_{s\sim p_{\pi^*}} [\ell(\pi_\theta(s),\pi^*(s))] }_{\text{behavioral cloning}} \quad\not\Rightarrow\quad \underbrace{ \min_\theta \mathbb E_{s\sim p_{\pi_\theta}} [\ell(\pi_\theta(s),\pi^*(s))] }_{\text{deployment objective}}.

DAgger follows the learner's state distribution, recovery data broadens it, memory restores missing context, and generative policies preserve the expert's distinct behavior modes.