Robot arm · imitation learning

Overview

The world, the four commands and the 60-command test stay exactly as in the Q-learning workshops. Only the learning signal changes: instead of a reward, every training pose comes with the command an expert chose there.

The mental model: the policy copies the expert where it has seen the expert, and guesses everywhere else.

Demonstrations

Our expert is the full-model reference from the Q-table workshop. It computed the shortest route from every pose using the whole world, so it always knows the way.

On the real SO-101 the expert is a person moving the leader arm. Here the computed reference stands in for that person, so these demonstrations are synthetic, not human recordings.

A demonstration is one episode: the expert drives the arm from a start to the cube, and every pose is saved with the command it chose there.

How to read the map:

• Each square is one pose, as in the Q-table workshop: shoulder across, elbow up.
• Dots mark every pose in the 16 training episodes.
• The orange line is the selected episode’s route.

def expert(state):
    """The expert's command at a pose. Ties go to the first command in the list, so labels never flip."""
    return int(np.argmax(EXPERT_Q[state]))


def demonstrate(start):
    """One demonstration episode: the expert's (pose, command) pairs from `start` to the cube."""
    state, pairs = tuple(start), []
    while not GOAL[state]:
        action = expert(state)
        pairs.append((state, action))
        state = step(state, action)[0]
    return pairs

"""Imitation learning on the same arm: copy an expert's commands with ordinary PyTorch.

Run: python3 classes/ml/labs/robot-arm/imitation.py

The expert is the full-model reference from world.py. It stands in for a person who knows the way, so its
demonstrations are synthetic. The policy network never sees a reward: it learns which command the expert
chose at each demonstrated pose (behavioral cloning). Evaluation is the same frozen 60-command test as the
Q-table and Q-network workshops.
"""
import numpy as np
import torch
from torch import nn
from world import STARTS, GOAL, HOME, step, encode, evaluate, reference

EXPERT_Q = reference()


def expert(state):
    """The expert's command at a pose. Ties go to the first command in the list, so labels never flip."""
    return int(np.argmax(EXPERT_Q[state]))


def demonstrate(start):
    """One demonstration episode: the expert's (pose, command) pairs from `start` to the cube."""
    state, pairs = tuple(start), []
    while not GOAL[state]:
        action = expert(state)
        pairs.append((state, action))
        state = step(state, action)[0]
    return pairs


def demonstrations(count, seed, near_home=False):
    """`count` episodes from random starts: anywhere (broad) or among the 40 starts nearest home (narrow)."""
    rng = np.random.default_rng(seed)
    pool = sorted(STARTS, key=lambda s: abs(s[0] - HOME[0]) + abs(s[1] - HOME[1]))[:40] if near_home else STARTS
    starts = [pool[i] for i in rng.choice(len(pool), count, replace=False)]
    return [demonstrate(s) for s in starts]


def make_policy(hidden=16):
    return nn.Sequential(nn.Linear(2, hidden), nn.ReLU(), nn.Linear(hidden, 4))   # 2 angles in, 4 logits out


def train(pairs, hidden=16, lr=1.0, epochs=300, seed=0, validation=(), observe=None):
    torch.set_num_threads(1)
    torch.manual_seed(seed)
    policy = make_policy(hidden)
    optimizer = torch.optim.SGD(policy.parameters(), lr=lr)
    loss_fn = nn.CrossEntropyLoss()
    poses = torch.tensor(encode([p for p, a in pairs]))
    labels = torch.tensor([a for p, a in pairs])
    order = torch.Generator().manual_seed(seed)
    for epoch in range(epochs):
        for batch in torch.randperm(len(labels), generator=order).split(32):
            logits = policy(poses[batch])
            loss = loss_fn(logits, labels[batch])     # cross-entropy against the expert's command
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        if observe:
            observe(epoch + 1, policy, accuracy(policy, pairs), accuracy(policy, validation))
    return policy


def logits_table(policy):
    """The policy's four logits at every one of the 19 × 19 poses."""
    with torch.no_grad():
        grid = torch.tensor(encode([[i, j] for i in range(19) for j in range(19)]))
        return policy(grid).numpy().reshape(19, 19, 4)


def accuracy(policy, pairs):
    if not pairs:
        return None
    table = logits_table(policy)
    return float(np.mean([np.argmax(table[p]) == a for p, a in pairs]))


