Reinforcement learning basics

Overview

A ball falls down a small screen. A paddle at the bottom can move left, stay or move right.

Catch the ball and the game says +1; miss it and it says −1. That is all the feedback there is.

Here is a network that learned to play from that alone:

Reinforcement learning

learning to act from rewards that come after acting, instead of from given answers.

We will play the game, look at the network that plays it, see how a score turns into labels, then watch 1,000 rounds of training.

Game

Play one drop yourself. Each button press is one move, and the ball falls one row per move.

Predict first: can you always reach the ball in time?

The game you just played has four parts, and every reinforcement learning problem has the same four:

State

everything the player can see right now; here, the screen.

Action

a choice the player can make; here, left, stay or right.

Reward

the number the game gives back; here, +1 or −1, and only at the end.

Episode

one attempt from start to finish; here, one drop from the top to the bottom.

Seven moves are always enough: the paddle can cross the whole screen in six.

paddle = min(COLS - 1, max(0, paddle + MOVES[move]))
reward = 1.0 if paddle == ball_col else -1.0  # the only feedback

"""Catch: a network learns to catch a falling ball from +1 / -1 alone.

Run: python3 site/ml/foundations/examples/reinforcement_catch.py
Nobody tells the network which move is right. After each drop, the moves it
made become its training labels: pushed up after a catch, down after a miss.
"""

import torch
from torch import nn

ROWS, COLS = 8, 7          # the screen: 8 rows of 7 pixels
MOVES = [-1, 0, 1]         # left, stay, right


def screen(ball_row, ball_col, paddle):
    """What the network sees: one bright pixel for the ball, one for the paddle."""
    pixels = torch.zeros(ROWS, COLS)
    pixels[ball_row, ball_col] = 1.0
    pixels[ROWS - 1, paddle] = 1.0
    return pixels.flatten()


def play(model, generator, greedy=False):
    """Drop one ball. Returns the screens seen, the moves made and the reward."""
    ball_col = int(torch.randint(COLS, (1,), generator=generator))
    paddle = int(torch.randint(COLS, (1,), generator=generator))
    return play_from(model, ball_col, paddle, generator, greedy)


def play_from(model, ball_col, paddle, generator=None, greedy=False):
    screens, moves = [], []
    for ball_row in range(ROWS - 1):              # 7 moves while the ball falls
        pixels = screen(ball_row, ball_col, paddle)
        with torch.no_grad():
            probabilities = torch.softmax(model(pixels), dim=0)
        if greedy:
            move = int(probabilities.argmax())
        else:
            move = int(torch.multinomial(probabilities, 1, generator=generator))
        screens.append(pixels)
        moves.append(move)
        paddle = min(COLS - 1, max(0, paddle + MOVES[move]))
    reward = 1.0 if paddle == ball_col else -1.0  # the only feedback
    return screens, moves, reward


def make_model():
    return nn.Sequential(nn.Linear(ROWS * COLS, 8), nn.ReLU(), nn.Linear(8, 3))


def catch_rate(model):
    """Frozen test: every ball column × every paddle start, best move each time."""
    caught = sum(play_from(model, b, p, greedy=True)[2] > 0
                 for b in range(COLS) for p in range(COLS))
    return caught / (COLS * COLS)


def train(seed=0, updates=1000, drops=128, lr=2.0, observe=None):
    torch.set_num_threads(1)                      # same numbers on every machine
    torch.manual_seed(seed)
    generator = torch.Generator().manual_seed(seed)
    model = make_model()
    optimizer = torch.optim.SGD(model.parameters(), lr=lr)
    for update in range(updates):
        # 1. Play a batch of drops with the current network.
        screens, moves, rewards = [], [], []
        for _ in range(drops):
            s, m, r = play(model, generator)
            screens += s
            moves += m
            rewards += [r] * len(m)               # every move in a drop shares its outcome
        screens = torch.stack(screens)
        moves = torch.tensor(moves).unsqueeze(1)
        rewards = torch.tensor(rewards)

        # 2. The moves it made are the labels; the reward says which way to push.
        log_probabilities = torch.log_softmax(model(screens), dim=1)
        taken = log_probabilities.gather(1, moves).squeeze(1)
        loss = -(rewards * taken).mean()

        # 3. The familiar PyTorch step.
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if observe:
            observe(update + 1, model, rewards)
    return model


