3.5 — Distribution Shift and Generalization

Learning objectives

After this lesson, you should be able to:

• Define distribution shift and covariate shift.
• Explain why models fail outside the range of their training data.
• Explain how varied training data helps, and its limit.
• Flag unfamiliar inputs with a simple statistic.

The one-sentence idea

Models are reliable where they had data, and unreliable elsewhere.

Plan for the situations the training data did not contain.

The test is not like the training

Machine learning quietly assumes that the data you meet later looks like the data you trained on. On a robot that assumption is broken all the time: different lighting, a new table, another object, a different robot. This is called distribution shift. The standard reference on the topic describes dataset shift as the situation where the joint distribution of inputs and outputs differs between training and test, and covariate shift as the special case where only the input distribution changes.^[1]

Robots add a second source that we met in lesson 3.3: the policy’s own mistakes change the states it visits, so the data it sees at run time differs from the demonstrations even in an unchanged room.^[2] And the simulator-to-real-world version is the reality gap: policies trained in simulation can perform worse on real hardware because simulated and real environments differ.^[3]

Outside the range it was trained on

The simplest picture: a model fitted on inputs between 0 and 1 and asked about inputs it never saw. The underlying rule here is a gentle curve. A flexible model fits it nicely where it has data, and goes wrong the moment it leaves:

import numpy as np
rng = np.random.default_rng(0)

x_train = rng.uniform(0.0, 1.0, 200)
y_train = np.sin(2 * x_train) + rng.normal(0, 0.02, 200)
coef = np.polyfit(x_train, y_train, 5)                 # a flexible model

for lo, hi in ((0.0, 1.0), (1.0, 1.5), (1.5, 2.5)):
    x = np.linspace(lo, hi, 100)
    err = np.abs(np.polyval(coef, x) - np.sin(2 * x)).mean()
    print(f"inputs in [{lo}, {hi}] -> average error {err:.3f}")

inputs in [0.0, 1.0] -> average error 0.003
inputs in [1.0, 1.5] -> average error 0.117
inputs in [1.5, 2.5] -> average error 9.691

Toy model. A curve fitted with a polynomial. Neural networks behave differently in detail, but the lesson holds: nothing in the training loss controls the model where there was no data.

The error is tiny inside the training range and grows quickly outside it. A robot gets there easily: a slightly different start position, an object in a new place.

A look-alike that breaks: lighting

Here is a visual version. A tiny classifier decides whether a patch of pixels is “object” or “background” using brightness. In training the room was always dimly lit. Then the lights come on:

import numpy as np
rng = np.random.default_rng(0)

def make_patches(n, light):
    label = rng.integers(0, 2, n)                              # 1 = object (darker than the background), 0 = background
    brightness = light * (0.6 + 0.3 * (1 - label)) + rng.normal(0, 0.03, n)
    return brightness, label

def accuracy(threshold, light):
    b, y = make_patches(2000, light)
    return float(np.mean((b < threshold) == (y == 1)))

b, y = make_patches(2000, light=0.5)                           # training: a dim room
threshold = (b[y == 1].mean() + b[y == 0].mean()) / 2          # the learned rule: darker than this means object
print("learned threshold:", round(float(threshold), 3))
print("accuracy in the dim room:    ", accuracy(threshold, 0.5))
print("accuracy in a bright room:   ", accuracy(threshold, 1.0))

learned threshold: 0.374
accuracy in the dim room:     0.9945
accuracy in a bright room:    0.508

Toy model. A one-feature classifier with made-up pixel values, chosen so that lighting scales every brightness.

In the bright room the whole picture is brighter, so nearly every patch is brighter than the learned threshold and gets called “background”; accuracy drops to about 0.5, no better than a coin. Nothing about the object changed. Only the conditions did.

Cheap defences: variety and a smoke alarm

The most common defence is to make training data more varied. Domain randomization, introduced in the context of sim-to-real, trains with randomised non-realistic textures in simulation; a detector trained only on such images localised real objects to about 1.5 cm, and the argument is that with enough variability in the simulator the real world looks like just another variation.^[4] The same idea works for real data: record in many rooms, many lights, many objects. Pooling data across robots is a version of it, and the Open X-Embodiment collaboration reported positive transfer between platforms.^[5] Even so, the RT-2 paper notes that physical skills remain limited to the distribution of skills seen in the robot data.^[6]

Applied to our brightness rule, variety means training under many lights. The rule learns its threshold as a fraction of the scene brightness (which, in this toy, it is told):

import numpy as np
rng = np.random.default_rng(0)

def make_patches(n, light):
    label = rng.integers(0, 2, n)
    brightness = light * (0.6 + 0.3 * (1 - label)) + rng.normal(0, 0.03, n)
    return brightness, label

def accuracy(rule, light):
    b, y = make_patches(2000, light)
    return float(np.mean(rule(b, light) == (y == 1)))

# Randomised training: the threshold is learned as a fraction of the scene brightness
ratios = []
for _ in range(50):
    light = rng.uniform(0.3, 1.2)                      # a different room each time
    b, y = make_patches(200, light)
    ratios.append(((b[y == 1].mean() + b[y == 0].mean()) / 2) / light)
ratio = float(np.mean(ratios))
rule = lambda b, light: b < ratio * light             # needs the room brightness as input (e.g. measured from the whole image)

print("accuracy in a dim room:   ", accuracy(rule, 0.5))
print("accuracy in a bright room:", accuracy(rule, 1.0))

accuracy in a dim room:    0.9925
accuracy in a bright room: 1.0

Variety only helps for kinds of variation that were present. A second defence is a smoke alarm: check whether new inputs look like the training inputs, and be careful when they do not. A simple version compares each new input’s feature with the training mean in units of standard deviation:

import numpy as np
rng = np.random.default_rng(0)

train_brightness = rng.normal(0.5, 0.1, 1000)
mean, std = train_brightness.mean(), train_brightness.std()

for name, value in (("dim room (as trained)", 0.52), ("slightly brighter", 0.7), ("lights fully on", 1.0)):
    z = (value - mean) / std
    print(f"{name:<22} z = {z:5.1f}  ->", "UNFAMILIAR, slow down or ask for help" if abs(z) > 3 else "looks familiar")

dim room (as trained)  z =   0.3  -> looks familiar
slightly brighter      z =   2.1  -> looks familiar
lights fully on        z =   5.2  -> UNFAMILIAR, slow down or ask for help

Real systems use richer checks, but the principle is the same: an alarm for “I have not seen this before” is cheap, and a policy that knows when it is out of its depth is safer (lesson 1.1).

Common misconceptions

Engineering takeaways

1. Shift = test differs from training.
2. Errors grow outside the training range.
3. Variety in training helps for the variation included.
4. An unfamiliarity alarm is cheap and useful.

A model knows only what its data showed it.

Knowledge check

Try each question first, then reveal the explanation.

What’s next

That completes the robot-learning foundations. Next the roadmap moves on to the VLA models themselves.

Back to the VLA roadmap