Balance bot on an Arduino Uno

// Balance bot: a two-wheel robot that stands up on its own.
//
//   robot SCL -> 12   robot SDA -> 11     the robot's own MPU6050, at 0x68
//   AIN1 -> 10  AIN2 -> 9   PWMA -> 6     left motor, on the robot's TB6612
//   BIN1 -> 5   BIN2 -> 4   PWMB -> 3     right motor
//   STBY -> 5 V                           the driver runs whenever it has power
//   A0, A1, A2   three knobs: Kp, Ki and Kd, live
//   A5 -> SCL   A4 -> SDA                 the OLED, angle and gains
//
// Every 5 ms the loop does three things.
//
// # 1. Work out the tilt
//
// The accelerometer knows which way is down, but it also feels every shove the
// wheels give the body, so on its own it is noisy and wrong while the robot
// moves. The gyroscope knows how fast the body is turning, smoothly, but adding
// up a rate drifts. The complementary filter takes the best of each:
//
//   angle = 0.98 * (angle + rate * dt) + 0.02 * accelAngle
//
// The gyro moves the angle from moment to moment; the accelerometer pulls it
// slowly back to the truth.
//
// # 2. PID
//
// The error is how far the body leans from upright. P drives the wheels under
// the lean, D (on the gyro rate) stops it overshooting, and I makes up the
// steady push a motor needs just to hold a speed, which is what stops this
// robot creeping off at full speed and falling. Turn the knobs while it runs.
//
// # 3. Drive
//
// Leaning forward, the wheels drive forward to get back under the body. The
// sign of the PID output picks AIN1/AIN2 and the size goes to analogWrite.
// Past 45 degrees it has fallen: the motors brake and the integral resets.
// Until a hand lets go of it (at power-up, and after standing it up) the
// motors stay braked, so the integral does not wind up while it is held.
// Press the stand-up button on the robot to try again.

#include "mokxi_i2c.h"
#include "mokxi_mpu6050.h"
#include "mokxi_ssd1306.h"

const uint8_t IMU_SCL = 12;
const uint8_t IMU_SDA = 11;
const uint8_t AIN1 = 10;
const uint8_t AIN2 = 9;
const uint8_t PWMA = 6;
const uint8_t BIN1 = 5;
const uint8_t BIN2 = 4;
const uint8_t PWMB = 3;
const uint8_t KNOB[3] = {A0, A1, A2};

// Full scale of each knob: the middle is a tune that balances.
const float KP_MAX = 80.0;
const float KI_MAX = 600.0;
const float KD_MAX = 2.0;

const unsigned long PERIOD_US = 5000;
const float ALPHA = 0.98;
const float FALLEN = 45.0;
const float UPRIGHT = 10.0;
// The gyro rate, deg/s, that says the hand has let go.
const float RELEASED = 1.0;
// The most PWM the I term may hold, so it cannot wind up without limit.
const float I_LIMIT = 150.0;

SoftI2c imuBus(IMU_SCL, IMU_SDA, 400);
SoftI2c oledBus(A5, A4, 400);
Mpu6050 imu(imuBus);
Ssd1306 oled(oledBus, SSD1306_ADDRESS);

float kp = KP_MAX / 2;
float ki = KI_MAX / 2;
float kd = KD_MAX / 2;
float angle = 0;
float integral = 0;
int output = 0;
bool fallen = false;
bool started = false;
// True from power-up, and from being stood up, until the hand lets go.
bool held = true;
bool reported = false;
bool disturbed = false;
unsigned long steadySince = 0;

unsigned long lastUs = 0;
unsigned long printedAt = 0;
unsigned long drawnAt = 0;
uint8_t knob = 0;
uint8_t row = 0;

// What each OLED row last showed, so an unchanged row costs no bus time.
long shown[5] = {-99999, -99999, -99999, -99999, -99999};

void drive(int pwm) {
  bool forward = pwm >= 0;
  int duty = abs(pwm);
  if (duty > 255) duty = 255;
  digitalWrite(AIN1, forward ? HIGH : LOW);
  digitalWrite(AIN2, forward ? LOW : HIGH);
  digitalWrite(BIN1, forward ? HIGH : LOW);
  digitalWrite(BIN2, forward ? LOW : HIGH);
  analogWrite(PWMA, duty);
  analogWrite(PWMB, duty);
}

void brake() {
  digitalWrite(AIN1, HIGH);
  digitalWrite(AIN2, HIGH);
  digitalWrite(BIN1, HIGH);
  digitalWrite(BIN2, HIGH);
  analogWrite(PWMA, 0);
  analogWrite(PWMB, 0);
}