if __name__ == "__main__":
    torch.manual_seed(0)
    print(f"Untrained network catches {catch_rate(make_model()):.0%} of the 49 test drops")
    def report(update, model, rewards):
        if update in (1, 100, 250, 500, 1000):
            practice = (rewards > 0).float().mean()   # share of this batch's drops caught
            print(f"update {update:4d}: practice catches {practice:.0%}, test {catch_rate(model):.0%}")
    train(observe=report)

• The first line moves the paddle and keeps it on the screen.
• The second runs once, after the seventh move.

Nothing in between says whether a move was good.

Policy network

The network sees what you saw: the screen, as 56 pixels (8 rows of 7). Every pixel is 0 except two: the ball and the paddle are 1.

It answers with three numbers, one per move. Click the screen to put the ball or the paddle somewhere else.

This is the same kind of network as in Neural network: several inputs, with more inputs: 56 pixels instead of height and weight. Its three scores, one per move, become probabilities through softmax; the table under the figure shows the steps for this screen.

The diagram reads as in the several-inputs lesson:

• blue lines are positive weights, magenta negative, and thicker is bigger;
• a darker hidden circle means a bigger output;
• faded lines leave a neuron that is switched off for this screen;
• only the lines from the two lit pixels are drawn, because the other 54 pixels are 0.

Switch between before and after training: everything the network learned is in those weights.

One thing is new. The network does not simply take its most likely move: it picks a move at random using the probabilities, like rolling a weighted die.

Before training all three are near ⅓, so the paddle wanders. That randomness is how it tries things.

Exploration

trying moves that are not the current favourite; here, built in by rolling the weighted die.

return nn.Sequential(nn.Linear(ROWS * COLS, 16), nn.ReLU(), nn.Linear(16, 3))

probabilities = torch.softmax(model(pixels), dim=0)
move = int(torch.multinomial(probabilities, 1, generator=generator))

The sizes:

• 56 inputs, one per pixel (\(8\times7\)).
• Each hidden neuron holds 56 weights, one per pixel, plus a bias. On any screen, only the weights of the ball's and the paddle's pixels count.
• 8 hidden neurons is room, chosen as in the one-input lesson. We tried 2, 4, 8, 16 and 32 (six seeds each): every size reached a median test catch rate between 94% and 100%, and 8 is the smallest whose worst seed (94%) matched the bigger sizes.
• 483 numbers to learn in all: \(56\times8+8+8\times3+3\).

torch.multinomial rolls the weighted die: it draws one move using the probabilities.

Policy

a rule that picks an action for each state.

Policy network

a neural network that outputs the probability of each action for a state; this one.

Labels

Here is the whole trick. After a drop, look at the seven moves the network made:

• If it caught the ball, treat each of those moves as a correct answer: make it more likely on that screen.
• If it missed, treat each as a wrong answer: make it less likely.

The network's own moves become the labels. The reward only decides the direction.

Both drops come from the same untrained network. In the caught drop, one training step raised all seven moves it made; in the missed drop, it lowered all seven.

Nobody said which move was the mistake. Every move of a missed drop gets a little of the blame, and over thousands of drops the moves that keep showing up in catches win.

Two things in the picture are worth a second look:

• Why the moves look random. The untrained network gives each move about a one-in-three chance, and the moves were drawn with the weighted die. Which moves appear in a drop is luck.
• Why the same move gets similar numbers on different screens. Each screen differs from the one before only by the ball being one row lower. To an untrained network these screens look almost alike, so its probabilities are almost the same, and one step moves them almost the same way.

