"""Learn to balance, on the four numbers the problem actually has.

Everything measured today says the search space, not the algorithm, is what
decided whether this worked:

  a 6594-parameter network fitted the teacher 17x more closely than a linear
  rule and then fell over in 1.95 seconds, while the linear rule balanced 10/10

  random guesses at just two gains balanced 28 percent of the time, and the
  first success came 2.6 seconds into the search

So this hands the policy only what a balancing robot needs to know:

    lean          how far it is tipped
    lean rate     how fast that is changing
    speed         how fast it is rolling
    odometry      how far it has strayed from where it started

and asks for one number back, applied to both wheels. Five parameters with the
bias. They are readable gains, not weights - you can print them, reason about
them, and type them into the robot's firmware.

The score is upright time AND staying put, because "stay upright" on its own is
satisfied perfectly well by rolling away at a constant speed, which is exactly
what this robot did for 112 metres before anyone noticed.

  python learn_balance.py --iters 40
"""

import argparse
import json
import math
import os
from pathlib import Path

import mujoco
import numpy as np

import controller as C
import firmware as FW
from env import RobotEnv

ROOT = Path(__file__).parent.parent
SPEC = str(ROOT / "robots/wheeled_biped.json")

# Effort is charged in EVERY stage score, at the same rate. Without it the
# search buys stability with the whole torque budget: the v19 balance bake
# scored 4.90 by railing the wheels at +-1.4 Nm for 66% of every episode - a
# 100 Hz bang-bang limit cycle that would cook a real driver and left no
# torque headroom for braking, turning or bracing. Charging mean|u| at this
# rate re-ranks quiet control above chatter while costing a good law almost
# nothing: the re-searched quiet gains score 4.88 on the old objective and
# hold still at 0.058 Nm, 0% railed. Rate is per second of full command, so
# episode length does not change the price.
EFFORT_RATE = 0.24
# The four numbers the working controller actually has. NOT five free linear
# gains: the real law has a CLAMP inside it, and the clamp is load-bearing.
#
#     lean_ref = clip( station * (0.5*speed + 0.1*odo),  +-clamp )
#     u        = -kp * (lean - lean_ref) + kd * rate
#
# Flattening that to a plain weighted sum drops the clamp, and then a 0.66 m/s
# roll asks for a 19 degree lean - four times what the robot survives. The
# hand-tuned controller scored 0.33 in that form, worse than all-zeros, because
# the form could not represent it. Search the structure, not a linear blur of
# it: these are gains you can read, argue with, and type into firmware.
SCORE_VERSION = 102  # v102: GET IN POSITION FAST (operator: "the
                    # robot rolls for a while before jumping - time
                    # to get into position needs to be penalised").
                    # The planted-feet gate made patience free, so
                    # patience is what the law learned. Now every
                    # second from episode start to liftoff beyond 1.0
                    # costs 1.5 points, and the dive no longer sits
                    # out a hardcoded 0.4s of dead air before
                    # starting - the entry brake and the crouch have
                    # to happen AT SPEED, together.
                    # v101: PLANT THE FEET BEFORE FIRING. Measured:
                    # liftoff vx tracks the messy entry's drift, and
                    # the roll-out IS that momentum coming back down -
                    # seed 7 fired drifting backward and rolled 555mm
                    # back, seed 3 fired drifting forward and rolled
                    # 496mm forward. Flight conserves what the fire
                    # gate lets through, and no in-air trick can shed
                    # linear momentum without traction (the fire-lean
                    # sweep measured WORSE at every deeper angle). The
                    # gate now also requires GROUND stillness (|speed|
                    # < 0.1 m/s) and its impatience timeout grows 0.6s
                    # -> 2.0s: the jump takes the beat it needs to
                    # stop rolling first, like any jumper planting
                    # their feet.
                    # v100: ARRIVE FORWARD, STAND TEN SECONDS
                    # (operator: "get it to lean forward into the
                    # jump for a stable landing that persists for 10
                    # seconds"). Slot 32 land lean: the in-flight
                    # attitude reference is no longer level - the
                    # wheels' air-attitude loop steers toward a
                    # searchable forward lean target, so the robot
                    # ARRIVES already leaning into the catch instead
                    # of level with backward momentum. And the episode
                    # no longer ends 3s after touchdown: it holds for
                    # a full 10 seconds, and "landed" means still
                    # standing at the end of them.
                    # v99: THE EXAM ENTERS MESSY + THE CATCH GETS LEGS
                    # (operator: "why is this auto baking when it's a
                    # fail?" - because it wasn't one: every exam
                    # episode entered from a perfect settled 455mm
                    # stance at zero speed, an entry the live bench
                    # never grants, so the exam literally could not
                    # see the fall the operator watched). Episodes now
                    # enter from a seeded random stance (400-460mm),
                    # pitch (+-3 deg) and rolling speed (+-0.25 m/s,
                    # wheels matching). And the catch stops being
                    # leg-frozen (operator: "why are we not opening
                    # the joints up on landing?"): slots 29-31 give
                    # the absorb a hip offset, a hip-answers-lean term
                    # and a knee-answers-lean term for the 3s window -
                    # the legs bend INTO the fall instead of riding
                    # the height gate.
                    # v98: THE BACKWARD LEAN IS THE FAIL (operator:
                    # "penalise lean back when we land as a fail, and
                    # make adjustments to roll forwards if rolling
                    # back - 3 second limit after touchdown"). The
                    # audit is now SIGNED: within 3s of touchdown a
                    # backward (nose-up) lean past 20 deg fails the
                    # episode outright, and every degree past 8 is
                    # charged on the way there, so the search feels
                    # the cliff before it falls off it. Slot 28 roll
                    # fix: a wheel term on measured ground speed for
                    # those same 3s - rolling backward becomes a
                    # command to drive back under the mass.
                    # v97: STICK THE LANDING (operator: "we're not
                    # scoring a stuck landing high, and a backwards
                    # landing is not treated as a fail"). The landed
                    # check only looked at the episode's LAST frame -
                    # a landing that dipped 60 deg backward and rolled
                    # half a metre before staggering up scored full
                    # marks. Now: the tilt audit runs through the
                    # catch, going past 45 deg after touchdown IS a
                    # fail (-1, same as any fall), and the roll-out
                    # after touchdown is charged (3.0/m past 10cm) -
                    # a stuck landing stops where it lands.
                    # v96: THE TUCK BALANCE (operator, watching v95 fly:
                    # "liftoff is perfect, all the force goes straight
                    # up - the backwards rotation starts when the shins
                    # pull up for the tuck; the hip needs to pull up as
                    # the shins pull up so they cancel"). Measured: the
                    # flight's +2.1 rad/s kick lands exactly on the
                    # tuck beat. Slot 25 ties the hip command to the
                    # shin's MEASURED fold rate through tuck and brace,
                    # so the two swings cancel inside the body. Slot 26
                    # land catch: a fading wheel-torque reflex 0.4s
                    # after touchdown. Slot 27 air spin: the wheels as
                    # reaction wheels in flight, torque tied to the
                    # shin fold rate - traction-free counter-rotation
                    # that fires on tuck and again on release, for the
                    # landing that LASTS (operator's brief).
                    # v95: rotation is charged ABSOLUTELY (operator,
                    # watching the v94 bake: "the feet come up, but the
                    # chassis rotates over backwards - incentivise not
                    # rotating at all"). The old leak charge was a
                    # FRACTION capped at 3 points, so a 6.4 rad/s
                    # backflip was affordable; and clearance is read at
                    # the wheels, so a backflip KICKS THE FEET UP and
                    # gets paid for altitude the chassis never earned.
                    # Now: every rad/s of in-flight pitch rate beyond
                    # 0.5 costs 1.2, every radian of true chassis tilt
                    # beyond 0.30 costs 8, and clearance pays 7.5/unit
                    # (was 5) so the honest route to points is a tall
                    # FLAT jump.
                    # v94: THE JUMP grows up (the M4 pays for it - one
                    # episode is 0.05s here, so the exam can afford to
                    # be harsh). Eight seeds instead of three, scored
                    # worst-half as before: the v93 bake landed its 3
                    # exam seeds and fell on 2 of the next 5 - "lands
                    # every time" now means every time. Run-out priced
                    # (1.5/m past 0.5m: a 2m landing slide scored 7.4).
                    # Five launch slots join the search - fire rate,
                    # launch top, fire gate, brace at, fire steer -
                    # all seeded 0 = the old hardcoded constants,
                    # because the fire is where the 43-52% rotation
                    # leak is born and the search could not touch it:
                    # the timing was frozen and the wheel drive flew
                    # the thrust open-loop, deaf to the gyro.
                    # v93: even-wheels price doubled again (1.2/unit,
                    # ~8.5 points for a held pivot - fatal, not a tax).
                    # v92: spin rate commands RAMP in (3 rad/s^2) -
                    # a stepped rate kicks a translation that smears
                    # into an elliptical orbit.
                    # v91: the bank sleeps at standstill (fades in
                    # with spin rate) - no bent leg when no turn is
                    # commanded.
                    # v90: the skater (slot 11) - working height
                    # rises with spin rate; max safe speed at every
                    # height becomes a built-in law.
                    # v89: even-wheels charge tripled (0.2 -> 0.6)
                    # - a held pivot must cost more than it earns;
                    # transient corrections stay cheap.
                    # v88: LEVEL is charged - roll beyond 5 deg
                    # bleeds like pitch wobble (the folded-leg grind
                    # was free until it fell).
                    # v87: direction-aware bank feedforward (slot
                    # 10) + the even-wheels charge (equal-magnitude
                    # counter-rotation or points deducted).
                    # v86: on-the-spot is the PATH - the mean
                    # displacement integral prices the one-wheel
                    # orbit the endpoint check let ride.
                    # v85: cent trim (slot 9) - equilibrium lean fed
                    # forward with yaw-rate squared; sustained-hold
                    # exam cases (the throttle's reality).
                    # v84: the reversal case - direction flips mid-
                    # episode (the demo's 270-degree sweep the exam
                    # never graded).
                    # v83: spin banks into itself - level-style leg
                    # roll answer (slots 7/8, sign measured), exam
                    # graded at 460/400/340 (operator: spin at
                    # various gate heights).
                    # v82: drift-null trim correction in the base
                    # curve (-0.75 deg at 340 -> +0.21 at 460) -
                    # affects every trim-consuming law; benches
                    # stop paying the phantom-drift charge.
                    # v81: payload exam carries 2kg ON TOP (CoM
                    # raised 130mm) - the operator's requirement
                    # for a robot this size.
                    # v80: level exam strengthened - flatness charge
                    # 0.9 -> 2.5, hold penalty 0.7 -> 1.2 (rigidity
                    # was cheaper than levelling; the baked law kept
                    # a 10mm leg cap where 7 deg needs 23mm).
                    # v79: tuck clearance cap raised to 192mm - the
                    # wheels-up height keeps paying (landed only).
                    # v78: a fall is a COMPLETE FAIL (-1 flat, no
                    # partial pay) and half the score is the WORST
                    # episode - landing every time is the product.
                    # v77: reseeded to the measured optimum (pike
                    # -20 + knee cancel -0.6 = 4.02); hip overdrive
                    # dropped (it only paid in the deep-fold show).
                    # v76: hip at FULL TORQUE in the tuck (command
                    # leads the joint's lag); knee cancel parked at
                    # zero on the operator's order.
                    # v75: knee cancel (slot 19) - the knee answers
                    # the gyro in flight, the shin-and-wheel mass as
                    # reaction wheel (measured 4.02); tuck hip
                    # deepened to -1.5 on the operator's order.
                    # v74: the tuck folds the HIP as well (slot 18,
                    # -20 deg seed) - the operator's fix that brings
                    # the full-power tuck to parity with v72 while
                    # flying cleaner.
                    # v73 EXPERIMENT (operator's): tuck fires at
                    # FULL POWER the instant contact is lost; brace
                    # by time-to-impact. Working v72 law tagged
                    # jump-v72-working + backed up on the Desktop.
                    # v72: the tuck moment is a searched slot (its
                    # reaction kick is the counter-rotation pulse -
                    # operator's insight).
                    # v71: chord tilt - the wheel->hip virtual leg
                    # holds a constant slight forward tilt through
                    # the launch (operator's triangle spec).
                    # v70: three-beat flight (operator's design):
                    # stretch -> tuck on REACHING extension (ground
                    # clearance) -> brace at the absorb stance while
                    # still falling, landing on travel not stops.
                    # v69: fire lean -12 deg (push perpendicular
                    # into the ground) + tuck rate as slot 15.
                    # v68: knee boost - the knee command leads its
                    # own measured lag (backdrive compensation on
                    # the knee side; hip side measured flat).
                    # v67: the hip gate reads the ACTUAL knee - the
                    # command sweep races the servo by ~100ms, so
                    # command-time gating released the hip before
                    # the real knee had moved (operator saw it).
                    # v66: hip hold - the hip stays pinned until a
                    # searched point in the sweep (70% seed), then
                    # releases (operator: it flipped back too early).
                    # v65: com set - feet planted aft so the CoM
                    # rides just ahead of the wheels (slot 12).
                    # v64: no pay for backflips - clearance only
                    # pays in full on a LANDED jump, and rotation
                    # is charged at its in-flight peak (the exam
                    # was scoring crashes 7.2; the operator caught
                    # it watching the demo flip).
                    # v63: the hip ratio - hip pinned at its true
                    # crouch angle, released at a fraction of the
                    # knee's pace (operator's experiment, peak at
                    # 0.45); kills the 12-deg fire snap.
                    # v62: knee lead - the knee opens ahead of the
                    # hip on the launch path (operator's diagnosis:
                    # the hip rotating first rolls the chassis
                    # back). Measured near-equivalent to the hip
                    # tilt; both levers kept for the search.
                    # v61: fire lean - the last beat ramps INTO a
                    # deliberate forward lean (6 deg nose-down at
                    # seed) so the jump spends its rotation budget
                    # getting back level instead of backflipping.
                    # v60: the fire is wheels-drive + full-power
                    # extension (operator's spec) - locked wheels
                    # were the backward topple. Feelings layer in
                    # (env.feel + FEEL_WEIGHTS, zero-weight).
                    # v59: launch damp - the hip corrects the felt
                    # pitch rate DURING the thrust (joint feedback,
                    # gain learned from the leak charge).
                    # v58: the leak is charged as an ENERGY FRACTION
                    # at liftoff (up vs fore-aft vs rotation), max 3
                    # points - a perfect jump loses nothing.
                    # v57: liftoff energy audit - rotation and
                    # fore-aft speed at liftoff are charged, and a
                    # thrust-tilt slot lets the search aim the push.
                    # v56: the jump is up-only, launched along the\n                    # measured tether trajectory.\n                    # v55: the tuck waits for the apex (extension
                    # held through the whole rise, legs in at the
                    # top). v54: the jump - crouch, fire, tuck,
                    # absorb, in
                    # three commanded modes (up / forward / back) with
                    # sustained-torque-saturation charged threshold-form.
                    # v53: the shutdown is bob's law verbatim in the
                    # MIRROR fold (operator's design: the gate in
                    # reverse, the knee going the other way) with the
                    # drift-null-measured mirror trim curve, descending
                    # until grounded, settling into the heap, and
                    # actually CUTTING POWER in the exam.
                    # (v52: # v52: the SHUTDOWN joins the firmware - sit is the
                    # rise run backwards (fold onto the rollers, un-tuck
                    # to the lay, settle into the heap), rewarded by the
                    # mirrored height integral, gently.
                    # (v51: # v51: back to v42 - the operator ran v43 (the
                    # asymmetric charge) and it did not find the stand;
                    # verdict recorded, v42's end-displacement charge is
                    # the working form.
                    # (v50: # v50: v43 installed at the operator's call - the
                    # asymmetric travel charge (max forward excursion
                    # billed, a metre of roll-back free), the truest form
                    # of the roll-back-for-a-straighter-lift intuition;
                    # one clean step from v42, nothing else changed.
                    # (v49: # v49: v42 restored at the operator's call - the
                    # height reward PLUS the forward-travel charge
                    # (1.5/m on the learned ascent), the combination
                    # whose search had recovered to 6.31 before the
                    # later tweaks buried it.
                    # (v48: # v48: RESTORATION at the operator's order - the
                    # rise exam is v41 verbatim, the version confirmed
                    # standing: pure time-averaged chassis height, fixed
                    # launch pose (53/-141), simple learned ascent with a
                    # height-gated balance handover. The five stacked
                    # refinements that crushed the stand to a 2.5 stall
                    # live on in git; any that return do so one at a
                    # time, measured.
                    # (v47: # v47: the catch is PITCH-triggered - balance takes
                    # the wheels as the chassis approaches level and
                    # arrests the rotation (the height gate watched the
                    # body blast through level onto its back).
                    # (v46: # v46: thrust-on-rotation - the trigger is the
                    # chassis STARTING TO ROTATE BACKWARDS (nose-up rate
                    # threshold, latched), not any angle; the operator's
                    # ride-the-rock physics.
                    # (v45: # v45: the thrust gets a TIMING lever - drive rocks
                    # the body with the fold held, the lift fires only
                    # past the searched nose-angle trigger (the operator
                    # watched premature thrusts face-plant).
                    # (v44: # v44: the launch pose joins the search - the
                    # deepest knee fold that can still lift, right at the
                    # backward-topple frontier (operator's directive).
                    # (v43: # v43: the travel charge goes ASYMMETRIC - forward
                    # excursion billed hard, a metre of roll-back free -
                    # so the search can discover the operator's
                    # rock-back-then-lift-straight strategy.
                    # (v42: # v42: the stand needs no run-up - forward travel
                    # during the learned ascent is charged (1.5/m), so
                    # height is earned steeply, in place.
                    # (v41: # v41: rise overhauled to the operator's directive -
                    # scripted choreography only to the LAUNCH pose
                    # (53/-141, the snapshot), then a LEARNED ascent
                    # (drive + taper + lift + a discovered handover
                    # height) rewarded purely by time-averaged CHASSIS
                    # HEIGHT ABOVE THE GROUND.
                    # (v40: # v40: the rise is the operator's COMPLETE verified
                    # sequence - lay, tuck, rock-over (wheels driving
                    # forward, knees extending to the straight leg), and
                    # the bench catch of the tallest pendulum; held 10.6s
                    # from a dead heap at the hand-found seeds.
                    # (v39: # v39: rise IS the operator's hand-discovered
                    # sequence - lay out flat (hip 48, knee 26), tuck the
                    # knees until the CoM topples the body into the squat,
                    # push up the instant the chassis approaches level;
                    # the search polishes the numbers, never the topology.
                    # (v38: # v38: the stand is tall and ABSOLUTE (455mm, the
                    # singularity finish) - the searched-stance loophole
                    # banked stand bonuses from a half-squat leg-snap.
                    # (v37: # v37: rise pays for TRYING - standing is the prize
                    # (height progress banks even through a topple, CoM-
                    # over-wheels is charged, the hold is the bonus) and
                    # upright-seconds pay nothing, because a slumped
                    # survivor was outscoring every bold attempt; wheels
                    # gain the tuck phase (signed, searched).
                    # (v36: # v36: rise starts DEAD - all motors off, the body
                    # collapses to its true heap, then wakes and stands on
                    # bench-interpolated gains at the measured hip height.
                    # (v35: # v35: the chatter meter (threshold form) - torque
                    # slew above 0.15/step is billed, so bang-bang laws
                    # pay while recoveries ride free; found via the kd
                    # noise-amplification limit cycle at the crouch.
                    # (v34: # v34: bob gains the realtime CoM hold - the squat
                    # holds the axle under the MEASURED mass through the
                    # journey, not just at the endpoints (4-slot layout).
                    # (v33: # v33: spin is learned STANDING (90, 180, 360 both
                    # ways at 460) - crouched spinning is composition via
                    # the gate slider, not a harder exam.
                    # (v32: # v32: spin joins the bench era - wheel gains
                    # interpolated at the commanded stance, kinematic trim
                    # direct, odometry out of its wheels, layout 10 -> 5.
                    # (v31: # v31: stop is aggressive balance WITH LEGS - wheel
                    # gains ride the benches at the live stance, a searched
                    # brake-crouch drops the CoM to brake harder, and the
                    # exam adds the operator's push case (shoved while
                    # standing, stop the uncommanded speed).
                    # (v30: # v30: look becomes ROTATE - full-circle chassis
                    # rotation about the freed hips (+/-360 hard-stopped),
                    # wheels on bench gains with wrapped lean error, legs
                    # holding stance with the realtime CoM hold, judged on
                    # gyro-integrated true rotation at all three stances.
                    # (v29: settled-gaze leg twitch is charged in look -
                    # wheel effort was priced but leg effort was FREE, so
                    # candidates could spasm their legs at no cost.
                    # (v28: look gains the realtime CoM-hold slot - legs
                    # walk the feet continuously to keep the axle under the
                    # measured total CoM through gaze transitions (sign
                    # measured negative; travel halved at 28 deg/s).
                    # (v27: look rebuilt on the benches - wheel gains
                    # interpolated at the gaze-commanded stance, no odometry
                    # in its wheels, exam symmetric to +/-90 degrees (sky to
                    # ground), 9-slot layout.
                    # (v26: the bench LAW loses its position memory - the
                    # odometry term in the lean reference leaned the robot
                    # home after every push, and station scaled it together
                    # with the essential speed damping so the search could
                    # never drop one without the other; benches now balance
                    # on lean and speed alone.
                    # (v25: settle meter is physical stillness, no
                    # return-to-origin charges. v24: moving entries.
                    # v23: settle meter. v22: three benches; bob
                    # interpolates wheel gains between the bakes)
NAMES = ["kp (lean)", "kd (lean rate)", "station", "lean clamp"]

# THE FEELINGS AS INCENTIVES (operator's design): a kid learns to stand
# by feeling its balance - it learns the legs first, then fine-tunes
# against the feel of pressure and lurch. The legs are learned; this is
# the fine-tuning layer. Every episode of every stage meters the
# vestibular sense in env.feel (world-frame g from the robot's own noisy
# IMU - see env._feel_update):
#   up_pk  peak vertical thrust (g)      fa_i/fa_pk    fore-aft lurch
#   side_i/side_pk  lateral wobble       jerk_i        roughness
#   rate_pk  peak rotation rate (rad/s)
# A stage prices any of them by adding weights here. Weights default to
# ABSENT = zero charge, so existing bakes stay valid; turning one on
# changes that stage's exam - bump SCORE_VERSION and expect to relearn
# that stage.
FEEL_WEIGHTS = {
    # "bob": {"jerk_i": 0.02, "fa_pk": 0.3},     # example shape
}


def feel_charge(env, stage):
    """The stage's feelings bill: weighted vestibular metrics."""
    w = FEEL_WEIGHTS.get(stage)
    if not w:
        return 0.0
    f = getattr(env, "feel", None) or {}
    return float(sum(wt * f.get(k, 0.0) for k, wt in w.items()))


