#!/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()