def dagger(pairs, rounds=4, starts_per_round=10, seed=0, **train_args):
    """DAgger: run the policy, ask the expert for the command at every new pose it visited, retrain."""
    rng = np.random.default_rng(1000 + seed)
    pairs = list(pairs)
    labelled = {p for p, a in pairs}
    history = []
    for r in range(rounds + 1):
        policy = train(pairs, seed=seed, **train_args)
        table = logits_table(policy)
        history.append({"round": r, "labels": len(pairs), "passed": evaluate(table)["successes"], "asked": []})
        if r == rounds:
            return policy, history
        visited = set()
        for i in rng.choice(len(STARTS), starts_per_round, replace=False):
            state = STARTS[i]
            for _ in range(60):                       # the policy drives; the expert only labels
                visited.add(state)
                state = step(state, int(np.argmax(table[state])))[0]
                if GOAL[state]:
                    break
        new = sorted(s for s in visited if s not in labelled)
        history[-1]["asked"] = new                    # the poses the expert was asked about this round
        pairs += [(s, expert(s)) for s in new]
        labelled.update(new)


if __name__ == "__main__":
    DRAW = 3        # the demonstration draw the page shows: the lower median of draws 0–9 (tools/run_imitation.py)
    episodes = demonstrations(20, seed=DRAW)
    train_pairs = [pair for e in episodes[:16] for pair in e]
    val_pairs = [pair for e in episodes[16:] for pair in e]
    print(f"Expert from home: {len(demonstrate(HOME))} commands. Demonstration draw {DRAW}: 20 episodes, "
          f"{len(train_pairs)} training pairs (16 episodes), {len(val_pairs)} validation pairs (4 episodes)")
    policy = train(train_pairs, seed=DRAW, validation=val_pairs)
    table = logits_table(policy)
    result = evaluate(table)
    print(f"Training accuracy {accuracy(policy, train_pairs):.3f}, validation accuracy {accuracy(policy, val_pairs):.3f}")
    print(f"Frozen test: home {'reached in ' + str(result['home']['moves']) + ' commands' if result['home']['success'] else 'failed'}; "
          f"starts passed {result['successes']} / {result['trials']}")

• expert(state) picks the command with the shortest remaining route. When two commands are equally short, it always takes the first in the list, so a pose never gets two different labels.
• demonstrate(start) follows the expert to the cube and returns one (pose, command) pair per step.

We record episodes from random starts. Whole episodes are split: for training ( pairs) and for validation ( pairs).

Why split whole episodes, not single pairs? Neighbouring poses on one route get similar labels. Splitting pairs would put near-copies on both sides, and validation would look better than it is.

Demonstration

one episode of an expert doing the task, saved as (observation, action) pairs.

Action label

the expert’s action at a saved observation: the answer the policy is trained to give.

Policy

The policy is a network that reads the two angles and gives one score per command. Softmax turns the four scores into probabilities, exactly as for the T-shirt sizes in Neural network: several inputs.

def make_policy(hidden=16):
    return nn.Sequential(nn.Linear(2, hidden), nn.ReLU(), nn.Linear(hidden, 4))   # 2 angles in, 4 logits out

Where each size comes from:

• 2 inputs: the shoulder and elbow angles, scaled to −1…1 as in the Q-network workshop.
• 16 hidden neurons: a choice, measured in Training below.
• 4 outputs: one per command.

That is \(2 \times 16 + 16 + 16 \times 4 + 4 = 116\) weights and biases, the same shape as the Q-network.

The outputs mean something different, though. A Q-network output estimates the reward still to come. These outputs only rank the commands: which one would the expert pick?

Logits

a classifier’s raw scores, one per class, before softmax turns them into probabilities.

The loss is cross-entropy, as for the T-shirts: minus the log of the probability given to the expert’s command. Here is one real label, worked by hand with the untrained network:

Behavioral cloning

training a policy to copy an expert’s actions with ordinary supervised learning: observations in, the expert’s actions as labels.

Training

Training is the PyTorch loop you know, with cross-entropy and plain SGD:

for epoch in range(epochs):
    for batch in torch.randperm(len(labels), generator=order).split(32):
        logits = policy(poses[batch])
        loss = loss_fn(logits, labels[batch])     # cross-entropy against the expert's command
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

• torch.randperm(...).split(32) shuffles the training pairs and cuts them into batches of 32.
• One epoch is one pass over all pairs; we run 300.
• The optimizer is torch.optim.SGD with learning rate 1.0.

• The loss is the mean cross-entropy over all training pairs after each epoch.
• Accuracy is the share of labels where the most likely command is the expert’s: blue dashed on training pairs, orange on validation pairs.
• The spikes come with rate 1.0: epochs jump to over 1.5 times the loss before. At rate 0.3 the same run has , but ends at loss instead of .

After 300 epochs the policy copies every training label, and of the validation labels.