class Sensed:
    """The four numbers, read from the robot's own noisy estimate.

    Speed is the TRUE rolling speed, not the raw wheel encoder: the encoder
    measures wheel-vs-shank, so when the shin rotates during a squat the
    raw reading includes the leg's own motion - the robot thinks it is
    rolling while it is kneeling (operator-diagnosed). All the axes are
    parallel, so the shank's absolute rotation is chassis pitch rate plus
    hip rate plus knee rate, and the compensation is a plain sum - the
    firmware does the same subtraction from its own encoders."""

    def __init__(self, env):
        self.e = env
        def dof(n):
            return env.model.jnt_dofadr[mujoco.mj_name2id(
                env.model, mujoco.mjtObj.mjOBJ_JOINT, n)]
        self.jd = [dof("l_wheel_j"), dof("r_wheel_j")]
        self.legd = [(dof("l_hip"), dof("l_knee")),
                     (dof("r_hip"), dof("r_knee"))]
        self.odo = 0.0

    def read(self):
        e = self.e
        up = e.est_up                       # estimator, never ground truth
        # FULL-QUADRANT lean: the old max(1e-6, up[2]) clamp saturated the
        # reading at +/-90, which made the far side of the rotation circle
        # INVISIBLE - the wheels were told to hold -170 while their sensor
        # swore the chassis was stuck at -90 (the operator caught it as
        # "the reverse kinematics look wrong past the centre"). atan2
        # handles every quadrant; other stages end their episodes by 60
        # degrees of pitch, so they never see the difference.
        lean = math.atan2(-up[0], up[2])
        rate = float(e.sensors.noisy("gyro", e.data.sensordata, e._raw)[1])
        w = 0.0
        # UNIT ratios, validated by the operator's static test: body pinned,
        # wheels planted, slow crouch cycle - the wheel encoder counter-
        # rotates the shank's motion at exactly -1.00 (measured -0.87 to
        # -1.15 across the range, -1.00 at slow speed), so true rolling =
        # encoder + hip + knee + pitch, coefficients of one. A dynamic
        # least-squares fit (9.3/5.1/7.4) predicted better on its own squat
        # trajectory but is a dynamics proxy, not kinematics - rejected for
        # the law in favour of the statically-proven identity.
        for (wj, (hj, kj)) in zip(self.jd, self.legd):
            w += (float(e.data.qvel[wj]) + float(e.data.qvel[hj])
                  + float(e.data.qvel[kj]) + rate)
        speed = 0.5 * w * 0.0625
        self.odo += speed * e.control_dt
        return lean, rate, speed, self.odo

    def command(self, g):
        """The control law, with its clamp, driven by four gains."""
        lean, rate, speed, odo = self.read()
        kp, kd, station, clamp = g
        ref = C.lean_reference(station, speed, odo, clamp)
        return C.wheel_command(kp, kd, lean, rate, ref)


def episode(g, seed, secs=10.0, lean0=0.0, shove=0.0, slope_deg=0.0,
            slope_step_deg=0.0, stance_mm=460.0, v0=0.0, pushes=None):
    """Balance at ONE held stance. Three benches share this episode -
    straight legs (460), half crouch (400), full crouch (340) - because one
    gain set does not fit every height: the crouch moves the CoM and the
    trim, and per-height searches give bob three proven points to
    interpolate between instead of one compromise. Legs are servo-held at
    the stance (the real robot's position loops, not a weld) and the
    kinematic trim for that stance is in the law, scale fixed at 1.0."""
    env = RobotEnv(SPEC, seed=seed, task="bob", randomise=False)
    env.reset()
    legs = [i for i, a in enumerate(env.spec["actuators"])
            if a.get("mode") == "position"]
    for i in legs:
        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])
    hip0, knee0 = C.stance_to_leg(stance_mm / 1000.0)
    pose = np.array([hip0, knee0, hip0, knee0])
    # spawn AT the stance, not falling into it: joints posed, base lowered
    for nm, v in zip(("l_hip", "l_knee", "r_hip", "r_knee"), pose):
        env.data.qpos[env.model.jnt_qposadr[mujoco.mj_name2id(
            env.model, mujoco.mjtObj.mjOBJ_JOINT, nm)]] = float(v)
    env.data.qpos[2] = stance_mm / 1000.0 + 0.0025
    if slope_deg:
        # A gradient, as a tilt of gravity - equivalent for a body on a plane,
        # and it needs no ramp geometry. Holding station on a slope means
        # holding a PERMANENT lean, which is a genuinely different problem from
        # balancing on the flat: the lean clamp has to be wide enough to allow it.
        th = math.radians(slope_deg)
        env.model.opt.gravity[:] = [9.81 * math.sin(th), 0.0, -9.81 * math.cos(th)]
    if lean0:
        env.data.qpos[3:7] = [math.cos(lean0 / 2), 0.0, math.sin(lean0 / 2), 0.0]
    if v0:
        # the robot ARRIVES moving - wheels spinning to match, a rolling
        # entry rather than a skid. The operator's point: an episode that
        # starts perfectly still never examines the damper from t=0, so a
        # law could score well having never stopped a swing it didn't
        # cause itself. Every case now opens with something to kill.
        env.data.qvel[0] = v0
        for nm in ("l_wheel_j", "r_wheel_j"):
            env.data.qvel[env.model.jnt_dofadr[mujoco.mj_name2id(
                env.model, mujoco.mjtObj.mjOBJ_JOINT, nm)]] = v0 / 0.0625
    mujoco.mj_forward(env.model, env.data)
    s = Sensed(env)
    rng = np.random.default_rng(seed)
    n = int(secs / env.control_dt)
    hits = set(rng.integers(10, max(11, n - 10), size=2)) if shove else set()
    up_s = 0.0
    effort = 0.0
    settle_i = 0.0
    chat = 0.0
    prev_u = 0.0
    kp, kd, station, clamp = g
    trim = math.radians(C.lean_trim_deg(stance_mm))
    # the operator's own hands as an exam case: pushes recorded from the
    # 3D view (mouse-drag forces on the live robot) are replayed here as
    # scheduled chassis forces, spaced through the episode
    push_bid = None
    push_gap = 0.0
    if pushes:
        push_bid = mujoco.mj_name2id(env.model, mujoco.mjtObj.mjOBJ_BODY,
                                     "chassis")
        push_gap = max(0.8, (secs - 1.0) / len(pushes))
    for k in range(n):
        # the ground changes under it mid-episode: the operator's point
        # that there is no virtue in balancing a spinning top on perfect
        # ground when the ground will not stay perfect
        if slope_step_deg and k == int(0.4 * n):
            th = math.radians(slope_step_deg)
            env.model.opt.gravity[:] = [9.81 * math.sin(th), 0.0,
                                        -9.81 * math.cos(th)]
        if push_bid is not None:
            t_now = k * env.control_dt
            f = 0.0
            for i2, ev in enumerate(pushes):
                t_s = 0.5 + i2 * push_gap
                if t_s <= t_now < t_s + min(1.5, float(ev.get("dur", 0.3))):
                    f = float(ev.get("n", 0.0))
            env.data.xfrc_applied[push_bid, 0] = f
        lean, rate, speed, odo = s.read()
        # NO position memory in the bench law: the odometry term leans the
        # robot back toward its start point, which is exactly the homing
        # the operator ruled out - and station scales it TOGETHER with the
        # essential speed damping, so the search could never drop one
        # without the other. The benches balance on lean and speed alone;
        # station-keeping belongs to the stages whose job it actually is.
        ref = C.lean_reference(station, speed, 0.0, clamp)
        u = C.wheel_command(kp, kd, lean, rate, ref, trim)
        # effort is priced when NOTHING is happening - quiet standing must
        # be cheap. During a disturbance response the burst is FREE: a
        # balancer that flinches from its own torque bill is a
        # fair-weather balancer. Quiet = small lean error, small rate,
        # AND the body actually still: a sway is small-lean with the body
        # translating, and the old two-condition gate charged for the very
        # torque that would kill it - the operator watched the robot sway
        # with almost no torque applied, because stinginess mid-sway was
        # literally what the exam paid for.
        quiet = (abs(lean - ref - trim) < math.radians(4.0) and abs(rate) < 0.8
                 and abs(speed) < 0.10)
        if quiet:
            effort += abs(u) * env.control_dt
        # the chatter meter: effort is only priced when QUIET, so a law
        # that never stops chattering is never quiet and pays nothing -
        # the crouch bench baked kd 8 for free that way (the kd term
        # amplifies gyro noise into a limit cycle at the crouch's higher
        # pendulum frequency; kd 2 holds 340mm at 3.9 degrees) and every
        # stage interpolating below 420mm inherited the oscillation.
        # Threshold form: dithers and honest recovery bursts under 0.15
        # per step ride free, bang-bang pays every step.
        chat += max(0.0, abs(u - prev_u) - 0.15)
        prev_u = u
        act = np.zeros(env.act_dim)
        for j, ai in enumerate(env.act_idx):
            if ai in legs:
                act[j] = np.clip((pose[legs.index(ai)] - env.act_center[ai])
                                 / max(1e-9, env.act_span[ai]), -1, 1)
            else:
                act[j] = u
        env.step(act)
        # the settle meter, defined as PHYSICAL STILLNESS: speed and pitch
        # rate, not position. The operator's clarification - a pushed robot
        # does not owe anyone its old GPS coordinate, it owes stillness as
        # fast as possible, wherever it ends up. Charging |speed| also
        # keeps the old loophole shut: rolling away at a steady lean pays
        # every second, while stabilised-300mm-downstream pays nothing.
        settle_i += (abs(speed) + 0.25 * abs(rate)) * env.control_dt
        if k in hits:
            env.data.qvel[0] += float(rng.choice([-1.0, 1.0])) * shove
        if env.up_z < 0.5:
            break
        up_s += env.control_dt
    # upright time is the goal; the settle meter (time-averaged) pays the
    # fastest damper and charges any sustained motion, so neither a limit
    # cycle nor a steady roll can hide; effort is priced only in the quiet.
    # Deliberately NO end-position or stray charge: where the robot settles
    # is its own business - that it settles FAST is the exam.
    _feel_stage = ("balance" if stance_mm >= 430 else
                   "balance_mid" if stance_mm >= 370 else "balance_crouch")
    return (- feel_charge(env, _feel_stage)
            + up_s - 3.0 * (settle_i / max(up_s, 1e-6))
            - 40.0 * (chat / max(k + 1, 1)) - EFFORT_RATE * effort)


# Conditions to be good at. Kept deliberately short: one round is
# (2 x directions + 1) x len(CASES) episodes, so six 10s cases cost over ten
# minutes of simulation per round and you cannot watch that happen. Four cases
# at 5s is about a minute a round, which is watchable, and early rounds are far
# quicker still because a bad rule falls over in half a second.
CASES = [dict(v0=0.4),                    # arrives rolling: brake and settle
         dict(lean0=0.08, v0=-0.3),       # tipped AND moving the wrong way
         dict(lean0=-0.08, v0=0.3),
         dict(shove=1.2, v0=0.3),
         dict(shove=2.2, v0=-0.4),        # the hard shove: max authority
         dict(slope_deg=6.0, v0=0.4), dict(slope_deg=-6.0, v0=-0.4),
         dict(slope_step_deg=6.0, v0=0.5)]  # ground changes, moving entry


def _push_case():
    """The operator's recorded mouse-pushes, as one extra exam case.

    Loaded fresh per score() call; the server only WRITES the bank while no
    search is running, so within a run the objective holds still. An empty
    bank means no extra case and unchanged scores."""
    try:
        evs = json.load(open(ROOT / "runs" / "push_bank.json")).get("events")
    except (OSError, ValueError):
        evs = None
    return [dict(pushes=evs)] if evs else []


def score(g, secs=10.0, seed0=0, stance_mm=460.0):
    cases = CASES + _push_case()
    return float(np.mean([episode(g, seed0 + i, secs, stance_mm=stance_mm,
                                  **c) for i, c in enumerate(cases)]))