The step even changes screens that were not in the drop at all. After the missed drop, which chose times, the average chance of “” over 49 other screens fell from to .

Every screen is computed with the same 483 numbers, so changing them for one screen changes them for all.

log_probabilities = torch.log_softmax(model(screens), dim=1)
taken = log_probabilities.gather(1, moves).squeeze(1)
loss = -(rewards * taken).mean()

In the several-inputs lesson the loss for a choice was \(-\log p\) of the right answer. Here nobody knows the right answer, so the move the network made stands in for it, and the reward sets the sign. For one move with probability \(p\) and reward \(r\):

• With \(r=+1\) this is exactly the several-inputs lesson's loss, so the move becomes more likely.
• With \(r=-1\) the sign flips, and the move becomes less likely.

gather picks, on each screen, the log-probability of the move that was made, just as it picked the right size in the several-inputs lesson.

Put the two training loops side by side:

log_probabilities = torch.log_softmax(scores, dim=1)
right = log_probabilities.gather(1, answers.unsqueeze(1))
loss = -right.mean()
optimizer.zero_grad()
loss.backward()
optimizer.step()

"""A network that chooses: T-shirt size S, M or L from height and weight.

Run: python3 site/ml/foundations/examples/pytorch_choices.py
Every person comes with the right size, so this is the PyTorch basics loop again,
with two inputs instead of one and three choices instead of one number.
"""

import torch
from torch import nn

SIZES = ["S", "M", "L"]


def size_chart(height, weight):
    """The shop's chart: the right answers we train on (made up for this lesson)."""
    if height > 180 or weight > 85:
        return 2                                   # L
    if height < 165 and weight < 60:
        return 0                                   # S
    return 1                                       # M


def make_people(count, generator):
    heights = 150 + 45 * torch.rand(count, generator=generator)    # 150-195 cm
    weights = 45 + 60 * torch.rand(count, generator=generator)     # 45-105 kg
    sizes = torch.tensor([size_chart(h, w) for h, w in zip(heights.tolist(), weights.tolist())])
    return torch.stack([heights, weights], dim=1), sizes


def scale(people):
    """Centre both inputs near 0 and give them a similar spread, so one learning rate suits both."""
    return (people - torch.tensor([170.0, 75.0])) / torch.tensor([10.0, 15.0])


def make_network():
    return nn.Sequential(nn.Linear(2, 4), nn.ReLU(), nn.Linear(4, 3))


def train(seed=0, steps=3000, lr=0.5, observe=None):
    torch.set_num_threads(1)                                       # same numbers on every machine
    torch.manual_seed(seed)
    generator = torch.Generator().manual_seed(seed)
    people, answers = make_people(80, generator)
    inputs = scale(people)
    network = make_network()
    optimizer = torch.optim.SGD(network.parameters(), lr=lr)
    for step in range(steps):
        scores = network(inputs)                                   # 80 × 3: one score per size
        log_probabilities = torch.log_softmax(scores, dim=1)       # softmax, then log
        right = log_probabilities.gather(1, answers.unsqueeze(1))  # log p of the right size
        loss = -right.mean()
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if observe:
            observe(step + 1, network, loss.item())
    return network, people, answers


if __name__ == "__main__":
    def report(step, network, loss):
        if step in (1, 100, 500, 3000):
            print(f"step {step:4d}: loss {loss:.3f}")

    network, people, answers = train(observe=report)
    with torch.no_grad():
        # nn.CrossEntropyLoss is the usual one-call name for the same loss.
        scores = network(scale(people))
        print("same loss from nn.CrossEntropyLoss:", round(nn.CrossEntropyLoss()(scores, answers).item(), 3))
        picked = scores.argmax(dim=1)
        print(f"training people sized correctly: {(picked == answers).sum().item()} of {len(answers)}")
        new_people, new_answers = make_people(200, torch.Generator().manual_seed(99))
        new_picked = network(scale(new_people)).argmax(dim=1)
        print(f"200 new people sized correctly: {(new_picked == new_answers).sum().item()}")
        person = torch.tensor([[178.0, 70.0]])
        probabilities = torch.softmax(network(scale(person)), dim=1)[0]
        print("178 cm, 70 kg:", ", ".join(f"{s} {p:.2f}" for s, p in zip(SIZES, probabilities.tolist())))

