Robot arm · Q-network

Overview

The Q-table learned one estimate per pose and command. A network can calculate the same four estimates from the two angles.

We keep the world, actions and rewards exactly the same, so we can see what replacing the table changes.

The new question is where training answers come from.

• In the ball workshop, MuJoCo supplied the measured next state.
• Here, the learner combines a reward it actually received with its own estimate of what can happen next.

Network

Select a pose: the arm and the four output bars show the same state.

Each output estimates the return from one joint command; the greedy policy chooses the largest.

network = nn.Sequential(
    nn.Linear(2, 16), nn.ReLU(), nn.Linear(16, 4)
)

"""Neural Q-learning on the same arm, using ordinary PyTorch.

Run: python3 classes/ml/labs/robot-arm/q_network.py
Actions collect transitions; optimizer steps learn from replay batches.
Evaluation freezes weights and uses the largest estimate, without exploration.
"""
import numpy as np
import torch
from torch import nn
from world import HOME, GAMMA, step, encode, evaluate


def make_network(initialization="custom"):
    model = nn.Sequential(nn.Linear(2, 16), nn.ReLU(), nn.Linear(16, 4))
    if initialization in ("custom", "scaled-hidden"):
        nn.init.normal_(model[0].weight, 0, 0.7)
        nn.init.zeros_(model[0].bias)
    if initialization in ("custom", "tiny-output"):
        nn.init.uniform_(model[2].weight, -0.003, 0.003)
        nn.init.zeros_(model[2].bias)
    return model


def train(seed=0, lr=0.3, episodes=300, observe=None,
          initialization="custom", target_period=100):
    torch.set_num_threads(1)
    torch.manual_seed(seed)
    rng = np.random.default_rng(seed)
    network = make_network(initialization)
    target_network = network if target_period == 1 else make_network(initialization)
    target_network.load_state_dict(network.state_dict())
    optimizer = torch.optim.SGD(network.parameters(), lr=lr)
    loss_fn = nn.MSELoss()
    memory, lengths, losses = [], [], []
    interactions, updates, first_success = 0, 0, None
    if observe:
        observe(0, network, interactions, updates)
    for episode in range(episodes):
        state = HOME
        total_loss = 0
        for move in range(400):
            with torch.no_grad():
                values = network(torch.tensor(encode(state))).numpy()
            action = int(rng.integers(4)) if rng.random() < 0.2 else int(np.argmax(values + 1e-9 * rng.random(4)))
            after, reward, done = step(state, action)
            memory.append([*state, action, reward, *after, done])
            interactions += 1
            if reward and first_success is None:
                first_success = {"episode": episode + 1, "move": interactions}
            for _ in range(2):
                indices = rng.integers(len(memory), size=min(32, len(memory)))
                batch = np.asarray([memory[k] for k in indices], dtype=np.float32)
                states = torch.tensor(encode(batch[:, :2]))
                actions = torch.tensor(batch[:, 2:3], dtype=torch.long)
                rewards = torch.tensor(batch[:, 3])
                next_states = torch.tensor(encode(batch[:, 4:6]))
                terminal = torch.tensor(batch[:, 6])
                with torch.no_grad():
                    next_values = target_network(next_states).max(dim=1).values
                    targets = rewards + GAMMA * (1 - terminal) * next_values
                predictions = network(states).gather(1, actions).squeeze(1)
                loss = loss_fn(predictions, targets)
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
                updates += 1
                total_loss += loss.item()
            if target_period > 1 and interactions % target_period == 0:
                target_network.load_state_dict(network.state_dict())
            state = after
            if done:
                break
        lengths.append(move + 1)
        losses.append(total_loss / (2 * (move + 1)))
        if observe:
            observe(episode + 1, network, interactions, updates)
    with torch.no_grad():
        all_states = torch.tensor(encode([[i, j] for i in range(19) for j in range(19)]))
        q = network(all_states).numpy().reshape(19, 19, 4)
    result = {"seed": seed, "lr": lr, "optimizer": "SGD", "initialization": initialization,
              "target_period": target_period, "episodes": episodes, "moves": interactions, "updates": updates,
              "first_success": first_success, "lengths": lengths, "losses": losses,
              "evaluation": evaluate(q)}
    return network, result