def search(iters=40, dirs=6, nu=0.35, alpha=0.6, secs=10.0, seed=0, g0=None,
           on_round=None, should_stop=None):
    """Augmented random search over five numbers.

    Direct search, not gradients. On a space this small it is both simpler and
    faster: jitter every gain, keep whichever direction helped, step that way.
    Always step from the BEST rule found so far - stepping from the current one
    lets a bad round walk the search downhill, which is how an earlier attempt
    at this produced 30 rounds that never beat their own starting point.
    """
    rng = np.random.default_rng(seed)
    # start from something sane rather than all-zeros: a robot that cannot
    # stand at all gives the search nothing to climb
    g = np.array([6.0, 1.0, 0.5, 0.06]) if g0 is None else np.array(g0, float)
    best = score(g, secs)
    print(f"  start {best:6.2f}")
    for it in range(iters):
        scale = np.maximum(np.abs(g), 0.05)      # relative jitter, with a floor
        deltas = [rng.normal(0, 1, 4) * scale for _ in range(dirs)]
        rp = np.array([score(g + nu * d, secs) for d in deltas])
        rm = np.array([score(g - nu * d, secs) for d in deltas])
        keep = np.argsort(-np.maximum(rp, rm))[:max(1, dirs // 2)]
        sd = np.concatenate([rp[keep], rm[keep]]).std() + 1e-6
        cand = g + (alpha / (len(keep) * sd)) * sum(
            (rp[i] - rm[i]) * deltas[i] for i in keep)
        s = score(cand, secs)
        if s > best:
            best, g = s, cand
        if on_round:
            on_round(it + 1, iters, float(s), float(best), g.tolist())
        if should_stop and should_stop():
            break
        if (it + 1) % 5 == 0 or it == 0:
            print(f"  round {it+1:3d}  {s:6.2f}   best {best:6.2f}   "
                  + " ".join(f"{v:+7.2f}" for v in g))
    return g, best


def report(g, label, secs=10.0):
    print(f"\n{label}")
    print(f"  {'condition':>16} {'upright':>9} {'strayed':>9}")
    rows = [("flat", {}), ("tipped 3 deg", dict(lean0=0.05)),
            ("tipped 7 deg", dict(lean0=0.12)), ("tipped -7 deg", dict(lean0=-0.12)),
            ("shoved 1.2 m/s", dict(shove=1.2)), ("shoved 2.0 m/s", dict(shove=2.0))]
    for name, c in rows:
        ups, strays = [], []
        for i in range(3):
            env = RobotEnv(SPEC, seed=9000 + i, task="balance", randomise=False)
            env.reset()
            if c.get("lean0"):
                th = c["lean0"]
                env.data.qpos[3:7] = [math.cos(th/2), 0, math.sin(th/2), 0]
                env.data.qpos[2] = 0.4625
            s = Sensed(env); rng = np.random.default_rng(9000 + i)
            n = int(secs / env.control_dt)
            hits = set(rng.integers(10, n - 10, size=2)) if c.get("shove") else set()
            for k in range(n):
                u = s.command(g)
                env.step(np.array([u, u]))
                if k in hits:
                    env.data.qvel[0] += float(rng.choice([-1.0, 1.0])) * c["shove"]
                if env.up_z < 0.5:
                    break
            ups.append(env.t); strays.append(abs(float(env.base_pos[0])) * 1000)
        print(f"  {name:>16} {np.mean(ups):8.2f}s {np.mean(strays):8.0f}mm")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--iters", type=int, default=40)
    ap.add_argument("--secs", type=float, default=10.0)
    ap.add_argument("--out", default=str(ROOT / "runs/balance_gains.json"))
    args = ap.parse_args()
    import time
    t0 = time.time()
    g, best = search(iters=args.iters, secs=args.secs)
    print(f"\n  {time.time()-t0:.0f}s of learning\n")
    print("  learned gains:")
    for n, v in zip(NAMES, g):
        print(f"    {n:11s} {v:+8.3f}")
    report(g, "how it does on conditions it never practised")
    Path(args.out).parent.mkdir(parents=True, exist_ok=True)
    json.dump({"gains": g.tolist(), "names": NAMES}, open(args.out, "w"))
    print(f"\nwrote {args.out}")


if __name__ == "__main__":
    main()


# ---------------------------------------------------------------------------
# Stages. Each one starts from what the previous stage learned, so the course
# accumulates instead of restarting - which is the whole point of doing them in
# order. Balance is four gains; bob adds two more for the legs and keeps the
# four it inherited, so it can refine them but does not have to rediscover them.

ROOT_RUNS = ROOT / "runs"
STAGES = {
    "balance": {
        "names": ["kp (lean)", "kd (lean rate)", "station", "lean clamp"],
        # The quiet family, found by re-running the search under the effort
        # charge, seeded FROM the old bang-bang bake so it had to climb out
        # of that optimum. Verified: 4.88 on the v19 objective (the bang-bang
        # set's 4.90, statistically the same) while holding still at 0.058 Nm
        # mean and 0% railed, against 1.171 Nm and 66% railed before.
        "seed": [3.65, 2.91, 1.02, 0.52],
    },
    # the operator's restructure: balance is THREE benches, one per posture,
    # because one gain set never fit every height - the crouch moves the CoM
    # and the trim, and the pendulum the wheels balance is a different length.
    # Each bench holds its stance for the whole exam; bob then interpolates
    # its wheel gains between the three bakes instead of learning its own.
    "balance_mid": {
        "names": ["kp (lean)", "kd (lean rate)", "station", "lean clamp"],
        "seed": [3.65, 2.91, 1.02, 0.52],
    },
    "balance_crouch": {
        "names": ["kp (lean)", "kd (lean rate)", "station", "lean clamp"],
        # reseeded from the mid bench's calm family with kd 1.6 - the
        # measured best at 340mm (4.25 vs the quarantined kd-8 bake's
        # -5.11 under the honest chatter-metered exam)
        "seed": [4.69, 1.6, 0.85, 0.52],
    },
    "bob": {
        # bob no longer owns wheel gains: at every commanded height the four
        # wheel gains are interpolated between the three balance bakes
        # (straight 460 / mid 400 / crouch 340), so the benches carry the
        # balancing and bob only learns the leg motion around them.
        # com hold: the squat's dynamic answer (operator: "when we squat,
        # the CoM change tips us - how do we get around that?"). The trim
        # curve is the STATIC answer; during the squat the feet now walk
        # the axle under the MEASURED CoM in realtime, same law rotate and
        # spin proved. Sign measured negative.
        "names": ["trim scale", "leg speed", "level shift", "com hold"],
        # level shift: the zero-travel mechanism. Gate travel was COMMANDED
        # - every height has a different trim lean (-0.8 tall, -8 crouched)
        # so each gate change ordered a 7-degree repitch, and a pendulum
        # repitches by driving its wheels (470mm per cycle, measured). The
        # feet now walk forward under the CoM as the crouch deepens (the
        # look stage's measured discovery: ~+50mm at 420) so the body keeps
        # ONE pitch at every height and no travel is ever ordered. Direct
        # wheel-braking was tried first and made travel WORSE (470 -> 806mm
        # - braking pitches the body; the old warning measured true again).
        # level shift seeds NEUTRAL: both hand-swept directions lost badly
        # against the wanderer's stale companion gains (forward 5-8m, back
        # 5-13m of travel) - the mapping only works co-tuned with the loop
        # gains around it, which is the search's job, judged by the drift
        # charges. The mechanism is here; the number is earned.
        "seed": [1.0, 1.0, 0.0, -0.5],
    },
}

# the three benches and their stances - the interpolation's x-axis
BALANCE_STANCES = [("balance", 460.0), ("balance_mid", 400.0),
                   ("balance_crouch", 340.0)]


def balance_points():
    """(stance_mm, [kp, kd, station, clamp]) for every baked bench."""
    pts = []
    for nm, h in BALANCE_STANCES:
        g = accepted(nm)
        if g:
            pts.append((h, np.array(g[:4], float)))
    return sorted(pts)


def interp_wheel_gains(pts, stance_mm):
    """Wheel gains at a height, interpolated between the balance bakes.

    Piecewise LINEAR through the baked points, clamped flat at the ends -
    a quadratic through three points can overshoot between them, and an
    overshot clamp or station is a law nobody examined. One point degrades
    to constant gains (the old single-bench behaviour)."""
    if not pts:
        raise RuntimeError("no balance bench is baked - bake balance / "
                           "balance_mid / balance_crouch first; bob "
                           "interpolates its wheel gains from them")
    h = float(stance_mm)
    if h <= pts[0][0]:
        return pts[0][1]
    if h >= pts[-1][0]:
        return pts[-1][1]
    for (h0, g0v), (h1, g1v) in zip(pts, pts[1:]):
        if h0 <= h <= h1:
            f = (h - h0) / max(1e-9, h1 - h0)
            return (1.0 - f) * g0v + f * g1v
    return pts[-1][1]


def accepted(stage):
    """What an earlier stage was signed off with, if anything."""
    try:
        return json.load(open(ROOT_RUNS / f"{stage}_gains.json"))["gains"]
    except Exception:
        return None


def bob_episode(g, seed, secs=8.0, heights=(460, 400, 340), yank=False,
                settle=False):
    """Hold a series of stance heights without falling over.

    The wheels balance on gains INTERPOLATED between the three balance
    bakes at the commanded height - the benches proved those numbers, bob
    only learns the leg motion around them. Scored on tracking as well as
    staying up, because a robot that ignores the command and stands still
    stays upright perfectly.
    """
    trim_scale, leg_speed, lvl_shift, com_hold = g
    bpts = balance_points()
    env = RobotEnv(SPEC, seed=seed, task="bob", randomise=False)
    env.reset()
    com_bid = mujoco.mj_name2id(env.model, mujoco.mjtObj.mjOBJ_BODY,
                                "chassis")
    com_wids = [mujoco.mj_name2id(env.model, mujoco.mjtObj.mjOBJ_BODY, n)
                for n in ("l_wheel", "r_wheel")]
    fdx = 0.0
    legs = [i for i, a in enumerate(env.spec["actuators"])
            if a.get("mode") == "position"]
    for i in legs:                       # full travel, not the policy's 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])
    s = Sensed(env)
    cur = np.array(C.stance_to_leg(0.460) * 2)
    rate = C.leg_rate(leg_speed)
    n = int(secs / env.control_dt)
    err = 0.0
    effort = 0.0
    drift_i = 0.0
    x0 = float(env.base_pos[0])
    if yank:
        # the operator's slider gesture as an exam case: square-wave gate
        # steps at 500 deg/s - survive the fast bob, not just the gentle one
        rate = math.radians(500.0)
    for k in range(n):
        t_now = k * env.control_dt
        if yank:
            target_mm = 460 if (int(t_now / 2.0) % 2 == 0) else 350
        elif settle:
            # the operator's continuous-watching reality: change gates,
            # then HOLD - the settle itself is what gets graded, because
            # fresh short episodes never had to finish calming down
            target_mm = (460 if t_now < 1.5 else 360 if t_now < 3.0
                         else 420)
        else:
            target_mm = heights[int(k / n * len(heights)) % len(heights)]
        # feet walk under the CoM as the crouch deepens: one body pitch at
        # every height, zero commanded travel - and the realtime CoM hold
        # keeps the axle under the MEASURED mass through the squat itself
        # (the trim curve is the static answer; this is the dynamic one)
        com_x = float(env.data.subtree_com[com_bid][0])
        ax_x = 0.5 * (float(env.data.xpos[com_wids[0]][0])
                      + float(env.data.xpos[com_wids[1]][0]))
        fdx += (env.control_dt / 0.04) * (com_hold * (com_x - ax_x) - fdx)
        dxb = float(np.clip(lvl_shift * (0.46 - target_mm / 1000.0)
                            + fdx, -0.06, 0.14))
        hip_t, knee_t = C.stance_to_leg(target_mm / 1000.0, dxb)
        want = np.array([hip_t, knee_t, hip_t, knee_t])
        cur = C.slew(cur, want, rate, env.control_dt)
        lean, lrate, speed, odo = s.read()
        # the benches' proven wheel gains at THIS height
        kp, kd, station, clamp = interp_wheel_gains(bpts, target_mm)
        ref = C.lean_reference(station, speed, odo, clamp)
        trim = math.radians(C.lean_trim_deg(target_mm)) * trim_scale
        u = C.wheel_command(kp, kd, lean, lrate, ref, trim)
        # effort priced only in the quiet (the fight-ready doctrine): a
        # burst that holds the floor position during a gate change is free.
        # Quiet requires actual stillness (speed too) - a sway is small-lean
        # with the body moving, and countering it must never be charged
        if (abs(lean - ref - trim) < math.radians(4.0) and abs(lrate) < 0.8
                and abs(speed) < 0.10):
            effort += abs(u) * env.control_dt
        act = np.zeros(env.act_dim)
        for j, ai in enumerate(env.act_idx):
            if ai in legs:
                act[j] = np.clip((cur[legs.index(ai)] - env.act_center[ai])
                                 / max(1e-9, env.act_span[ai]), -1, 1)
            else:
                act[j] = u
        env.step(act)
        err += abs(env.height - target_mm / 1000.0) * env.control_dt
        drift_i += abs(float(env.base_pos[0]) - x0) * env.control_dt
        if env.up_z < 0.5:
            break
    # upright, tracked, AND PLANTED: the operator watched bob wander the
    # floor with idle wheels - the wheels' job during gate changes is to
    # keep the robot exactly where it stood, and now that is what pays
    x_end = abs(float(env.base_pos[0]) - x0)
    return (- feel_charge(env, "bob")
            + env.t - 2.0 * err - 1.5 * (drift_i / max(env.t, 1e-6))
            - 1.0 * min(x_end, 2.0) - EFFORT_RATE * effort)



def stage_max(stage, secs):
    """The ceiling a perfect run could score, for honest display.

    Every stage's score is 'time upright minus its own penalty', except stop,
    which is a fixed 10-point formula - so raw numbers are not comparable
    across stages (5.88 bob vs 2.56 level says nothing until you know one is
    out of 6 and the other out of 5). Bakes record this so the UI can always
    show 'score / ceiling'.
    """
    if stage == "stop":
        return 10.0
    if stage == "spin":
        return float(secs) + 4.5   # completed-quarter bonuses sit on top
    if stage == "bob":
        # the settle case runs a full 10s regardless of secs, so the
        # honest ceiling is the mean of the actual episode lengths
        return (4.0 * float(secs) + max(float(secs), 10.0)) / 5.0
    if stage == "rotate":
        # completion bonus (2) + speed bonus (3) sit on top of upright time
        return float(secs) + 5.0
    if stage == "rise":
        return 12.0    # the height-integral reward's ceiling
    if stage == "sit":
        return 11.0    # low-height (6) + still landing (3) + quick (2)
    if stage == "jump":
        return 17.0    # clearance (12) + landing (3) + stay-put (2)
    return float(secs)

def stage_score(stage, g, secs=8.0):
    if stage == "balance_mid":
        return score(g, secs, stance_mm=400.0)
    if stage == "balance_crouch":
        return score(g, secs, stance_mm=340.0)
    if stage == "bob":
        return float(np.mean([bob_episode(g, i, secs) for i in range(3)]
                             + [bob_episode(g, 3, secs, yank=True),
                                bob_episode(g, 4, max(secs, 10.0),
                                            settle=True)]))
    if stage == "slope":
        return slope_score(g, secs)
    if stage == "level":
        return level_score(g, secs)
    if stage == "stop":
        return stop_score(g, secs)
    if stage == "cruise":
        return cruise_score(g, secs)
    if stage == "turn":
        return turn_score(g, secs)
    if stage == "bump":
        return bump_score(g, secs)
    if stage == "slip":
        return slip_score(g, secs)
    if stage == "payload":
        return payload_score(g, secs)
    if stage == "spin":
        return spin_score(g, secs)
    if stage == "rotate":
        return rotate_score(g, secs)
    if stage == "jump":
        return jump_score(g, secs)
    if stage == "rise":
        return rise_score(g, secs)
    if stage == "sit":
        return sit_score(g, secs)
    if stage == "recover":
        # 0.0 first: the null case. A recovery stage whose base pipeline
        # cannot stand still UNSHOVED scores well on nothing, and that exact
        # failure went unseen because no case ever exercised it.
        return float(np.mean([recover_episode(g, i, secs, shove=sh)
                              for i, sh in enumerate((0.0, 2.0, 2.5, 3.0))]))
    return score(g, secs)


def base_reflex():
    """The recovery reflex as a BASE LAYER under every stage.

    Learned once, in recover (brace gain and fade, slots 8/9 of its bake),
    then active during cruise, turn, spin, bump, slip and payload - not just
    during recover's own exam. Skills used to inherit only NUMBERS; the
    operator's question "can it recover itself while spinning?" exposed that
    the catching behaviour itself never rode along. Now it does: a stage
    that gets shoved (or shakes itself apart - the post-spin station
    oscillation) braces with whatever recover has learned so far. Unbaked
    recover means no reflex, which is one more reason to keep recover
    strong.
    """
    b = accepted("recover")
    if b and len(b) >= 10:
        return C.ShoveBrace(b[8], b[9])
    return C.ShoveBrace(0.0, 0.3)


def _recover_seed():
    """Recover stands on balance's STATION-CAPABLE slots plus bob's legs.

    It used to copy bob's first six wholesale - but bob's objective never
    pays for holding position, so bob's baked station gain is ~0.08, and a
    recover seeded from it drifts away after every catch (the away penalty
    ate the 46% bake). Balance owns slots 0-3, bob contributes only what it
    actually learned: trim scale and leg speed."""
    g = list(STAGES["recover"]["seed"])
    bal = accepted("balance")
    if bal:
        g[:4] = bal[:4]
    prev = accepted("bob")
    if prev and len(prev) >= 2:
        # bob's NEW 3-slot layout: 0 trim scale, 1 leg speed
        g[4], g[5] = prev[0], prev[1]
    # NOTE deliberately no inheritance from "stop" here. Slots 8 and 9 are
    # brace gain and brace fade; an earlier version wrote stop's brake gain
    # (0.5) into the brace slot (seeded 0.02), a 25x overdose that fired on
    # the first noise spike and put the robot down in under a second. If the
    # brake handover comes back it gets its OWN slots, never someone else's.
    return g


def _nice_worker():
    """Search workers yield the CPU to the live view: the demo stutters
    into chunky substeps whenever the pool saturates the cores, which is
    exactly when the operator is watching a candidate learn."""
    try:
        os.nice(12)
    except OSError:
        pass


def _eval_stage(args):
    """One pool job: score a candidate. Module-level so spawn can find it."""
    stage, g, secs = args
    return stage_score(stage, np.array(g, dtype=float), secs)


# --- the polish pass: flatten the waves after the bake ---------------------
# The operator's design: movements are wave signals, and the flatter the
# balance wave, the more efficient the robot. After a stage bakes (the exam
# certifies it WORKS), this second, gentler search holds the exam score and
# minimizes the wave height - fore-aft lurch, lateral wobble, vertical
# shudder - read from the same IMU feelings layer every episode meters.

def measure_wave(stage, g, secs=8.0):
    """Exam score plus the ROUGHNESS behind it - tick-to-tick jaggedness,
    not amplitude: a sweeping turn is a slow wave and stays cheap, ripple
    and chatter are expensive (operator: smooth the lines, don't flatten
    the move)."""
    import env as _envmod
    _envmod.FEEL_ACCUM = {}
    try:
        raw = stage_score(stage, list(g), secs)
        wave = _envmod.FEEL_ACCUM.get("rough", 0.0)
    finally:
        _envmod.FEEL_ACCUM = None
    return float(raw), float(wave)


def _eval_polish(args):
    stage, g, secs = args
    return measure_wave(stage, g, secs)


def polish_stage(stage, rounds=15, secs=8.0, dirs=8, nu=0.06, tol=0.03,
                 on_round=None, should_stop=None, on_eval=None):
    """Local search from the baked gains: minimize wave height, HARD floor
    on the exam score (a flatter law that fails the exam is no law)."""
    base = accepted(stage)
    if base is None:
        return None
    g = np.array(base, dtype=float)
    import multiprocessing as _mp
    import os as _os
    pool = _mp.get_context("spawn").Pool(
        max(2, (_os.cpu_count() or 4) - 2),
        initializer=_nice_worker)
    try:
        raw0, wave0 = pool.apply(_eval_polish, ((stage, g.tolist(), secs),))
        floor = raw0 - abs(raw0) * tol - 0.05
        best_wave = wave0
        rng = np.random.default_rng(0)
        for it in range(rounds):
            if should_stop and should_stop():
                break
            deltas = rng.standard_normal((dirs, len(g)))
            cands = []
            for dlt in deltas:
                for sgn in (1.0, -1.0):
                    cands.append(g + sgn * nu * dlt * np.maximum(
                        0.05, np.abs(g)))
            if on_eval:
                on_eval(0, len(cands))
            outs = []
            for i, cg in enumerate(cands):
                outs.append(pool.apply_async(
                    _eval_polish, ((stage, cg.tolist(), secs),)))
            done = 0
            scored = []
            for cg, o in zip(cands, outs):
                raw, wave = o.get()
                done += 1
                if on_eval:
                    on_eval(done, len(cands))
                # the exam floor is HARD: flat but failing = rejected
                scored.append((wave if raw >= floor else 1e9, cg))
            scored.sort(key=lambda t: t[0])
            if scored[0][0] < best_wave - 1e-6:
                g = scored[0][1]
                best_wave = scored[0][0]
            if on_round:
                pct = 100.0 * (1.0 - best_wave / max(wave0, 1e-9))
                on_round(it + 1, rounds, pct, pct, g.tolist())
        raw1, wave1 = pool.apply(_eval_polish, ((stage, g.tolist(), secs),))
        return {"gains": g.tolist(), "raw_before": raw0, "raw_after": raw1,
                "wave_before": wave0, "wave_after": wave1,
                "kept": bool(raw1 >= floor and wave1 < wave0)}
    finally:
        pool.close()
        pool.join()


def search_stage(stage="balance", iters=40, secs=8.0, dirs=6, nu=0.3,
                 alpha=0.6, seed=0, on_round=None, should_stop=None, g0=None,
                 on_eval=None, scale_floor=0.05):
    # scale_floor: perturbation steps scale with each slot's magnitude,
    # floored here - near-zero slots (a reseeded trim, a brand-new v94
    # slot) explore in 0.015-size steps at the default and can take
    # hundreds of rounds to reach a distant optimum. A wide pass
    # (0.15-0.25) lets every slot range; the default keeps the polishing
    # behaviour every previous bake was learned under.
    cfg = STAGES[stage]
    g = np.array(cfg["seed"], dtype=float)
    # Each stage stands on the one before it - but inherits ONLY the slots
    # that mean the same thing in its own layout (the INHERIT map, defined
    # with the ladder below). Copying whole gain vectors across layouts is
    # how stop's brake gain once landed in recover's brace slot, and how
    # bob's trim scale quietly seeded level's level-gain slot.
    if stage == "recover":
        g = np.array(_recover_seed(), dtype=float)
    else:
        for src, spec in INHERIT.get(stage, []):
            prev = accepted(src)
            if not prev:
                continue
            if isinstance(spec, int):
                k = min(spec, len(prev), len(g))
                g[:k] = prev[:k]
            else:
                for si, di in spec:
                    if si < len(prev) and di < len(g):
                        g[di] = prev[si]
    # A stage that already has a bake seeds from ITS OWN proven gains -
    # inheritance is for first-time learning. Without this, relearning
    # balance_mid inherited the straight bench's kd 4.71, a MEASURED
    # chatterer at crouch heights (the limit-cycle discovery), and the
    # search started at -4.3 climbing out of a known hole.
    own = accepted(stage)
    if own is not None and len(own) == len(g):
        g = np.array(own, dtype=float)
    if g0 is not None and len(g0) == len(g):
        g = np.array(g0, dtype=float)    # continue from a previous run's best
    rng = np.random.default_rng(seed)
    # The round's 12 perturbation scores are independent - farm them across
    # cores. A round used to be 13 sequential score calls (~10 minutes of
    # simulated episodes on ONE core) with nothing on screen but "round 0",
    # which read as a hang; now the trials run in parallel and each
    # completion ticks the on_eval counter the UI shows.
    import multiprocessing as _mp
    import os as _os
    pool = _mp.get_context("spawn").Pool(
        max(2, (_os.cpu_count() or 4) - 2),
        initializer=_nice_worker)
    try:
        if on_eval:
            on_eval(0, 2 * dirs + 1)
        # EVERY episode runs in a worker - the seed and candidate evals
        # used to run in the server process itself, holding the GIL for
        # 30-60s of simulation per round while live-view requests queued
        # behind them: the operator's "still jerky" was those stalls
        best = pool.apply(_eval_stage, ((stage, g.tolist(), secs),))
        print(f"  [{stage}] start {best:6.2f}")
        for it in range(iters):
            scale = np.maximum(np.abs(g), scale_floor)
            deltas = [rng.normal(0, 1, len(g)) * scale for _ in range(dirs)]
            jobs = ([(stage, (g + nu * d).tolist(), secs) for d in deltas]
                    + [(stage, (g - nu * d).tolist(), secs) for d in deltas])
            res = []
            for r in pool.imap(_eval_stage, jobs):
                res.append(r)
                if on_eval:
                    on_eval(len(res), 2 * dirs + 1)
            rp = np.array(res[:dirs])
            rm = np.array(res[dirs:])
            keep = np.argsort(-np.maximum(rp, rm))[:max(1, dirs // 2)]
            sd = np.concatenate([rp[keep], rm[keep]]).std() + 1e-6
            cand = g + (alpha / (len(keep) * sd)) * sum(
                (rp[i] - rm[i]) * deltas[i] for i in keep)
            sc = pool.apply(_eval_stage, ((stage, cand.tolist(), secs),))
            if on_eval:
                on_eval(2 * dirs + 1, 2 * dirs + 1)
            if sc > best:
                best, g = sc, cand
            if on_round:
                on_round(it + 1, iters, float(sc), float(best), g.tolist())
            if should_stop and should_stop():
                break
    finally:
        pool.terminate()
        pool.join()
    return g, best


# --- stage 3: survive a shove ----------------------------------------------
# Inherits everything bob was signed off with and adds the two numbers that
# control the recovery reflex: how far the feet slide toward a fall, and the
# limit on that. Measured by hand earlier, sliding the feet cut peak lean on a
# 2 m/s shove from 24.9 degrees to 6.7 - this stage is asking the search to
# find the gain rather than me guessing it.
# Recover now stands on the stop stage. Being shoved is, after the first
# instant, a stopping problem: something gave you a speed you did not ask for
# and the job is to lose it quickly without going over. So the brace handles
# the impact, then the braking law it learned in "stop" takes over until the
# speed is gone. It inherits those two numbers rather than rediscovering them.
STAGES["recover"] = {
    "names": ["kp (lean)", "kd (lean rate)", "station", "lean clamp",
              "trim scale", "leg speed", "foot shift", "shift limit",
              "brace gain", "brace fade"],
    "seed": [12.0, 2.0, 1.0, 0.08, 1.0, 1.0, 0.30, 0.10, 0.02, 0.30]
}


def recover_episode(g, seed, secs=8.0, shove=2.5, stance_mm=460):
    """Take a shove and get back to standing still.

    This is the version that measured well (7.50 / 7.15 / 6.27 at 2.0 / 2.5 /
    3.0 m/s). Several attempts to improve it - handing over to the braking law
    from the stop stage, then collapsing both into a single lean-into-the-motion
    term - all scored worse, so it is back to what was measured rather than
    what sounded better.

    Two mechanisms, both continuous, neither switching modes:
      - the feet slide toward the fall, putting the contact patch back under
        the mass faster than the wheels could drive it there
      - a brace leans into a detected push, sized by how hard it was, fading
        because it answers an event rather than holding a posture
    """
    (kp, kd, station, clamp, trim_scale, leg_speed, shift_k, shift_max,
     brace_k, brace_fade) = g
    env = RobotEnv(SPEC, seed=seed, task="bob", randomise=False)
    env.reset()
    legs = [i for i, a in enumerate(env.spec["actuators"])
            if a.get("mode") == "position"]
    for i in legs:
        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])
    s = Sensed(env)
    cur = np.array(C.stance_to_leg(stance_mm / 1000.0) * 2)
    # THE LAW LIVES IN firmware.StanceLaw (the standing law the whole
    # firmware rests on) - this exam searches its numbers
    law = FW.StanceLaw((kp, kd, station, clamp),
                       math.radians(C.lean_trim_deg(stance_mm))
                       * trim_scale,
                       shift_k, shift_max, leg_speed,
                       C.ShoveBrace(brace_k, brace_fade),
                       stance_mm, env.control_dt)
    rate = law.leg_rate
    n = int(secs / env.control_dt)
    rng = np.random.default_rng(seed)
    # Both shoves land in the first 60% of the episode, leaving at least 3s
    # to watch the recovery. Drawn from the full episode, seed 0 put both hits
    # AFTER the point the robot had already fallen - so every sweep tonight
    # measured an unshoved robot, and shove strength changed nothing, twice,
    # without being noticed. Flat scores across conditions are the tell.
    hits = set(rng.integers(int(0.15 * n), int(0.60 * n), size=2))
    away = 0.0
    effort = 0.0
    for k in range(n):
        lean, lrate, speed, odo = s.read()
        want, _, u = law.tick(lean, lrate, speed, odo)
        cur = C.slew(cur, want, rate, env.control_dt)
        effort += abs(u) * env.control_dt
        act = np.zeros(env.act_dim)
        for j, ai in enumerate(env.act_idx):
            act[j] = (np.clip((cur[legs.index(ai)] - env.act_center[ai])
                              / max(1e-9, env.act_span[ai]), -1, 1)
                      if ai in legs else u)
        env.step(act)
        if k in hits:
            env.data.qvel[0] += float(rng.choice([-1.0, 1.0])) * shove
        away += abs(float(env.base_pos[0])) * env.control_dt
        if env.up_z < 0.5:
            break
    # time spent away from where it started, so getting back quickly is what pays
    return env.t - 1.2 * (away / max(env.t, 1e-6)) - EFFORT_RATE * effort


STAGES["stop"] = {
    # the operator's reframe: stopping IS aggressive balance, and the LEGS
    # are part of the brake - crouching drops the CoM, which raises the
    # tipping threshold and lets the wheels bite harder before the body
    # goes over. Wheel gains are not learned here: they ride the benches,
    # interpolated at whatever stance the brake law commands, with the
    # kinematic trim in the loop. Brake stance is an order-1 offset
    # (x100mm from 400) - the level stage measured raw millimetres getting
    # +-125mm lurches from the search's proportional jitter.
    "names": ["brake gain", "brake clamp", "brake stance", "crouch speed"],
    # brake clamp seeded WIDE (the old 0.20 was the binding constraint,
    # measured 5.69/10 rising to ~7.3 at 0.60); brake stance seeds at a
    # modest crouch (-0.2 = 380mm), crouch speed 1.0 = 300mm/s of stance
    # travel - all four are the search's to move.
    "seed": [0.5, 0.45, -0.2, 1.0],
}


def stop_episode(g, seed, secs=8.0, cruise=1.5, settle=2.5, shove=0.0):
    """Get up to speed (or take a shove while standing), then stop as fast
    as possible without falling - legs in play: the brake crouch drops the
    CoM while the wheels haul the speed down."""
    brake_gain, brake_clamp, brake_st, crouch_speed = g
    bpts = balance_points()
    env = RobotEnv(SPEC, seed=seed, task="bob", randomise=False)
    env.reset()
    legs = [i for i, a in enumerate(env.spec["actuators"])
            if a.get("mode") == "position"]
    for i in legs:
        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])
    s = Sensed(env)
    CRUISE_ST = 430.0
    stance_now = CRUISE_ST
    brake_mm = float(np.clip(400.0 + 100.0 * brake_st, 330.0, 462.0))
    rate_mm = 300.0 * max(0.2, abs(crouch_speed))    # stance travel, mm/s
    n = int(secs / env.control_dt)
    n_run = int(settle / env.control_dt)
    v_cmd = 0.0
    x_at_stop = None
    t_stop = None
    t_brake = 0.0
    effort = 0.0
    v_peak = 0.1
    for k in range(n):
        lean, lrate, speed, odo = s.read()
        if shove and k == n_run:
            # the push arrives mid-stand: stopping an uncommanded speed is
            # the same skill as braking a commanded one (operator's point)
            env.data.qvel[0] += shove
        braking = k >= n_run
        # the legs' part of the brake: crouch toward the brake stance
        tgt_mm = brake_mm if braking else CRUISE_ST
        stance_now += float(np.clip(tgt_mm - stance_now,
                                    -rate_mm * env.control_dt,
                                    rate_mm * env.control_dt))
        hip_t, knee_t = C.stance_to_leg(stance_now / 1000.0)
        # the benches' proven wheel gains AT THIS STANCE, trim included
        kp, kd, station, clamp = interp_wheel_gains(bpts, stance_now)
        trim = math.radians(C.lean_trim_deg(stance_now))
        if not braking:
            v_cmd = min(cruise, v_cmd + 1.5 * env.control_dt)  # wind it up
            ref = max(-abs(clamp), min(abs(clamp),
                      station * (0.5 * (speed - v_cmd))))
        else:
            if x_at_stop is None:
                x_at_stop = float(env.base_pos[0])
                t_brake = env.t
            # braking gets its own gain and its own, wider, limit
            ref = max(-abs(brake_clamp), min(abs(brake_clamp),
                      brake_gain * speed))
        u = C.wheel_command(kp, kd, lean, lrate, ref, trim)
        effort += abs(u) * env.control_dt
        act = np.zeros(env.act_dim)
        for j, ai in enumerate(env.act_idx):
            if ai in legs:
                pose = (hip_t, knee_t, hip_t, knee_t)[legs.index(ai)]
                act[j] = np.clip((pose - env.act_center[ai])
                                 / max(1e-9, env.act_span[ai]), -1, 1)
            else:
                act[j] = u
        env.step(act)
        if env.up_z < 0.5:
            # falling is not stopping - but the fall credit is SURVIVAL
            # TIME after the brake began, not speed killed: the base
            # ACCELERATES during a nose-down fall (wheels driving under
            # it), so every velocity-based credit measured a flat -5.0.
            # Lasting longer into the catch is the gradient that exists.
            alive = env.t - t_brake if x_at_stop is not None else 0.0
            return -5.0 + 3.0 * min(1.0, alive / 1.5)
        if x_at_stop is not None and t_stop is None and abs(speed) < 0.1:
            t_stop = env.t - t_brake
            dist = abs(float(env.base_pos[0]) - x_at_stop)
            # quick AND short, still upright - and not by burning the wheels
            return 10.0 - 3.0 * t_stop - 2.0 * dist - EFFORT_RATE * effort
    return -2.0                                 # never stopped


def stop_score(g, secs=8.0):
    # three commanded speeds plus the operator's push case: shoved while
    # standing, kill the uncommanded speed the same way. The push sits AT
    # the measured frontier (1.2 catchable, 1.6 not, with fall credit for
    # speed killed) so the search has a slope to climb; harder shoves are
    # recover's exam, where the brace lives. Measured and rejected here:
    # fast-deep crouching and foot-shift DURING the catch both fell where
    # the mild config caught - leg motion mid-catch disturbs the recovery.
    return float(np.mean([stop_episode(g, i, secs, cruise=c)
                          for i, c in enumerate((1.0, 1.5, 2.0))]
                         + [stop_episode(g, 3, secs, cruise=0.0,
                                         shove=1.4)]))


