aleksa24-balance/balance.py

#!/usr/bin/env python3
"""
  sudo python3 balance.py imu                   # find the balance point
  sudo python3 balance.py balance --offset X    # balance
"""
import argparse, math, os, sys, threading, time
import numpy as np
import imufusion
import serial
from simple_pid import PID
from smbus2 import SMBus

# ---------------- config ----------------
LOOP_HZ     = 200
FALL_ANGLE  = 45.0                  # deg: motors off beyond this
START_ANGLE = 5.0                   # deg: arms when held this close to upright and still
PWM_MAX     = 2000                  # power-bank limit
PWM_SLEW    = 300                   # max PWM change per tick (no instant reversals)
ENC_TIMEOUT = 0.05                  # s without encoder data -> motors off, re-init driver board
LEFT, RIGHT = 2, 3                  # M3 = left, M4 = right
MOTOR_SIGN  = (-1, +1)              # +PWM on M3 rolls the robot backward, so flip it
ENC_SIGN    = (-1, +1)              # M3 counts down when rolling forward, so flip it
MM_PER_TICK = math.pi * 67.0 / 1320 # 67 mm wheel, 1320 counts per wheel turn


def clamp(x, limit):
    return max(-limit, min(limit, x))


class IMU:
    """MPU-6500. read() gives body-frame data: x forward, y left, z up."""

    def __init__(self):
        self.bus = SMBus(1)
        # reset, wake, gyro filter 92 Hz, gyro +-500 deg/s, accel +-2 g, accel filter 41 Hz
        for reg, val in [(0x6B, 0x80), (0x6B, 0x01), (0x1A, 0x02), (0x1B, 0x08), (0x1C, 0x00), (0x1D, 0x03)]:
            self.bus.write_byte_data(0x68, reg, val)
            time.sleep(0.1)
        print("Calibrating gyro - keep the robot STILL...", flush=True)
        self.bias = np.zeros(3)
        samples = []
        for _ in range(300):
            samples.append(self.read()[1])
            time.sleep(0.003)
        self.bias = np.mean(samples, axis=0)
        print("gyro bias (deg/s):", np.round(self.bias, 2))

    def read(self):
        d = self.bus.read_i2c_block_data(0x68, 0x3B, 14)
        raw = np.frombuffer(bytes(d), dtype=">i2").astype(float)   # ax ay az temp gx gy gz
        # sensor is rotated 180 deg about the axle: body = (-x, y, -z)
        acc = imufusion.remap(raw[0:3] / 16384.0, imufusion.ALIGNMENT_NXPYNZ)              # g
        gyr = imufusion.remap(raw[4:7] / 65.5, imufusion.ALIGNMENT_NXPYNZ) - self.bias     # deg/s
        return acc, gyr


class Tilt:
    def __init__(self):
        self.ahrs = imufusion.Ahrs()
        self.ahrs.set_settings(imufusion.AhrsSettings(
            sample_rate=LOOP_HZ,
            convention=imufusion.CONVENTION_NWU,   # x forward, y left, z up
            gain=0.5,
            gyroscope_range=500,
            acceleration_rejection=10,             # ignore accel while it's 10+ deg off (wheel pushes)
            rejection_timeout=5))

    def update(self, acc, gyr, dt):
        """Returns pitch in degrees, + = leaning forward."""
        self.ahrs.set_sample_period(dt)
        self.ahrs.update_no_magnetometer(gyr, acc)
        return imufusion.quaternion_to_euler(self.ahrs.get_quaternion())[1]

    def ready(self):
        return not self.ahrs.get_flags().startup   # Fusion converges for ~3 s at startup


class Motors:
    def __init__(self):
        self.ser = serial.Serial("/dev/serial0", 115200, timeout=0.02)
        self.counts, self.stamp, self.error = None, 0.0, None
        threading.Thread(target=self._reader, daemon=True).start()
        self.configure()

    def configure(self):
        """Stop, set motor type (1 = 520 motor), stream encoder totals. Also used after a board reset."""
        for cmd in (b"$pwm:0,0,0,0#", b"$mtype:1#", b"$upload:1,0,0#"):
            self.ser.write(cmd)
            time.sleep(0.05)

    def pwm(self, left, right):
        v = [0, 0, 0, 0]
        v[LEFT], v[RIGHT] = int(MOTOR_SIGN[0] * left), int(MOTOR_SIGN[1] * right)
        self.ser.write(b"$pwm:%d,%d,%d,%d#" % tuple(v))

    def wheels(self, now):
        """((left, right) total counts with forward = +, seconds since last update)."""
        c = self.counts
        if c is None:
            return None, 1e9
        return (ENC_SIGN[0] * c[LEFT], ENC_SIGN[1] * c[RIGHT]), now - self.stamp

    def _reader(self):
        buf = b""
        while True:
            try:
                buf += self.ser.read(self.ser.in_waiting or 1)
            except serial.SerialException as e:
                self.error = e
                return
            *lines, buf = buf.split(b"#")
            for line in lines:
                if b"$MAll:" in line:
                    try:
                        self.counts = [int(x) for x in line.split(b"$MAll:")[-1].split(b",")]
                        self.stamp = time.perf_counter()
                    except ValueError:
                        pass

    def stop(self):
        for _ in range(3):
            self.ser.write(b"$pwm:0,0,0,0#")
            time.sleep(0.02)