if __name__ == "__main__":
    network, result = train()
    print("Training moves:", result["moves"])
    print("Home test:", result["evaluation"]["home"])
    print("Starts passed:", result["evaluation"]["successes"], "/", result["evaluation"]["trials"])

The network:

• 2 inputs: the two angles;
• 16 ReLU neurons in the hidden layer;
• 4 linear outputs: one Q estimate per command.

That is \(2\times16+16+16\times4+4=116\) weights and biases. As in the Neural network lessons, each line is a weight and each input circle holds a number.

The grid indexes range from 0 to 18. Dividing by nine and subtracting one scales each input to −1…1, so home becomes about (0.333, −0.333).

All poses share the same 116 parameters, so one update can change the estimates at many poses at once.

We start output weights near zero and hidden weights wider; Experiment shows why.

Experience

A command gives us one record: the old state, the chosen action, the reward, the next state, and whether the episode ended.

We store those records and study small random batches of them.

Replay memory

the store of past records that training samples its batches from.

after, reward, done = step(state, action)
memory.append([*state, action, reward, *after, done])

That is the first successful record from the Q-table workshop. Its target is one, because the episode ends.

For an earlier zero-reward move whose next-state estimate is 0.5, the target would be \(0+0.9\times0.5=0.45\).

The record’s reward is observed, but its continuation is estimated. As those estimates improve, the training targets change too.

Bootstrapping

building a training target partly from another estimate.

Targets

The earlier action’s target uses the next pose’s estimated return. If training changes that estimate immediately, it also changes the next training answer. Compare using the latest weights with holding a copy steady for 100 commands:

python3 classes/ml/labs/robot-arm/tools/train_network.py --seed 0 --lr 0.3 --target-period 1
python3 classes/ml/labs/robot-arm/tools/train_network.py --seed 0 --lr 0.3 --target-period 100

Holding the copy raised median coverage from 103.5/251 to 181.5/251, with home reached in 8/10 seeds in both settings.

So: hold the continuation answers steady between refreshes.

Target network

that separate copy of the network, used only to compute continuation values and refreshed every 100 commands.

target_network = network if target_period == 1 else make_network(initialization)

# after every robot command:
if target_period > 1 and interactions % target_period == 0:
    target_network.load_state_dict(network.state_dict())

• The first line makes the copy. With --target-period 1 there is no copy: continuation values come from the latest weights.
• The refresh copies weights and biases every 100 commands. It does not collect an action or run an optimizer step, and the copy stays fixed in between.

with torch.no_grad():
    next_values = target_network(next_states).max(dim=1).values
    targets = rewards + GAMMA * (1 - terminal) * next_values

• max(dim=1) chooses the largest of the four action values in each batch row.
• terminal is one for a successful ending and zero otherwise, so 1 - terminal removes the continuation at the cube.
• no_grad() treats the target as a fixed answer during this update.

Training update

Only score the output for the action that was actually recorded. If the robot tried the elbow command, that transition tells us about the elbow output.

optimizer = torch.optim.SGD(network.parameters(), lr=lr)

predictions = network(states).gather(1, actions).squeeze(1)
loss = loss_fn(predictions, targets)
optimizer.zero_grad()
loss.backward()
optimizer.step()

gather(1, actions) picks the recorded action’s column from each row; squeeze(1) removes the one-column dimension, matching the target shape. The rest is the familiar PyTorch loop: mean squared error, backward() and plain SGD.

One command adds one record. We then make two optimizer updates, each studying up to 32 randomly sampled records.

Training

The runs:

• seeds 0–9, 300 episodes each;
• every attempt starts at home and lasts at most 400 commands;
• the same ε-greedy policy explores 20% of the time.

The default is the lower median run: sort the ten runs by starts passed, then by seed, and show the fifth.

Choose any seed, including the failed runs, and scrub its saved snapshots.

