40 min

Fixing covariate shift

Every fix for covariate shift is the same move - get expert labels on states the policy will actually visit - and they differ only in how well they guess where those are.

Where you are. You can prove that a cloned policy fails because it is asked about states it never trained on. This lesson fixes that, and the interesting part is which fixes do not work.

Move the shove

The clone from the last lesson fails 40 out of 40 episodes. The obvious repair is to show it how to recover, so: take the fifty demonstrations and start half of them four centimetres off the groove, letting the expert drive back to centre. Now the training data contains being-off-course and the correction for it. Same expert, same 3,000 labelled frames, no other change.

Zero out of forty. The median episode survives 125 steps instead of 101.

Now change one thing. Keep the shove, keep its size, keep the number of labels exactly the same, and move when it happens from the first step of the episode to a random step in the middle.

Forty out of forty.

Two datasets of identical size, produced by the identical expert, using the identical trick. One is worthless and one is a complete fix. The difference is nothing but where on the path the corrective frames sit.

The idea in one paragraph

There is only one fix for covariate shift and everything in this lesson is a version of it: get expert labels on the states the policy will actually visit. Behaviour cloning fails because the training states come from the expert and the test states come from the policy, so the repair is to move the training states toward the policy’s. The three standard methods differ only in how they find those states. Noise injection guesses by disturbing the expert until it visits a wider region. Recovery demonstrations guess by putting the robot somewhere wrong on purpose. DAgger stops guessing: it runs the trained policy, sees where it actually went, and asks the expert what to do there. Guessing is cheap and works when your guess is right. Measuring is expensive and always right.

Noise injection disturb what actually executes, so the expert corrects all episode Recovery demos start it somewhere wrong on purpose, demonstrate the way back DAgger let the trained policy drift, then label exactly where it ended up the demonstrated path The first two guess where the policy will go wrong. The third measures it, which is why it needs a trained policy in the loop and an expert willing to answer questions about poses it never chose.
Three ways to add data away from the demonstrated path

Wider than the screen; scroll it sideways.

Noise injection

Disturb what actually executes. The expert computes its action, you add a small random vector to it before applying it, and the expert then has to correct from wherever that put the arm. It corrects continuously, all episode, so every recorded frame is a slightly-off-course state paired with the expert’s answer for it.

It costs nothing beyond a single line in the collection loop, and it produced a complete fix here at both settings tried:

noise added to the executed actionsuccess at 240 steps
none0.00 ± 0.00
0.02 rad1.00 ± 0.00
0.05 rad1.00 ± 0.00

Recovery demonstrations, and why the obvious version does nothing

Start the robot somewhere wrong and demonstrate the way back. It is the first thing everyone suggests, it is what the hook did, and on this task it does essentially nothing: 0.00 success, median survival 101 to 125.

The reason is visible the moment you look at where the frames landed. Counting rows more than one centimetre off the groove centre, and asking where along the path they sit as a fraction from A to B:

datasetrows off centrewhere they are on the path
plain demonstrations590.01 average, 0.03 at the 90th percentile
recovery at the start850.02 average, 0.04 at the 90th percentile
recovery mid-episode1010.25 average, 0.87 at the 90th percentile
noise injection1,9410.62 average, 1.00 at the 90th percentile

Starting the arm off centre buys 26 extra corrective frames, all of them in the first four percent of the path, because a good expert has fixed the problem within a handful of steps. And the clone does not fail in the first four percent of the path. It fails after step 80, parked at the far end. Moving the same shove into the middle of the episode buys a similar number of frames and spreads them over the whole path instead.

Measured directly: take the state each clone rollout is in when it jams, and find the distance to the nearest row in each training set. python fix_shift.py --where prints both of these tables.

training setdistance from the failure states to its nearest row
plain demonstrations0.0391
recovery at the start0.0391
recovery mid-episode0.0134
noise injection0.0014

The recovery dataset is, to four decimal places, exactly as far from the failure as the dataset it was supposed to improve on. It added corrective data to a place that was never in trouble. The two fixes that work are the two that closed the distance.

DAgger

Stop guessing. Train the clone, run it, write down the states it visits, ask the expert what it would have done at each of those states, add those pairs to the dataset, retrain. Repeat.

Dataset states and the actions the expert wants there fit Clone the policy so far run it Where it went its own states, including the ones it drifted into Expert "here is what I would do" aggregate, then do it all again the states belong to the policy. The actions belong to the expert. That swap is the whole algorithm. it also means the expert must answer questions about poses it would never have chosen to visit
Train, roll out, label where it went, aggregate, repeat