// A value in hundredths, drawn with the given number of decimals.
void field(uint8_t r, const char *label, long hundredths, uint8_t decimals) {
  if (shown[r] == hundredths) return;
  shown[r] = hundredths;
  int16_t y = 16 + r * 10;
  oled.fillRect(0, y, 128, 8, false);
  oled.text(0, y, label);
  char buf[12];
  uint8_t n = 0;
  long v = hundredths;
  if (v < 0) {
    buf[n++] = '-';
    v = -v;
  }
  long whole = v / 100;
  long frac = v % 100;
  char digits[8];
  uint8_t d = 0;
  do {
    digits[d++] = (char)('0' + whole % 10);
    whole /= 10;
  } while (whole > 0 && d < 7);
  while (d > 0) buf[n++] = digits[--d];
  if (decimals > 0) {
    buf[n++] = '.';
    buf[n++] = (char)('0' + frac / 10);
    if (decimals > 1) buf[n++] = (char)('0' + frac % 10);
  }
  buf[n] = 0;
  oled.text(48, y, buf);
}

// One row of the screen per call, so the bus is never busy for long.
void drawRow() {
  switch (row) {
    case 0: field(0, "angle", (long)(angle * 100), 1); break;
    case 1: field(1, "Kp", (long)(kp * 100), 1); break;
    case 2: field(2, "Ki", (long)(ki * 100), 0); break;
    case 3: field(3, "Kd", (long)(kd * 100), 2); break;
    default:
      if (shown[4] != (long)fallen) {
        shown[4] = fallen;
        oled.fillRect(0, 56, 128, 8, false);
        oled.text(0, 56, fallen ? "FALLEN: stand me up" : "balancing");
      }
      break;
  }
  oled.display();
  row = (row + 1) % 5;
}

void readKnob() {
  float x = analogRead(KNOB[knob]) / 1023.0;
  if (knob == 0) kp = x * KP_MAX;
  if (knob == 1) ki = x * KI_MAX;
  if (knob == 2) kd = x * KD_MAX;
  knob = (knob + 1) % 3;
}

void setup() {
  pinMode(AIN1, OUTPUT);
  pinMode(AIN2, OUTPUT);
  pinMode(BIN1, OUTPUT);
  pinMode(BIN2, OUTPUT);
  pinMode(PWMA, OUTPUT);
  pinMode(PWMB, OUTPUT);
  brake();
  Serial.begin(115200);
  imuBus.begin();
  oledBus.begin();
  if (!imu.begin()) {
    Serial.println("No MPU6050 at 0x68: check SCL and SDA");
  }
  imu.setRanges(0, 0);
  oled.begin();
  oled.clear();
  oled.text(0, 0, "BALANCE BOT");
  oled.hLine(0, 10, 128);
  oled.display();
  for (uint8_t i = 0; i < 3; i++) readKnob();
  Serial.println("Balance bot ready");
  lastUs = micros();
}

void loop() {
  unsigned long now = micros();
  if (now - lastUs < PERIOD_US) return;
  float dt = (now - lastUs) / 1000000.0;
  lastUs = now;

  Mpu6050Reading r;
  if (!imu.read(r)) return;
  float accelAngle = atan2(-(float)r.ax, (float)r.az) * 57.29578;
  float rate = r.gy / 131.0;
  if (!started) {
    angle = accelAngle;
    started = true;
  }
  angle = ALPHA * (angle + rate * dt) + (1 - ALPHA) * accelAngle;

  if (fallen) {
    brake();
    // Lying still, gravity is the whole story: trust the accelerometer, so
    // the moment it is stood up the angle is right.
    angle = accelAngle;
    if (fabs(angle) < UPRIGHT) {
      fallen = false;
      held = true;
      Serial.println("Upright again");
    }
  } else if (held) {
    // Held still by a hand: wait for it to let go before balancing. Driving
    // now would only fill the integral with a lean the wheels cannot fix,
    // and that stored push would throw it over the moment it was let go.
    // The gyro tells: held, the body cannot turn.
    brake();
    angle = accelAngle;
    if (fabs(rate) > RELEASED) {
      held = false;
      integral = 0;
      steadySince = millis();
      Serial.println("Let go: balancing");
    }
  } else if (fabs(angle) > FALLEN) {
    fallen = true;
    integral = 0;
    output = 0;
    reported = false;
    disturbed = false;
    brake();
    Serial.println("FALLEN");
  } else {
    integral += angle * dt;
    if (ki > 0.01) {
      float limit = I_LIMIT / ki;
      if (integral > limit) integral = limit;
      if (integral < -limit) integral = -limit;
    }
    float pid = kp * angle + ki * integral + kd * rate;
    output = (int)constrain(pid, -255, 255);
    drive(output);
    // A shove that tips it past 3 degrees, and the moment it is back.
    if (fabs(angle) > 3) {
      disturbed = true;
    } else if (disturbed && fabs(angle) < 1) {
      disturbed = false;
      Serial.println("Caught it: upright again");
    }
    if (fabs(angle) > 2) {
      steadySince = millis();
    } else if (!reported && millis() - steadySince >= 3000) {
      reported = true;
      Serial.println("Balanced for 3 s");
    }
  }

  readKnob();
  if (millis() - drawnAt >= 40) {
    drawnAt = millis();
    drawRow();
  }
  if (millis() - printedAt >= 250) {
    printedAt = millis();
    Serial.print("t=");
    Serial.print(millis());
    Serial.print(" angle=");
    Serial.print(angle, 2);
    Serial.print(" out=");
    Serial.println(output);
  }
}