7.5 — Training a Small Policy from Scratch

Learning objectives

After this lesson, you should be able to:

• Record demonstrations and split them by episode.
• Train a small policy by behavior cloning in PyTorch.
• Evaluate it in closed loop rather than by loss alone.
• Use ablations to check that the policy uses both its inputs.

The one-sentence idea

Build the pipeline once, small, and every large system becomes recognisable.

Loss is for training; closed-loop success is for judging.

The smallest complete project

Everything in the course so far appears here in miniature: demonstrations (lesson 3.2), behavior cloning (3.3), a loss and a validation split (3.4), and closed-loop evaluation (1.5, 3.4). The task: a hand on a table, a red block and a blue block, and an instruction “red” or “blue”. The hand must reach the instructed block. The policy sees the hand position, both block positions and the instruction, and outputs a small movement. It is a toy, but you will build, train and test every part yourself in PyTorch (these cells run on the course’s run server).

Because the cells below build on each other like a notebook, run them in order.

Step 1: record demonstrations from an expert

An expert rule plays the part of the human demonstrator: move 30% of the way to the target block, with a little noise. We record episodes, and keep whole episodes together so we can split by episode:

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

def expert_action(hand, red, blue, instruction):                 # instruction: 0 = red, 1 = blue
    target = red if instruction == 0 else blue
    return 0.3 * (target - hand)

def record_episode(steps=12):
    red, blue = rng.uniform(-1, 1, 2), rng.uniform(-1, 1, 2)
    hand, instruction = rng.uniform(-1, 1, 2), int(rng.integers(2))
    rows = []
    for _ in range(steps):
        a = expert_action(hand, red, blue, instruction) + rng.normal(0, 0.01, 2)
        rows.append((np.concatenate([hand, red, blue, np.eye(2)[instruction]]), a))
        hand = hand + a
    return rows

episodes = [record_episode() for _ in range(300)]
train_eps, val_eps = episodes[:240], episodes[240:]               # split by episode
def stack(eps):
    X = np.array([o for ep in eps for o, a in ep], dtype=np.float32)
    Y = np.array([a for ep in eps for o, a in ep], dtype=np.float32)
    return torch.tensor(X), torch.tensor(Y)
Xtr, Ytr = stack(train_eps); Xva, Yva = stack(val_eps)
print("training steps:", len(Xtr), "| validation steps:", len(Xva), "| observation size:", Xtr.shape[1], "| action size:", Ytr.shape[1])

training steps: 2880 | validation steps: 720 | observation size: 8 | action size: 2

Toy model. A synthetic expert and a table with two blocks. The observation here is a list of positions, not camera images.

Step 2: define and train the policy

The policy is a small multilayer network. The loss is mean squared error between the predicted and the demonstrated action. We print training and validation loss, as in lesson 3.4:

import torch.nn as nn

policy = nn.Sequential(nn.Linear(8, 64), nn.ReLU(), nn.Linear(64, 64), nn.ReLU(), nn.Linear(64, 2))
opt = torch.optim.Adam(policy.parameters(), lr=3e-3)

for epoch in range(1, 41):
    perm = torch.randperm(len(Xtr))
    for i in range(0, len(perm), 128):
        idx = perm[i:i + 128]
        opt.zero_grad()
        loss = ((policy(Xtr[idx]) - Ytr[idx]) ** 2).mean()
        loss.backward(); opt.step()
    if epoch in (1, 10, 40):
        with torch.no_grad():
            tr = float(((policy(Xtr) - Ytr) ** 2).mean()); va = float(((policy(Xva) - Yva) ** 2).mean())
        print(f"epoch {epoch:>2}: training loss {tr:.1e} | validation loss {va:.1e}")

epoch  1: training loss 5.2e-03 | validation loss 6.5e-03
epoch 10: training loss 2.0e-04 | validation loss 3.4e-04
epoch 40: training loss 1.4e-04 | validation loss 2.6e-04

The validation loss tracks the training loss, so the model is not just memorising the 240 training episodes. But remember lesson 3.4: a low loss on demonstrated states is not yet success.

Step 3: run it on its own, in closed loop

Now the real test. For new random scenes, let the policy drive the hand for 20 steps and count how often it ends within 0.1 of the instructed block. We also test two ablations from lesson 5.5: tell it the wrong instruction, and hide the block positions:

def evaluate(trials=300, wrong_instruction=False, hide_blocks=False):
    wins = 0
    for _ in range(trials):
        red, blue = rng.uniform(-1, 1, 2), rng.uniform(-1, 1, 2)
        hand, instruction = rng.uniform(-1, 1, 2), int(rng.integers(2))
        told = 1 - instruction if wrong_instruction else instruction
        for _ in range(20):
            r, b = (np.zeros(2), np.zeros(2)) if hide_blocks else (red, blue)
            obs = torch.tensor(np.concatenate([hand, r, b, np.eye(2)[told]]), dtype=torch.float32)
            with torch.no_grad(): a = policy(obs).numpy()
            hand = hand + a
        target = red if instruction == 0 else blue
        wins += np.linalg.norm(hand - target) < 0.1
    return round(wins / trials, 2)

print("normal:                      ", evaluate())
print("given the wrong instruction: ", evaluate(wrong_instruction=True))
print("block positions hidden:      ", evaluate(hide_blocks=True))

normal:                       0.96
given the wrong instruction:  0.0
block positions hidden:       0.0

The policy reaches the instructed block in about 96% of episodes, and the two ablations show it genuinely uses both inputs: told the wrong colour, it goes to the other block (so it fails the original task); with positions hidden, it cannot find anything. These are the same checks as the ablation in lesson 2.5, now on a policy you trained.

Where to go from here

You now have the full loop: record, train, evaluate. Real VLAs replace each piece with something heavier: a vision encoder instead of position numbers, a pretrained language model instead of a one-hot instruction, a chunked or generative action output instead of a single step, and robot data instead of a script. The structure of the work, and the discipline (split by episode, closed-loop evaluation, ablations), stay the same.

Try changing things and re-running the cells. Make the demonstrations noisier (raise the 0.01) and watch the closed-loop success; reduce the number of training episodes from 240 to 20 and watch the gap between training and validation loss; make the expert move 60% of the way instead of 30% and see whether the learned policy copies it. Each experiment teaches one of the lessons of this roadmap in a minute, and a result that surprises you is worth following up, because it usually points at an assumption you had not noticed.

Common misconceptions

Engineering takeaways

1. Record, split, train, evaluate.
2. Split by episode.
3. Evaluate in closed loop.
4. Ablate.

Build it small enough to see every part.

Knowledge check

Try each question first, then reveal the explanation.

What’s next

Next: measuring results and sorting failures into causes.

Back to the VLA roadmap