Wider than the screen; scroll it sideways.

The whole algorithm is one swap: the states belong to the policy, the actions belong to the expert. That is what makes the training distribution converge on the deployment distribution, and it is what buys back the linear error bound from the last lesson.

Starting from twenty demonstrations and spending the remaining budget on rounds of three 120-step rollouts:

expert labelssuccess at 240 steps
1,200 (demonstrations only)0.28
1,560 (one round)1.00
3,000 (five rounds)1.00

One round. Three hundred and sixty labelled frames, chosen by running the policy rather than by guessing, and the problem is gone.

The comparison, on one budget

Every row below gets 3,000 expert-labelled frames, three training seeds, forty evaluation episodes each.

image/svg+xml Matplotlib v3.11.1, https://matplotlib.org/ 0.0 0.2 0.4 0.6 0.8 1.0 success rate over 240 steps (3 seeds, 40 episodes each) plain BC recovery noise 0.02 noise 0.05 DAgger the scripted expert started at home started 4 cm off centre
Success over 240 steps, every method on 3,000 expert labels

methodsuccess at 240 stepsstarted 4 cm off centre
the scripted expert1.001.00
plain behaviour cloning0.00 ± 0.000.00
recovery demonstrations, at the start0.00 ± 0.000.00
recovery demonstrations, mid-episode1.00 ± 0.001.00
noise injection, 0.02 rad1.00 ± 0.001.00
noise injection, 0.05 rad1.00 ± 0.001.00
DAgger1.00 ± 0.001.00

image/svg+xml Matplotlib v3.11.1, https://matplotlib.org/ joint space plain BC recovery noise 0.02 DAgger
The same budget of labels, spent in three different places

Four of the six work perfectly and two fail completely, and the split is not between clever and simple methods. It is between the ones whose extra frames landed where the policy fails and the ones whose did not.

Is noise injection just smoothing the network?

Worth ruling out, because there is a much duller explanation available: maybe the noisy dataset simply regularises the fit, and the expert’s relabelling has nothing to do with it.

The control is to add the same noise to the network’s input during training, on the original clean demonstrations, so the network is smoothed but no new expert label is ever collected:

what was disturbedsuccess at 240 steps
the executed action, expert relabels (0.02)1.00 ± 0.00
the network input at training time (0.02)0.58 ± 0.41
the network input at training time (0.05)0.00 ± 0.00

Smoothing alone helps unreliably at one setting and actively hurts at another. The thing that works is the expert answering a question about a state it would not otherwise have visited. That is the difference between regularisation and supervision, and it is why all of this is expensive.

What this costs when the expert is a person

Everything above is easy here for one reason: the expert is a program, so it can be asked about any pose, instantly, for free.

Replace it with a human teleoperator and DAgger’s requirement becomes a strange demand. The human must look at a pose their own hands would never have chosen, one produced by a half-trained policy, and answer “what should you have done here” - not drive the task, just label the moment. It is slow, it is unnatural, and the labels get worse as the poses get stranger. This is the honest reason pure DAgger is rare in production manipulation pipelines, and it is reasoning from the cost model rather than a result anyone has published.

What the field does instead keeps DAgger’s insight and drops its ergonomics: let the policy run, and have the human take over when it is about to fail. The states are still the policy’s and the actions are still the expert’s, but the human acts rather than annotates, which is a thing humans are good at.

The other two answers

Fixing the data is one of three responses to the last lesson, and the module takes all three.

The second is to give the policy fewer chances to drift. If it commits to fifty actions from one observation instead of one, the number of places error can enter drops by a factor of fifty, and that is action chunking, the next lesson.

The third is to stop the drift starting. A policy trained with squared error on demonstrations that contain two valid ways to do something will output the average of them, which is often a third thing that works at all - and that averaged action is itself a step off the demonstrated distribution. Fixing it needs a policy that samples rather than averages, which is multimodal demonstrations and the diffusion lesson after it.

Check yourself

1. Recovery demonstrations and mid-episode shoves use the same expert, the same disturbance size and the same number of labels. Why does one score 0.00 and the other 1.00?

Because the corrective frames land in different places. A recovery demonstration starts the arm off course and a competent expert has it back on the path within a few steps, so all the corrective frames sit in the first few percent of the path. The clone does not fail there; it fails after step 80, near the end. Moving the shove into the middle of the episode puts corrective frames along the whole path, including where the failure actually happens. Measured: the recovery dataset is 0.0391 from the failure states, exactly the same distance as the plain dataset it was supposed to improve on, while the mid-episode version is 0.0134.