def realtime():
    """Pin to core 3 with real-time priority (needs sudo). Run Doom with: taskset -c 0-2 ..."""
    try:
        os.sched_setaffinity(0, {3})
        os.sched_setscheduler(0, os.SCHED_FIFO, os.sched_param(50))
    except OSError as e:
        print(f"No real-time priority ({e}) - run with sudo. Expect more timing jitter.")


def ticker(hz):
    """Yields (dt, now) at a steady rate."""
    period = 1.0 / hz
    nxt = last = time.perf_counter()
    while True:
        nxt += period
        time.sleep(max(0.0, nxt - time.perf_counter()))
        now = time.perf_counter()
        nxt = max(nxt, now - period)    # don't try to catch up after a stall
        yield now - last, now
        last = now


def show_imu(imu, tilt, args):
    worst = 0.0
    for i, (dt, now) in enumerate(ticker(LOOP_HZ)):
        acc, gyr = imu.read()
        pitch = tilt.update(acc, gyr, dt)
        worst = max(worst, dt)
        if i % 20 == 0:
            acc_pitch = math.degrees(math.atan2(-acc[0], acc[2]))
            print(f"accel-only {acc_pitch:+7.2f}   FUSION {pitch:+7.2f}   minus offset {pitch - args.offset:+7.2f}   "
                  f"rate {gyr[1]:+7.1f}   {'ready' if tilt.ready() else 'settling'}   worst dt {worst * 1000:4.1f} ms",
                  flush=True)
            worst = 0.0


def balance(imu, tilt, motors, args):
    # outer loop: wheel speed -> lean target. Error = 0 - speed, so P leans back against speed
    # and I (integral of speed = distance) leans back toward the start position.
    speed_pid = PID(args.kv, args.ki, 0, setpoint=0, sample_time=None, output_limits=(-8, 8))
    armed = False
    out = target = v = 0.0
    prev = prev_t = None
    last_reinit = time.perf_counter()
    print(f"kp={args.kp} kd={args.kd} kv={args.kv} ki={args.ki} offset={args.offset}")
    print("Hold it upright and still to arm. Ctrl-C to quit.", flush=True)

    for tick, (dt, now) in enumerate(ticker(LOOP_HZ)):
        if motors.error:
            sys.exit(f"Serial error: {motors.error}")
        acc, gyr = imu.read()
        theta = tilt.update(acc, gyr, dt) - args.offset   # deg, + = leaning forward
        rate = gyr[1]                                      # deg/s, + = tipping forward

        # watchdog: the driver board went quiet (power dip resets it)
        wheels, age = motors.wheels(now)
        if wheels is None or age > ENC_TIMEOUT:
            if armed:
                print("ENCODER DATA LOST - motors off, re-initialising driver board", flush=True)
            armed, out = False, 0.0
            motors.pwm(0, 0)
            if now - last_reinit > 0.5:
                motors.configure()
                last_reinit = time.perf_counter()
            continue

        if not armed:
            if tilt.ready() and abs(theta) < START_ANGLE and abs(rate) < 30:
                armed, out, target, v = True, 0.0, 0.0, 0.0
                speed_pid.reset()
                prev, prev_t = wheels, now
                print("ARMED", flush=True)
            elif tick % 20 == 0:
                motors.pwm(0, 0)
            continue

        if abs(theta) > FALL_ANGLE:
            armed = False
            motors.pwm(0, 0)
            print(f"FELL ({theta:+.1f} deg) - motors off", flush=True)
            continue

        # outer loop, 50 Hz: speed PI -> target lean angle
        if tick % 4 == 0:
            v_now = (wheels[0] - prev[0] + wheels[1] - prev[1]) / 2 * MM_PER_TICK / (now - prev_t)
            v = 0.7 * v + 0.3 * v_now                      # mm/s, smoothed
            target = speed_pid(v, dt=now - prev_t)
            prev, prev_t = wheels, now

        # inner loop, 200 Hz: angle PD (D straight from the gyro, no differentiating)
        u = clamp(args.kp * (theta - target) + args.kd * rate, PWM_MAX)
        out += clamp(u - out, PWM_SLEW)                   # ramp, never jump
        motors.pwm(out, out)

        if tick % 20 == 0:
            print(f"th {theta:+6.2f}  tgt {target:+5.2f}  rate {rate:+7.1f}  v {v:+6.0f} mm/s  pwm {out:+6.0f}",
                  flush=True)


def main():
    p = argparse.ArgumentParser()
    p.add_argument("mode", choices=["imu", "balance"])
    p.add_argument("--offset", type=float, default=0.0, help="pitch at the balance point, deg")
    p.add_argument("--kp", type=float, default=100.0, help="PWM per deg of lean")
    p.add_argument("--kd", type=float, default=2.5, help="PWM per deg/s of tipping")
    p.add_argument("--kv", type=float, default=0.0, help="deg of lean per mm/s of speed")
    p.add_argument("--ki", type=float, default=0.0, help="deg of lean per mm of drift")
    args = p.parse_args()

    realtime()
    imu, tilt = IMU(), Tilt()
    motors = Motors() if args.mode == "balance" else None
    try:
        if motors:
            balance(imu, tilt, motors, args)
        else:
            show_imu(imu, tilt, args)
    except KeyboardInterrupt:
        pass
    finally:
        if motors:
            motors.stop()
            print("\nmotors stopped")


if __name__ == "__main__":
    main()