gamma = 0.95
env = gym.make('FrozenLake-v1', is_slippery=True)
mdp = d2l.TabularMDP.from_gym(env, gamma)
V_star = d2l.value_iteration(mdp, num_iters=1000)[-1]
pi_star = mdp.backup(V_star).argmax(axis=1)
def demonstrations(num_episodes):
"""Roll the expert; record only what it saw and what it did."""
states, actions = [], []
for _ in range(num_episodes):
s, done = env.reset()[0], False
while not done:
states.append(s)
actions.append(int(pi_star[s]))
s, reward, terminated, truncated, _ = env.step(actions[-1])
done = terminated or truncated
return np.array(states), np.array(actions)
env.reset(seed=0)
demo_s, demo_a = demonstrations(3)
print(f'{demo_s.size} state-action pairs from 3 expert episodes, '
f'covering {np.unique(demo_s).size} of 11 reachable states')Learning from Demonstrations
Dive into Deep Learning · §14.3
Learning from demonstrations
behavior cloning as classification · deployment distribution shift · compounding error \varepsilon T^2 · DAgger data collection
Behavior Cloning Is Supervised Action Prediction
Behavior cloning consumes two columns of a demonstration: the states visited and the actions taken. Fitting \pi_\theta(a \mid s) to the pairs is softmax regression.
96 state-action pairs from 3 expert episodes, covering 7 of 11 reachable states96 labeled pairs; 4 of 11 reachable states never appear.
One Policy Object for Two Chapters
class ActorCritic(nnx.Module):
"""A policy and a value function, each with its own optimizer."""
def __init__(self, policy, value, lr=1e-2):
self.policy, self.value = policy, value
self.opt_pi = nnx.Optimizer(policy, optax.adam(lr), wrt=nnx.Param)
self.opt_v = nnx.Optimizer(value, optax.adam(lr), wrt=nnx.Param)
def log_prob(self, obs, act, policy=None):
"""log pi(a|s). Gradients flow w.r.t. the module you differentiate;
the update functions pass that module back in as `policy`."""
policy = self.policy if policy is None else policy
logp = jax.nn.log_softmax(policy(obs), axis=-1)
return jnp.take_along_axis(logp, act[:, None], axis=-1).squeeze(-1)
def V(self, obs, value=None):
value = self.value if value is None else value
return value(obs).squeeze(-1)
@classmethod
def tabular(cls, num_states, num_actions, lr=0.1, rngs=None):
"""One preference theta_{s,a} per state-action pair: an embedding."""
rngs = nnx.Rngs(d2l.get_key()) if rngs is None else rngs
zeros = nnx.initializers.zeros_init()
return cls(nnx.Embed(num_states, num_actions,
embedding_init=zeros, rngs=rngs),
nnx.Embed(num_states, 1, embedding_init=zeros, rngs=rngs),
lr)nn.Embedding(16, 4) represents the preference table \theta_{s,a}; zero initialization produces a uniform policy. The value head is first used in :numref:sec_policygradient.
Training Fit Does Not Cover Unvisited States
def clone(states, actions, num_steps=200):
"""Behavior cloning: cross-entropy fit of pi(a|s) to expert choices."""
ac = ActorCritic.tabular(16, 4)
obs, act = jnp.asarray(states), jnp.asarray(actions)
def nll_fn(policy):
return -ac.log_prob(obs, act, policy).mean()
for _ in range(num_steps):
loss, grads = nnx.value_and_grad(nll_fn)(ac.policy)
ac.opt_pi.update(ac.policy, grads)
return ac, float(loss)
bc, nll = clone(demo_s, demo_a)
print(f'cross-entropy on the demonstrations after the fit: {nll:.3f}')
for s in (9, 3):
probs = np.exp(bc.log_prob_np(np.repeat(s, 4), np.arange(4)))
print(f'clone pi(.|s={s}): {np.round(probs, 3)}')cross-entropy on the demonstrations after the fit: 0.004
clone pi(.|s=9): [0.001 0.996 0.001 0.001]
clone pi(.|s=3): [0.25 0.25 0.25 0.25]At an unobserved state, \pi(\cdot \mid s = 3) remains exactly uniform, and greedy tie-breaking selects left.
Distribution Shift Reduces Policy Return
env.reset(seed=1)
expert_rate = d2l.evaluate(env, lambda s, rng: int(pi_star[s]),
num_episodes=1000)
clone_rate = d2l.evaluate(env, bc.act_greedy, num_episodes=1000)
mistakes = sum(bc.act_greedy(s) != a for s, a in zip(demo_s, demo_a))
print(f'mistakes on the {demo_s.size} demonstration pairs: {mistakes}')
print(f'success rate: expert {expert_rate:.1%}, clone {clone_rate:.1%}')mistakes on the 96 demonstration pairs: 0
success rate: expert 73.4%, clone 17.5%The classifier is certified on the expert’s states. The agent is tested on the states its own actions produce.
Compounding Error
Proposition. Per-step error \varepsilon under the expert’s distribution can cost \Theta(\varepsilon T^2) return; the same \varepsilon under the learner’s own distribution costs O(\varepsilon T).
After the first deviation, the supervised guarantee no longer controls later states; a deviation at step t can lose all T - t remaining rewards.
lost return at T=10: cloned 2.40 (eps T^2/2 = 2.50), recovering 0.51 (eps T = 0.50)Not a Defect of the Fit
after 3 steps: total variation 0.000
after 5 steps: total variation 0.004
after 10 steps: total variation 0.063
after 20 steps: total variation 0.227
mass in the hole at s=12 after 20 steps: expert 0.000, clone 0.221The distributions are identical for three steps. By step 20, the clone assigns 22% probability to a hole that the expert reaches with probability zero.
DAgger: Relabel the Learner’s States
Roll the learner, keep its states, ask the expert what it would have done, aggregate, refit.
round 0: trained on 96 pairs, success rate 18.0%
round 1: trained on 188 pairs, success rate 72.0%
round 2: trained on 273 pairs, success rate 72.2%
round 3: trained on 426 pairs, success rate 71.9%The added labels include states where the clone differs from the expert; this procedure requires continued expert access rather than only a fixed dataset.
Recap
- Behavior cloning = cross-entropy on (s, a) pairs: no kernel, no reward.
- The guarantee holds on the expert’s distribution; acting moves the test distribution.
- \Theta(\varepsilon T^2) under the expert’s states, O(\varepsilon T) under your own: the gap is the missing off-distribution guarantee.
- DAgger moves training onto the learner’s states with a relabeling loop; the guarantee needs iteration and a no-regret learner.
- SFT of a language model is behavior cloning (:numref:
sec_rl_sequences); BC is the offline baseline (:numref:sec_offline). - :numref:
sec_policygradientreusesActorCriticandpolicy_step, replacing expert labels with sampled rewards.