Foundations — Imitation Learning, Covariate Shift & the Multimodality Problem

Pure PyTorch + NumPy + Matplotlib. No robot, no simulator, runs on CPU in seconds. The point is to see the theory above with your own eyes before we touch hardware.

Setup

# Python 3.10+ — versions are forgiving here, this is a toy.
pip install torch numpy matplotlib

1. Generate a multimodal expert dataset. A point starts on the left, must reach a goal on the right, with an obstacle in the middle. The scripted expert picks — per episode, at random — the upper or lower detour. That coin-flip is what injects multimodality: from the same start state, the expert action is sometimes "go up," sometimes "go down."

   import numpy as np, torch, torch.nn as nn
   
   GOAL = np.array([9., 0.]); OBST = np.array([4.5, 0.]); R = 1.4
   
   def expert_action(s, route):
       # route=+1 detour up, route=-1 detour down. Returns a velocity step.
       way = GOAL if s[0] > OBST[0] else OBST + np.array([0., route*2.2])
       d = way - s; n = np.linalg.norm(d) + 1e-8
       return (d / n) * 0.6            # fixed-speed step toward the waypoint
   
   def rollout_expert(route, steps=40):
       s = np.array([0., np.random.uniform(-.5, .5)]); S, A = [], []
       for _ in range(steps):
           a = expert_action(s, route); S.append(s.copy()); A.append(a); s = s + a
       return np.array(S), np.array(A)
   
   S, A = [], []
   for ep in range(120):                  # ~half go up, half go down
       s, a = rollout_expert(route=1 if ep % 2 else -1)
       S.append(s); A.append(a)
   Sd = torch.tensor(np.concatenate(S), dtype=torch.float32)
   Ad = torch.tensor(np.concatenate(A), dtype=torch.float32)
   print("dataset:", Sd.shape, Ad.shape)   # ~4800 (s,a) pairs, bimodal near the obstacle

2. Train an MSE behavioral-cloning MLP and roll it out. Standard regression. Then run the policy in closed loop from a fresh start and watch two things: near the obstacle it predicts the average of "up" and "down" (≈ straight = into the obstacle — mode-averaging), and any small deviation puts it in states the demos never covered (covariate shift).

   def mlp(out=2):
       return nn.Sequential(nn.Linear(2,128), nn.ReLU(),
                            nn.Linear(128,128), nn.ReLU(), nn.Linear(128,out))
   
   bc = mlp(); opt = torch.optim.Adam(bc.parameters(), 1e-3)
   for it in range(3000):                  # pure MSE = the L2 objective from Section D
       pred = bc(Sd); loss = ((pred - Ad)**2).mean()
       opt.zero_grad(); loss.backward(); opt.step()
   
   def rollout_policy(net, steps=60):
       s = np.array([0., 0.]); traj = [s.copy()]
       for _ in range(steps):
           a = net(torch.tensor(s, dtype=torch.float32)).detach().numpy()
           s = s + a; traj.append(s.copy())
       return np.array(traj)
   
   traj = rollout_policy(bc)
   hit = ((traj - OBST)[:, :]**2).sum(1).min()**.5 < R
   print("closest approach to obstacle:", round(float(((traj-OBST)**2).sum(1).min()**.5),2),
         "| COLLISION" if hit else "| clear")
   # Expected: the MSE policy drives ~straight, collides, and/or drifts off into
   # states no demo covered. That is mode-averaging + covariate shift, live.

3. One DAgger iteration. We have a scripted expert, so we can do the real thing: roll out the current policy, relabel the states it actually visited with the expert's action, aggregate, retrain. Error on the policy's own visitation drops.

   def dagger_relabel(net, n=40):
       Sn, An = [], []
       for _ in range(n):
           s = np.array([0., np.random.uniform(-.5,.5)])
           route = np.random.choice([-1, 1])     # expert still has a preferred mode
           for _ in range(60):
               a_star = expert_action(s, route)       # label = what the EXPERT would do here
               Sn.append(s.copy()); An.append(a_star)
               a = net(torch.tensor(s,dtype=torch.float32)).detach().numpy()  # move by LEARNER
               s = s + a
       return torch.tensor(np.array(Sn),dtype=torch.float32), torch.tensor(np.array(An),dtype=torch.float32)
   
   Sx, Ax = dagger_relabel(bc)
   Sd2 = torch.cat([Sd, Sx]); Ad2 = torch.cat([Ad, Ax])     # AGGREGATE
   for it in range(3000):                       # RETRAIN on the union
       loss = ((bc(Sd2)-Ad2)**2).mean()
       opt.zero_grad(); loss.backward(); opt.step()
   print("after 1 DAgger round, on-policy MSE:",
         round(float(((bc(Sx)-Ax)**2).mean()),4))
   # Expected: on-policy error drops — the model has now seen its own drift states.
   # Note it still mode-averages near the obstacle: DAgger fixes covariate shift,
   # NOT multimodality. Those are two different diseases. Hence step 4.

4. (Optional) Swap the MSE head for a 2-mode mixture head. Predict two candidate actions plus a logit choosing between them, and train with the "min-over-modes" loss so each mode specializes (this is the toy ancestor of the CVAE/diffusion idea on Page 2). At rollout, sample a mode — the policy now commits to a real detour instead of the dead-center average.

   K = 2
   head = mlp(out=K*2 + K)     # K action means (2D each) + K mixing logits
   opt2 = torch.optim.Adam(head.parameters(), 1e-3)
   for it in range(4000):
       o = head(Sd); mus = o[:, :K*2].view(-1, K, 2); logit = o[:, K*2:]
       se = ((mus - Ad[:, None, :])**2).sum(-1)       # (N,K) error to each mode
       win = se.argmin(1)                                # assign each sample to its closest mode
       reg = se[torch.arange(len(se)), win].mean()       # "min-over-modes" regression loss
       cls = nn.functional.cross_entropy(logit, win)        # learn the mixing weights
       loss = reg + cls
       opt2.zero_grad(); loss.backward(); opt2.step()
   
   def sample_action(s):
       o = head(torch.tensor(s,dtype=torch.float32))
       mus = o[:K*2].view(K,2); k = torch.distributions.Categorical(logits=o[K*2:]).sample()
       return mus[k].detach().numpy()                      # COMMIT to one mode, never the mean
   # Expected: rolling out sample_action follows a clean upper OR lower detour and
   # reaches the goal. Re-run it a few times — you'll see BOTH routes recovered.

Expected observations

Done when

• You can state in one sentence why MSE fails on multimodal data — it outputs 𝔼[a|s], the mean of the modes, which is itself not a valid action.
• You can explain why DAgger reduces compounding error — it trains on the learner's own visited states, shrinking the train/test distribution gap from O(εT2) toward O(εT).
• You observed empirically that DAgger and the 2-mode head fix different problems (covariate shift vs multimodality) — one does not substitute for the other.
• The 2-mode rollout recovers both detour routes across re-samples, not the average of them.