# --- stage: stay level on tilted ground ------------------------------------
# The wheels share an axle, so they can do nothing about roll. The only handle
# is leg LENGTH: shorten the uphill leg, extend the downhill one, and the body
# sits level while the ground does not.
#
# This is deliberately a learned stage. A plain proportional gain on measured
# roll was tried by hand three separate times today - for turning, for shove
# recovery, and for levelling - and fell over every time, including on flat
# ground. Three failures by guesswork is enough.
STAGES["level"] = {
    # Three leg degrees of freedom, all searchable - the earlier law only used
    # the first, so "full use of the legs" was mechanically true but
    # structurally false:
    #   level limit    left/right length DIFFERENCE (rolls the body level)
    #   stance height  common-mode leg length (posture on the slope)
    #   foot shift     fore-aft placement, feet toward the lean - the same
    #                  reflex that cut shove peak lean 24.9 -> 6.7 deg, here
    #                  as active suspension against fore-aft slosh
    # stance is an ORDER-1 OFFSET (x100mm from 400), not raw millimetres.
    # As a raw 420 among order-1 gains, the search's proportional jitter threw
    # +-125mm stance lurches into every probe - both sides nulled, every
    # direction read as zero gradient, and the search froze at its seed while
    # the rounds ticked by looking busy.
    "names": ["kp (lean)", "kd (lean rate)", "station", "lean clamp",
              "level gain", "level damping", "level limit", "leg speed",
              "stance offset", "foot shift"],
    # Leveller seeded OFF and stance at 460, both from measurement (12 Aug):
    # under the peak-torque legs the old seed (level gain -1.0) is an
    # unstable loop that falls on a FLAT plate (leveller off stands 4.79,
    # leveller on -3.53, sign verified corrective - it is gain margin, not
    # sign), and stance 420 sits on the station-hold floor. Rigid survives
    # the whole feasible tilt band, so the search starts from rigid and must
    # DISCOVER any leveling that pays. The 8-12 degree holds are withdrawn
    # from the score: real lateral bank is back (the table is a real tilt)
    # and the crowned tyre cannot statically hold it - NEXT.md's tyre
    # contact investigation gates those cases, exactly as it prescribes.
    # reseeded under the strengthened exam (operator: the legs must
    # ABSORB the gradient): leveller gain measured monotone-better
    # to 7.0 (6.62 vs the rigid bake at -4.64) with the 60mm cap
    # restored - the search refines from a law that already levels
    "seed": [3.051, 3.478, 0.961, 0.118, 7.0,
             0.05, 0.06, 0.054, 0.967, 0.037],
}

_BANKED = {}


def banked_spec(bank_deg, pitch_deg=0.0):
    """A spec with the floor tilted. Baked in, because MuJoCo precomputes the
    frame of static world geoms and setting it at runtime silently does
    nothing - which produced identical results at 3, 6 and 10 degrees."""
    key = (round(bank_deg, 2), round(pitch_deg, 2))
    if key not in _BANKED:
        import copy
        s = json.load(open(SPEC))
        s.setdefault("world", {}).setdefault("floor", {})["bank_deg"] = bank_deg
        s["world"]["floor"]["pitch_deg"] = pitch_deg
        path = ROOT_RUNS / f"_bank_{key[0]}_{key[1]}.json"
        ROOT_RUNS.mkdir(exist_ok=True)
        json.dump(s, open(path, "w"))
        _BANKED[key] = str(path)
    return _BANKED[key]


def level_episode(g, seed, secs=8.0, vary=8.0, stance_mm=None, **_ignored):
    """The level stage on the tilting table, sensed the way the real robot
    will be: a BNO086 flat on the chassis top reporting absolute tilt vs
    gravity - leg bounce or not. Control and score share that one reference.

    The plate steps and holds (lean, STAY leaned, return); the legs must keep
    the module level to 1 degree, position is secondary, chatter is charged,
    and the joints may move as fast as torque allows - on this reference the
    old accel-pollution feedback loop does not exist.
    """
    (kp, kd, station, clamp, lvl_k, lvl_d, lvl_max, leg_speed,
     stance_g, dx_k) = (list(g) + [420.0, 0.0])[:10]
    stance = 455.0            # the tall finish, absolute, not searchable
    env = platform_env(seed=seed)
    legs = [i for i, a in enumerate(env.spec["actuators"])
            if a.get("mode") == "position"]
    for i in legs:
        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])
    # placed on the table in the working pose - a step command from a spawn
    # mismatch is a torque slam once the ramp is open
    hp0, kn0 = C.stance_to_leg(stance / 1000.0)
    for n, v in zip(("l_hip", "l_knee", "r_hip", "r_knee"),
                    (hp0, kn0, hp0, kn0)):
        env.data.qpos[env.model.jnt_qposadr[
            mujoco.mj_name2id(env.model, mujoco.mjtObj.mjOBJ_JOINT, n)]] = v
    # the FULL base pose, orientation included. Placement used to rely on
    # the tethered reset zeroing qpos; once reset went back to the normal
    # spawn path its random base-orientation noise survived this placement,
    # and every episode started up to ~5 deg tilted on a flat plate
    # (measured: up_z 0.997 at t=0) - which is why every v20 level score
    # sat near the fall floor whatever the gains were.
    env.data.qpos[0:3] = [0.0, 0.0, 0.73 + stance / 1000.0 + 0.003]
    env.data.qpos[3:7] = [1.0, 0.0, 0.0, 0.0]
    env.data.qvel[:] = 0.0
    mujoco.mj_forward(env.model, env.data)
    s = Sensed(env)
    cur = np.array([hp0, kn0, hp0, kn0])
    # above |leg speed| 3 the ramp opens to servo speed: on the BNO reference
    # the old accel-pollution feedback loop does not exist
    rate = C.leg_rate(leg_speed, floor=0.3, unlimited_at=3.0)
    trim = math.radians(C.lean_trim_deg(stance))
    roll_err = away = settle_pen = chatter = effort = 0.0
    prev_u = 0.0
    x0, y0 = float(env.base_pos[0]), float(env.base_pos[1])
    for k in range(int(secs / env.control_dt)):
        t = k * env.control_dt
        fr = (0.0 if t < 1.0 else (t - 1.0) if t < 2.0 else 1.0 if t < 4.5
              else (1.0 - 2.0 * (t - 4.5)) if t < 5.5 else -1.0)
        env.data.ctrl[-1] = math.radians(vary) * fr
        # the BNO086: absolute tilt of the top plate, gyro-fast, unshaken by
        # leg acceleration - one reference for the leveller AND the score
        roll = math.atan2(env.est_up[1], max(1e-6, env.est_up[2]))
        roll_rate = float(env.sensors.noisy(
            "gyro", env.data.sensordata, env._raw)[0])
        dh = float(np.clip(lvl_k * roll + lvl_d * roll_rate,
                           -abs(lvl_max), abs(lvl_max)))
        lean_now = math.atan2(-env.est_up[0], max(1e-6, env.est_up[2]))
        dx = float(np.clip(dx_k * lean_now, -0.08, 0.08))
        hl, kl = C.stance_to_leg(max(0.20, min(0.462, stance / 1000.0 + dh)), dx)
        hr, kr = C.stance_to_leg(max(0.20, min(0.462, stance / 1000.0 - dh)), dx)
        cur = C.slew(cur, np.array([hl, kl, hr, kr]), rate, env.control_dt)
        lean, lrate, speed, odo = s.read()
        ref = C.lean_reference(station, speed, odo, clamp)
        u = C.wheel_command(kp, kd, lean, lrate, ref, trim)
        chatter += abs(u - prev_u)
        effort += abs(u) * env.control_dt
        prev_u = u
        act = np.zeros(env.act_dim)
        for j, ai in enumerate(env.act_idx):
            act[j] = (np.clip((cur[legs.index(ai)] - env.act_center[ai])
                              / max(1e-9, env.act_span[ai]), -1, 1)
                      if ai in legs else u)
        env.step(act)
        roll_deg = abs(math.degrees(math.atan2(env.est_up[1],
                                               max(1e-6, env.est_up[2]))))
        dist = math.hypot(float(env.base_pos[0]) - x0,
                          float(env.base_pos[1]) - y0)
        if dist > 2.0 or env.up_z < 0.5:
            return -5.0 + 0.3 * env.t
        if t > 1.0:
            roll_err += max(0.0, roll_deg - 1.0) * env.control_dt
            if (2.0 <= t < 4.5 or t >= 5.5) and roll_deg > 1.0:
                settle_pen += env.control_dt
        away += dist * env.control_dt
    n_steps = max(1, int(env.t / env.control_dt))
    # STRENGTHENED (operator: the legs just don't absorb the gradient) -
    # flatness now dominates: the old 0.9 weight priced rigidity cheaper
    # than levelling, and the baked law settled at a 10mm leg cap where
    # the 7-degree case needs 23mm
    return (env.t - 2.5 * (roll_err / max(env.t - 1.0, 1e-6))
                  - 0.6 * (away / max(env.t, 1e-6))
                  - 1.2 * settle_pen
                  - 3.0 * (chatter / n_steps)
                  - EFFORT_RATE * effort)


def level_score(g, secs=8.0):
    # Graded MOVING slopes only. The static banked-floor cases are withdrawn:
    # measured with legs frozen adapted and wheels dead, the model cannot
    # statically stand on an 8 deg lateral bank AT ALL - the crowned tyre's
    # roll-plane curvature at the contact is b^2/a = 6.4mm, close to a knife
    # edge, and the contact migrates downhill on a tilted plane. Scoring the
    # search on a physically impossible condition poisoned every result and
    # every observation made about the legs. Bank statics are a contact-model
    # investigation (NEXT.md), not a levelling-control one. Gravity-slew is
    # dynamically equivalent for the levelling skill, and works.
    # No flat case. Flat standing is balance's and bob's banked skill; here it
    # acted as a subsidy for rigidity - a straight robot banked +2.1 on flat,
    # outweighing its loss on the 12 deg hold, so the COMBINED optimum was
    # legs-off even though bending won every tilted case. (Checked per-case
    # last time, declared victory, never summed. Sum first.)
    # Tilts capped at 7 degrees: the feasible band under the current tyre.
    # The table is REAL lateral bank, and the crowned tyre cannot statically
    # hold 8+ degrees (NEXT.md) - scoring the search on a physically
    # impossible case poisoned every v20 level result until measured.
    cases = [dict(vary=3.0), dict(vary=5.0), dict(vary=7.0)]
    return float(np.mean([level_episode(g, i, secs, **c)
                          for i, c in enumerate(cases)]))


# --- stage: hold station on a fore-aft slope -------------------------------
# Legs stay the SAME length. On a hill the robot has to lean into the slope by
# roughly the slope angle, so its weight falls over the contact patch, and then
# balance about that leaning attitude instead of about vertical. A 10% grade is
# 5.7 degrees and needs 0.19 Nm a wheel held continuously just to stay put.
#
# The robot is not told the gradient. It has to find the lean from the only
# evidence it has - that it is sliding downhill - which is what the odometry
# term is for. That is why this stage learns an integral gain and a wide enough
# clamp to let the lean get there.
STAGES["slope"] = {
    "names": ["kp (lean)", "kd (lean rate)", "station", "lean clamp",
              "slope integral", "integral limit"],
    "seed": [12.0, 2.0, 1.0, 0.08, 0.30, 0.25],
}


def slope_episode(g, seed, secs=10.0, grade_pct=10.0, stance_mm=420):
    kp, kd, station, clamp, ki, ki_max = g
    pitch_deg = math.degrees(math.atan(grade_pct / 100.0))
    env = RobotEnv(banked_spec(0.0, pitch_deg), seed=seed, task="bob",
                   randomise=False)
    env.reset()
    legs = [i for i, a in enumerate(env.spec["actuators"])
            if a.get("mode") == "position"]
    for i in legs:
        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])
    # both legs the same, held there: this stage is about leaning, not levelling
    hp, kn = C.stance_to_leg(stance_mm / 1000.0)
    hold = np.array([hp, kn, hp, kn])
    s = Sensed(env)
    trim = math.radians(C.lean_trim_deg(stance_mm))
    slide = 0.0
    lean_i = 0.0
    effort = 0.0
    x0 = float(env.base_pos[0])
    for k in range(int(secs / env.control_dt)):
        lean, lrate, speed, odo = s.read()
        # it is not told the gradient - it works the lean out from sliding
        lean_i += ki * odo * env.control_dt
        lean_i = float(np.clip(lean_i, -abs(ki_max), abs(ki_max)))
        ref = (C.lean_reference(station, speed, odo, clamp)
               + brx.update(speed, 0.0, env.control_dt))
        u = C.wheel_command(kp, kd, lean, lrate, ref, trim + lean_i)
        effort += abs(u) * env.control_dt
        act = np.zeros(env.act_dim)
        for j, ai in enumerate(env.act_idx):
            act[j] = (np.clip((hold[legs.index(ai)] - env.act_center[ai])
                              / max(1e-9, env.act_span[ai]), -1, 1)
                      if ai in legs else u)
        env.step(act)
        slide += abs(float(env.base_pos[0]) - x0) * env.control_dt
        if env.up_z < 0.5:
            break
    # staying UP and staying PUT: sliding down the hill is the failure mode.
    # NOTE the slope's steady lean costs real held torque, and that is the
    # point - the effort charge prices honest work, not just chatter.
    return env.t - 1.5 * (slide / max(env.t, 1e-6)) - EFFORT_RATE * effort


def slope_score(g, secs=10.0):
    cases = [dict(grade_pct=0.0), dict(grade_pct=10.0), dict(grade_pct=-10.0),
             dict(grade_pct=18.0), dict(grade_pct=-18.0)]
    return float(np.mean([slope_episode(g, i, secs, **c)
                          for i, c in enumerate(cases)]))


# --- the tilting platform -----------------------------------------------
# A real rotating plate under the robot, driven by its own actuator. The
# gravity-slew world could not represent the BNO086: a gyro never sees
# gravity move, but it sees a PLATFORM rotate instantly - which is what the
# real module, mounted flat on the chassis top, reports as absolute tilt
# regardless of leg bounce. On this world the sim's fused estimator IS the
# BNO: orientation w.r.t. true vertical, gyro-fast, immune to leg jolts.
PLATFORM_XML = '''
    <body name="platform" pos="0 0 0.65">
      <joint name="plat_j" type="hinge" axis="1 0 0" damping="200"/>
      <geom name="plat" type="box" size="2.5 2.5 0.08" mass="400"
            friction="1.5 0.005 0.0001" rgba="0.35 0.37 0.42 1"/>
    </body>'''
PLATFORM_ACT = '''
    <position name="plat_a" joint="plat_j" kp="2000000" kv="200000"
              ctrlrange="-0.35 0.35" forcerange="-1e6 1e6"/>'''


def platform_env(seed=0):
    """A RobotEnv standing on the actuated plate. The platform body is
    appended AFTER the robot so every existing qpos/actuator index is
    untouched; its hinge is the last dof and its actuator the last ctrl."""
    import mjcf
    env = RobotEnv(SPEC, seed=seed, task="bob", randomise=False)
    xml = mjcf.spec_to_mjcf(env.spec, props=env.props)
    xml = xml.replace("</worldbody>", PLATFORM_XML + "\n  </worldbody>")
    xml = xml.replace("</actuator>", PLATFORM_ACT + "\n  </actuator>")
    env.model = mujoco.MjModel.from_xml_string(xml)
    env.data = mujoco.MjData(env.model)
    env._rebind()
    # the env drives only the robot's actuators; the platform's is ours.
    # _rebind re-read the ranges from the 7-actuator model, which broke the
    # 6-wide command math in step()
    n = len(env.spec["actuators"])
    env.ctrl_lo = env.ctrl_lo[:n]
    env.ctrl_hi = env.ctrl_hi[:n]
    env.tau_limit = env.tau_limit[:n]
    env.reset()
    return env


# ===========================================================================
# The ladder. Survival first (stand, stop, take a shove), then locomotion
# (go somewhere, turn, cross a bump, keep grip), then living in the body
# (carry a load). Each stage inherits ONLY the slots that mean the same
# thing in its own layout - copying whole gain vectors across layouts is how
# stop's brake gain once landed in recover's brace slot (25x overdose), and
# how bob's trim scale was quietly seeding level's level-gain slot.
#
# Sit-to-stand and falling-well are deliberately NOT here: they are contact-
# rich, sequence-shaped skills a handful of gains cannot express - they
# belong to the PPO track, seeded by demonstrations from these controllers.
# Level comes BEFORE turn: the leveller's skill - leg-differential control
# of body roll - is exactly the mechanism turn's lean-in reuses, so turn
# inherits level's measured roll gains instead of rediscovering them.
STAGE_ORDER = ["balance", "balance_mid", "balance_crouch", "bob", "stop",
               "recover", "cruise", "level", "turn", "bump", "slip",
               "payload"]
# slope exists and is scoreable but is withdrawn from the ladder until the
# tilted-floor lean question is understood (AUDIT.md standing items)

# Entries are applied IN ORDER, later sources overwriting earlier for the
# slots they cover, and a source is skipped when unbaked. An entry is either
# (src, k) - copy the first k slots - or (src, [(si, di), ...]) - copy
# specific slots across DIFFERENT layouts, which prefix-copying cannot say
# (whole-vector inheritance is how stop's brake gain landed in recover's
# brace slot).
INHERIT = {
    # each bench seeds the next posture down - the crouch prefers the mid
    # bench's answer over the straight one when both exist
    "balance_mid":    [("balance", 4)],
    "balance_crouch": [("balance", 4), ("balance_mid", 4)],
    # bob inherits NOTHING: its wheel gains are interpolated live from the
    # balance bakes and its own three slots (trim scale, leg speed, level
    # shift) have no earlier owner
    # stop inherits nothing: its wheel gains ride the benches live
    "level":   [("balance", 4), ("balance_mid", 4)],
    "slope":   [("balance", 4)],
    "cruise":  [("balance", 4)],
    "turn":    [("cruise", 6),
                # level gain -> roll gain, level damping -> roll damping:
                # the same dh convention, the same roll reference
                ("level", [(4, 9), (5, 10)])],
    "bump":    [("cruise", 6)],
    "slip":    [("cruise", 6)],
    # bob's NEW layout: slot 0 trim scale, slot 1 leg speed
    "payload": [("balance", 4), ("bob", [(0, 4), (1, 5)])],
}


# --- stage: cruise ----------------------------------------------------------
# Hold a commanded road speed through a whole profile: pull away, cruise,
# slow to half, reverse, stop. This is the single most-used skill in the
# robot's life and it was hand-coded assist until now. The two new numbers
# are the speed loop's own: how hard the measured speed is filtered before
# it drives the lean target (raw speed feedback is the drive limit cycle),
# and how fast the speed command may ramp.
STAGES["cruise"] = {
    "names": ["kp (lean)", "kd (lean rate)", "station", "lean clamp",
              "speed filter", "accel"],
    "seed": [12.0, 2.0, 1.0, 0.08, 0.10, 1.2],
}


def cruise_profile(f, v_top):
    """Fractional-time speed profile: out, cruise, half, reverse, home."""
    return (0.0 if f < 0.10 else v_top if f < 0.45 else 0.5 * v_top
            if f < 0.60 else -0.5 * v_top if f < 0.85 else 0.0)


def cruise_episode(g, seed, secs=8.0, v_top=1.0):
    # THE LAW LIVES IN firmware.DriveLaw - this episode is the exam
    # around it (profile, scoring), not a second copy of the control
    env = RobotEnv(SPEC, seed=seed, task="balance", randomise=False)
    env.reset()
    s = Sensed(env)
    law = FW.DriveLaw(g, base_reflex(), env.control_dt)
    err = eff = 0.0
    n = int(secs / env.control_dt)
    for k in range(n):
        want = cruise_profile(k / n, v_top)
        lean, lrate, speed, _ = s.read()
        u, v_cmd = law.tick(want, lean, lrate, speed)
        eff += abs(u) * env.control_dt
        env.step(np.array([u, u]))
        # scored on TRUE speed - the filter is the controller's business,
        # the score's business is whether the body actually tracks
        err += abs(float(env.vel_local[0]) - v_cmd) * env.control_dt
        if env.up_z < 0.5:
            break
    return env.t - 1.5 * err - EFFORT_RATE * eff


def cruise_score(g, secs=8.0):
    # null first: a cruiser that cannot simply STAND has no business moving
    cases = [dict(v_top=0.0), dict(v_top=0.5), dict(v_top=1.0),
             dict(v_top=1.5)]
    return float(np.mean([cruise_episode(g, i, secs, **c)
                          for i, c in enumerate(cases)]))


# --- stage: turn ------------------------------------------------------------
# Track a commanded yaw RATE while holding a road speed - with the LEGS in
# the loop. The first version ran legs-locked-straight: the tallest possible
# pose, zero roll authority, no way to lean into the corner - and it fell
# after turns exactly the way an uncompensated inverted pendulum under
# centrifugal load must. The wheels share one axle and can do nothing about
# roll, so the only anti-centrifugal handle is a leg-length difference.
#
# SIGN, measured three ways (the first two disagreed and the label was the
# bug): +dh (left leg longer) raises the left chassis edge 17mm above the
# right - the body leans RIGHT, away from the longer leg (world-coordinate
# proof, not an angle formula). Leaning INTO a left turn therefore means
# the INSIDE leg squats: dh negative for +wz - so the feedforward seed is
# NEGATIVE. An earlier stance-420 sweep that preferred the outward lean was
# confounded: 420 is below the documented station-hold floor and its falls
# were stance falls, not turn falls.
#
# Deeper finding (turn_mechanism probe, stance 460): at working stance the
# turn deaths are NOT roll deaths - peak roll stays under 3 deg while the
# robot winds up forward speed (2.3 m/s measured in a "spot" turn) and
# pitches over at ~5s. Centrifugal lean matters at speed (1.5 m/s at
# 1 rad/s needs ~9 deg), but the pitch/speed slots are what decide whether
# turning survives at all - the same family as NEXT.md's old drive limit
# cycle. If the search cannot hold this, the coupling goes back on the
# blocker list with the evidence above.
STAGES["turn"] = {
    "names": ["kp (lean)", "kd (lean rate)", "station", "lean clamp",
              "speed filter", "accel", "yaw gain", "yaw clamp",
              "lean-in ff", "roll gain", "roll damping", "load balance",
              "stance offset", "curve budget"],
    # lean-in ff seeded NEGATIVE (inside leg squats - the world-coordinate
    # proof above). roll gain/damping seed near zero and inherit level's
    # measured values once level is baked.
    #
    # "load balance" closes the loop the operator described: measure the
    # per-wheel contact force (the touch sensors - on hardware, motor
    # current), and when the OUTSIDE wheel loads up in a corner, drop the
    # inside leg until the load rebalances. Left turn -> right wheel loads
    # -> imbalance positive -> lean left -> dh negative, so the seed is
    # negative. This finds the coordinated lean from measurement instead of
    # from v*wz/g, which also makes it honest under payload and slope.
    #
    # stance seeded 460: MEASURED (null case) - 440 scores 0.25 and 450
    # scores 0.44 where 460 scores 5.96, so under the peak-torque legs the
    # practical station-hold floor sits just below 460, far above the
    # documented 420 (a finding in its own right - the hold boundary moved
    # with the leg physics). Lean authority is therefore one-sided: the
    # corner lean comes from DROPPING THE INSIDE LEG, which is precisely
    # the operator's scheme, and lowers the body in corners as a bonus.
    # "curve budget" is the lateral-acceleration allowance (m/s^2) that
    # ties speed to turn radius: a = v*w = v^2/R, so the budget IS the
    # ideal-radius-for-speed rule, learned rather than guessed. A commanded
    # corner tighter than the budget allows is widened by the governor to
    # the radius the body can hold, at a mild score charge for the forgone
    # tightness so "never turn" cannot win.
    "seed": [12.0, 2.0, 1.0, 0.08, 0.10, 1.2, 0.35, 0.30, -0.03, 0.0,
             0.02, -0.05, 0.6, 1.0],
}


