273 lines
11 KiB
Python
273 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
sudo python3 balance.py imu # find the balance point
|
|
sudo python3 balance.py balance --offset X # balance
|
|
sudo python3 balance.py balance --offset X --speed 200 # drive forward at 200 mm/s
|
|
"""
|
|
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
|
|
ACCEL_MAX = 150 # mm/s per second: how fast the speed setpoint ramps
|
|
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()
|
|
stats = [] # (wheel position mm, tilt, pwm) for the stillness readout
|
|
print(f"kp={args.kp} kd={args.kd} kv={args.kv} ki={args.ki} offset={args.offset} "
|
|
f"slew={args.slew} deadband={args.deadband} speed={args.speed}")
|
|
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()
|
|
speed_pid.setpoint = 0.0 # always start standing still, then ramp
|
|
stats.clear()
|
|
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
|
|
speed_pid.setpoint += clamp(args.speed - speed_pid.setpoint, ACCEL_MAX * (now - prev_t))
|
|
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 = args.kp * (theta - target) + args.kd * rate
|
|
if u:
|
|
u += math.copysign(args.deadband, u) # help the wheels past static friction
|
|
u = clamp(u, PWM_MAX)
|
|
out += clamp(u - out, args.slew) # ramp, never jump
|
|
motors.pwm(out, out)
|
|
|
|
# stillness readout every 5 s: lower = stiller
|
|
stats.append(((wheels[0] + wheels[1]) / 2 * MM_PER_TICK, theta, out))
|
|
if len(stats) == 5 * LOOP_HZ:
|
|
pos, th, pw = (np.array(c) for c in zip(*stats))
|
|
print(f"STILLNESS 5s: wheels p-p {np.ptp(pos):5.1f} mm rms {pos.std():4.1f} mm "
|
|
f"tilt rms {th.std():5.2f} deg pwm rms {pw.std():4.0f}", flush=True)
|
|
stats.clear()
|
|
|
|
if tick % 20 == 0:
|
|
print(f"th {theta:+6.2f} tgt {target:+5.2f} rate {rate:+7.1f} v {v:+6.0f} mm/s "
|
|
f"sp {speed_pid.setpoint:+4.0f} 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=120.0, help="PWM per deg of lean")
|
|
p.add_argument("--kd", type=float, default=3.0, help="PWM per deg/s of tipping")
|
|
p.add_argument("--kv", type=float, default=0.02, help="deg of lean per mm/s of speed")
|
|
p.add_argument("--ki", type=float, default=0.003, help="deg of lean per mm of drift")
|
|
p.add_argument("--speed", type=float, default=0.0, help="drive speed, mm/s, + = forward")
|
|
p.add_argument("--slew", type=float, default=PWM_SLEW, help="max PWM change per tick")
|
|
p.add_argument("--deadband", type=float, default=0.0, help="PWM added to beat static friction")
|
|
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() |