adding balance
This commit is contained in:
+391
@@ -0,0 +1,391 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
balance.py - two-wheel balancing robot
|
||||
Pi Zero 2 W
|
||||
MPU-6500 (GY-6500 / GY-9250 board) on I2C 0x68, mounted upside down
|
||||
Yahboom 4-ch driver over UART (/dev/serial0), 520 motors: left = M3, right = M4
|
||||
|
||||
sudo python3 balance.py imu # watch tilt, find the balance point (--offset)
|
||||
sudo python3 balance.py balance --offset X # balance (tune kp/kd first, kv/ki = 0)
|
||||
sudo python3 balance.py balance --offset X --kp 120 --kd 3 --kv 0.02 --ki 0.003
|
||||
|
||||
Tilt comes from a 6-axis Madgwick filter. Pitch is + when leaning FORWARD (toward the tape).
|
||||
"""
|
||||
import argparse, gc, math, os, sys, threading, time
|
||||
import serial
|
||||
from smbus2 import SMBus
|
||||
|
||||
# ============================ CONFIG ============================
|
||||
I2C_BUS, IMU_ADDR = 1, 0x68
|
||||
UART_PORT, UART_BAUD = "/dev/serial0", 115200
|
||||
|
||||
LOOP_HZ = 200
|
||||
RT_CPU = 3 # core for the control loop (run Doom with: taskset -c 0-2 ...)
|
||||
BETA = 0.05 # Madgwick gain (rad/s): how hard the accelerometer pulls out gyro drift
|
||||
ANGLE_OFFSET = 0.0 # deg: pitch reading at the true balance point (find it in 'imu' mode)
|
||||
|
||||
# motors / encoders (measured with motor_test.py)
|
||||
LEFT_CH, RIGHT_CH = 2, 3 # M3 = left, M4 = right
|
||||
LEFT_DIR, RIGHT_DIR = -1, +1 # +PWM on M3 rolls the robot backward, so flip it
|
||||
ENC_LEFT_DIR, ENC_RIGHT_DIR = -1, +1 # M3 counts down when rolling forward, so flip it
|
||||
TICKS_PER_REV = 30 * 11 * 4 # 1320 (measured ~1300 by hand)
|
||||
WHEEL_D_MM = 67.0
|
||||
|
||||
# safety
|
||||
PWM_MAX = 2000 # power-bank limit; barely any extra speed above this anyway
|
||||
PWM_SLEW = 150 # max PWM change per tick: full reverse takes ~27 ms instead of 0
|
||||
DEADBAND = 0 # PWM added to beat static friction (wheels moved at 800 in the air)
|
||||
FALL_ANGLE = 45.0 # deg: motors off beyond this
|
||||
START_ANGLE = 5.0 # deg: arms when held this close to upright and still
|
||||
ENC_TIMEOUT = 0.05 # s without encoder data -> motors off, re-init the driver board
|
||||
TARGET_LIMIT = 8.0 # deg: max lean the speed loop may ask for
|
||||
SPEED_DIV = 4 # speed loop every 4th tick = 50 Hz
|
||||
# ================================================================
|
||||
|
||||
|
||||
def clamp(x, lo, hi):
|
||||
return lo if x < lo else hi if x > hi else x
|
||||
|
||||
|
||||
class IMU:
|
||||
"""MPU-6500 family. read() returns BODY-frame data: x forward, y left, z up."""
|
||||
|
||||
def __init__(self, bus):
|
||||
self.bus = bus
|
||||
who = bus.read_byte_data(IMU_ADDR, 0x75)
|
||||
if who not in (0x70, 0x71, 0x73, 0x68):
|
||||
sys.exit(f"Unexpected WHO_AM_I 0x{who:02X} at 0x{IMU_ADDR:02X}")
|
||||
w = lambda reg, val: bus.write_byte_data(IMU_ADDR, reg, val)
|
||||
w(0x6B, 0x80); time.sleep(0.1) # reset
|
||||
w(0x6B, 0x01); time.sleep(0.05) # wake, PLL clock
|
||||
w(0x1A, 0x02) # gyro low-pass ~92 Hz (less delay than 41 Hz)
|
||||
w(0x1B, 0x08) # gyro +-500 deg/s
|
||||
w(0x1C, 0x00) # accel +-2 g
|
||||
w(0x1D, 0x03) # accel low-pass ~41 Hz
|
||||
time.sleep(0.05)
|
||||
self.bias = (0.0, 0.0, 0.0) # gyro bias, sensor frame, deg/s
|
||||
|
||||
def _raw(self):
|
||||
d = self.bus.read_i2c_block_data(IMU_ADDR, 0x3B, 14)
|
||||
s = lambda i: (((d[i] << 8) | d[i + 1]) ^ 0x8000) - 0x8000
|
||||
acc = (s(0) / 16384.0, s(2) / 16384.0, s(4) / 16384.0)
|
||||
gyr = (s(8) / 65.5, s(10) / 65.5, s(12) / 65.5)
|
||||
return acc, gyr
|
||||
|
||||
def read(self):
|
||||
"""(accel in g, gyro in deg/s) in body frame.
|
||||
The sensor sits rotated 180 deg about the axle (y), as measured:
|
||||
upright az = -1 g, forward tilt makes ax + and gy +.
|
||||
So body x = -sensor x, body y = sensor y, body z = -sensor z."""
|
||||
(ax, ay, az), (gx, gy, gz) = self._raw()
|
||||
bx, by, bz = self.bias
|
||||
return (-ax, ay, -az), (-(gx - bx), gy - by, -(gz - bz))
|
||||
|
||||
def calibrate(self, seconds=1.0):
|
||||
for attempt in range(5):
|
||||
print("Calibrating gyro - keep the robot STILL...", flush=True)
|
||||
samples = []
|
||||
end = time.perf_counter() + seconds
|
||||
while time.perf_counter() < end:
|
||||
samples.append(self._raw()[1])
|
||||
time.sleep(0.003)
|
||||
spread = max(max(s[k] for s in samples) - min(s[k] for s in samples) for k in range(3))
|
||||
if spread < 4.0:
|
||||
n = len(samples)
|
||||
self.bias = tuple(sum(s[k] for s in samples) / n for k in range(3))
|
||||
print("Gyro bias (deg/s): " + ", ".join(f"{b:+.2f}" for b in self.bias))
|
||||
return
|
||||
print(f" moved during calibration (spread {spread:.1f} deg/s), retrying")
|
||||
sys.exit("Could not calibrate - robot kept moving.")
|
||||
|
||||
|
||||
class Madgwick:
|
||||
"""6-axis Madgwick filter (gyro + accel), quaternion q = [w, x, y, z]."""
|
||||
|
||||
def __init__(self, beta):
|
||||
self.beta = beta
|
||||
self.q = [1.0, 0.0, 0.0, 0.0]
|
||||
|
||||
def init_from_accel(self, a):
|
||||
ax, ay, az = a
|
||||
roll = math.atan2(ay, az)
|
||||
pitch = math.atan2(-ax, math.hypot(ay, az))
|
||||
cr, sr = math.cos(roll / 2), math.sin(roll / 2)
|
||||
cp, sp = math.cos(pitch / 2), math.sin(pitch / 2)
|
||||
self.q = [cr * cp, sr * cp, cr * sp, -sr * sp]
|
||||
|
||||
def update(self, g, a, dt):
|
||||
"""g: body gyro in rad/s, a: body accel (any unit), dt in s."""
|
||||
q0, q1, q2, q3 = self.q
|
||||
gx, gy, gz = g
|
||||
ax, ay, az = a
|
||||
# rate of change of the quaternion from the gyro
|
||||
qd0 = 0.5 * (-q1 * gx - q2 * gy - q3 * gz)
|
||||
qd1 = 0.5 * ( q0 * gx + q2 * gz - q3 * gy)
|
||||
qd2 = 0.5 * ( q0 * gy - q1 * gz + q3 * gx)
|
||||
qd3 = 0.5 * ( q0 * gz + q1 * gy - q2 * gx)
|
||||
# accelerometer correction: step toward the orientation that matches gravity
|
||||
n = math.sqrt(ax * ax + ay * ay + az * az)
|
||||
if n > 0:
|
||||
ax /= n; ay /= n; az /= n
|
||||
_2q0, _2q1, _2q2, _2q3 = 2 * q0, 2 * q1, 2 * q2, 2 * q3
|
||||
_4q0, _4q1, _4q2 = 4 * q0, 4 * q1, 4 * q2
|
||||
_8q1, _8q2 = 8 * q1, 8 * q2
|
||||
q0q0, q1q1, q2q2, q3q3 = q0 * q0, q1 * q1, q2 * q2, q3 * q3
|
||||
s0 = _4q0 * q2q2 + _2q2 * ax + _4q0 * q1q1 - _2q1 * ay
|
||||
s1 = (_4q1 * q3q3 - _2q3 * ax + 4 * q0q0 * q1 - _2q0 * ay - _4q1
|
||||
+ _8q1 * q1q1 + _8q1 * q2q2 + _4q1 * az)
|
||||
s2 = (4 * q0q0 * q2 + _2q0 * ax + _4q2 * q3q3 - _2q3 * ay - _4q2
|
||||
+ _8q2 * q1q1 + _8q2 * q2q2 + _4q2 * az)
|
||||
s3 = 4 * q1q1 * q3 - _2q1 * ax + 4 * q2q2 * q3 - _2q2 * ay
|
||||
sn = math.sqrt(s0 * s0 + s1 * s1 + s2 * s2 + s3 * s3)
|
||||
if sn > 0:
|
||||
b = self.beta / sn
|
||||
qd0 -= b * s0; qd1 -= b * s1; qd2 -= b * s2; qd3 -= b * s3
|
||||
q0 += qd0 * dt; q1 += qd1 * dt; q2 += qd2 * dt; q3 += qd3 * dt
|
||||
n = math.sqrt(q0 * q0 + q1 * q1 + q2 * q2 + q3 * q3)
|
||||
self.q = [q0 / n, q1 / n, q2 / n, q3 / n]
|
||||
|
||||
def pitch(self):
|
||||
"""Forward lean in degrees, measured in the plane the robot falls in."""
|
||||
q0, q1, q2, q3 = self.q
|
||||
vx = 2 * (q1 * q3 - q0 * q2) # 'up' direction seen from the body
|
||||
vz = q0 * q0 - q1 * q1 - q2 * q2 + q3 * q3
|
||||
return math.degrees(math.atan2(-vx, vz))
|
||||
|
||||
|
||||
def accel_pitch(a):
|
||||
return math.degrees(math.atan2(-a[0], a[2]))
|
||||
|
||||
|
||||
class Motors:
|
||||
def __init__(self):
|
||||
self.ser = serial.Serial(UART_PORT, UART_BAUD, timeout=0.02)
|
||||
self.lock = threading.Lock()
|
||||
self.totals = None
|
||||
self.last_rx = 0.0
|
||||
self.error = None
|
||||
self.running = True
|
||||
threading.Thread(target=self._reader, daemon=True).start()
|
||||
self.configure()
|
||||
|
||||
def configure(self):
|
||||
"""(Re)send settings. Also used after the board resets from a power dip."""
|
||||
self.ser.write(b"$pwm:0,0,0,0#")
|
||||
self.ser.write(b"$mtype:1#"); time.sleep(0.05) # 1 = Yahboom 520 motor
|
||||
self.ser.write(b"$upload:1,0,0#"); time.sleep(0.05) # stream totals as $MAll
|
||||
|
||||
def pwm(self, left, right):
|
||||
v = [0, 0, 0, 0]
|
||||
v[LEFT_CH] = int(LEFT_DIR * left)
|
||||
v[RIGHT_CH] = int(RIGHT_DIR * right)
|
||||
self.ser.write(f"$pwm:{v[0]},{v[1]},{v[2]},{v[3]}#".encode())
|
||||
|
||||
def wheels(self):
|
||||
"""((left, right) total counts with forward = +, time of last update)."""
|
||||
with self.lock:
|
||||
t, ts = self.totals, self.last_rx
|
||||
if t is None:
|
||||
return None, ts
|
||||
return (ENC_LEFT_DIR * t[LEFT_CH], ENC_RIGHT_DIR * t[RIGHT_CH]), ts
|
||||
|
||||
def _reader(self):
|
||||
buf = b""
|
||||
while self.running:
|
||||
try:
|
||||
buf += self.ser.read(self.ser.in_waiting or 1)
|
||||
except serial.SerialException as e:
|
||||
self.error = str(e)
|
||||
return
|
||||
while b"#" in buf:
|
||||
line, buf = buf.split(b"#", 1)
|
||||
i = line.rfind(b"$")
|
||||
if i < 0:
|
||||
continue
|
||||
msg = line[i + 1:].decode(errors="ignore")
|
||||
if msg.startswith("MAll:"):
|
||||
try:
|
||||
vals = [int(x) for x in msg[5:].split(",")]
|
||||
except ValueError:
|
||||
continue
|
||||
with self.lock:
|
||||
self.totals = vals
|
||||
self.last_rx = time.perf_counter()
|
||||
|
||||
def stop(self):
|
||||
for _ in range(3):
|
||||
try:
|
||||
self.ser.write(b"$pwm:0,0,0,0#")
|
||||
except serial.SerialException:
|
||||
pass
|
||||
time.sleep(0.02)
|
||||
self.running = False
|
||||
self.ser.close()
|
||||
|
||||
|
||||
def go_realtime():
|
||||
try:
|
||||
os.sched_setaffinity(0, {RT_CPU})
|
||||
except OSError as e:
|
||||
print("affinity failed:", e)
|
||||
try:
|
||||
os.sched_setscheduler(0, os.SCHED_FIFO, os.sched_param(50))
|
||||
except PermissionError:
|
||||
print("Not root: no real-time priority, expect more timing jitter.")
|
||||
|
||||
|
||||
def ticker(hz):
|
||||
"""Yields the real dt of each tick while holding the loop rate."""
|
||||
period = 1.0 / hz
|
||||
nxt = last = time.perf_counter()
|
||||
while True:
|
||||
nxt += period
|
||||
delay = nxt - time.perf_counter()
|
||||
if delay > 0:
|
||||
time.sleep(delay)
|
||||
else:
|
||||
nxt = time.perf_counter()
|
||||
now = time.perf_counter()
|
||||
yield now - last
|
||||
last = now
|
||||
|
||||
|
||||
def run_imu(imu, args):
|
||||
f = Madgwick(args.beta)
|
||||
a, _ = imu.read()
|
||||
f.init_from_accel(a)
|
||||
comp = accel_pitch(a)
|
||||
worst, n = 0.0, 0
|
||||
print("\n accel compl. MADGWICK rel.offset rate worst dt")
|
||||
for dt in ticker(LOOP_HZ):
|
||||
a, g = imu.read()
|
||||
f.update((math.radians(g[0]), math.radians(g[1]), math.radians(g[2])), a, dt)
|
||||
acc = accel_pitch(a)
|
||||
comp = 0.98 * (comp + g[1] * dt) + 0.02 * acc
|
||||
mad = f.pitch()
|
||||
worst = max(worst, dt); n += 1
|
||||
if n % (LOOP_HZ // 10) == 0:
|
||||
print(f"{acc:+7.2f} {comp:+7.2f} {mad:+8.2f} {mad - args.offset:+8.2f} "
|
||||
f"{g[1]:+7.1f} {worst * 1000:5.1f} ms", flush=True)
|
||||
worst = 0.0
|
||||
|
||||
|
||||
def run_balance(imu, motors, args):
|
||||
f = Madgwick(args.beta)
|
||||
a, _ = imu.read()
|
||||
f.init_from_accel(a)
|
||||
mm_per_tick = math.pi * WHEEL_D_MM / TICKS_PER_REV
|
||||
|
||||
armed = False
|
||||
out = 0.0 # PWM actually sent, after the slew limit
|
||||
tick = 0
|
||||
target = v_filt = dist = 0.0
|
||||
prev_w, prev_t = None, 0.0
|
||||
last_reinit = 0.0
|
||||
|
||||
print(f"\nkp={args.kp} kd={args.kd} kv={args.kv} ki={args.ki} offset={args.offset} "
|
||||
f"beta={args.beta}\nHold it upright and still to arm. Ctrl-C to quit.\n")
|
||||
|
||||
for dt in ticker(LOOP_HZ):
|
||||
if motors.error:
|
||||
sys.exit(f"Serial error: {motors.error}")
|
||||
|
||||
a, g = imu.read()
|
||||
f.update((math.radians(g[0]), math.radians(g[1]), math.radians(g[2])), a, dt)
|
||||
theta = f.pitch() - args.offset # deg, + = leaning forward
|
||||
rate = g[1] # deg/s, + = tipping forward
|
||||
tick += 1
|
||||
now = time.perf_counter()
|
||||
|
||||
# --- watchdog: encoder data must keep flowing ---
|
||||
w, w_time = motors.wheels()
|
||||
if w is None or now - w_time > ENC_TIMEOUT:
|
||||
if armed:
|
||||
armed = False
|
||||
out = 0.0
|
||||
motors.pwm(0, 0)
|
||||
print("ENCODER DATA LOST - motors off, re-initialising the driver board", flush=True)
|
||||
if now - last_reinit > 0.5:
|
||||
motors.configure()
|
||||
last_reinit = time.perf_counter()
|
||||
continue
|
||||
|
||||
# --- waiting to arm ---
|
||||
if not armed:
|
||||
if tick % 20 == 0:
|
||||
motors.pwm(0, 0)
|
||||
if abs(theta) < START_ANGLE and abs(rate) < 30:
|
||||
armed = True
|
||||
out = 0.0
|
||||
target = v_filt = dist = 0.0
|
||||
prev_w, prev_t = w, now
|
||||
print("ARMED", flush=True)
|
||||
continue
|
||||
|
||||
# --- fall cut-off ---
|
||||
if abs(theta) > FALL_ANGLE:
|
||||
armed = False
|
||||
out = 0.0
|
||||
motors.pwm(0, 0)
|
||||
print(f"FELL ({theta:+.1f} deg) - motors off", flush=True)
|
||||
continue
|
||||
|
||||
# --- outer loop, 50 Hz: wheel speed -> lean target ---
|
||||
if tick % SPEED_DIV == 0:
|
||||
dts = now - prev_t
|
||||
v = ((w[0] - prev_w[0]) + (w[1] - prev_w[1])) * 0.5 * mm_per_tick / dts # mm/s
|
||||
prev_w, prev_t = w, now
|
||||
v_filt = 0.7 * v_filt + 0.3 * v
|
||||
dist = clamp(dist + v_filt * dts, -500.0, 500.0) # mm
|
||||
target = clamp(-(args.kv * v_filt + args.ki * dist), -TARGET_LIMIT, TARGET_LIMIT)
|
||||
|
||||
# --- inner loop, 200 Hz: angle PD ---
|
||||
u = args.kp * (theta - target) + args.kd * rate
|
||||
if u > 0:
|
||||
u += args.deadband
|
||||
elif u < 0:
|
||||
u -= args.deadband
|
||||
u = clamp(u, -PWM_MAX, PWM_MAX)
|
||||
out += clamp(u - out, -PWM_SLEW, PWM_SLEW) # never jump, ramp
|
||||
turn = 0.0 # hook for steering later
|
||||
motors.pwm(out + turn, out - turn)
|
||||
|
||||
if tick % (LOOP_HZ // 10) == 0:
|
||||
print(f"th {theta:+6.2f} tgt {target:+5.2f} rate {rate:+7.1f} "
|
||||
f"v {v_filt:+6.0f} mm/s pwm {out:+6.0f}", flush=True)
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("mode", choices=["imu", "balance"])
|
||||
ap.add_argument("--kp", type=float, default=100.0) # PWM per deg
|
||||
ap.add_argument("--kd", type=float, default=2.5) # PWM per deg/s
|
||||
ap.add_argument("--kv", type=float, default=0.0) # deg of lean per mm/s
|
||||
ap.add_argument("--ki", type=float, default=0.0) # deg of lean per mm
|
||||
ap.add_argument("--offset", type=float, default=ANGLE_OFFSET)
|
||||
ap.add_argument("--beta", type=float, default=BETA)
|
||||
ap.add_argument("--deadband", type=float, default=DEADBAND)
|
||||
args = ap.parse_args()
|
||||
|
||||
go_realtime()
|
||||
gc.collect(); gc.disable()
|
||||
bus = SMBus(I2C_BUS)
|
||||
imu = IMU(bus)
|
||||
imu.calibrate()
|
||||
motors = Motors() if args.mode == "balance" else None
|
||||
try:
|
||||
if args.mode == "imu":
|
||||
run_imu(imu, args)
|
||||
else:
|
||||
run_balance(imu, motors, args)
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
finally:
|
||||
if motors:
|
||||
motors.stop()
|
||||
bus.close()
|
||||
print("\nmotors stopped")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user