def turn_episode(g, seed, secs=8.0, v_top=0.5, wz=1.0):
    (kp, kd, station, clamp, v_tau, accel, k_yaw, t_max,
     rff, rk, rd, kl, stance_g, a_bud) = g
    env = RobotEnv(SPEC, seed=seed, task="bob", randomise=False)
    env.reset()
    legs = [i for i, a in enumerate(env.spec["actuators"])
            if a.get("mode") == "position"]
    for i in legs:
        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])
    touch_adr = [env.model.sensor_adr[mujoco.mj_name2id(
        env.model, mujoco.mjtObj.mjOBJ_SENSOR, nm)]
        for nm in ("l_touch", "r_touch")]
    stance = 455.0            # the tall finish, absolute, not searchable
    trim = math.radians(C.lean_trim_deg(stance))
    s = Sensed(env)
    filt = C.SpeedFilter(abs(v_tau))
    accel = max(0.3, abs(accel))
    k_yaw = abs(k_yaw)
    t_max = min(0.5, max(0.02, abs(t_max)))
    cur = np.array([float(env.data.qpos[env.model.jnt_qposadr[
        mujoco.mj_name2id(env.model, mujoco.mjtObj.mjOBJ_JOINT, n)]])
        for n in ("l_hip", "l_knee", "r_hip", "r_knee")])
    brx = base_reflex()
    v_cmd = 0.0
    oerr = 0.0
    ff_lat = 0.0
    gate = 0.0
    fb_slow = 0.0
    yerr = verr = rerr = sliperr = settle = eff = 0.0
    n = int(secs / env.control_dt)
    for k in range(n):
        f = k / n
        want = 0.0 if f < 0.15 else v_top
        # one corner each way, then a 30% settle window: "stumbling after
        # the turn" is the exit transient, so the exit is a scored phase of
        # its own, not an afterthought
        yaw_cmd = 0.0 if f < 0.25 else wz if f < 0.50 else -wz if f < 0.70 else 0.0
        v_cmd += float(np.clip(want - v_cmd,
                               -accel * env.control_dt, accel * env.control_dt))
        lean, lrate, speed, _ = s.read()
        v = filt.update(speed, env.control_dt)
        oerr += (v - v_cmd) * env.control_dt
        # THE CURVATURE GOVERNOR: the ideal radius for the speed. Lateral
        # acceleration is v*w (= v^2/R), so the learned budget caps the
        # executed yaw rate at what the body can lean against; a tighter
        # command is widened to the holdable radius. Everything downstream
        # (yaw loop, lean, roll reference) follows the GOVERNED curve so
        # radius, speed and lean stay one consistent manoeuvre.
        yaw_eff = yaw_cmd
        if abs(v) > 0.15 and abs(yaw_cmd) > 1e-6:
            w_cap = max(0.1, abs(a_bud)) / abs(v)
            yaw_eff = float(np.clip(yaw_cmd, -w_cap, w_cap))
        # the commanded lateral acceleration, low-passed (tau 0.25s) so the
        # lean walks in with the corner and RECOVERS GRADUALLY after it -
        # the published controllers recover the roll gradually post-turn
        # rather than snapping it back while the yaw is still decaying
        ff_lat += (env.control_dt / max(env.control_dt, 0.25)) * (
            v * yaw_eff - ff_lat)
        # lean INTO the turn (feedforward), plus the leveller's loop:
        # proportional roll feedback and rate damping, inherited from level
        roll_rate = float(env.sensors.noisy("gyro", env.data.sensordata,
                                            env._raw)[0])
        roll_now = math.atan2(env.est_up[1], max(1e-6, env.est_up[2]))
        # the measured-force loop: normalised left/right wheel load
        # imbalance (outside wheel loads up in a corner); drop the inside
        # leg until it rebalances. Guarded low so a momentary unload (a
        # bump, a hop) cannot command a lurch.
        fl = float(env.data.sensordata[touch_adr[0]])
        fr = float(env.data.sensordata[touch_adr[1]])
        imb = (fr - fl) / max(25.0, fr + fl)
        # The corner-lean machinery ENGAGES WITH THE TURN and stands down at
        # rest (fade over 0.3s). Ungated, the measured-feedback terms close
        # a lagged loop on their own consequences at standstill - drop a
        # leg, that wheel loads, the loop commands the opposite lean - and
        # the robot bobs side to side on its knees while doing nothing,
        # exactly as the operator observed. Standing means still legs.
        gate += (env.control_dt / 0.3) * (
            (1.0 if abs(yaw_cmd) > 0.05 else 0.0) - gate)
        # the measured-feedback terms TRIM, they do not servo: through a
        # 0.4s low-pass they settle the lean toward the force sweet spot,
        # while the feedforward carries the fast dip. Fast lagged feedback
        # through the leg servos is an oscillator - the alternating
        # left-right dip the operator saw mid-corner.
        fb = rk * roll_now + rd * roll_rate + kl * imb
        fb_slow += (env.control_dt / 0.4) * (fb - fb_slow)
        dh = float(np.clip(gate * (rff * ff_lat + fb_slow), -0.05, 0.05))
        hl, kl = C.stance_to_leg(max(0.20, min(0.462, stance / 1000.0 + dh)))
        hr, kr = C.stance_to_leg(max(0.20, min(0.462, stance / 1000.0 - dh)))
        cur = C.slew(cur, np.array([hl, kl, hr, kr]), math.radians(120),
                     env.control_dt)
        ref = (C.lean_reference(station, v - v_cmd, oerr, clamp)
               + brx.update(speed, v_cmd, env.control_dt))
        u = C.wheel_command(kp, kd, lean, lrate, ref, trim)
        yaw_rate = float(env.sensors.noisy("gyro", env.data.sensordata,
                                           env._raw)[2])
        # sign measured: +differential on (l = u+t, r = u-t) yaws NEGATIVE
        t = float(np.clip(-k_yaw * (yaw_eff - yaw_rate), -t_max, t_max))
        # balance first, turning with what is left
        t = float(np.clip(t, -(1.0 - abs(u)), 1.0 - abs(u)))
        eff += abs(u) * env.control_dt
        act = np.zeros(env.act_dim)
        for j, ai in enumerate(env.act_idx):
            if ai in legs:
                act[j] = np.clip((cur[legs.index(ai)] - env.act_center[ai])
                                 / max(1e-9, env.act_span[ai]), -1, 1)
            else:
                nm = env.spec["actuators"][ai]["name"]
                act[j] = float(np.clip(u + (t if nm.startswith("l") else -t),
                                       -1, 1))
        env.step(act)
        yerr += abs(float(env.gyro_true[2]) - yaw_eff) * env.control_dt
        # the governor's forgone tightness costs, mildly - shrinking the
        # budget to zero and never turning cannot win
        yerr += 0.3 * abs(yaw_cmd - yaw_eff) * env.control_dt
        verr += abs(float(env.vel_local[0]) - v_cmd) * env.control_dt
        # roll is scored against the COORDINATED-TURN reference, not zero:
        # the correct cornering lean (roll balancing the centripetal force,
        # negative = into a left turn in this convention) is what the
        # published controllers command - a flat-roll charge fought the very
        # lean the corner requires
        roll = math.atan2(env.est_up[1], max(1e-6, env.est_up[2]))
        roll_ref = -math.atan(ff_lat / 9.81)
        rerr += max(0.0, abs(roll - roll_ref) - 0.025) * env.control_dt
        # lateral slip: sideways velocity is the stumble building - the one
        # signal every published reward charges that this score did not
        sliperr += abs(float(env.vel_local[1])) * env.control_dt
        if f >= 0.70:
            # the settle window: residual yaw and roll after the corner
            settle += (abs(float(env.gyro_true[2]))
                       + 2.0 * abs(roll)) * env.control_dt
        if env.up_z < 0.5:
            break
    return (env.t - 0.8 * yerr - 0.8 * verr - 1.2 * rerr - 0.8 * sliperr
            - 0.5 * settle - EFFORT_RATE * eff)


def turn_score(g, secs=8.0):
    # GRADED corners, like bump's heights and slip's grades: at full rate
    # (1.0 rad/s at speed) every seed candidate fell, every score sat at the
    # fall floor, and the search had no hill - round 19 of one run produced
    # a best of 0.33/10 and a roll gain of 12 flapping the robot from wheel
    # to wheel. Gentle corners first; raise the rates once a bake holds.
    cases = [dict(v_top=0.0, wz=0.0),          # null: stand still
             dict(v_top=0.0, wz=0.6),          # gentle spot turn
             dict(v_top=0.4, wz=0.8),          # moderate arc
             dict(v_top=0.8, wz=0.6)]          # faster, wider arc
    return float(np.mean([turn_episode(g, i, secs, **c)
                          for i, c in enumerate(cases)]))


# --- stage: bump ------------------------------------------------------------
# Cross door-threshold ridges at speed without falling. This is where the
# legs stop being posture and start being suspension - and where the peak-
# torque leg bound earns its keep. Ridges are PROPS (drawn by the bench and
# present in the physics), heights chosen around real thresholds.
BUMP_RIDGE_X = (0.8, 1.7, 2.6)


def bump_props(h):
    if not h:
        return []
    return [{"name": f"ridge{i}", "type": "box",
             "size": [0.02, 2.5, h / 2.0], "pos": [x, 0.0, h / 2.0],
             "rgba": [0.30, 0.32, 0.38, 1.0],
             "friction": [1.0, 0.005, 0.0001]}
            for i, x in enumerate(BUMP_RIDGE_X)]


def bump_env(seed=0, h=0.010):
    """A bob-task env whose world carries the ridge props."""
    import mjcf
    env = RobotEnv(SPEC, seed=seed, task="bob", randomise=False)
    if h:
        env.props = list(env.props or []) + bump_props(h)
        xml = mjcf.spec_to_mjcf(env.spec, props=env.props)
        env.model = mujoco.MjModel.from_xml_string(xml)
        env.data = mujoco.MjData(env.model)
        env._rebind()
    env.reset()
    return env


STAGES["bump"] = {
    "names": ["kp (lean)", "kd (lean rate)", "station", "lean clamp",
              "speed filter", "accel", "stance offset", "leg speed"],
    "seed": [12.0, 2.0, 1.0, 0.08, 0.10, 1.2, 0.2, 1.0],
}


def bump_episode(g, seed, secs=8.0, h=0.010, v_cruise=0.8):
    (kp, kd, station, clamp, v_tau, accel, stance_g, leg_speed) = g
    env = bump_env(seed, h)
    legs = [i for i, a in enumerate(env.spec["actuators"])
            if a.get("mode") == "position"]
    for i in legs:
        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])
    stance = 455.0            # the tall finish, absolute, not searchable
    rate = C.leg_rate(leg_speed)
    trim = math.radians(C.lean_trim_deg(stance))
    s = Sensed(env)
    cur = np.array([float(env.data.qpos[env.model.jnt_qposadr[
        mujoco.mj_name2id(env.model, mujoco.mjtObj.mjOBJ_JOINT, n)]])
        for n in ("l_hip", "l_knee", "r_hip", "r_knee")])
    filt = C.SpeedFilter(abs(v_tau))
    accel = max(0.3, abs(accel))
    brx = base_reflex()
    v_cmd = 0.0
    oerr = 0.0
    err = eff = 0.0
    n = int(secs / env.control_dt)
    hp, kn = C.stance_to_leg(stance / 1000.0)
    want = np.array([hp, kn, hp, kn])
    for k in range(n):
        v_want = 0.0 if k * env.control_dt < 1.0 else v_cruise
        v_cmd += float(np.clip(v_want - v_cmd,
                               -accel * env.control_dt, accel * env.control_dt))
        cur = C.slew(cur, want, rate, env.control_dt)
        lean, lrate, speed, _ = s.read()
        v = filt.update(speed, env.control_dt)
        oerr += (v - v_cmd) * env.control_dt
        ref = (C.lean_reference(station, v - v_cmd, oerr, clamp)
               + brx.update(speed, v_cmd, env.control_dt))
        u = C.wheel_command(kp, kd, lean, lrate, ref, trim)
        eff += abs(u) * env.control_dt
        act = np.zeros(env.act_dim)
        for j, ai in enumerate(env.act_idx):
            act[j] = (np.clip((cur[legs.index(ai)] - env.act_center[ai])
                              / max(1e-9, env.act_span[ai]), -1, 1)
                      if ai in legs else u)
        env.step(act)
        err += abs(float(env.vel_local[0]) - v_cmd) * env.control_dt
        if env.up_z < 0.5:
            break
    return env.t - 1.5 * err - EFFORT_RATE * eff


def bump_score(g, secs=8.0):
    # null is FLAT ground: crossing nothing must not regress plain driving
    cases = [dict(h=0.0), dict(h=0.005), dict(h=0.010), dict(h=0.015)]
    return float(np.mean([bump_episode(g, i, secs, **c)
                          for i, c in enumerate(cases)]))


# --- stage: slip ------------------------------------------------------------
# Keep tracking through a slick patch. Mid-episode every slide-friction
# coefficient is scaled down (MuJoCo pairs take the element-wise max of the
# two geoms, so lowering the floor alone does nothing while the tyre says
# 1.5 - the scale hits everything, which for a robot whose only contacts
# are wheel-on-floor is exactly "the ground got slick"). The one new number
# is a grip ceiling on the wheel command - on ice you must ask for less.
STAGES["slip"] = {
    "names": ["kp (lean)", "kd (lean rate)", "station", "lean clamp",
              "speed filter", "accel", "grip ceiling"],
    "seed": [12.0, 2.0, 1.0, 0.08, 0.10, 1.2, 1.0],
}


def slip_episode(g, seed, secs=8.0, mu_scale=1.0, v_top=0.8):
    kp, kd, station, clamp, v_tau, accel, grip = g
    env = RobotEnv(SPEC, seed=seed, task="balance", randomise=False)
    env.reset()
    mu0 = env.model.geom_friction[:, 0].copy()
    u_cap = min(1.0, max(0.2, abs(grip)))
    s = Sensed(env)
    brx = base_reflex()
    filt = C.SpeedFilter(abs(v_tau))
    accel = max(0.3, abs(accel))
    v_cmd = 0.0
    oerr = 0.0
    err = eff = 0.0
    n = int(secs / env.control_dt)
    for k in range(n):
        t = k * env.control_dt
        on_ice = 0.40 * secs <= t < 0.65 * secs
        env.model.geom_friction[:, 0] = mu0 * (mu_scale if on_ice else 1.0)
        v_want = 0.0 if t < 1.0 else v_top
        v_cmd += float(np.clip(v_want - v_cmd,
                               -accel * env.control_dt, accel * env.control_dt))
        lean, lrate, speed, _ = s.read()
        v = filt.update(speed, env.control_dt)
        oerr += (v - v_cmd) * env.control_dt
        ref = (C.lean_reference(station, v - v_cmd, oerr, clamp)
               + brx.update(speed, v_cmd, env.control_dt))
        u = float(np.clip(C.wheel_command(kp, kd, lean, lrate, ref),
                          -u_cap, u_cap))
        eff += abs(u) * env.control_dt
        env.step(np.array([u, u]))
        err += abs(float(env.vel_local[0]) - v_cmd) * env.control_dt
        if env.up_z < 0.5:
            break
    env.model.geom_friction[:, 0] = mu0
    return env.t - 1.5 * err - EFFORT_RATE * eff


def slip_score(g, secs=8.0):
    cases = [dict(mu_scale=1.0), dict(mu_scale=0.4), dict(mu_scale=0.25),
             dict(mu_scale=0.15)]
    return float(np.mean([slip_episode(g, i, secs, **c)
                          for i, c in enumerate(cases)]))


# --- stage: payload ---------------------------------------------------------
# A robot that lives with people carries things. Mid-episode a mass appears
# on the chassis (sometimes off-centre), shifting the centre of mass; the
# stage learns the integral that re-finds the trim from the only evidence
# available - that it is drifting - exactly the mechanism slope uses.
STAGES["payload"] = {
    "names": ["kp (lean)", "kd (lean rate)", "station", "lean clamp",
              "trim scale", "leg speed", "load integral", "integral limit"],
    "seed": [12.0, 2.0, 1.0, 0.08, 1.0, 1.0, 0.30, 0.25],
}


def payload_episode(g, seed, secs=8.0, kg=0.5, at_x=0.0, at_z=0.13,
                    stance_mm=440):
    (kp, kd, station, clamp, trim_scale, leg_speed, ki, ki_max) = g
    env = RobotEnv(SPEC, seed=seed, task="bob", randomise=False)
    env.reset()
    legs = [i for i, a in enumerate(env.spec["actuators"])
            if a.get("mode") == "position"]
    for i in legs:
        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])
    s = Sensed(env)
    wheel_bids = [mujoco.mj_name2id(env.model, mujoco.mjtObj.mjOBJ_BODY, n)
                  for n in ("l_wheel", "r_wheel")]
    root_bid = mujoco.mj_name2id(env.model, mujoco.mjtObj.mjOBJ_BODY,
                                 "chassis")
    fdx = 0.0
    hp, kn = C.stance_to_leg(stance_mm / 1000.0)
    want = np.array([hp, kn, hp, kn])
    cur = np.array([float(env.data.qpos[env.model.jnt_qposadr[
        mujoco.mj_name2id(env.model, mujoco.mjtObj.mjOBJ_JOINT, n)]])
        for n in ("l_hip", "l_knee", "r_hip", "r_knee")])
    rate = C.leg_rate(leg_speed)
    trim = math.radians(C.lean_trim_deg(stance_mm)) * trim_scale
    brx = base_reflex()
    bid = mujoco.mj_name2id(env.model, mujoco.mjtObj.mjOBJ_BODY, "chassis")
    m0 = float(env.model.body_mass[bid])
    x0 = float(env.model.body_ipos[bid][0])
    z0 = float(env.model.body_ipos[bid][2])
    loaded = False
    lean_i = 0.0
    away = eff = 0.0
    n = int(secs / env.control_dt)
    for k in range(n):
        if not loaded and k * env.control_dt >= 2.0 and kg:
            # the load lands ON TOP (operator's spec: 2kg carried on the
            # chassis roof) - mass up, CoM toward where it sits, and
            # UP, which is what actually changes the balance problem
            env.model.body_mass[bid] = m0 + kg
            env.model.body_ipos[bid][0] = (m0 * x0 + kg * at_x) / (m0 + kg)
            env.model.body_ipos[bid][2] = (m0 * z0 + kg * at_z) / (m0 + kg)
            loaded = True
        cur = C.slew(cur, want, rate, env.control_dt)
        lean, lrate, speed, odo = s.read()
        lean_i += ki * odo * env.control_dt
        lean_i = float(np.clip(lean_i, -abs(ki_max), abs(ki_max)))
        ref = (C.lean_reference(station, speed, odo, clamp)
               + brx.update(speed, 0.0, env.control_dt))
        u = C.wheel_command(kp, kd, lean, lrate, ref, trim + lean_i)
        eff += abs(u) * env.control_dt
        act = np.zeros(env.act_dim)
        for j, ai in enumerate(env.act_idx):
            act[j] = (np.clip((cur[legs.index(ai)] - env.act_center[ai])
                              / max(1e-9, env.act_span[ai]), -1, 1)
                      if ai in legs else u)
        env.step(act)
        away += abs(float(env.base_pos[0])) * env.control_dt
        if env.up_z < 0.5:
            break
    env.model.body_mass[bid] = m0
    env.model.body_ipos[bid][0] = x0
    return env.t - 1.2 * (away / max(env.t, 1e-6)) - EFFORT_RATE * eff


def payload_score(g, secs=8.0):
    # the operator's requirement: at least 2kg carried on top. Graded
    # up to it, with one offset case - a parcel never sits centred.
    cases = [dict(kg=0.0), dict(kg=1.0), dict(kg=2.0),
             dict(kg=2.0, at_x=0.06)]
    return float(np.mean([payload_episode(g, i, secs, **c)
                          for i, c in enumerate(cases)]))


# --- stage: spin ------------------------------------------------------------
# Spin on the spot in indexed 90-degree sections, at every gate height. The
# operator's specification, literally: from normal balance, power the wheels
# in opposition and find the spin rate that sits a safe margin inside
# toppling - so the rate is a SEARCHED gain, pushed up by the reward for
# completed quarters and pushed down by wobble and falls; where it settles
# is the safe margin. Each quarter is judged on TRUE heading after settling
# (the judge may use ground truth; the controller must not) - the
# controller integrates the gyro, drift and all, exactly as the firmware
# would.
STAGES["spin"] = {
    # spin was the last stage still built pre-bench: fixed straight-leg
    # wheel gains at every stance (the crouch bench's own numbers existed
    # and were never used) and the odometry homing still in its wheels -
    # the exact underdamped loop exorcised from everything else in v26.
    # The operator asked why turn works and spin does not: the yaw law was
    # never the problem (a single crouched quarter traced clean), the
    # balancing underneath it was. Wheel gains now interpolate from the
    # benches at the commanded stance, kinematic trim direct, no odometry.
    # com hold: the operator's axis - a spin should turn ABOUT the centre
    # of mass, which the robot measures at all times; the feet walk
    # continuously to keep the axle under it (rotate's proven law), so the
    # spin axis passes through the CoM by construction instead of swinging
    # any residual offset in a circle every revolution.
    "names": ["leg speed", "yaw gain", "yaw clamp",
              "heading gain", "spin rate", "com hold", "bank gain",
              "bank damp", "cent trim", "bank ff", "skater rise"],
    # reseeded from the v83 wide search's discovery (15.41: every
    # height, both directions)
    "seed": [0.211, 0.3577, 0.1642, 3.2114, 6.6958,
             -1.6137, 0.1684, -0.0335, 0.005, -0.01, 1.5],
}
STAGE_ORDER.insert(STAGE_ORDER.index("turn") + 1, "spin")
# turn's yaw GAIN transfers; its yaw CLAMP does not - 0.05 is sized for
# gentle driving arcs and cannot brake a 1.5 rad/s spin (measured: the
# overshoot and the long exposed pull-back both trace to it)
INHERIT["spin"] = [("bob", [(1, 0)]), ("turn", [(6, 1)])]