2. State the one sentence that covers noise injection, recovery demonstrations and DAgger at once.

Get expert labels on states the policy will actually visit. They differ only in how those states are found: noise injection disturbs the expert until it wanders into them, recovery demonstrations place the robot in guessed ones, and DAgger runs the trained policy and observes the real ones. The first two guess and are cheap; the third measures and is expensive.

3. Why does DAgger keep all previous data rather than training each round on the newest rollouts only?

Because a policy trained only on the latest round’s states forgets how to behave on the earlier ones, and next round’s rollouts start from those. You would get a policy that oscillates between competences rather than accumulating one, and the training distribution would chase the policy instead of converging with it. The aggregation is what makes the two distributions meet, which is the entire mechanism, and it is why the algorithm is named after it.

4. Noise injection and training-time input jitter both blur the function the network fits. One works and one does not. What separates them?

Whether an expert was consulted. Injecting noise into the executed action puts the robot in a genuinely new state and records what the expert does about it, which is new supervision. Jittering the network’s input at training time asks the network to give the same old answer for nearby inputs, which is regularisation and contains no information about what the expert would actually do out there. Measured: 1.00 against 0.58 at the same noise scale, and at a larger scale the jittered version collapses to 0.00 while the relabelled one stays at 1.00.

5. Your policy fails by releasing the gripper slightly too early. Which of the three fixes can find that failure, and why can the other two not?

Only DAgger, reliably. Noise injection perturbs joint commands around the expert’s trajectory, and the expert never releases early, so no amount of joint noise produces the “released early, object falling” state. Recovery demonstrations require you to already know that this is the failure and to stage it deliberately, which is possible but is you doing the diagnosis rather than the method. DAgger finds it because it runs the actual policy and the actual policy does it.

6. A human demonstrator makes DAgger much harder than the version in this lesson. Precisely what becomes hard, and what does the field do instead?

Labelling off-policy states. The human must be shown a pose their own hands would never have produced and asked what the correct action is there, without driving through it, which is slow and unnatural and gets worse as the poses get stranger. The practical descendant keeps the on-policy states and drops the annotation: let the policy run and have the operator take over when it is about to fail, so the human acts instead of annotating. The states are still the policy’s; the actions are still the expert’s.

Do this

code/fix_shift.py has three TODO(you) markers: the mid-episode shove and the two lines of DAgger.

python fix_shift.py            # all five methods, one seed, a few minutes
python fix_shift.py --trace    # DAgger round by round
python fix_shift.py --where    # where each fix put its corrective frames
python fix_shift.py --seeds 3  # the table in this lesson, slower

--seeds 3 retrains around thirty times and takes several minutes on a laptop CPU. When the TODOs are right you will reproduce this lesson’s table: plain behaviour cloning and start-of-episode recovery at 0.00, everything else at 1.00, and DAgger reaching 1.00 after a single round.

Then:

  1. Find the cheapest DAgger that works. --trace shows one round is enough here. Cut the seed demonstrations from twenty to ten, then five, and find where it stops recovering. You are measuring how much of your label budget is genuinely better spent on-policy, which is the number that decides whether this is worth doing on real hardware.

  2. Break noise injection deliberately. Add a failure mode that joint noise cannot reach - the simplest is to make the groove branch, so a correct policy must pick a side. Then check whether noise injection still fixes it and whether DAgger still does. Predict both answers first. This is the experiment behind the “do not read this table as noise injection is as good as DAgger” warning, and running it once is worth more than the warning.

  3. Do the honest version on your Module 2 recording. You have a scripted expert there too, which means you can run genuine DAgger on it: roll your clone out in the simulator, label the states it visits with your script, aggregate, retrain. Then write down what you would have to do differently if the expert were you, holding a leader arm. That difference is the whole reason the rest of this module exists.

What you can now do

You can name the single move that every fix for covariate shift makes, and place noise injection, recovery demonstrations and DAgger on the spectrum from guessing where the policy fails to measuring it. You can take a policy that fails every episode to one that fails none, without spending a single extra expert label. You can run the control that separates a real fix from mere smoothing. And you can explain why the method with the strongest guarantee is the one people avoid on real robots, and what they run instead.

What you can now do

You can turn a policy that fails every episode into one that never fails, on the same number of expert labels, and say why the obvious version of the fix does nothing.