runtime/controllers/gait_controller.py

#!/usr/bin/env python3
"""
ARGOS-Roach gait controller module.

GaitController is the SWAPPABLE interface: any neuro-controller can replace the
CPG implementation without touching the simulation harness, as long as it
implements:
    step(t, state) -> np.ndarray[18]   # joint angle targets, rad
    set_command(cmd, t, ttl=None)      # 'GO' | 'STOP' | 'HOLD'
    reset(rng=None)

state dict keys: dt (s), t (s), qpos (18,), qvel (18,), body_pos (3,),
body_quat (4, wxyz), body_angvel (3,), foot_contact (6, 0/1 per leg).

Command / stop policy (fixed rule, test D):
  * 'GO'        : CPG amplitudes ramp in from 0 with smoothstep over T_RAMP s.
  * 'STOP'      : amplitudes decay exponentially (tau_amp = 0.15 s), oscillator
                  frequency decays (tau_om = 0.10 s); targets converge to the
                  standing pose q0. Below amp < 0.005 controller holds q0
                  (standing) until the next command.
  * 'GO' + ttl  : command expires at t_ttl = t + ttl -> automatic smooth stop
                  (same rule as STOP). No unbounded execution of a stale command.

Yaw fix (26.09.2026): the vestibular bias was applied with the per-leg side
sign s, which shifts BOTH sides' feet the same way (a speed trim, not steering).
A yaw moment needs DIFFERENTIAL coxa action: the bias is now added WITHOUT s
(right coxa += bias, left coxa += bias -> right foot back, left foot forward ->
clockwise moment for bias > 0). bias = k_psi*yaw + k_psid*yaw_rate gives
heading stabilization (proportional + damping).

Retracted sweep (26.09.2026): sweep(phi) = cos(phi) + k_ret*sin(2*phi).
The raw +cos profile has zero foot velocity at touchdown (phi=pi), so the
planted foot of the new tripod is dragged forward against the body's motion
(slip) and jerks the body backward. With k_ret > 0 the foot already moves
BACKWARD at touchdown (d/dphi sweep = +2*k_ret at phi=pi, and the foot
x-offset is -L*A*sweep, so x-velocity is negative) and at lift-off, matching
the body's forward speed -> no touchdown jerk, plus a push-off retraction at
the end of stance. sweep and its derivative are continuous across the
phi=0/2pi wrap.

Units: radians, seconds. Joint order per leg: [coxa, femur, tibia], legs 0..5.
"""

import math
import numpy as np

LEG_COUNT = 6
JOINTS_PER_LEG = 3
N_ACT = LEG_COUNT * JOINTS_PER_LEG

# leg numbering: 0 RF, 1 LF, 2 LM, 3 RM, 4 RR, 5 LR ; tripod A={0,2,4}, B={1,3,5}
SIDE = np.array([+1, -1, -1, +1, +1, -1], dtype=float)  # +1 right, -1 left


class GaitController:
    """Abstract interface (swappable module)."""
    def set_command(self, cmd, t, ttl=None):
        raise NotImplementedError

    def step(self, t, state):
        raise NotImplementedError

    def reset(self, rng=None):
        raise NotImplementedError