Network size and learning rate were chosen on validation accuracy only, over demonstration draws 0–9 (each draw is a different set of 20 random starts, and also seeds the network):

The rule: the highest median validation accuracy, then the highest worst draw, then the smaller network. Rate 0.1 scores lowest at every size. 16 neurons at rate 1.0 wins; 8 neurons came close, with a worse worst draw.

Expert and policy

Training only scored the policy on demonstrated poses. What does it do everywhere else? Click any pose on the map to see the expert’s move and the policy’s move side by side.

Test

Agreeing pose by pose is not the task. The task is reaching the cube, so we freeze the policy and run the same test as the Q-learning workshops: always the most likely command, from every one of the 251 starts, for at most 60 commands.

Distribution shift

Where are the failing starts? Select one of them:

All sit along the edge of the forbidden region, where the arm is squeezed between the barrier and the table. Some are right beside the cube.

None of their routes ever touches a demonstrated pose. Each guessed command leads to another pose nobody demonstrated, until the policy picks a command the world rejects, again and again.

Distribution shift

the situations met at test time differ from the ones in the training data.

Compounding error

each action chooses the next observation, so one wrong guess leads to situations even less like the training data, and the errors pile up.

So: a cloned policy is only as good as the places its demonstrations covered.

Coverage

If coverage is what matters, where the demonstrations start should matter more than how many there are. We test that with two kinds of demonstration sets:

• Broad: episodes from random starts anywhere on the map.
• Narrow: episodes from random starts among the 40 starts nearest home, like a person who always begins in the same place.

Each kind is trained with 4, 8, 16 and 32 episodes, all on demonstration draws 0–9, with the network and training above. Predict: will 32 narrow episodes beat 4 broad ones?

Narrow demonstrations stop at starts, however many episodes there are: 32 narrow episodes, about labels, do worse than 4 broad ones. All the extra labels repeat the same patch near home.

So: to clone a skill, demonstrate it from the situations the robot will actually meet. For the SO-101, that means varying the object’s position across the whole workspace while recording.

DAgger

But which situations will the robot meet? The policy itself decides that, through its own mistakes. DAgger asks the expert exactly there.

Each round:

1. Drive: the current policy runs from 10 random starts, and every pose it visits is recorded.
2. Ask: the expert labels each visited pose that has no label yet. It only answers; it never drives.
3. Retrain: the policy is trained again on all labels so far.

visited = set()
for i in rng.choice(len(STARTS), starts_per_round, replace=False):
    state = STARTS[i]
    for _ in range(60):                       # the policy drives; the expert only labels
        visited.add(state)
        state = step(state, int(np.argmax(table[state])))[0]
        if GOAL[state]:
            break
new = sorted(s for s in visited if s not in labelled)
history[-1]["asked"] = new                    # the poses the expert was asked about this round
pairs += [(s, expert(s)) for s in new]
labelled.update(new)

We start from the narrow set, 16 episodes near home, which passes 116 starts. Step through the rounds:

How to read the lab:

• On the map, dots are labelled poses and orange rings are the poses the expert is asked about next.
• In the chart, the blue line is DAgger and the orange line spends the same number of labels on more expert episodes from random starts, both for this draw.

Over all ten draws, after each round:

After four rounds DAgger passes all 251 starts in of 10 draws. The same number of labels spent on more expert episodes never does: median .

The difference is where the labels go. Expert episodes repeat the expert’s own routes; DAgger’s labels land exactly where the policy goes wrong.

The path is not smooth: in some draws one round passes fewer starts than the round before, because each retrained network fills in unlabelled poses differently. Every draw is in the table below.

DAgger

dataset aggregation: let the policy act, have the expert label the situations it reached, add them to the data and retrain.

DAgger needs an expert who can answer at any moment. With a person and a real arm, that means corrections during practice, which is what methods such as HIL-SERL build on.

Recap

• expert(state) is the labeller: the full-model reference, standing in for a person. The demonstrations are synthetic.
• demonstrate(start) records one episode of (pose, command) pairs; whole episodes are split into training and validation.
• make_policy() builds the 2 → 16 → 4 network: two angles in, four logits out.
• train(pairs) is behavioral cloning: cross-entropy against the expert’s command, plain SGD, 300 epochs.
• evaluate(logits_table(policy)) freezes the policy and tests reaching from all 251 starts.
• dagger(pairs) lets the policy drive, asks the expert at every new pose it visits, and retrains.

An expert’s labels replace the reward, so the arm learns fast and from far less data. But it only knows the places the demonstrations covered. Broad demonstrations, or DAgger’s targeted questions, carry it everywhere.