log_probabilities = torch.log_softmax(model(screens), dim=1)
taken = log_probabilities.gather(1, moves).squeeze(1)
loss = -(rewards * taken).mean()
optimizer.zero_grad()
loss.backward()
optimizer.step()

The first line does the same job, and the last three are identical. Two things changed:

• answers, typed in before training, became moves, made by the network playing the game a moment ago;
• each term is multiplied by rewards, so a missed drop pushes its moves down instead of up.

Policy gradient

training a policy network by making rewarded actions more likely and punished actions less likely. The version on this page is known as REINFORCE.

Training

Training repeats three steps 1,000 times:

• play 128 drops with the current network;
• label every move: pushed up after a catch, down after a miss;
• step: one SGD update at learning rate 2.0.

Scrub through the saved run and predict: when does the paddle start heading the right way?

What the two views show:

• The left screen plays drops with the network as it was at that update, always taking its most likely move.
• The map reads its rule: one square for every ball column and paddle column, with the move it prefers when the ball is at the chosen row.

Before training the arrows are a jumble. By update 250 a pattern appears, but not the one you might write.

While the ball is high, the network ignores where the paddle is and runs toward the wall on the ball's side.

Choose a lower row and some arrows flip: it steps back out from the wall just in time, using the ball's height as a clock.

Nobody wrote that rule, and it is not the obvious one. It is simply a rule that +1 and −1 kept rewarding.

The curve shows how often the practice drops were caught, while the network was still rolling its die. It starts near one in seven, which is what wandering gets you, and climbs to about 93%.

Test

Now freeze the trained network, turn the die off (always take the most likely move), and try every start: 7 ball columns × 7 paddle positions. Click a square to watch that drop.

Seed 0 catches 46 of the 49 starts.

The three misses are the same mistake: the paddle runs past the ball on its way to the wall, turns back, and ends one column off.

We trained ten seeds with the same settings. All ten land between 92% and 100%, with a median of 94%. Seed 0 reaches exactly the median, so it is the one shown throughout the page.

Scale

A fair question: is this really machine learning?

• You could write “move toward the ball” as one if statement.
• The screen only has 343 possible pictures, so a table could store a move for each.

That is true of every example small enough to see in full.

What this page shows is the mechanism: a rule that nobody wrote, appearing from +1 and −1 alone, in a world small enough to check.

The need shows up when the world grows. In Atari Pong the screen is 210 × 160 pixels and practically never repeats, so there is nothing to look up, and nobody can write the rule down.

Andrej Karpathy trained a network on Pong with exactly this loop: play, use the network's own moves as labels, push them up after a won point and down after a lost one.

His summary: “Run a policy for a while. See what actions led to high rewards. Increase their probability.”

Recap

• model = make_model() creates the 56 → 8 → 3 policy network once.
• play(model, generator) drops a ball; on each of seven screens the network samples a move from its probabilities.
• reward arrives after the last move: +1 caught, −1 missed.
• rewards += [r] * len(m) gives every move in that drop the drop's reward.
• taken = log_probabilities.gather(1, moves) reads the log-probability of each move that was made.
• loss = -(rewards * taken).mean() asks for caught moves to become more likely and missed moves less likely.
• optimizer.step() updates the 483 weights and biases; repeat with fresh drops.

The network was never given an answer. It produced its own labels by playing, and the score sorted them into “do more of this” and “do less of this”.

The robot-arm workshops use the same idea on a harder task, where a reward can be hundreds of moves away.