class CPGGaitController(GaitController):
    """
    CPG: 6 coupled phase oscillators (Kuramoto), tripod gait.
    Oscillator i: dphi_i/dt = omega + K * sum_j sin(phi_j - phi_i - psi_ij)
    Coupling: same-side neighbours psi=0 (in-phase), opposite tripod pairs psi=pi.
    Joint targets (per leg, side sign s):
        coxa  = 0 + s*A_c*sweep(phi)  + bias        # fore-aft sweep (retracted)
        femur = s * (0.8 + A_f * lift(phi))         # lift
        tibia = -s * (1.45 + A_t * lift(phi))       # knee tuck
        lift(phi) = max(0, sin(phi))^2 > 0 during swing (phi in (0, pi)).
    Gait cycle: phi=0 lift-off at back (with backward retraction) -> swing
    back->front while lifted -> touchdown at front with the foot already
    moving backward (retraction, no slip jerk) -> stance sweep front->back.
    Subclass hooks: _update_phases(dt) (rhythm generator) and
    _apply_reflex(state, dt, prev_phases) (feedback; no-op in this class).
    """

    def __init__(self, freq=1.0, a_coxa=0.45, a_femur=0.40, a_tibia=0.42,
                 coupling_k=2.0, t_ramp=1.0, tau_amp=0.15, tau_om=0.10,
                 amp_floor=0.005, k_psi=0.0, k_psid=0.0, k_ret=0.30,
                 tau_yf=0.0, bias_stance=True, a_coxa_legs=None):
        self.omega0 = 2.0 * math.pi * freq
        self.a_coxa, self.a_femur, self.a_tibia = a_coxa, a_femur, a_tibia
        # per-leg coxa amplitude scale (None = 1.0 for all legs)
        self.a_coxa_legs = (np.asarray(a_coxa_legs, dtype=float)
                            if a_coxa_legs is not None else np.ones(LEG_COUNT))
        self.coupling_k = coupling_k
        self.t_ramp, self.tau_amp, self.tau_om = t_ramp, tau_amp, tau_om
        self.amp_floor = amp_floor
        # vestibular yaw stabilization: coxa bias = k_psi * yaw + k_psid * yaw_rate
        self.k_psi, self.k_psid = k_psi, k_psid
        self.k_ret = k_ret
        self.tau_yf = tau_yf          # low-pass filter time constant (0 = off)
        self.bias_stance = bias_stance  # apply yaw bias to stance feet too?
        self._yaw = 0.0
        self._yaw_rate = 0.0
        self._yaw_f = 0.0
        self._yawrate_f = 0.0
        self.F0, self.T0 = 0.8, 1.45  # standing femur / tibia magnitude (rad)
        # coupling pairs: (i, j, psi). Within-tripod pairs in-phase (psi=0),
        # cross-tripod pairs anti-phase (psi=pi). NOTE (26.09.2026 fix): the
        # original list used psi=0 for same-SIDE pairs (0,3),(3,4),(1,2),(2,5),
        # which connect legs in OPPOSITE tripods — a frustrated coupling whose
        # tripod equilibrium is unstable; the phases collapse in ~10-25 s into
        # an all-right vs all-left pattern, which is the root cause of the
        # late-run yaw drift and jerk regime.
        self.pairs = [(0, 2, 0.0), (2, 4, 0.0), (4, 0, 0.0),
                      (1, 3, 0.0), (3, 5, 0.0), (5, 1, 0.0),
                      (0, 1, math.pi), (2, 3, math.pi), (4, 5, math.pi)]
        self.reset()

    # ---------------- interface ----------------
    def reset(self, rng=None):
        self.phases = np.array([0.0, math.pi, 0.0, math.pi, 0.0, math.pi])
        self.omega_eff = self.omega0
        self.amp = 0.0
        self.mode = "hold"          # hold | go | stopping
        self._t_ref = 0.0           # reference time of last command/ramp state
        self._ttl_exp = None
        self._last_targets = np.zeros(N_ACT)
        self._yaw = 0.0
        self._yaw_rate = 0.0
        self._yaw_f = 0.0
        self._yawrate_f = 0.0

    @staticmethod
    def _yaw_from_quat(q):
        w, x, y, z = q[0], q[1], q[2], q[3]
        siny = 2.0 * (w * z + x * y)
        cosy = 1.0 - 2.0 * (y * y + z * z)
        return math.atan2(siny, cosy)

    def set_command(self, cmd, t, ttl=None):
        cmd = str(cmd).upper()
        if cmd == "GO":
            self.mode = "go"
            self._t_ref = t
            self._ttl_exp = (t + ttl) if ttl is not None else None
            self.omega_eff = self.omega0
        elif cmd == "STOP":
            self._begin_stop(t)
        elif cmd == "HOLD":
            self.mode = "hold"
            self.amp = 0.0
            self._ttl_exp = None
            self._t_ref = t
        else:
            raise ValueError(f"unknown command: {cmd}")

    def _begin_stop(self, t):
        self.mode = "stopping"
        self._t_ref = t
        self._ttl_exp = None

    # ---------------- core ----------------
    @staticmethod
    def standing_targets():
        """Standing pose q0 for all 18 joints (rad)."""
        q = np.zeros(N_ACT)
        for leg in range(LEG_COUNT):
            s = SIDE[leg]
            q[3 * leg + 0] = 0.0
            q[3 * leg + 1] = s * 0.8
            q[3 * leg + 2] = -s * 1.45
        return q

    def _update_phases(self, dt):
        """Rhythm generator. Base: Kuramoto-coupled phase oscillators."""
        dphi = np.full(LEG_COUNT, self.omega_eff * dt)
        for i, j, psi in self.pairs:
            d = self.phases[j] - self.phases[i] - psi
            dphi[i] += self.coupling_k * math.sin(d) * dt
        self.phases = (self.phases + dphi) % (2.0 * math.pi)

    def _apply_reflex(self, state, dt, prev_phases):
        """Feedback hook (no-op in the open-loop base controller)."""

    def _post_targets(self, state, targets):
        """Post-target hook for additive limb corrections (no-op in base)."""
        return targets

    def step(self, t, state):
        dt = float(state["dt"])

        # --- command expiry (ttl) -> same smooth stop as explicit STOP ---
        if self.mode == "go" and self._ttl_exp is not None and t >= self._ttl_exp:
            self._begin_stop(t)

        # --- mode state machine ---
        if self.mode == "hold":
            self._last_targets = self.standing_targets()
            return self._last_targets.copy()

        if self.mode == "go":
            # smoothstep ramp-in of amplitude
            x = max(0.0, min(1.0, (t - self._t_ref) / self.t_ramp))
            self.amp = x * x * (3.0 - 2.0 * x)
            self.omega_eff = self.omega0
        elif self.mode == "stopping":
            self.amp *= math.exp(-dt / self.tau_amp)
            self.omega_eff *= math.exp(-dt / self.tau_om)
            if self.amp < self.amp_floor:
                self.mode = "hold"
                self.amp = 0.0
                self._last_targets = self.standing_targets()
                return self._last_targets.copy()

        # --- rhythm update (overridable) + feedback reflex (overridable) ---
        prev_phases = self.phases.copy()
        self._update_phases(dt)
        self._apply_reflex(state, dt, prev_phases)

        # --- vestibular yaw feedback (steering via DIFFERENTIAL coxa bias) ---
        if "body_quat" in state:
            self._yaw = self._yaw_from_quat(state["body_quat"])
            if "body_angvel" in state:
                self._yaw_rate = float(state["body_angvel"][2])
        if self.tau_yf > 0:
            a = dt / self.tau_yf
            self._yaw_f += a * (self._yaw - self._yaw_f)
            self._yawrate_f += a * (self._yaw_rate - self._yawrate_f)
        else:
            self._yaw_f, self._yawrate_f = self._yaw, self._yaw_rate
        bias = self.k_psi * self._yaw_f + self.k_psid * self._yawrate_f

        # --- joint targets ---
        sinp = np.sin(self.phases)
        lift = np.clip(sinp, 0.0, None) ** 2
        sweep = np.cos(self.phases) + self.k_ret * np.sin(2.0 * self.phases)
        q0 = self.standing_targets()
        targets = q0.copy()
        for leg in range(LEG_COUNT):
            s = SIDE[leg]
            # bias NOT multiplied by s: right coxa += bias, left coxa += bias
            # -> differential fore-aft foot shift -> yaw moment (heading control).
            # Optional: apply the bias only while the leg is in SWING (the
            # touchdown point is steered without disturbing planted feet).
            if not self.bias_stance and (self.phases[leg] % (2.0 * math.pi)) >= math.pi:
                bias_leg = 0.0
            else:
                bias_leg = bias
            targets[3 * leg + 0] = (q0[3 * leg + 0]
                                    + s * self.a_coxa * self.a_coxa_legs[leg] * sweep[leg] * self.amp
                                    + bias_leg * self.amp)
            targets[3 * leg + 1] = s * (self.F0 + self.a_femur * lift[leg] * self.amp)
            targets[3 * leg + 2] = -s * (self.T0 + self.a_tibia * lift[leg] * self.amp)
        targets = self._post_targets(state, targets)  # hook: limb reflexes
        self._last_targets = targets
        return targets.copy()

    # ---------------- introspection (for CSV/report) ----------------
    def state_report(self):
        return {
            "mode": self.mode,
            "amp": self.amp,
            "omega_eff": self.omega_eff,
            "phases": np.round(self.phases, 3).tolist(),
        }