a lɀj'@sdZddlZddlZdZdZdddZd d Zd d Zdd dZ ddZ d ddZ d!ddZ ddZ GdddZGdddZGdddZdS)"a)The one implementation of Johnny 6's control laws. Five copies of the balance law used to exist - three in serve.py, one in learn_balance.py's episodes, one in assist_lab.py - and they drifted: the live demo negated the foot-shift reflex the trainer had measured and fixed, the demo dropped recover's brace slots, and drive mode re-derived balance with enough small differences to produce a 1.9 Nm limit cycle while standing still. The fix recorded in AUDIT.md and NEXT.md is structural: one module owns the law, and the trainer, the scorer, the live demo and the drive mode all call it. A divergence is now an import error, not a Saturday. Everything here is deliberately plain: readable gains in, one command out, state carried explicitly. These are the functions that would be transcribed into firmware. Ng?g?c Cstd|t}t||}t|dtd}tdtd||dtddtd}|t|}t||ttt|tdt |}||fS)a_Two-link leg: put the axle at height h, and dx ahead of the hip. dx=0 reproduces the simple case - hip +q/2, knee -q, axle directly under the hip. A non-zero dx moves the wheel fore or aft WITHOUT changing height, which is how the robot shifts its mass over the contact patch: sliding the feet forward under a forward dive is the same as throwing the body back. branch picks which way the knee folds. Both postures reach the same foot position, but each can only recover a fall one way, and swapping is a large, slow movement - choose by direction of travel, not mid-fall. g{Gz?g@g-C6?r?) maxR_WHEELmathhypotminL_SEGacosatan2sincos)hdxbranchHrcZkneeZhipr./home/ghrups/robot-bench/trainer/controller.py stance_to_legs  ,rc Cst|}d|dd|dd|d }d}||ddkrP|dd }n||d dkrn|d d }nl|dd }t||d d D]L\\}}\}}||kr|krnq||||} || ||}qq||S) aThe lean the robot must hold at a given stance height to stay put. REFITTED FOR THE R2 BODY (13 Aug): the r1 polynomial was not just stale, it was BACKWARDS - at a 380 crouch it commanded -6.1 degrees where the r2 body truly needs +7.0, opposite in sign at every height, because the knee-mounted drive motors flipped the crouch CoM geometry. That 13-degree error is why bob wandered (a wrong trim can only be "balanced" by accelerating), why gate changes rocked, and why look measured +50mm of feet-forward compensation (~6.6 degrees - the same error, absorbed geometrically). Solved kinematically on the r2 model: pose the body at each stance, find the pitch that puts the total CoM over the axle. g1] ʾg=  [m?rga?g4Op1Ui@))g@u@g)gy@gܿ)g|@gzG?rN)floatzip) stance_mmsZkin_fixrs0Zc0s1c1frrr lean_trim_deg1s$   "r%cCs*tt| tt||d|d|S)aThe lean target: speed and odometry pick a small lean, clamped. The clamp is load-bearing, not incidental - flattening the law into a plain weighted sum was measured at 0.33, worse than all-zeros, because a 0.66 m/s roll then asks for a 19 degree lean. g?皙?)rabsr )stationspeedodoclamprrrlean_reference\sr,cCs(tt| |||||ddS)zmThe inner balance law. Returns a normalised command in [-1, 1]; the actuator maps 1.0 to its full 1.4 Nm.rr)rnpclip)kpkdleanratereftrimrrr wheel_commandgsr5cCsB||}ttt|}|dkr>||t||||}|S)zMove all leg joints TOGETHER: one scale factor for the whole move so hip and knee arrive in ratio and the axle stays under the hip.g& .>)rr-rr'r )curwantr2dtdfarrrrslewms r;r&cCs6t|t|}|dur(||kr(tdStd|S)aLeg slew rate in rad/s from a stage's leg-speed gain (sign-blind, floored - a negative half-line was an inescapable dead zone that froze the legs). Stages that may use the servos flat out (level, on the BNO reference) pass unlimited_at: above it the ramp opens to servo speed. Base is 150 deg/s per unit. The old 60 was the accidental-stabiliser era: the accel-polluted roll reference punished fast legs for their own motion, and the crawl was load-bearing. On the BNO reference with the r2 peak-torque legs it is only a brake - the searches now choose their speed on a scale where 1.0 means genuinely quick.Ngj@gb@)rr'rradians)Z leg_speedfloor unlimited_atZlsrrrleg_ratews  r?cCsVtd|t}|t||t|}|t||t|}t|t||S)aLeg IK for a PITCHED chassis: hips and knees move together so the wheel stays planted at dx_w (world) while the chassis rotates by phi. The naive version - adding the gaze to the hip joint alone - swings the whole leg and translates the wheel (the leg is the long lever, the chassis the heavy part); measured at 5m of drift. Rotating the leg TARGET into the pitched chassis frame is the coordinated pose the operator specified: gaze through hips AND knees, wheels balancing only. 皙?)rrrrrr)Zstance_mZdx_wphirZHhZdx_cZHh_crrr gaze_poses rBcCs tt||t| t|S)aFeet slide TOWARD the fall - dx has the SAME sign as lean. Measured in the assist lab: toward the fall cuts a 2 m/s shove from 24.9 deg peak lean to 6.7; the inverted sign (feet fleeing) closed an unstable pitch loop that fell in 5s unshoved. That inversion has now happened twice, once in the trainer and once in serve.py's copies - which is why the sign lives here and nowhere else.)rr-r.r')shift_kr1 shift_maxrrr foot_shiftsrEc@s(eZdZdZdZdZddZddZdS) ShoveBracea\Detect a shove and lean into it, fading afterwards. Fires on acceleration the robot did not command for itself: shove_felt = filtered_accel - (commanded_speed_change / dt) Raw wheel-speed derivative carries the balance loop's own activity (measured 17 m/s^2 peak just standing, against a threshold of 3 - the brace became a random 12 degree injection). Through a 0.10s low-pass the same standing noise peaks at 3.7 while a real 2.5 m/s shove reads 10.5, so 5.0 separates cleanly. A driving stage must feed its own v_cmd or it will brace against its own throttle. r&g@cCs(||_||_d|_d|_d|_d|_dSNr)gainfadev_filt prev_speed prev_v_cmdbrace)selfrHrIrrr__init__s zShoveBrace.__init__c Cs|j|t||j||j7_|j|j|}|j|_||j|}||_||}t||jkrr|j||_|jt | tdt|j 9_|jS)Nr@) rJrTAUrKrLr' THRESHOLDrHrMrexprI)rNr)v_cmdr8accelexpectedZfeltrrrupdates$ &zShoveBrace.updateN)__name__ __module__ __qualname____doc__rPrQrOrVrrrrrFs  rFc@s"eZdZdZdddZddZdS) SpeedFilteraLow-pass on wheel 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 stands perfectly still - measured at 2.22 Nm mean where holding still needs about 0.05. tau=0 disables (raw).r&cCs||_d|_dSrG)tauvrNr\rrrrOszSpeedFilter.__init__cCs.|t||j}|j|||j7_|jS)N)rr\r])rNv_rawr8arrrrVszSpeedFilter.updateN)r&rWrXrYrZrOrVrrrrr[s r[c@s"eZdZdZdddZddZdS) ComplementaryTiltaGyro-fast tilt with slow accel correction, for the REAL robot. In the sim the fused estimator (est_up) already plays the BNO086 and is the reference every stage controls and scores on; this class is the firmware-portable equivalent for hardware without a fused IMU. The gyro carries the fast path (immune to the linear acceleration the legs inject - the loop that made accel-only references punish fast legs); the accelerometer corrects drift with time constant tau.rcCs||_d|_dSrGr\angler^rrrrOszComplementaryTilt.__init__cCs6|j|j|}||j||d|||_|jS)Nrrc)rNZ gyro_rateZ accel_angler8r`rrrrVs zComplementaryTilt.updateN)rrarrrrrbs rb)rr)rr)r&N)r)rZrnumpyr-r rrr%r,r5r;r?rBrErFr[rbrrrrs +    %