def spin_episode(g, seed, secs=12.0, stance_mm=460,
                 quarters=(90.0, -90.0, 90.0), hold_wz=None):
    bpts = balance_points()
    env = RobotEnv(SPEC, seed=seed, task="bob", randomise=False)
    env.reset()
    legs = [i for i, a in enumerate(env.spec["actuators"])
            if a.get("mode") == "position"]
    for i in legs:
        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])
    s = Sensed(env)
    wheel_bids = [mujoco.mj_name2id(env.model, mujoco.mjtObj.mjOBJ_BODY, n)
                  for n in ("l_wheel", "r_wheel")]
    root_bid = mujoco.mj_name2id(env.model, mujoco.mjtObj.mjOBJ_BODY,
                                 "chassis")
    cur = np.array([float(env.data.qpos[env.model.jnt_qposadr[
        mujoco.mj_name2id(env.model, mujoco.mjtObj.mjOBJ_JOINT, n)]])
        for n in ("l_hip", "l_knee", "r_hip", "r_knee")])
    # THE LAW LIVES IN firmware.SpinLaw - the exam is the goal logic
    # (quarters, settling, the judged charges) around the same tick the
    # firmware runs in free roam
    law = FW.SpinLaw(g, stance_mm, interp_wheel_gains(bpts, stance_mm),
                     base_reflex(), env.control_dt)
    rate = law.leg_rate
    trim = law.trim

    def true_heading():
        q = env.base_quat
        return math.atan2(2 * (q[0] * q[3] + q[1] * q[2]),
                          1 - 2 * (q[2] * q[2] + q[3] * q[3]))

    h0_true = true_heading()
    h_est = 0.0                 # the controller's own gyro-integrated heading
    target = 0.0
    qi = -1                     # current quarter (starts after 1s of settle)
    settled_for = 0.0
    completed = 0
    miss = wobble = eff = away = asym = 0.0
    wheel_dofs = [env.model.jnt_dofadr[mujoco.mj_name2id(
        env.model, mujoco.mjtObj.mjOBJ_JOINT, n)]
        for n in ("l_wheel_j", "r_wheel_j")]
    x0, y0 = float(env.base_pos[0]), float(env.base_pos[1])
    n = int(secs / env.control_dt)
    for k in range(n):
        t_now = k * env.control_dt
        lean, lrate, speed, odo = s.read()
        gz = float(env.sensors.noisy("gyro", env.data.sensordata, env._raw)[2])
        h_est += gz * env.control_dt
        remaining = target - h_est
        # advance to the next quarter once this one is landed and quiet
        if abs(remaining) < math.radians(3.0) and abs(gz) < 0.25:
            settled_for += env.control_dt
        else:
            settled_for = 0.0
        # the first quarter waits for the FULL settle second: the
        # settled-advance path used to fire at 0.4s while the legs were
        # still descending into a crouch stance, spinning mid-transition
        if (qi < 0 and t_now >= 1.0) or (qi >= 0 and settled_for >= 0.4
                                         and qi < len(quarters) - 1):
            qi += 1
            if qi < len(quarters):
                target += math.radians(quarters[qi])
            settled_for = 0.0
        # spin toward the target: P on remaining angle, capped at the
        # learned safe rate - "add power until the correct spin speed".
        # THE OPERATOR'S HANDOVER: once the quarter has LANDED this is
        # just balance - no pull-back chase at spin rates (the falls were
        # at the stop, not the spin). A landed spin trims heading at a
        # gentle 0.3 rad/s cap; full authority re-engages only when a new
        # quarter (or a real disturbance) opens the error again.
        xm = env.data.xmat[root_bid].reshape(3, 3)
        _fx, _fy = float(xm[0][0]), float(xm[1][0])
        _fn = math.hypot(_fx, _fy) or 1.0
        com = env.data.subtree_com[root_bid]
        axp = 0.5 * (env.data.xpos[wheel_bids[0]]
                     + env.data.xpos[wheel_bids[1]])
        off = ((float(com[0]) - float(axp[0])) * _fx
               + (float(com[1]) - float(axp[1])) * _fy) / _fn
        _rollr = math.atan2(env.est_up[1], max(1e-6, env.est_up[2]))
        _rrate = float(env.sensors.noisy("gyro", env.data.sensordata,
                                         env._raw)[0])
        want, _, u, t, landed = law.tick(
            lean, lrate, speed, gz, _rollr, _rrate, off,
            hold_wz=hold_wz, remaining=remaining)
        cur = C.slew(cur, want, rate, env.control_dt)
        eff += abs(u) * env.control_dt
        act = np.zeros(env.act_dim)
        for j, ai in enumerate(env.act_idx):
            if ai in legs:
                act[j] = np.clip((cur[legs.index(ai)] - env.act_center[ai])
                                 / max(1e-9, env.act_span[ai]), -1, 1)
            else:
                nm = env.spec["actuators"][ai]["name"]
                act[j] = float(np.clip(u + (t if nm.startswith("l") else -t),
                                       -1, 1))
        env.step(act)
        # wobble: pitch beyond 4 degrees of the trim line while spinning is
        # the topple margin being spent
        wobble += max(0.0, abs(lean - trim) - math.radians(4.0)) * env.control_dt
        # LEVEL is paid for (operator: the dip in the chassis is what
        # causes the issues - and roll was never charged, only pitch).
        # 5 degrees of grace covers the deliberate lean-in; a folded-
        # leg grind beyond it bleeds points every tick instead of only
        # at the fall it always ends in
        _rollw = math.atan2(env.est_up[1], max(1e-6, env.est_up[2]))
        wobble += max(0.0, abs(_rollw) - math.radians(5.0)) * env.control_dt
        # ON THE SPOT means the whole path, not the endpoint (operator
        # caught the cheat): a one-wheel pivot orbits the planted wheel
        # and returns home after each revolution, so end-point drift
        # rode free - the integral of displacement prices the orbit
        away += math.hypot(float(env.base_pos[0]) - x0,
                           float(env.base_pos[1]) - y0) * env.control_dt
        # EVEN WHEELS (operator's rule): a true spot spin counter-
        # rotates at equal magnitudes - dramatic speed difference is
        # the pivot, and it is deducted directly
        _wl = abs(float(env.data.qvel[wheel_dofs[0]]))
        _wr = abs(float(env.data.qvel[wheel_dofs[1]]))
        asym += abs(_wl - _wr) * env.control_dt
        if env.up_z < 0.5:
            break
    # judged on TRUE heading: each landed quarter pays, missing the mark
    # charges by how far out the final heading is from the commanded total
    turned_true = true_heading() - h0_true
    want_total = math.radians(sum(quarters[:max(qi, 0) + 1] if qi >= 0
                                  else []))
    completed = max(0, qi + (1 if settled_for >= 0.4 else 0))
    miss = abs(turned_true - want_total)
    drift = math.hypot(float(env.base_pos[0]) - x0, float(env.base_pos[1]) - y0)
    return (env.t + 1.5 * completed - 2.0 * min(miss, 1.5)
            - 1.0 * wobble - 0.8 * min(drift, 2.0)
            - 3.0 * min(away / max(env.t, 1e-6), 1.0)
            - 1.2 * min(asym / max(env.t, 1e-6), 8.0)
            - EFFORT_RATE * eff)


def spin_score(g, secs=16.0):
    # the operator's simplification: spin is learned STANDING - the
    # benches already own balancing at every height, so spinning at a
    # crouch is composition (gate slider + this skill), not a harder
    # exam. The progression is theirs too: turn 90, then 180, then a full
    # 360, both ways, plus the null (a spinner that cannot simply stand
    # scores on nothing).
    # graded at every gate height now (operator's goal): the standing
    # quarters both ways, the null, and the mid/crouch spins the bank
    # makes possible
    # SCREENING ORDER + EARLY ABORT (operator: probing a known-dead
    # candidate through all nine cases is wasted learning): the null
    # and the two sustained holds run first - they are where hopeless
    # candidates die in seconds - and a candidate averaging under 3.0
    # after them is dismissed on that evidence instead of billed for
    # the six long quarter cases it cannot pass. ARS only ranks, so a
    # consistently pessimistic partial mean preserves the ordering;
    # any candidate good enough to matter always runs the full exam.
    cases = [dict(quarters=()),
             dict(quarters=(), hold_wz=1.2),
             dict(quarters=(), hold_wz=0.8, stance_mm=340),
             dict(quarters=(90.0, 180.0, 360.0)),
             dict(quarters=(-90.0, -180.0, -360.0)),
             dict(quarters=(90.0, 180.0), stance_mm=400),
             dict(quarters=(90.0, 180.0), stance_mm=340),
             dict(quarters=(-90.0, -180.0), stance_mm=340),
             dict(quarters=(180.0, -270.0))]
    scores = []
    for i, c in enumerate(cases):
        scores.append(spin_episode(g, i, max(secs, 16.0),
                                   **{"stance_mm": 460, **c}))
        if i == 2 and float(np.mean(scores)) < 3.0:
            return float(np.mean(scores))
    return float(np.mean(scores))


# --- stage: look ------------------------------------------------------------
# Point the chassis up or down ON COMMAND while holding a gate height and
# holding station - the camera-platform skill. The operator's spec: hip
# movement points the front of the chassis, the rest compensates while
# balancing. Mechanically the compensation is FOOT PLACEMENT: looking up
# moves the mass back, so the feet walk back underneath it until the
# commanded pitch is the sustainable equilibrium; the balance loop then
# holds that pitch instead of the trim-down slouch. The null case (look
# zero) is the operator's floor-gazing complaint: hold the chassis LEVEL at
# a crouched gate, which stock trim never did.
STAGES["rotate"] = {
    # the operator's restructure of "look": with the hips freed to a full
    # turn each way, pointing the camera is not a pitch trick - it is
    # CHASSIS ROTATION about the hip axis, anywhere on the circle, at any
    # bench stance. The legs hold stance and walk the feet under the
    # measured CoM; the wheels balance on bench-interpolated gains; the
    # hip angle IS the gaze axis. Conventions MEASURED (13 Aug): est lean
    # positive = nose UP; positive rot command = nose DOWN (ground side,
    # matching the old look slider); nose-down theta needs hip = hip0 -
    # theta; base qvel[4] positive = nose-down rate. Floor clearance
    # measured CLEAR through the full circle at all three stances.
    "names": ["rot rate", "lever slow", "com hold", "foot place",
              "static trim", "leg speed"],
    # rot rate: rad/s of commanded rotation. lever slow: fraction of the
    # rate given back near the +/-90 lever arcs (|sin| shape) - the
    # chassis is heaviest to carry horizontal. com hold: the measured
    # realtime CoM law (sign measured NEGATIVE, the naive derivation was
    # backwards). foot place: the measured near-constant +25mm. static
    # trim: position creep feeds the FEET, wheels keep no position memory.
    # rate reseeded 0.8 after the far side became visible (the est-lean
    # clamp had saturated at +/-90 and blinded everything beyond - fixing
    # it took the 180 case from 6.1 to 13.9 and the full wrap from 6.0 to
    # 13.7): 1.0 completes the circle rougher, 1.5 falls - the speed
    # reward pushes the search up this measured slope.
    "seed": [0.8, 0.5, -1.2, 0.025, 0.2, 2.0],
}
STAGE_ORDER.insert(STAGE_ORDER.index("bob") + 1, "rotate")
INHERIT["rotate"] = [("bob", [(1, 5)])]   # bob's leg speed -> rotate's


def rotate_episode(g, seed, secs=12.0, target_deg=90.0, stance_mm=460.0):
    """Rotate the chassis about the hip axis to a commanded angle, then
    hold it STILL - wheels balancing throughout, legs holding stance with
    the feet walking under the measured CoM. Positive = nose down; the
    circle is open to +/-360 with hard stops (no continuous spinning)."""
    (rot_rate, lever_slow, com_hold, foot_place, static_ki, leg_speed) = g
    bpts = balance_points()
    env = RobotEnv(SPEC, seed=seed, task="bob", randomise=False)
    env.reset()
    legs = [i for i, a in enumerate(env.spec["actuators"])
            if a.get("mode") == "position"]
    for i in legs:
        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])
    s = Sensed(env)
    brx = base_reflex()
    cur = np.array([float(env.data.qpos[env.model.jnt_qposadr[
        mujoco.mj_name2id(env.model, mujoco.mjtObj.mjOBJ_JOINT, n)]])
        for n in ("l_hip", "l_knee", "r_hip", "r_knee")])
    wheel_bids = [mujoco.mj_name2id(env.model, mujoco.mjtObj.mjOBJ_BODY, n)
                  for n in ("l_wheel", "r_wheel")]
    root_bid = mujoco.mj_name2id(env.model, mujoco.mjtObj.mjOBJ_BODY,
                                 "chassis")
    hip_jids = [mujoco.mj_name2id(env.model, mujoco.mjtObj.mjOBJ_JOINT, n)
                for n in ("l_hip", "r_hip")]
    kp, kd, station, clamp = [float(v) for v in
                              interp_wheel_gains(bpts, stance_mm)]
    tgt = math.radians(target_deg)
    rot_cmd = 0.0
    fdx = 0.0
    branch = -1.0          # knee back; mirrors when the face flips
    theta = 0.0            # continuous true rotation, gyro-integrated
    reach = 0.0
    t_arrive = None        # first moment the true angle reaches the target
    perr = drift_i = com_i = eff = jit = 0.0
    prev_cur = cur.copy()
    x0 = float(env.base_pos[0])
    n = int(secs / env.control_dt)
    for k in range(n):
        # slow through the lever arcs: the chassis is heaviest to carry
        # near horizontal (|sin| peaks at the +/-90 family)
        # lever-slow capped at 0.75: railing it to 0.9 was half of the
        # crawl exploit (crossings at 2 deg/s) - a crossing may be careful
        # but never stalled
        rr = max(0.05, abs(rot_rate)) * (
            1.0 - min(0.75, abs(lever_slow)) * abs(math.sin(rot_cmd)))
        rot_cmd += float(np.clip(tgt - rot_cmd, -rr * env.control_dt,
                                 rr * env.control_dt))
        lean, lrate_g, speed, odo = s.read()
        # realtime CoM hold: feet walk to keep the axle under the mass
        com_x = float(env.data.subtree_com[root_bid][0])
        ax_x = 0.5 * (float(env.data.xpos[wheel_bids[0]][0])
                      + float(env.data.xpos[wheel_bids[1]][0]))
        fdx += (env.control_dt / 0.04) * (com_hold * (com_x - ax_x) - fdx)
        dx = float(np.clip(foot_place + fdx
                           - np.clip(abs(static_ki) * odo, -0.08, 0.08),
                           -0.14, 0.14))
        # the operator's mirror rule: when the face points the other way
        # (past +/-90), the knee folds the other way too, so it stays
        # toward the BACK of whichever way the robot now faces. Hysteresis
        # (swap past ~95, revert under ~85) so a hold at exactly 90 never
        # chatters.
        _c = math.cos(rot_cmd)
        if branch < 0 and _c < -0.087:
            branch = +1.0
        elif branch > 0 and _c > 0.087:
            branch = -1.0
        # the operator's straight-leg crossing: the leg extends to full
        # height exactly as the rotation CROSSES the +/-90 family, where
        # the two knee folds MEET at the straight-leg singularity - the
        # mirror swap becomes continuous, no swing-through. Straighten
        # ONLY when a crossing actually lies ahead on the path (measured:
        # blending during a plain 90-degree HOLD cost 1.5 points - the
        # tallest pose is the hardest hold, and a hold needs no swap).
        _m = math.pi / 2.0
        _lo, _hi = sorted((rot_cmd - 0.3 * np.sign(tgt - rot_cmd), tgt))
        # a boundary AT the target is a destination, not a crossing - a
        # plain 90-degree hold must not straighten (measured: it cost 1.5)
        _cross_ahead = any(
            nn % 2 != 0 and abs(nn * _m - tgt) > 0.12
            for nn in range(math.ceil(_lo / _m - 1e-9),
                            math.floor(_hi / _m + 1e-9) + 1))
        _bl = (abs(math.sin(rot_cmd)) ** 4) if _cross_ahead else 0.0
        stance_eff = stance_mm + (462.0 - stance_mm) * _bl
        hip0, knee0 = C.stance_to_leg(stance_eff / 1000.0, dx, branch)
        want = np.array([hip0 - rot_cmd, knee0, hip0 - rot_cmd, knee0])
        cur = C.slew(cur, want, C.leg_rate(leg_speed), env.control_dt)
        # settled gaze must mean still legs (the paid-spasm lesson)
        if abs(tgt - rot_cmd) < 0.02:
            jit += float(np.mean(np.abs(cur - prev_cur)))
        prev_cur = cur.copy()
        ref = (C.lean_reference(station, speed, 0.0, clamp)
               + brx.update(speed, 0.0, env.control_dt))
        # the chassis lives on a circle: wrap the lean error so 350 and
        # -10 are neighbours, not a full-turn correction
        e = lean - ref - (-rot_cmd)
        e = math.atan2(math.sin(e), math.cos(e))
        u = float(np.clip(-kp * e + kd * lrate_g, -1.0, 1.0))
        eff += abs(u) * env.control_dt
        act = np.zeros(env.act_dim)
        for j, ai in enumerate(env.act_idx):
            act[j] = (np.clip((cur[legs.index(ai)] - env.act_center[ai])
                              / max(1e-9, env.act_span[ai]), -1, 1)
                      if ai in legs else u)
        env.step(act)
        # continuous true rotation (the up vector cannot tell 270 from
        # -90); qvel[4] positive = nose-down rate, same as the command
        theta += float(env.data.qvel[4]) * env.control_dt
        if abs(tgt) > 1e-6:
            reach = max(reach, float(np.clip(theta / tgt, 0.0, 1.0)))
            if t_arrive is None and abs(theta - tgt) < 0.09:
                t_arrive = env.t
        # judged against the TARGET, not the moving command: tracking a
        # crawling command is free, and the search found that exploit in
        # three rounds (rot rate 0.39, lever railed - 2 deg/s through the
        # crossings). Distance-to-target integrated over time IS the
        # journey clock: every second spent not-there is charged, so
        # speed is the score, wobble is the score, and a crawl pays for
        # its whole leisurely arc.
        perr += abs(theta - tgt) * env.control_dt
        drift_i += abs(float(env.base_pos[0]) - x0) * env.control_dt
        com_i += abs(com_x - ax_x) * env.control_dt
        # fallen = the LEGS collapsed, judged by hip-axis height (the
        # chassis pointing somewhere deliberate is not a fall - up_z is
        # meaningless on the far side of the circle)
        hip_z = 0.5 * (float(env.data.xanchor[hip_jids[0]][2])
                       + float(env.data.xanchor[hip_jids[1]][2]))
        if hip_z < 0.55 * (stance_mm / 1000.0):
            break
    x_end = abs(float(env.base_pos[0]) - x0)
    prog = reach if abs(tgt) > 1e-6 else 1.0
    # SPEED PAYS (operator's directive): arriving at the target early
    # earns up to 3 points, scaled by how much of the episode was saved -
    # a rotation that lands in a third of the time banks 2. Falls already
    # forfeit upright time, so speed is never free of its risk.
    quick = (3.0 * max(0.0, 1.0 - t_arrive / secs)
             if t_arrive is not None else 0.0)
    return (env.t + 2.0 * prog + quick
            - 2.0 * (perr / max(env.t, 1e-6))
            - 2.0 * (drift_i / max(env.t, 1e-6))
            - 1.5 * min(x_end, 2.0)
            - 10.0 * (com_i / max(env.t, 1e-6))
            - 2.0 * jit - EFFORT_RATE * eff)


def rotate_score(g, secs=12.0):
    # the circle, graded, at all three bench stances: the working arcs,
    # the +/-90 lever holds (the hardest), the inverted hang (past 90 the
    # CoM drops and life gets easier), and the full wrap. Continuous
    # credit: an unreached target still pays per degree of progress.
    cases = [dict(target_deg=45, stance_mm=460),
             dict(target_deg=-45, stance_mm=400),
             dict(target_deg=90, stance_mm=460),
             dict(target_deg=-90, stance_mm=340),
             dict(target_deg=180, stance_mm=400),
             dict(target_deg=-180, stance_mm=460),
             dict(target_deg=360, stance_mm=400)]
    return float(np.mean([rotate_episode(g, i, secs, **c)
                          for i, c in enumerate(cases)]))


# --- stage: rise ------------------------------------------------------------
# Stand up from rest and hand over to balance. The rest posture is the
# operator's: collapsed onto the knee rollers in FRONT, main wheels grounded
# BEHIND - which is a statically stable two-line base, so the rise is
# quasi-static until the rollers lift and the balance loop catches the
# handover. The sit pose lives in the knee-FORWARD fold, so the whole rise
# stays inside that fold family: no swing-through anywhere.
STAGES["rise"] = {
    # COMPLETE OVERHAUL (operator's directive): the scripted rock-over is
    # gone. The approach to the LAUNCH POSE is fixed choreography (heap ->
    # lay 48/26 -> launch 53/-141, slow, wheels locked - the operator's
    # snapshot, reached smoothly); everything after is LEARNED - the
    # search experiments with forward drive and lift to reach full height
    # as quickly as safely possible, and discovers its own balance
    # handover height. The reward is the operator's, verbatim: CHASSIS
    # HEIGHT ABOVE THE GROUND, time-averaged - slowness costs area under
    # the height curve, falling collapses it, nothing else is scored.
    "names": ["drive", "drive taper", "lift rate", "balance-on mm"],
    # seed = the measured best of the launch grid: hard drive + fast
    # lift + an EARLY balance handover (360mm) scored 7.12 where every
    # timid combination sat at 3-4 - the operator's "as quick as safely
    # possible" emerging from the height integral by itself
    "seed": [0.5, 0.0, 0.7, 360.0],
}
STAGE_ORDER.insert(STAGE_ORDER.index("bob") + 1, "rise")
# rise inherits nothing: its wheel gains ride the benches live


def limp_settle(env, secs=1.2):
    """ALL motors off: zero every actuator's gain and bias, let the body
    collapse under gravity to its true rest, then power back on."""
    m, d = env.model, env.data
    d.qvel[:] = 0.0            # collapse from stillness, not mid-bounce
    gain0 = m.actuator_gainprm.copy()
    bias0 = m.actuator_biasprm.copy()
    m.actuator_gainprm[:, 0] = 0.0
    m.actuator_biasprm[:, 1:3] = 0.0
    nst = max(1, int(round(env.control_dt / m.opt.timestep)))
    for _ in range(int(secs / env.control_dt)):
        mujoco.mj_step(m, d, nst)
    m.actuator_gainprm[:] = gain0
    m.actuator_biasprm[:] = bias0


SIT_POSE = {"l_hip": -0.5235987755982988, "l_knee": 2.356194490192345,
            "r_hip": -0.5235987755982988, "r_knee": 2.356194490192345}


def place_sitting(env):
    """Put the robot in the operator's rest: rollers front, wheels behind,
    chassis low and level."""
    for n, v in SIT_POSE.items():
        env.data.qpos[env.model.jnt_qposadr[
            mujoco.mj_name2id(env.model, mujoco.mjtObj.mjOBJ_JOINT, n)]] = v
    env.data.qpos[0:3] = [0.0, 0.0, 0.27]   # slightly high; it settles on
    env.data.qpos[3:7] = [1.0, 0.0, 0.0, 0.0]
    env.data.qvel[:] = 0.0
    mujoco.mj_forward(env.model, env.data)


RISE_LAY = (48.0, 26.0)       # the operator's lay-out
RISE_LAUNCH = (53.0, -141.0)  # the operator's launch pose (snapshot 13:41)


