"""The bench's drive-assist controller, in its own sandbox. The server keeps ONE live simulation. While the bench page is open it posts to it about eleven times a second, so any experiment run through the API is interleaved with whatever the browser is sending - different stance, different turn, sometimes a different model entirely. Two sweeps were quietly ruined that way before I noticed. So this is the same controller against a private RobotEnv. Nothing here touches the live sim, which means experiments are repeatable and you can carry on driving while they run. python assist_lab.py --sweep roll """ import argparse import math from pathlib import Path import mujoco import numpy as np from env import RobotEnv # ONE implementation. These used to be defined here and copied into serve.py, # where the copy drifted (an inverted foot-shift sign, among others). The # law, the IK and the trim now live in controller.py and everyone imports it; # the re-exports below keep older callers working. from controller import (L_SEG, R_WHEEL, stance_to_leg, # noqa: F401 lean_trim_deg) ROOT = Path(__file__).parent.parent SPEC = str(ROOT / "robots/wheeled_biped.json") class Assist: """kp 12, kd 2, station 1.95 - the gains confirmed by hand in the bench.""" def __init__(self, env, kp=12.0, kd=2.0, station=1.95, lean_max=0.08, k_roll=0.0, k_rolld=0.0, roll_ff=0.0, leg_rate_deg=60.0, k_yaw=0.25, ki_yaw=0.30, yaw_max=1.5, trim=True, reserve=True, v_tau=0.0, duck_k=0.0, duck_start=4.0, duck_floor=300.0, shift_k=0.0, shift_max=0.10, branch=-1.0, branch_band=0.15, duck_rate_deg=60.0): self.e = env self.kp, self.kd, self.station, self.lean_max = kp, kd, station, lean_max self.k_roll, self.k_rolld, self.roll_ff = k_roll, k_rolld, roll_ff self.leg_rate = math.radians(leg_rate_deg) self.legs = [i for i, a in enumerate(env.spec["actuators"]) if a.get("mode") == "position"] self.whl = [i for i, a in enumerate(env.spec["actuators"]) if a.get("mode") == "torque"] self.names = [a["name"] for a in env.spec["actuators"]] for i in self.legs: # full travel, not the policy band env.act_center[i] = 0.5 * (env.ctrl_lo[i] + env.ctrl_hi[i]) env.act_span[i] = 0.5 * (env.ctrl_hi[i] - env.ctrl_lo[i]) self.jd = [env.model.jnt_dofadr[mujoco.mj_name2id( env.model, mujoco.mjtObj.mjOBJ_JOINT, n)] for n in ("l_wheel_j", "r_wheel_j")] # Start the ramp from where the robot ACTUALLY is. env.reset() puts it # in the task's crouch pose, so assuming a straight leg here is a step # change at t=0 that knocks it over before any test begins - which is # what made every run in two sweeps identical. self.leg = np.array([float(env.data.qpos[env.model.jnt_qposadr[ mujoco.mj_name2id(env.model, mujoco.mjtObj.mjOBJ_JOINT, self.names[i][:-2])]]) for i in self.legs]) self.odo = 0.0 self.v_cmd = 0.0 self.v_f = 0.0 self.yaw_i = 0.0 self.k_yaw, self.ki_yaw, self.yaw_max = k_yaw, ki_yaw, yaw_max self.trim = trim self.v_tau = v_tau self.reserve = reserve self.duck_k, self.duck_start = duck_k, duck_start self.duck_floor = duck_floor self.shift_k, self.shift_max = shift_k, shift_max self.branch = branch self.branch_band = branch_band self.duck_rate = math.radians(duck_rate_deg) def act(self, stance_mm=460.0, drive=0.0, turn=0.0, lean_deg=0.0): e = self.e dt = e.control_dt up = e.est_up pitch = math.atan2(-up[0], max(1e-6, up[2])) roll = math.atan2(up[1], max(1e-6, up[2])) # Recovery reflex: when it starts to go, get low and get the feet # under the mass. Lowering shortens the pendulum so there is less # toppling torque for the wheels to fight; sliding the feet toward the # fall puts the contact patch back under the centre of mass, which the # wheels alone can only do by driving (and they may not have the grip # or the room). The legs can do it directly and much faster. dx = 0.0 if self.duck_k > 0.0 or self.shift_k != 0.0: over = max(0.0, abs(math.degrees(pitch)) - self.duck_start) stance_mm = max(self.duck_floor, stance_mm - self.duck_k * over) # Forward only. The knee bends one way, so the leg can slide the # foot AHEAD of the hip easily but fights being dragged behind it. # Measured: allowing both directions turned a recovered backward # shove (7.2 deg peak) into a fall, while the forward case improved # from a fall to a 10.5 deg recovery. So catch forward dives with # the legs and leave backward ones to the wheels, which handle them # perfectly well already. dx = max(-self.shift_max, min(self.shift_max, self.shift_k * pitch)) # Pick the knee branch to suit the way it is going. Both postures # put the foot in the same place, but one of them can drive into # the fall and the other fights it: knee-back recovers a forward # dive (6.7 deg peak vs 24.9 unaided) and drops a backward one, # knee-forward does exactly the reverse. Hysteresis on the sign, # because swapping branch is a big leg movement and doing it every # time the estimate crosses zero would be worse than not doing it. # Choose the posture by DIRECTION OF TRAVEL, not by the fall. # The branch decides which way it can recover, and swapping is a # big slow leg movement - far too slow once it is already going. # But the direction you are driving is known in advance and changes # slowly, and it is also the way you are most likely to lose it. # So commit to the posture that suits where you are heading. if self.v_cmd > self.branch_band: self.branch = -1.0 # heading forward: knee back elif self.v_cmd < -self.branch_band: self.branch = +1.0 # reversing: knee forward g = e.sensors.noisy("gyro", e.data.sensordata, e._raw) rate, roll_rate = float(g[1]), float(g[0]) # roll: the wheels share an axle and can do nothing about it, so the # only handle is making one leg longer than the other dh = self.k_roll * roll + self.k_rolld * roll_rate + self.roll_ff * turn / 100.0 dh = max(-0.06, min(0.06, dh)) h = stance_mm / 1000.0 tgt = [] for side in (+1.0, -1.0): a, b = stance_to_leg(max(0.20, min(0.462, h + side * dh)), dx, self.branch) tgt += [a, b] tgt = np.array(tgt) # one scale factor for the whole move, so hip and knee arrive together d = tgt - self.leg far = float(np.max(np.abs(d))) if far > 1e-9: # NOT "rate" - that name already holds the gyro pitch rate, and # clobbering it turned the balance loop's damping term into a # constant. The robot fell over in one second and the cause was a # reused variable name, not the control design. lim = max(self.leg_rate, self.duck_rate) if self.duck_k > 0 else self.leg_rate self.leg = self.leg + d * (min(far, lim * dt) / far) self.v_cmd += float(np.clip(np.clip(drive / 100.0, -1, 1) * 1.0 - self.v_cmd, -0.6 * dt, 0.6 * dt)) v_raw = 0.5 * sum(float(e.data.qvel[i]) for i in self.jd) * 0.0625 # Low-pass the speed before it drives the lean target. Raw wheel speed # contains the limit cycle itself, so feeding it back unfiltered closes # a positive loop: the lean target saturates at its clamp, flips sign # every step, and the wheels chatter at full torque while the robot is # standing perfectly still - 2.22 Nm mean when holding upright needs # about 0.05. a_lp = dt / max(dt, self.v_tau) self.v_f += a_lp * (v_raw - self.v_f) v = self.v_f self.odo += (v - self.v_cmd) * dt lean = max(-self.lean_max, min(self.lean_max, self.station * ( 0.5 * (v - self.v_cmd) + 0.1 * self.odo))) trim = math.radians(lean_trim_deg(stance_mm)) if self.trim else 0.0 u = max(-1.0, min(1.0, -self.kp * (pitch - lean - math.radians(lean_deg) - trim) + self.kd * rate)) # Turn is a RATE command, closed on the gyro - not a torque # difference. An open-loop difference is a constant yaw acceleration # with nothing to damp it, so the spin rate winds up without limit: # even a 2% turn eventually reached the speed that topples the robot, # it just took 7.7s instead of 1.8s to get there. # Proportional only. An integral term is worse than useless here: with # no turn commanded it still integrates gyro noise, winds up to the # differential limit and spins the robot up on its own - every test run # was already on the floor before the turn even started. The P term is # the part that matters anyway, because it provides the yaw DAMPING # that an open-loop torque difference completely lacks. yaw_rate = float(g[2]) yaw_cmd = float(np.clip(turn / 100.0, -1, 1)) * self.yaw_max # NOTE THE SIGN. A differential of t on (left = u+t, right = u-t) # yaws the robot the OTHER way: t = -0.30 measured +11.1 rad/s. My # first check saw +0.015 rad/s against a gyro noise sd of 0.041 and I # took it as proof - it was noise, and the wrong sign turned this loop # into positive feedback that span the robot up and threw it over. t = (float(np.clip(-self.k_yaw * (yaw_cmd - yaw_rate), -0.30, 0.30)) if self.k_yaw > 0 else float(np.clip(turn / 100.0, -1, 1)) * 0.30) # Balance first, turning with whatever is left. left = u+t and # right = u-t both clip at 1, so when the balancer is asking for a big # forward surge to get under a lean, the turn differential clips one # wheel and the AVERAGE forward push comes out lower than asked for. # The robot then cannot drive forward hard enough to catch itself - it # leans further and carries on. Reserving the headroom means a turn can # only ever use torque the balance loop is not already using. if self.reserve: t = float(np.clip(t, -(1.0 - abs(u)), 1.0 - abs(u))) a = np.zeros(e.act_dim) for j, ai in enumerate(e.act_idx): if ai in self.whl: a[j] = np.clip(u + (t if self.names[ai].startswith("l") else -t), -1, 1) else: k = self.legs.index(ai) a[j] = np.clip((self.leg[k] - e.act_center[ai]) / max(1e-9, e.act_span[ai]), -1, 1) return a, roll, pitch def run(secs=12.0, settle=2.0, stance_mm=460.0, drive=0.0, turn=0.0, seed=0, **kw): env = RobotEnv(SPEC, seed=seed, task="stand", randomise=False) env.reset() c = Assist(env, **kw) for _ in range(int(settle / env.control_dt)): a, _, _ = c.act(stance_mm=stance_mm) env.step(a) peak_roll = peak_pitch = 0.0 for k in range(int(secs / env.control_dt)): a, roll, pitch = c.act(stance_mm=stance_mm, drive=drive, turn=turn) env.step(a) peak_roll = max(peak_roll, abs(math.degrees(roll))) peak_pitch = max(peak_pitch, abs(math.degrees(pitch))) if env.up_z < 0.5: return k * env.control_dt, peak_roll, peak_pitch return None, peak_roll, peak_pitch def main(): ap = argparse.ArgumentParser() ap.add_argument("--sweep", default="roll") ap.add_argument("--turn", type=float, default=40.0) ap.add_argument("--stance", type=float, default=460.0) ap.add_argument("--secs", type=float, default=12.0) args = ap.parse_args() if args.sweep == "baseline": print(f"No roll control at all, {args.secs:.0f}s:\n") print(f" {'stance':>8} {'turn':>6} {'result':>16} {'peak roll':>11}") for st in (460, 420, 380): for tn in (0, 20, 40, 60, 80): t, pr, pp = run(secs=args.secs, stance_mm=st, turn=tn) print(f" {st:6.0f}mm {tn:5.0f}% " f"{('fell %.1fs' % t if t else 'survived'):>16} {pr:9.1f}d") else: print(f"Roll control by leg difference. turn {args.turn:.0f}%, " f"stance {args.stance:.0f}mm, {args.secs:.0f}s\n") print(f" {'k_roll':>8} {'k_rolld':>9} {'result':>16} {'peak roll':>11}") for kr in (0.0, 0.3, 0.6, 1.0, 1.5): for kd in (0.0, 0.1): t, pr, pp = run(secs=args.secs, stance_mm=args.stance, turn=args.turn, k_roll=kr, k_rolld=kd) print(f" {kr:8.2f} {kd:9.2f} " f"{('fell %.1fs' % t if t else 'SURVIVED'):>16} {pr:9.1f}d") if __name__ == "__main__": main()