robots/spot/controllers/trot_gait.py

"""Trot Gait Controller for Spot robot."""
import numpy as np
import os
import importlib.util

# Import BaseController using importlib to work when loaded via importlib
_controllers_dir = os.path.dirname(os.path.abspath(__file__))
_base_path = os.path.join(_controllers_dir, "base.py")
_spec = importlib.util.spec_from_file_location("spot_base", _base_path)
_base_module = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(_base_module)
BaseController = _base_module.BaseController


class TrotGaitController(BaseController):
    """
    Classical trot gait controller for Spot quadruped.

    Trot gait: Diagonal legs move together (FL+HR, FR+HL).
    Uses simple sinusoidal trajectories for leg motion.
    """

    # Standing pose - from model's "home" keyframe
    STANDING_POSE = np.array([
        0.0, 1.04, -1.8,  # FL: hx, hy, kn
        0.0, 1.04, -1.8,  # FR
        0.0, 1.04, -1.8,  # HL
        0.0, 1.04, -1.8,  # HR
    ], dtype=np.float32)

    def __init__(self, num_joints: int = 12):
        super().__init__(num_joints)

        # Velocity command [vx, vy, vyaw]
        self.cmd = np.array([0.0, 0.0, 0.0], dtype=np.float32)

        # Gait parameters - tuned for visible motion
        self.gait_freq = 2.0  # Hz - gait cycles per second
        self.swing_height = 0.3  # How much to bend knee during swing (radians)
        self.stride_length = 0.25  # Hip pitch amplitude (radians)

        # Phase offset for trot (diagonal pairs in sync)
        # FL and HR at phase 0, FR and HL at phase 0.5
        self.phase_offsets = np.array([0.0, 0.5, 0.5, 0.0])  # FL, FR, HL, HR

    def set_command(self, vx: float = 0.0, vy: float = 0.0, vyaw: float = 0.0):
        """Set velocity command for the robot."""
        self.cmd[0] = np.clip(vx, -1.0, 1.0)
        self.cmd[1] = np.clip(vy, -0.5, 0.5)
        self.cmd[2] = np.clip(vyaw, -1.5, 1.5)

    def get_command(self):
        """Get current velocity command."""
        return self.cmd.copy()

    def compute_action(self, obs: np.ndarray, data) -> np.ndarray:
        """Compute target joint positions for trot gait."""
        vx, vy, vyaw = self.cmd

        # Check if we should be walking
        speed = np.sqrt(vx**2 + vy**2 + vyaw**2)
        if speed < 0.05:
            # Just stand
            return self.STANDING_POSE.copy()

        # Gait phase (0 to 1)
        phase = (self.time * self.gait_freq) % 1.0

        # Base standing angles (from model's home keyframe)
        hy_stand = 1.04   # Hip pitch
        kn_stand = -1.8   # Knee

        # Output array
        target = np.zeros(12, dtype=np.float32)

        for leg_idx in range(4):
            # Get leg phase with offset
            leg_phase = (phase + self.phase_offsets[leg_idx]) % 1.0

            # Determine if in swing (0-0.5) or stance (0.5-1.0)
            in_swing = leg_phase < 0.5

            # Compute phase within swing or stance (0 to 1)
            if in_swing:
                p = leg_phase / 0.5
            else:
                p = (leg_phase - 0.5) / 0.5

            # Hip pitch: sinusoidal forward/back motion
            # During swing: leg moves forward
            # During stance: leg moves backward (pushes body forward)
            if in_swing:
                # Swing forward: start back, end front
                hy_offset = self.stride_length * vx * (p - 0.5)
            else:
                # Stance backward: start front, end back
                hy_offset = self.stride_length * vx * (0.5 - p)

            # Knee: lift during swing phase
            if in_swing:
                # Sinusoidal lift
                kn_offset = -self.swing_height * np.sin(np.pi * p)
            else:
                kn_offset = 0.0

            # Hip roll for turning
            is_left = leg_idx in [0, 2]
            hx_offset = 0.0
            if abs(vyaw) > 0.01:
                # Differential motion for turning
                sign = 1.0 if is_left else -1.0
                if in_swing:
                    hx_offset = sign * 0.1 * vyaw * np.sin(np.pi * p)

            # Lateral motion
            if abs(vy) > 0.01:
                sign = 1.0 if is_left else -1.0
                if in_swing:
                    hx_offset += sign * 0.2 * vy * (p - 0.5)
                else:
                    hx_offset += sign * 0.2 * vy * (0.5 - p)

            # Set joint targets
            base_idx = leg_idx * 3
            target[base_idx + 0] = hx_offset                    # hx (hip roll)
            target[base_idx + 1] = hy_stand + hy_offset         # hy (hip pitch)
            target[base_idx + 2] = kn_stand + kn_offset         # kn (knee)

        return target

    def reset(self):
        super().reset()
        self.cmd = np.array([0.0, 0.0, 0.0], dtype=np.float32)

    @property
    def name(self) -> str:
        cmd = self.cmd
        return f"Trot Gait (cmd: {cmd[0]:.1f}, {cmd[1]:.1f}, {cmd[2]:.1f})"