A small loss can still fit poor estimates. Select the failing seed with the button below: its loss can be low while its frozen choices fail to reach.

So: test what the arm does, not only how well the network fits its targets.

Test

Choose the learned network or the table from the previous workshop, freeze it, and start the arm somewhere else.

Select any allowed nonterminal square on the map, or use the start menu. The test makes no training updates.

The coverage map is the payoff:

• the table passes 59/251 starts;
• the default network, seed 4, passes 176/251;
• across all ten network seeds, the median is 181.5/251.

Both learners trained from home. The shared parameters also produced useful choices at starting poses the arm never trained from.

Choose a policy and click a blue passing square or a red failing one. Each policy gets the same 251 starts and 60-command limit.

Experiment

Try smaller parameter updates while keeping the world, seeds, initialization and 300-episode budget fixed:

optimizer = torch.optim.SGD(network.parameters(), lr=lr)

python3 classes/ml/labs/robot-arm/tools/train_network.py --seed 0 --lr 0.03
python3 classes/ml/labs/robot-arm/tools/train_network.py --seed 0 --lr 0.3

The two rates across ten seeds:

• at SGD 0.03, home was reached in 4/10 seeds, median coverage 28/251;
• at SGD 0.3, home was reached in 8/10 seeds, median coverage 181.5/251.

So: choose the learning rate by testing reaching across seeds, then keep it fixed for the initialization and target-copy comparisons.

Starting weights

Output weights. Before any experience, PyTorch’s default starting weights already give the arm preferences. Averaged across seeds 0–9 on the 361-pose grid:

• the largest absolute Q is 0.658, bigger than the real home return of 0.430;
• the best output leads the second by 0.169 on average.

With our tiny output weights and zero output biases, that lead shrinks to 0.002 and the largest absolute Q to 0.011.

Starting near zero keeps these random preferences small, so the reward records can shape the choices.

Hidden weights. With PyTorch’s default hidden weights and biases, an average of 3.9 of the 16 ReLU neurons are always on or always off across the whole map.

• an always-on neuron stays on one straight piece;
• an always-off neuron contributes zero.

Neither bends on this map. With wider hidden weights (standard deviation 0.7) and zero hidden biases, all 16 cross zero somewhere on the grid, so each can bend as the pose changes.

Measure these starting values without training. The script prints every seed and their averages; “largest” means each seed’s maximum over all poses and outputs, then averaged across seeds.

python3 classes/ml/labs/robot-arm/tools/probe_initialization.py

Open the measurement script ↗

Now train with each initialization at the same SGD rate of 0.3, using seeds 0–9 and the same 300-episode budget. Using both changes raised median coverage from PyTorch defaults’ 24/251 to 181.5/251 starts:

nn.init.normal_(model[0].weight, 0, 0.7)
nn.init.zeros_(model[0].bias)
nn.init.uniform_(model[2].weight, -0.003, 0.003)
nn.init.zeros_(model[2].bias)

The first two lines give hidden weights a standard deviation of 0.7 and zero their biases; the last two start output weights near zero and zero their biases. So: keep both changes for this lesson’s run.

Run the comparison with --init default, --init tiny-output, --init scaled-hidden or --init custom. The last combines both changes.

python3 classes/ml/labs/robot-arm/tools/train_network.py --seed 0 --lr 0.3 --init default
python3 classes/ml/labs/robot-arm/tools/train_network.py --seed 0 --lr 0.3 --init custom

Recap

• make_network() creates the 2 → 16 → 4 model; a second copy supplies continuation estimates.
• step(state, action) moves the arm and supplies reward and an ending flag.
• memory.append(...) keeps the observed transition.
• targets combine rewards with the target copy’s best next-state estimates.
• gather(...) selects each recorded action’s output; loss, backward and optimizer step update shared parameters.
• load_state_dict(...) periodically refreshes the target copy.
• evaluate(...) freezes decisions and tests reaching from multiple starts.

The world supplies rewards, the learner constructs training targets, and the network learns estimates used to choose actions.

In the imitation workshop, after Robot basics, demonstrations supply the action labels directly.