Week 3 · embedded programming

A neural network that learns on a $5 board

What the datasheet told me

Building the board

Exactly what I did at the bench

Reproduce the firmware

// touch_test_rp2040.ino: how a pad is read, no hardware touch sensing needed
const uint8_t PADS[6] = { 3, 4, 2, 27, 1, 26 };   // Q0..Q5
const int SAMPLES = 32;   // averaged reads per pad
const int MARGIN  = 12;   // counts above baseline that count as a touch

uint16_t readPad(uint8_t pin) {
  pinMode(pin, OUTPUT); digitalWrite(pin, LOW); delayMicroseconds(50);  // discharge
  pinMode(pin, INPUT_PULLUP);
  uint16_t n = 0;
  while (digitalRead(pin) == LOW && n < 2000) n++;                        // charge time
  return n;
}

What it does

One step of the simulation

And backwards

Controls

The bug that made it too good

void loop1() {
  if (!trainingEnabled) {
    delay(5);
    return;
  }

  // One BPTT step at a time so learning is actually visible on the OLED.
  float l = trainOneBPTTStep();

  mutex_enter_blocking(&mtx);
  trainStep += 1;
  trainLoss = l;
  memcpy(visibleTheta, theta, sizeof(theta));
  mutex_exit(&mtx);

  // Deliberately pace training. Without this, the RP2040 can converge
  // before you have time to see the random policy behave badly.
  delay(180);
}

Code

Forward pass through the network

void networkForward(
  const float *th,
  const float in[N_IN],
  float hidden[N_H],
  float out[N_OUT]
) {
  for (int h = 0; h < N_H; h++) {
    float z = th[idxB1(h)];
    for (int i = 0; i < N_IN; i++) {
      z += th[idxW1(h, i)] * in[i];
    }
    hidden[h] = tanhf(z);
  }

  for (int o = 0; o < N_OUT; o++) {
    float z = th[idxB2(o)];
    for (int h = 0; h < N_H; h++) {
      z += th[idxW2(o, h)] * hidden[h];
    }
    out[o] = tanhf(z);
  }
}

Start of the training step

float trainOneBPTTStep() {
  float tx, ty, cx, cy;

  mutex_enter_blocking(&mtx);
  tx = targetX;
  ty = targetY;
  cx = currentX;
  cy = currentY;
  mutex_exit(&mtx);

  // Small randomized starting condition makes the learned controller useful
  // from more than one exact initial state.
  st[0][0] = 0.20f + 0.10f * (urand(rngState) - 0.5f);
  st[0][1] = 0.45f + 0.10f * (urand(rngState) - 0.5f);
  st[0][2] = 0.08f * (urand(rngState) - 0.5f);
  st[0][3] = 0.08f * (urand(rngState) - 0.5f);

  float loss = 0.0f;

  // ---------------- forward rollout ----------------
  for (int t = 0; t < HORIZON; t++) {
    float x  = st[t][0];
    float y  = st[t][1];
    float vx = st[t][2];
    float vy = st[t][3];

    ins[t][0] = tx - x;
    ins[t][1] = ty - y;
    ins[t][2] = vx;
    ins[t][3] = vy;

    networkForward(theta, ins[t], hid[t], act[t]);

    float nvx = DRAG * vx + THRUST * act[t][0] + cx;
    float nvy = DRAG * vy + THRUST * act[t][1] + cy;

    float nx = x + DT * nvx;
    float ny = y + DT * nvy;

    st[t + 1][0] = nx;
    st[t + 1][1] = ny;
    st[t + 1][2] = nvx;
    st[t + 1][3] = nvy;

    float dx = x - tx;
    float dy = y - ty;
  // ... full sketch continues with the backward pass

A true reset

void resetLearning() {
  // B means a *true* reset:
  //   1. stop learning
  //   2. randomize all weights
  //   3. freeze those random weights
  //   4. let you watch the untrained swimmer
  //
  // Press A when you want training to begin.
  trainingEnabled = false;
  delay(10);

  randomizePolicy();

  mutex_enter_blocking(&mtx);
  memcpy(frozenTheta, theta, sizeof(frozenTheta));
  memcpy(visibleTheta, theta, sizeof(visibleTheta));
  mutex_exit(&mtx);

  randomDemoMode = true;
  testMode = true;

  resetVisible();
}

Before the swimmer

Files