def rise_episode(g, seed, secs=15.0):
    (drive, taper, lift_rate, bal_mm) = g
    bpts = balance_points()
    env = RobotEnv(SPEC, seed=seed, task="bob", randomise=False)
    env.reset()
    env.min_h = 0.02
    place_sitting(env)
    limp_settle(env, 1.2)               # the dead start
    legs = [i for i, a in enumerate(env.spec["actuators"])
            if a.get("mode") == "position"]
    for i in legs:
        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])
    s = Sensed(env)
    m, d = env.model, env.data
    jq = lambda nm: m.jnt_qposadr[mujoco.mj_name2id(
        m, mujoco.mjtObj.mjOBJ_JOINT, nm)]
    hip_jids = [mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_JOINT, nm)
                for nm in ("l_hip", "r_hip")]
    chassis_bid = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_BODY, "chassis")
    wq_adr = [jq(n) for n in ("l_wheel_j", "r_wheel_j")]
    wd_adr = [m.jnt_dofadr[mujoco.mj_name2id(
        m, mujoco.mjtObj.mjOBJ_JOINT, n)]
        for n in ("l_wheel_j", "r_wheel_j")]
    lock_q = [float(d.qpos[a]) for a in wq_adr]
    cur = np.array([float(d.qpos[jq(nm)])
                    for nm in ("l_hip", "l_knee", "r_hip", "r_knee")])
    lay = np.array([math.radians(RISE_LAY[0]),
                    math.radians(RISE_LAY[1])] * 2)
    launch = np.array([math.radians(RISE_LAUNCH[0]),
                       math.radians(RISE_LAUNCH[1])] * 2)
    st_h, st_k = C.stance_to_leg(0.455)
    stand = np.array([st_h, st_k] * 2)
    slow = math.radians(35.0)
    liftr = math.radians(150.0 * min(1.5, max(0.1, abs(lift_rate))))
    LAUNCH_HIP_Z = 0.165                  # snapshot height, for the taper
    phase = 0                             # 0 lay, 1 launch, 2 LEARNED
    h_sum = 0.0
    eff = 0.0
    x_ascent = None                       # where the learned ascent began
    n = int(secs / env.control_dt)
    for k in range(n):
        t_now = k * env.control_dt
        lean, lrate_g, speed, odo = s.read()
        hip_z = 0.5 * (float(d.xanchor[hip_jids[0]][2])
                       + float(d.xanchor[hip_jids[1]][2]))
        uw = None
        u = 0.0
        if phase == 0:
            tgt, rate_now = lay, slow
            uw = [float(np.clip(
                -3.0 * (float(d.qpos[wq_adr[i]]) - lock_q[i])
                - 0.3 * float(d.qvel[wd_adr[i]]), -1, 1)) for i in (0, 1)]
            if (t_now > 0.3 and float(np.max(np.abs(cur - lay)))
                    < math.radians(5.0)):
                phase = 1
        elif phase == 1:
            tgt, rate_now = launch, slow
            uw = [float(np.clip(
                -3.0 * (float(d.qpos[wq_adr[i]]) - lock_q[i])
                - 0.3 * float(d.qvel[wd_adr[i]]), -1, 1)) for i in (0, 1)]
            if float(np.max(np.abs(cur - launch))) < math.radians(5.0):
                phase = 2
                x_ascent = float(env.base_pos[0])
        else:
            # THE LEARNED ASCENT: forward drive (tapering with height)
            # and lift, until the learned handover height, then the
            # benches balance the rest of the way up
            tgt, rate_now = stand, liftr
            if 1000.0 * hip_z < max(200.0, min(462.0, abs(bal_mm))):
                u = float(np.clip(
                    drive + taper * (hip_z - LAUNCH_HIP_Z) / 0.29,
                    -1.0, 1.0))
            else:
                eff_st = float(np.clip(1000.0 * hip_z, 340.0, 462.0))
                kp, kd, station, clamp = interp_wheel_gains(bpts, eff_st)
                trim = math.radians(C.lean_trim_deg(eff_st))
                ref = C.lean_reference(station, speed, 0.0, clamp)
                u = C.wheel_command(kp, kd, lean, lrate_g, ref, trim)
        cur = C.slew(cur, tgt, rate_now, env.control_dt)
        eff += abs(u) * env.control_dt
        act = np.zeros(env.act_dim)
        wheel_i = 0
        for j, ai in enumerate(env.act_idx):
            if ai in legs:
                act[j] = np.clip(
                    (cur[legs.index(ai)] - env.act_center[ai])
                    / max(1e-9, env.act_span[ai]), -1, 1)
            else:
                act[j] = uw[wheel_i] if uw is not None else u
                wheel_i += 1
        env.step(act)
        # THE REWARD, the operator's verbatim: chassis height above the
        # ground, every single step - integrated over the whole episode
        # so slowness costs area and falling collapses the curve. No
        # phase bonuses, no position terms, nothing to game.
        h_sum += float(d.xpos[chassis_bid][2]) * env.control_dt
    # the operator's refinement: the stand should not need a run-up -
    # forward travel during the LEARNED ascent is charged (the scripted
    # approach's slide is choreography, not the law's choice), so the
    # search trades drive for a steeper, tighter rise
    travel = (abs(float(env.base_pos[0]) - x_ascent)
              if x_ascent is not None else 0.0)
    return (- feel_charge(env, "rise")
            + 12.0 * (h_sum / max(secs, 1e-6)) / 0.46
            - 1.5 * min(travel, 3.0)
            - EFFORT_RATE * eff)


def rise_score(g, secs=15.0):
    # three rises to different stands, three seeds of contact noise
    cases = [dict(), dict(), dict()]
    return float(np.mean([rise_episode(g, i, secs) for i in range(3)]))


# --- stage: sit (the shutdown) ----------------------------------------
MIRROR_TRIM = ((455.0, 1.0), (420.0, 3.0), (380.0, 5.0), (340.0, 8.0),
               (300.0, 10.0), (260.0, 14.0), (220.0, 18.0),
               (180.0, 22.0))


def mirror_trim_deg(stance_mm):
    """The MIRROR fold's trim, measured by drift-null (14 Aug): almost
    exactly the negated normal curve - the off pose lives on the other
    side of the straight-leg singularity and leans the other way."""
    pts = MIRROR_TRIM
    s = float(stance_mm)
    if s >= pts[0][0]:
        return pts[0][1]
    if s <= pts[-1][0]:
        return pts[-1][1]
    for (h1, t1), (h0, t0) in zip(pts, pts[1:]):
        if h0 <= s <= h1:
            f = (s - h0) / (h1 - h0)
            return t0 + f * (t1 - t0)
    return pts[-1][1]


STAGES["sit"] = {
    # THE SHUTDOWN, the operator's final design: 'literally copy the
    # gate and balance and just make it go lower, then turn off - the
    # knee goes the other way'. This is BOB'S LAW VERBATIM (its baked
    # gains, its CoM hold, its level shift) running in the MIRROR knee
    # fold - the off pose lives across the straight-leg singularity
    # (knees forward, wheels behind) - with the mirror trim curve
    # measured by drift-null and the gate simply ramping down past the
    # bench floor until the body grounds; then the wheels lock, the
    # legs settle into the rest heap, and the POWER CUTS.
    "names": ["down rate", "settle rate"],
    "seed": [0.9, 0.7],
}
STAGE_ORDER.insert(STAGE_ORDER.index("rise") + 1, "sit")


def sit_episode(g, seed, secs=12.0):
    (down_rate, settle_rate) = g
    bobg = accepted("bob") or [1.0, 1.0, 0.0, -0.5]
    (trim_scale, leg_speed, lvl_shift, com_hold) = (
        list(bobg) + [1.0, 1.0, 0.0, -0.5])[:4]
    bpts = balance_points()
    env = RobotEnv(SPEC, seed=seed, task="bob", randomise=False)
    env.reset()
    env.min_h = 0.02
    m, d = env.model, env.data
    jq = lambda nm: m.jnt_qposadr[mujoco.mj_name2id(
        m, mujoco.mjtObj.mjOBJ_JOINT, nm)]
    st_h, st_k = C.stance_to_leg(0.455, 0.0, branch=+1.0)
    for nm, v in (("l_hip", st_h), ("l_knee", st_k),
                  ("r_hip", st_h), ("r_knee", st_k)):
        d.qpos[jq(nm)] = v
    d.qpos[0:3] = [0.0, 0.0, 0.4575]
    d.qpos[3:7] = [1.0, 0.0, 0.0, 0.0]
    d.qvel[:] = 0.0
    mujoco.mj_forward(m, d)
    legs = [i for i, a in enumerate(env.spec["actuators"])
            if a.get("mode") == "position"]
    for i in legs:
        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])
    s = Sensed(env)
    com_bid = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_BODY, "chassis")
    com_wids = [mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_BODY, n)
                for n in ("l_wheel", "r_wheel")]
    floor_g = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_GEOM, "floor")
    tyre_g = {gi for gi in range(m.ngeom)
              if int(m.geom_bodyid[gi]) in com_wids}
    wq_adr = [jq(n) for n in ("l_wheel_j", "r_wheel_j")]
    wd_adr = [m.jnt_dofadr[mujoco.mj_name2id(
        m, mujoco.mjtObj.mjOBJ_JOINT, n)]
        for n in ("l_wheel_j", "r_wheel_j")]
    cur = np.array([st_h, st_k, st_h, st_k])
    heap = np.array([SIT_POSE["l_hip"], SIT_POSE["l_knee"]] * 2)
    rate = C.leg_rate(leg_speed)
    settle_r = math.radians(150.0 * min(1.0, max(0.08, abs(settle_rate))))
    target_mm = 455.0
    dmm = 300.0 * min(1.0, max(0.08, abs(down_rate)))
    fdx = 0.0
    phase = 0             # 0 the descending gate, 1 settle + power off
    lock_q = None
    off = False
    gain0 = bias0 = None
    h_sum = 0.0
    peak_rate = 0.0
    eff = 0.0
    t_done = None
    n = int(secs / env.control_dt)
    for k in range(n):
        lean, lrate, speed, odo = s.read()
        grounded = False
        for ci in range(d.ncon):
            c = d.contact[ci]
            if floor_g == c.geom1:
                other = c.geom2
            elif floor_g == c.geom2:
                other = c.geom1
            else:
                continue
            if other not in tyre_g:
                grounded = True
                break
        u = 0.0
        uw = None
        if phase == 0:
            # BOB'S LAW, VERBATIM, in the mirror fold - the gate just
            # keeps going lower
            target_mm = max(170.0, target_mm - dmm * env.control_dt)
            com_x = float(d.subtree_com[com_bid][0])
            ax_x = 0.5 * (float(d.xpos[com_wids[0]][0])
                          + float(d.xpos[com_wids[1]][0]))
            fdx += (env.control_dt / 0.04) * (com_hold * (com_x - ax_x)
                                              - fdx)
            dxb = float(np.clip(
                lvl_shift * (0.46 - target_mm / 1000.0) + fdx,
                -0.06, 0.14))
            hip_t, knee_t = C.stance_to_leg(target_mm / 1000.0, dxb,
                                            branch=+1.0)
            cur = C.slew(cur, np.array([hip_t, knee_t, hip_t, knee_t]),
                         rate, env.control_dt)
            kp, kd, station, clamp = interp_wheel_gains(bpts, target_mm)
            ref = C.lean_reference(station, speed, odo, clamp)
            trim = (math.radians(mirror_trim_deg(target_mm))
                    * trim_scale)
            u = C.wheel_command(kp, kd, lean, lrate, ref, trim)
            # the balance law will hover the chassis a few mm off the
            # floor forever - the end of the gate IS the trigger, with
            # early grounding as the backup
            if target_mm <= 171.0 or (grounded and target_mm < 340.0):
                phase = 1
                lock_q = [float(d.qpos[a]) for a in wq_adr]
        else:
            cur = C.slew(cur, heap, settle_r, env.control_dt)
            uw = [float(np.clip(
                -3.0 * (float(d.qpos[wq_adr[i]]) - lock_q[i])
                - 0.3 * float(d.qvel[wd_adr[i]]), -1, 1))
                for i in (0, 1)]
            if (not off and float(np.max(np.abs(cur - heap)))
                    < math.radians(10.0)):
                # THE POWER CUTS - the shutdown ends as the real one
                # will, limp in its heap
                off = True
                t_done = env.t
                gain0 = m.actuator_gainprm.copy()
                bias0 = m.actuator_biasprm.copy()
                m.actuator_gainprm[:, 0] = 0.0
                m.actuator_biasprm[:, 1:3] = 0.0
        eff += abs(u) * env.control_dt
        act = np.zeros(env.act_dim)
        wheel_i = 0
        for j, ai in enumerate(env.act_idx):
            if ai in legs:
                act[j] = np.clip(
                    (cur[legs.index(ai)] - env.act_center[ai])
                    / max(1e-9, env.act_span[ai]), -1, 1)
            else:
                act[j] = uw[wheel_i] if uw is not None else u
                wheel_i += 1
        env.step(act)
        h_sum += float(d.xpos[com_bid][2]) * env.control_dt
        peak_rate = max(peak_rate, abs(lrate))
        roll = math.degrees(math.atan2(env.est_up[1],
                                       max(1e-6, env.est_up[2])))
        if abs(roll) > 50.0:
            break
    if gain0 is not None:
        m.actuator_gainprm[:] = gain0
        m.actuator_biasprm[:] = bias0
    still = (off and abs(float(d.qvel[0])) < 0.15 and env.up_z > 0.85)
    quick = (2.0 * max(0.0, 1.0 - t_done / secs)
             if t_done is not None else 0.0)
    return (- feel_charge(env, "sit")
            + 6.0 * (1.0 - (h_sum / max(env.t, 1e-6)) / 0.46)
            + 3.0 * (1.0 if still else 0.0) + quick
            - 1.0 * max(0.0, peak_rate - 2.5)
            - EFFORT_RATE * eff)


def sit_score(g, secs=12.0):
    return float(np.mean([sit_episode(g, i, secs) for i in range(3)]))


STAGES["jump"] = {
    # THE JUMP, v56, up only (forward/back removed - they complicated
    # it). The launch follows the TETHER TRAJECTORY: the operator's
    # design - hang the crouched robot from its centre, hold the
    # chassis level, wheels rim-locked, legs limp, and winch it up;
    # the angles the legs fall into on the way ARE the path. Measured
    # 15 Aug: the hip runs 10-13 deg lower than straight IK the whole
    # way - the IK path is what shoved the CoM aft and pitched every
    # bang-launch nose-up. Firing along the measured path lifts
    # straight. Joint limits may be hit briefly; sustained torque
    # saturation is charged threshold-form, like the chatter meter.
    "names": ["crouch", "dive rate", "launch drive", "tuck", "absorb",
              "recover rate", "air gain", "hip trim", "launch damp",
              "fire lean", "hip ratio", "com set", "hip hold",
              "knee boost", "tuck rate", "chord tilt", "tuck at",
              "tuck hip", "knee cancel",
              "fire rate", "launch top", "fire gate", "brace at",
              "fire steer", "tuck balance", "land catch", "air spin",
              "roll fix", "catch hip", "catch flex", "catch knee",
              "land lean"],
    # thrust tilt seeded from the measured sweep: +2.0 (hip +16 deg
    # into the launch) cancels most of the back-rotation leak
    # launch damp is the JOINT FEEDBACK (operator's design): the hip
    # steers against the pitch rate it feels mid-thrust, gain learned
    # from the leak charge
    # drive +0.4 and tilt +2.0 from the measured joint grid (6.52);
    # fire lean = the deliberate forward lean INTO the jump (operator:
    # the CoM backflips it - spend the rotation budget before launch)
    # THE HIP RATIO (operator's experiment): the hip starts PINNED at
    # its true static crouch angle and releases toward extension at a
    # fraction of the knee's pace - measured peak at 0.45 (7.24, the
    # best the jump has scored; 0 = pinned throughout 5.3, 1 = full
    # pace 6.3, 1.5 = hip leading 2.8). This also removed a 12-deg
    # hip snap at fire: the old law aimed the hip at the TETHER
    # path's crouch angle (42 deg) but the robot crouches in the
    # BALANCE pose (54 deg). Fire lean -1.0 = 6 deg nose-down; hip
    # trim and launch damp are fine-trim on top, reseeded 0.
    # com set (slot 12, operator's diagnosis): feet planted aft so
    # the CoM rides just ahead of the wheels through crouch and fire
    # (+20mm best measured with drive +0.4); the slip the operator saw
    # in slow motion is the drive slot's job - it spans both signs.
    # hip hold (slot 13, operator's diagnosis: the hip flips back too
    # EARLY in the launch): the hip stays pinned until the sweep
    # passes this fraction, then releases - 70% measured best (1.80)
    # knee boost (slot 14, from the operator's backdrive question):
    # the knee command is pushed further ahead in proportion to its
    # own measured lag - more torque sooner; the hip-side equivalent
    # measured zero-best both signs
    # fire lean deepened to -12 deg (operator: push perpendicular
    # into the ground; measured 2.23); tuck rate searchable (operator:
    # tuck faster once rotation shrinks - 500 measures best today,
    # the slot lets the search revisit as the launch cleans up)
    # launch drive re-measured under the v69 law at the operator's
    # prompt: a LIGHT BACKWARD drive (-0.15) now beats forward (2.51
    # vs 2.23) - the reaction torque of gently braking the wheels
    # counters the residual back-rotation. Narrow peak (-0.22 falls
    # to 0.4, -0.05 to 0.5): exactly the kind of ridge the search
    # walks better than a hand.
    # brace stance 380mm measured best (2.82) - land on long, ready
    # legs with the full travel beneath
    # chord tilt (slot 16, operator's triangle spec): the virtual leg
    # (wheel->hip chord) holds a slight FORWARD tilt through the whole
    # launch - the foot offset scales with the commanded leg length so
    # the angle stays constant as the leg grows. -1.7 deg measured
    # best (3.15; vertical 0.76, past -3 deg collapses).
    # tuck at (slot 17, operator's counter-rotation): the tuck's leg
    # swing kicks the chassis nose-down in reaction - WHEN it fires is
    # an attitude control input. Knee-angle threshold, searchable;
    # seed 0 = the measured best trigger (~-20 deg, earlier forced
    # tucks measured 2.94-2.99 vs 3.15 today - the search re-places
    # the pulse as the launch's residual spin changes).
    # tuck hip (slot 18, operator's fix for the full-power tuck): the
    # stance-IK tuck pose barely moves the hip - the thigh must fold
    # too (clockwise on the right side, toward the chassis front).
    # -20 deg measured best: 3.10 vs 1.51 without it, catching the
    # v72 law's 3.15 while keeping the stable flight attitude.
    # tuck hip -1.5 = the thigh swings ~85 deg further round on the
    # tuck (OPERATOR'S ORDER: see the hips work, even if it breaks -
    # backups exist). knee cancel -0.6 (slot 19, operator's design):
    # the knee answers the gyro live in flight, swinging the
    # shin-and-wheel mass against the body's pitch rate - measured
    # 4.02, the best mean the jump has ever scored.
    # knee cancel parked at 0 (operator: forget the knee on tuck,
    # concentrate on the hip - full torque once airborne)
    # tuck hip FLIPPED positive on the operator's eye: negative swung
    # the thigh forward-down (a pike); the true knees-to-chest tuck
    # folds the thigh back and UP past the crouch angle
    # restored to the measured optimum: shallow PIKE fold (-20 deg -
    # the fold direction IS the counter-rotation: this robot leaves
    # the ground spinning backward, so its correct mid-air move
    # mirrors a human tuck) + the operator's knee-gyro cancel (-0.6).
    # Together: 4.02, the best mean ever measured. The chest-tuck
    # direction lost at every depth and rate (grid 15 Aug).
    # THE EVERY-TIME LANDER: found by the 20-direction search (7.00
    # under v77), certified under v78 (6.8 - mathematically all three
    # episodes land, since any crash caps the score at 5.0)
    # v94 slots (M4 relearn): fire rate = the extension command ramp
    # (was hardcoded 2000 mm/s), launch top = how far the command
    # rides the path before liftoff (was 458mm), fire gate = the
    # settled-crouch stillness threshold that times the fire against
    # the crouch's own sway (was 0.15 m/s), brace at = the
    # time-to-impact that ends the tuck (was 0.1s). fire steer
    # (operator's question: "the leaks come from the gforce signals -
    # can the adjustment put them into the jump?"): the wheel drive
    # during the thrust was a CONSTANT - the strongest actuator flew
    # the whole fire open-loop while the spin was being born. It now
    # answers the FELT pitch rate (the gyro, same signal launch damp
    # gives the hip), so reaction torque spends the leak as it
    # appears instead of pricing it post-mortem. All zero = the old
    # constants. Seeded from the v93 8.08 bake, which landed its
    # three exam seeds and fell on two of the next five - the 8-seed
    # exam makes it climb out of that.
    # v96 seed = the v95 9.95 bake (clean straight-up liftoff, launch
    # drive railed at -1) + two zero slots. Slot 25 tuck balance
    # (operator's design): the hip target tracks the shin's measured
    # fold rate through tuck AND brace - the thigh swing cancels the
    # shin swing's reaction instead of the chassis eating it. Slot 26
    # land catch (operator: "it falls back when it lands - torque the
    # wheels to straighten us up"): a fading wheel-torque reflex for
    # the first 0.4s after touchdown, driving the wheels back under
    # the mass; the standing benches' quiet gains were never sized for
    # catching a robot arriving backward out of the sky.
    # Slot 27 air spin (operator: "spin the wheels slightly backwards
    # in the air - we cannot rely on traction, they may bounce; the
    # wheel torque itself must counter the shin tucking up"): wheel
    # torque proportional to the shin's fold rate through the whole
    # flight - reaction-wheel control, traction-free. Tied to knee
    # rate it fires on the tuck AND (sign-flipped) on the tuck
    # release, arriving at touchdown with the momentum that makes the
    # landing LAST. Engaging it fades the liftoff position-lock,
    # which otherwise springs any deliberate spin straight back.
    # Slot 28 roll fix (operator: "make adjustments to roll forwards
    # if we are rolling back"): wheel command on the MEASURED ground
    # speed for 3s after touchdown - the landing that rolls backward
    # drives itself back under the mass instead of riding it out.
    # Slots 29-31, the catch's legs (operator: "open the joints up on
    # landing"): catch hip = a held hip offset through the 3s window,
    # catch flex = hip answers the felt lean, catch knee = knee
    # answers the felt lean - the one beat where the legs were still
    # frozen to the symmetric stance curve.
    # Slot 32 land lean (operator: "lean forward into the jump for a
    # landing that persists"): the flight attitude reference - the
    # wheels' air loop steers toward 6deg * slot of FORWARD lean, so
    # touchdown happens already leaning into the catch.
    "seed": [-0.0529, 5.6698, -3.4865, -0.4280, 0.5925,
             0.4717, 0.0867, -0.4442, 1.3041, -0.7674,
             0.0775, 0.3786, 0.0291, -0.0675, -0.1401,
             -0.4180, -0.2299, -0.0932, -0.7294, -0.0861,
             -0.4116, 0.1136, -0.1797, 0.1431, 0.0, 0.0, 0.0, 0.0,
             0.0, 0.0, 0.0, 0.0],
}
STAGE_ORDER.insert(STAGE_ORDER.index("spin") + 1, "jump")

# the tether measurement, verbatim (height mm, hip deg, knee deg):
# limp legs, level chassis, locked wheels, winched from the centre
JUMP_TRAJ = ((218.0, 67.6, -135.2), (238.1, 57.5, -127.2),
             (258.1, 46.8, -119.6), (278.1, 45.9, -113.4),
             (298.1, 42.4, -106.3), (318.2, 38.8, -98.9),
             (338.2, 35.0, -91.0), (358.2, 30.8, -82.5),
             (378.2, 26.2, -73.2), (398.3, 21.0, -62.8),
             (418.3, 14.7, -50.5), (438.3, 6.4, -34.2),
             (458.0, -3.4, -10.0))


def jump_traj(h_mm):
    """Hip/knee radians on the measured tether path, clamped ends."""
    pts = JUMP_TRAJ
    if h_mm <= pts[0][0]:
        hh, kk = pts[0][1], pts[0][2]
    elif h_mm >= pts[-1][0]:
        hh, kk = pts[-1][1], pts[-1][2]
    else:
        for (h0, p0, k0), (h1, p1, k1) in zip(pts, pts[1:]):
            if h0 <= h_mm <= h1:
                f = (h_mm - h0) / max(1e-9, h1 - h0)
                hh = p0 + f * (p1 - p0)
                kk = k0 + f * (k1 - k0)
                break
    return math.radians(hh), math.radians(kk)


def jump_episode(g, seed, secs=6.0, out=None):
    (g_crouch, g_dive, g_launch, g_tuck, g_absorb, g_rec,
     g_air, g_tilt, g_damp, g_flean,
     g_ratio, g_dxs, g_hold, g_kboost,
     g_trate, g_chord, g_tuckat, g_thip,
     g_kcancel, g_frate, g_ltop, g_fgate,
     g_brace, g_fsteer, g_tbal, g_catch,
     g_aspin, g_rollfix, g_chip, g_cflex,
     g_cknee, g_llean) = (list(g) + [0.0] * 32)[:32]
    # v94 launch timing, all zero = the pre-v94 hardcoded constants
    ext_rate = 2000.0 * (1.0 + float(np.clip(g_frate, -0.7, 1.0)))
    top_mm = float(np.clip(458.0 + 40.0 * g_ltop, 380.0, 461.0))
    vz_gate = 0.15 * (1.0 + float(np.clip(g_fgate, -0.9, 2.0)))
    brace_tti = 0.1 * (1.0 + float(np.clip(g_brace, -0.8, 3.0)))
    chord_k = float(np.clip(0.1 * g_chord, -0.15, 0.15))
    tuck_r = 1000.0 * min(2.0, max(0.3, abs(g_trate)))
    dxs = float(np.clip(0.05 * g_dxs, -0.06, 0.06))
    crouch_mm = float(np.clip(300.0 + 80.0 * g_crouch, 218.0, 380.0))
    tuck_mm = float(np.clip(340.0 + 80.0 * g_tuck, 250.0, 430.0))
    absorb_mm = float(np.clip(300.0 + 80.0 * g_absorb, 230.0, 400.0))
    dive_r = 300.0 * min(1.0, max(0.08, abs(g_dive)))
    launch_r = 2500.0 * min(1.0, max(0.1, abs(g_launch)))
    rec_r = 300.0 * min(1.0, max(0.08, abs(g_rec)))
    bobg = (list(accepted("bob") or [1.0]) + [1.0])[:1]
    trim_scale = bobg[0] or 1.0
    bpts = balance_points()
    env = RobotEnv(SPEC, seed=seed, task="bob", randomise=False)
    env.reset()
    env.min_h = 0.02
    m, d = env.model, env.data
    jq = lambda nm: m.jnt_qposadr[mujoco.mj_name2id(
        m, mujoco.mjtObj.mjOBJ_JOINT, nm)]
    # v99: the exam enters MESSY, the way the live bench actually
    # arrives at a jump - a seeded random stance, a few degrees of
    # pitch, and real rolling speed with the wheels matching. The old
    # perfect settled 455mm zero-speed entry was an entry the robot
    # never gets, and the exam could not see the falls the operator
    # watched from live entries.
    erng = np.random.default_rng(1000 + seed)
    stance0 = float(erng.uniform(0.40, 0.46))
    pitch0 = math.radians(float(erng.uniform(-3.0, 3.0)))
    vx0 = float(erng.uniform(-0.25, 0.25))
    st_h, st_k = C.stance_to_leg(stance0, 0.0)
    for nm, v in (("l_hip", st_h), ("l_knee", st_k),
                  ("r_hip", st_h), ("r_knee", st_k)):
        d.qpos[jq(nm)] = v
    d.qpos[0:3] = [0.0, 0.0, stance0 + 0.0025]
    d.qpos[3:7] = [math.cos(pitch0 / 2.0), 0.0,
                   math.sin(pitch0 / 2.0), 0.0]
    d.qvel[:] = 0.0
    d.qvel[0] = vx0
    mujoco.mj_forward(m, d)
    legs = [i for i, a in enumerate(env.spec["actuators"])
            if a.get("mode") == "position"]
    for i in legs:
        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])
    s = Sensed(env)
    floor_g = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_GEOM, "floor")
    wids = [mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_BODY, n)
            for n in ("l_wheel", "r_wheel")]
    wq_adr = [jq(n) for n in ("l_wheel_j", "r_wheel_j")]
    wd_adr = [m.jnt_dofadr[mujoco.mj_name2id(
        m, mujoco.mjtObj.mjOBJ_JOINT, n)]
        for n in ("l_wheel_j", "r_wheel_j")]
    kn_dof = m.jnt_dofadr[mujoco.mj_name2id(
        m, mujoco.mjtObj.mjOBJ_JOINT, "l_knee")]
    # the entry's rolling speed reaches the wheels too - a chassis
    # moving over stationary wheels is a skid, not an entry
    for a in wd_adr:
        d.qvel[a] = vx0 / 0.0625
    frange = np.abs(m.actuator_forcerange[:, 1])
    cur = np.array([st_h, st_k, st_h, st_k])
    target = 1000.0 * stance0
    hcmd = None            # the launch's height parameter on the path
    # phases: 0 settle+dive, 1 settle-at-crouch, 2 launch (ride the
    # tether path), 3 air, 4 absorb
    phase = 0
    ph_t = 0.0
    air_t = 0.0
    free_t = 0.0
    air_mode = 0
    lock_q = None
    x_launch = None
    vx_lo = 0.0
    vz_lo = 0.0
    wy_lo = 0.0
    wy_pk = 0.0
    tilt_pk = 0.0
    catch_tilt = 0.0
    back_tilt = 0.0
    x_land = None
    t_lift = None
    peak = 0.0
    ch_bid = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_BODY, "chassis")
    sat_t = 0.0
    eff = 0.0
    landed_t = None
    n = int(secs / env.control_dt)
    trace = out.get("trace") if out is not None else None
    for k in range(n):
        t = k * env.control_dt
        lean, lrate, speed, odo = s.read()
        if trace is not None:
            # the energy audit, live: pitch rate vs climb rate at every
            # tick, so the leak can be SEEN being born instead of read
            # off the liftoff post-mortem
            trace.append((round(t, 3), phase, float(d.qvel[4]),
                          float(d.qvel[2]),
                          min(float(d.xpos[wids[0]][2]),
                              float(d.xpos[wids[1]][2])) - 0.0625))
        contact = False
        for ci in range(d.ncon):
            c = d.contact[ci]
            if floor_g in (c.geom1, c.geom2):
                contact = True
                break
        u = 0.0
        uw = None
        if phase in (0, 1):
            if phase == 0:
                # v102: no dead air - the dive starts as soon as the
                # sim is warm; braking the entry drift and crouching
                # happen together, at speed
                if t > 0.1:
                    target = max(crouch_mm, target
                                 - dive_r * env.control_dt)
                if target <= crouch_mm + 0.5:
                    phase = 1
                    ph_t = 0.0
            else:
                ph_t += env.control_dt
                # fire from a SETTLED crouch at the FIRE LEAN (slot 9,
                # operator's call): the launch backflips, so the last
                # beat ramps INTO a deliberate forward lean and the
                # jump spends its rotation budget getting back level.
                # Sign measured by sweep, like every sign here.
                fire_lean = math.radians(6.0 * g_flean)
                # v101: the gate demands PLANTED FEET - ground speed
                # near zero, not just a settled crouch. Fired while
                # drifting, the flight conserves the drift and the
                # landing inherits it as the roll-out (measured:
                # 555mm backward from one backward-drifting fire).
                # The impatience fallback gets the time stopping
                # actually takes from a 0.25 m/s entry.
                if (ph_t >= 0.1 and abs(float(d.qvel[2])) < vz_gate
                        and (abs(speed) < 0.1 or ph_t > 2.5)
                        and (abs(lean - fire_lean) < math.radians(5.0)
                             or ph_t > 2.0)):
                    phase = 2
                    ph_t = 0.0
                    hcmd = crouch_mm
                    x_launch = float(env.base_pos[0])
                    lock_q = [float(d.qpos[a]) for a in wq_adr]
            hip_t, knee_t = C.stance_to_leg(
                target / 1000.0, dxs if phase == 1 else 0.0)
            cur = C.slew(cur, np.array([hip_t, knee_t] * 2),
                         math.radians(400.0), env.control_dt)
            kp, kd, station, clamp = interp_wheel_gains(bpts, target)
            ref = C.lean_reference(station, speed, 0.0, clamp)
            trim = math.radians(C.lean_trim_deg(target)) * trim_scale
            if phase == 1:
                # blend from the balance trim to the fire lean
                _f = max(0.0, 1.0 - ph_t / 0.15)
                trim = trim * _f + math.radians(6.0 * g_flean) * (1 - _f)
            u = C.wheel_command(kp, kd, lean, lrate, ref, trim)
        elif phase == 2:
            # THE FIRE (operator's spec): the wheels DRIVE forward
            # while knee and hip open at full power along the measured
            # tether path - the contact point chases the centre of
            # mass instead of pinning behind it. (Locked wheels were
            # the backward topple; direct extension without the path
            # measured 4.8 vs 6.5 with it - the path sequences the
            # joints, the drive keeps the feet under the load.)
            # Slot 2 is the drive, slot 7 aims the line, slot 8 steers
            # it live against the felt pitch rate - signs swept.
            hcmd = min(top_mm, hcmd + ext_rate * env.control_dt)
            # the operator's triangle: the wheel->hip chord holds a
            # constant slight forward tilt as the leg grows - the foot
            # offset scales with commanded length
            _dxl = chord_k * (hcmd / 1000.0)
            # the knee drives the launch along the path; the hip
            # starts at its TRUE crouch angle (the balance pose it is
            # actually in - the tether path's angle there is 12 deg
            # different and used to snap at fire) and releases toward
            # extension at the searched fraction of the knee's pace
            knee_t = C.stance_to_leg(hcmd / 1000.0, _dxl)[1]
            _ch = C.stance_to_leg(crouch_mm / 1000.0,
                                  chord_k * crouch_mm / 1000.0)[0]
            _eh = C.stance_to_leg(0.462, chord_k * 0.462)[0]
            # the hip is gated on the ACTUAL knee, not the command
            # sweep (operator watched the thigh move before the knee:
            # the command races the servo by ~100ms, so a command-time
            # gate released the hip while the real knee had opened 5
            # degrees). Progress = how far the measured knee has
            # actually travelled from crouch to extension.
            _ck = C.stance_to_leg(crouch_mm / 1000.0,
                                  chord_k * crouch_mm / 1000.0)[1]
            _ek = C.stance_to_leg(0.462, chord_k * 0.462)[1]
            _ka = float(d.qpos[jq("l_knee")])
            _prog = float(np.clip((_ka - _ck) / (_ek - _ck), 0.0, 1.0))
            _hld = float(np.clip(g_hold, 0.0, 0.95))
            _rel = max(0.0, (_prog - _hld) / max(0.05, 1.0 - _hld))
            hip_t = _ch + min(1.0, (g_ratio / 0.45) * _rel) * (_eh - _ch)
            # knee boost: the command leads the lagging joint so the
            # servo pulls harder through the loaded part of the thrust
            knee_t = knee_t + float(np.clip(
                g_kboost * (knee_t - _ka), -1.2, 1.2))
            hip_t += math.radians(8.0 * g_tilt) + float(np.clip(
                g_damp * lrate, -0.35, 0.35))
            cur = np.array([hip_t, knee_t] * 2)
            # fire steer (v94): the drive answers the FELT pitch rate
            # mid-thrust - the same gyro signal launch damp gives the
            # hip - so wheel reaction torque spends the rotation leak
            # as it is born; zero = the old constant-drive fire
            uw = [float(np.clip(3.0 * g_launch + g_fsteer * lrate,
                                -1, 1))] * 2
            ph_t += env.control_dt
            wz = min(float(d.xpos[wids[0]][2]),
                     float(d.xpos[wids[1]][2]))
            peak = max(peak, wz - 0.0625)
            if not contact and ph_t > 0.02:
                phase = 3
                air_t = 0.0
                # OPERATOR'S EXPERIMENT (v73): tuck begins the instant
                # contact is lost, at full power
                air_mode = 1
                # LIFTOFF ENERGY AUDIT: whatever is not going straight
                # up here is the leak the score charges
                vx_lo = float(d.qvel[0])
                vz_lo = float(d.qvel[2])
                wy_lo = float(d.qvel[4])
                t_lift = t
        elif phase == 3:
            # airborne, the operator's three beats: finish the stretch,
            # TUCK the moment it is reached (ground clearance), then
            # BRACE - legs pre-set to the absorb stance while still
            # falling, so touchdown lands on travel, not on stops.
            # The wheels double as reaction attitude control (slot 6)
            air_t = air_t + env.control_dt if not contact else 0.0
            free_t = max(free_t, air_t)
            wy_pk = max(wy_pk, abs(float(d.qvel[4])))
            _ka3 = float(d.qpos[jq("l_knee")])
            _ek3 = C.stance_to_leg(0.462, 0.0)[1]
            if air_mode == 0:
                hip_t, knee_t = C.stance_to_leg(0.462, 0.0)
                cur = C.slew(cur, np.array([hip_t, knee_t] * 2),
                             math.radians(900.0), env.control_dt)
                # the tuck fires at a searched knee angle (its leg
                # swing is the nose-down counter-rotation pulse -
                # operator's call: WHEN is a control input), or at
                # apex as the fallback
                _t_at = _ek3 - math.radians(15.0 - 15.0 * g_tuckat)
                if _ka3 > _t_at or float(d.qvel[2]) < 0.0:
                    air_mode = 1
            elif air_mode == 1:
                # the thigh folds too (slot 18) - the stance-IK pose
                # alone left the hip almost still and the shins did
                # all the tucking (operator's diagnosis)
                hip_t, knee_t = C.stance_to_leg(tuck_mm / 1000.0, 0.0)
                hip_t += g_thip
                # tuck balance (slot 25, operator's design): the hip
                # pulls up WITH the shins - proportional to the knee's
                # measured fold rate, so the thigh swing cancels the
                # shin swing's reaction instead of the chassis pitching
                # backward to conserve the momentum
                hip_t += float(np.clip(
                    g_tbal * float(d.qvel[kn_dof]), -0.7, 0.7))
                knee_t += float(np.clip(g_kcancel * lrate, -0.6, 0.6))
                cur = C.slew(cur, np.array([hip_t, knee_t] * 2),
                             math.radians(3000.0), env.control_dt)
                # brace by time-to-impact so the tuck owns the fall
                _vz3 = float(d.qvel[2])
                _hw3 = min(float(d.xpos[wids[0]][2]),
                           float(d.xpos[wids[1]][2])) - 0.0625
                if _vz3 < -0.05 and _hw3 / max(0.05, -_vz3) < brace_tti:
                    air_mode = 2
            else:
                # the brace: meet the ground with bent, ready legs -
                # the knee keeps cancelling rotation all the way down
                hip_t, knee_t = C.stance_to_leg(absorb_mm / 1000.0,
                                                0.0)
                # the shins swing back OUT here - the same balance
                # term cancels the reverse reaction too
                hip_t += float(np.clip(
                    g_tbal * float(d.qvel[kn_dof]), -0.7, 0.7))
                knee_t += float(np.clip(g_kcancel * lrate, -0.6, 0.6))
                cur = C.slew(cur, np.array([hip_t, knee_t] * 2),
                             math.radians(tuck_r), env.control_dt)
            # land lean (slot 32): the flight attitude reference is a
            # searchable FORWARD lean, not level - arrive leaning into
            # the catch (operator's design)
            att = float(np.clip(g_air * (
                (lean - math.radians(6.0 * g_llean)) * 8.0 + lrate),
                -1.0, 1.0))
            # air spin (slot 27, operator's design): the wheels are
            # reaction wheels in flight - torque proportional to the
            # shin's fold rate counters the tuck's kick WITHOUT
            # traction, fires sign-flipped again on the tuck release,
            # and lands with the momentum that makes the catch stick.
            # Engaging it fades the liftoff position-lock, which
            # otherwise springs any deliberate spin straight back.
            _wspin = float(np.clip(
                g_aspin * float(d.qvel[kn_dof]), -1.0, 1.0))
            _wlock = max(0.0, 1.0 - min(1.0, abs(g_aspin)))
            uw = [float(np.clip(
                _wlock * (-1.5 * (float(d.qpos[wq_adr[i]]) - lock_q[i])
                          - 0.15 * float(d.qvel[wd_adr[i]]))
                + _wspin + att, -1, 1))
                for i in (0, 1)]
            wz = min(float(d.xpos[wids[0]][2]),
                     float(d.xpos[wids[1]][2]))
            peak = max(peak, wz - 0.0625)
            # DEBOUNCED landing: the launch grazes the floor for a
            # while - a touch only counts after 60ms of CONSECUTIVE
            # clean air (grazes reset the clock)
            if contact and free_t > 0.06:
                phase = 4
                target = absorb_mm
                landed_t = t
                x_land = float(env.base_pos[0])
        else:
            # absorb: the gate rides from the yield depth back up to
            # stand, the benches balancing the whole way
            target = min(455.0, target + rec_r * env.control_dt)
            hip_t, knee_t = C.stance_to_leg(target / 1000.0, 0.0)
            # the catch's legs (slots 29-31, operator: "open the
            # joints up on landing"): a held hip offset plus hip and
            # knee answers to the FELT lean, fading over the 3s
            # window - the legs bend into the fall instead of riding
            # the symmetric height gate through the whole catch
            if landed_t is not None:
                _cfad = max(0.0, 1.0 - (t - landed_t) / 3.0)
                hip_t += float(np.clip(
                    0.5 * g_chip + g_cflex * lean, -0.5, 0.5)) * _cfad
                knee_t += float(np.clip(
                    g_cknee * lean, -0.5, 0.5)) * _cfad
            cur = C.slew(cur, np.array([hip_t, knee_t] * 2),
                         math.radians(400.0), env.control_dt)
            kp, kd, station, clamp = interp_wheel_gains(bpts, target)
            ref = C.lean_reference(station, speed, 0.0, clamp)
            trim = math.radians(C.lean_trim_deg(target)) * trim_scale
            u = C.wheel_command(kp, kd, lean, lrate, ref, trim)
            # land catch (slot 26, operator's design): the benches'
            # quiet standing gains were never sized for catching a
            # robot arriving backward out of the sky - for the first
            # 0.4s after touchdown an extra wheel torque drives the
            # wheels back under the mass, fading as the absorb settles
            if landed_t is not None:
                _fade = max(0.0, 1.0 - (t - landed_t) / 0.4)
                u += float(np.clip(g_catch * (8.0 * lean + lrate),
                                   -1.0, 1.0)) * _fade
                # roll fix (slot 28, operator's design): for the 3s
                # window, measured ground speed feeds the wheels -
                # rolling backward becomes a drive back under the mass
                u += (float(np.clip(g_rollfix * speed, -1.0, 1.0))
                      * max(0.0, 1.0 - (t - landed_t) / 3.0))
        eff += abs(u) * env.control_dt
        act = np.zeros(env.act_dim)
        wi = 0
        for j, ai in enumerate(env.act_idx):
            if ai in legs:
                act[j] = np.clip(
                    (cur[legs.index(ai)] - env.act_center[ai])
                    / max(1e-9, env.act_span[ai]), -1, 1)
            else:
                act[j] = uw[wi] if uw is not None else u
                wi += 1
        env.step(act)
        # sustained torque saturation is the charged sin - brief limit
        # hits are free (the operator allowed the launch its violence)
        fsat = np.abs(np.array(d.actuator_force))[:4] >= 0.97 * frange[:4]
        if fsat.any():
            sat_t += env.control_dt
        # v95: TRUE chassis tilt through launch and flight (not the
        # sensed lean - the accel misreads while rotating). This is
        # what the operator SEES: the body going over backwards.
        # v97: the audit continues through the CATCH - the backwards
        # dip after touchdown was invisible to the score and the
        # staggering recovery collected full marks.
        _tilt_now = math.acos(min(1.0, max(-1.0, env.up_z)))
        if phase >= 2 and landed_t is None:
            tilt_pk = max(tilt_pk, _tilt_now)
        elif landed_t is not None:
            catch_tilt = max(catch_tilt, _tilt_now)
            # v98: the SIGNED backward lean in the 3s catch window -
            # the chassis x-axis tipping nose-up (z-component
            # positive, measured) is the roll-back the operator
            # watches lose the landing
            if t - landed_t < 3.0:
                back_tilt = max(back_tilt, math.asin(min(1.0, max(
                    0.0, float(d.xmat[ch_bid][6])))))
        if phase >= 2 and env.up_z < 0.35:
            break
        # v100: the landing must PERSIST - ten full seconds standing
        # after touchdown before the episode may end (operator's
        # standard; the 3.2s window let a slow topple graduate)
        if landed_t is not None and t - landed_t > 10.0 \
                and abs(speed) < 0.15 and target >= 454.0:
            break
    stance_now = 1000.0 * float(env.base_pos[2])
    # v97: a catch that goes past 45 deg backward IS a fall - the
    # last-frame check alone let a 60-deg dip-and-stagger pass
    landed = (landed_t is not None and env.up_z > 0.9
              and abs(float(d.qvel[0])) < 0.4 and stance_now > 420.0
              and catch_tilt < math.radians(45.0)
              # v98: a backward lean past 20 deg in the catch window
              # IS the fail the operator called
              and back_tilt < math.radians(20.0))
    dx = (float(env.base_pos[0]) - x_launch) if x_launch is not None \
        else 0.0
    land_roll = (abs(float(env.base_pos[0]) - x_land)
                 if x_land is not None else 0.0)
    # THE ENERGY-DIRECTION CHARGE (operator's brief: humans aim the
    # whole push). Rotation is priced at its IN-FLIGHT PEAK, not the
    # liftoff sample - the liftoff read came before the launch had
    # finished grazing, and the tuck multiplies whatever spin remains
    # (angular momentum conserved, inertia drops). And a jump that
    # does not LAND pays only 35 cents on its clearance dollar: the
    # exam was scoring backflips 7.2 while the operator watched them
    # crash - height only counts when you come back from it.
    e_up = vz_lo * vz_lo
    e_leak = vx_lo * vx_lo + (0.25 * wy_pk) ** 2
    leak_frac = e_leak / (e_leak + e_up + 1e-6)
    if out is not None:
        out.update(peak_mm=1000.0 * peak, vx_lo=vx_lo, vz_lo=vz_lo,
                   wy_lo=wy_lo, wy_pk=wy_pk, tilt_pk=tilt_pk,
                   catch_tilt=catch_tilt, back_tilt=back_tilt,
                   land_roll=land_roll, t_lift=t_lift,
                   leak_frac=leak_frac,
                   sat_t=sat_t, eff=eff, dx=dx, free_t=free_t,
                   landed=landed, landed_t=landed_t,
                   feel=feel_charge(env, "jump"))
    # A FALL IS A COMPLETE FAIL (operator's rule): an episode that
    # does not land pays nothing - flat -1, no partial clearance, no
    # stay-put consolation. Landing every time is the product.
    if not landed:
        return -1.0
    # clearance uncapped further (operator: the tuck's wheel clearance
    # is the incentive) - pays to 192mm now, and only on landed jumps,
    # so it cannot be farmed by leaping into a crash
    return (- feel_charge(env, "jump")
            # v95: clearance pays 7.5/unit (was 5) - height stays worth
            # chasing once the rotation charges below make the backflip
            # route to "feet up" unprofitable
            + 7.5 * min(peak / 0.12, 1.6)
            + 3.0
            + 2.0 * max(0.0, 1.0 - abs(dx) / 0.5)
            # v94: run-out is CHARGED past the bonus's edge, not just
            # unpaid - a 2m landing slide scored 7.4 under v93 (the
            # bonus bottomed out at 0.5m and stopped caring)
            - 1.5 * max(0.0, abs(dx) - 0.5)
            # v97: the roll-out AFTER touchdown is the stagger the
            # operator watched - a stuck landing stops where it lands
            - 3.0 * max(0.0, land_roll - 0.1)
            # v98: the backward lean is charged from 8 deg so the
            # search feels the 20-deg fail cliff before falling off it
            - 5.0 * max(0.0, back_tilt - math.radians(8.0))
            # v102: rolling around before the jump is charged - every
            # second from start to liftoff past 1.0 costs 1.5 (the
            # operator watched it dawdle; the stillness gate had made
            # patience free)
            - 1.5 * max(0.0, (t_lift or 0.0) - 1.0)
            - 3.0 * leak_frac
            # v95, the operator's order: not rotating AT ALL is the
            # incentive. Absolute charges, not fractions - the v94 bake
            # pays ~7 points for its 6.4 rad/s backflip and ~2 for the
            # body going 55 deg over, so killing the rotation is worth
            # more than every other term combined.
            - 1.2 * max(0.0, wy_pk - 0.5)
            - 8.0 * max(0.0, tilt_pk - 0.30)
            - 4.0 * max(0.0, sat_t - 0.25)
            - EFFORT_RATE * eff)


def jump_score(g, secs=12.0):
    # half the score is the WORST episode: 'lands every time' is a
    # worst-case standard, and a mean alone lets one crash hide
    # behind two good jumps (operator caught it at 6.2).
    # v94: eight seeds, not three - the v93 bake landed its three
    # exam seeds and fell on two of the next five (different sensor
    # noise draws, nothing else). One episode costs 0.05s on the M4;
    # there is no excuse for a three-noise-history definition of
    # "every time".
    s = [jump_episode(g, i, secs) for i in range(8)]
    return float(0.5 * np.mean(s) + 0.5 * np.min(s))
