[real] 整理 Python Sim2Real v2

05_software/real/sim2real_v2/startup/stand_balance.py

@@ -0,0 +1,141 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict
+
+import numpy as np
+
+
+@dataclass
+class StandBalanceDebug:
+    roll: float
+    pitch: float
+    roll_rate: float
+    pitch_rate: float
+    hip_base: float
+    knee_base: float
+    roll_corr: float
+    pitch_corr: float
+    pitch_compensation_enabled: bool
+    target: list[float]
+    stable: bool
+
+
+class StandBalanceController:
+    def __init__(self, cfg: Dict[str, Any], control_dt: float):
+        self.enabled = bool(cfg.get("enabled", True))
+        self.control_dt = float(control_dt)
+        self.height = float(cfg.get("height", 0.33))
+        self.kp_roll = float(cfg.get("kp_roll", 0.85))
+        self.pitch_compensation_enabled = bool(cfg.get("pitch_compensation_enabled", True))
+        self.kp_pitch = float(cfg.get("kp_pitch", 0.70))
+        self.kd_roll_rate = float(cfg.get("kd_roll_rate", 0.03))
+        self.kd_pitch_rate = float(cfg.get("kd_pitch_rate", 0.025))
+        self.pitch_deadband = float(np.radians(cfg.get("pitch_deadband_deg", 0.0)))
+        self.pitch_corr_clip = float(cfg.get("pitch_corr_clip", 0.12))
+        self.pitch_corr_filter_alpha = float(np.clip(cfg.get("pitch_corr_filter_alpha", 1.0), 0.0, 1.0))
+        self.pitch_front_sign = float(cfg.get("pitch_front_sign", -1.0))
+        self.lateral_lean_gain = float(cfg.get("lateral_lean_gain", 0.0))
+        self.hip_abduction_clip = float(cfg.get("hip_abduction_clip", 0.45))
+        self.hip_pitch_clip = tuple(cfg.get("hip_pitch_clip", [-1.0, 2.5]))
+        self.knee_clip = tuple(cfg.get("knee_clip", [-2.6, -0.3]))
+        self.stable_roll_deg = float(cfg.get("stable_roll_deg", 6.0))
+        self.stable_pitch_deg = float(cfg.get("stable_pitch_deg", 8.0))
+        self.stable_gyro_deg_s = float(cfg.get("stable_gyro_deg_s", 45.0))
+        self.enter_hold_s = float(cfg.get("enter_hold_s", 1.0))
+
+        self.profile_h = np.asarray(
+            cfg.get("profile_h", [0.157, 0.248, 0.311, 0.366, 0.411, 0.448]),
+            dtype=np.float32,
+        )
+        self.profile_hip = np.asarray(
+            cfg.get("profile_hip", [1.5, 1.2, 1.0, 0.8, 0.6, 0.4]),
+            dtype=np.float32,
+        )
+        self.profile_knee = np.asarray(
+            cfg.get("profile_knee", [-2.5, -2.1, -1.8, -1.5, -1.2, -0.9]),
+            dtype=np.float32,
+        )
+        self._stable_time = 0.0
+        self._pitch_corr_filtered = 0.0
+        self._last_debug = StandBalanceDebug(0.0, 0.0, 0.0, 0.0, 0.9, -1.8, 0.0, 0.0, False, [], False)
+
+    @property
+    def last_debug(self) -> StandBalanceDebug:
+        return self._last_debug
+
+    def reset(self) -> None:
+        self._stable_time = 0.0
+        self._pitch_corr_filtered = 0.0
+
+    def _estimate_roll_pitch(self, projected_gravity: np.ndarray) -> tuple[float, float]:
+        gx, gy, gz = [float(v) for v in projected_gravity]
+        roll = float(np.arctan2(-gy, max(1e-6, -gz)))
+        pitch = float(np.arctan2(gx, np.sqrt(max(1e-6, gy * gy + gz * gz))))
+        return roll, pitch
+
+    def _base_leg_pose(self) -> tuple[float, float]:
+        h_clamp = float(np.clip(self.height, float(self.profile_h[0]), float(self.profile_h[-1])))
+        hip = float(np.interp(h_clamp, self.profile_h, self.profile_hip))
+        knee = float(np.interp(h_clamp, self.profile_h, self.profile_knee))
+        return hip, knee
+
+    def compute_target(self, state: Dict[str, Any], command: np.ndarray | None = None) -> np.ndarray:
+        projected_gravity = np.asarray(state["projected_gravity"], dtype=np.float32)
+        imu_gyro = np.asarray(state["imu_gyro"], dtype=np.float32)
+        cmd = np.zeros(3, dtype=np.float32) if command is None else np.asarray(command, dtype=np.float32)
+
+        hip_base, knee_base = self._base_leg_pose()
+        roll, pitch = self._estimate_roll_pitch(projected_gravity)
+        roll_rate = float(imu_gyro[0])
+        pitch_rate = float(imu_gyro[1])
+
+        roll_corr = -self.kp_roll * roll - self.kd_roll_rate * roll_rate
+        if self.pitch_compensation_enabled:
+            pitch_for_ctrl = 0.0 if abs(pitch) < self.pitch_deadband else pitch
+            pitch_corr_raw = -self.kp_pitch * pitch_for_ctrl - self.kd_pitch_rate * pitch_rate
+            pitch_corr_raw = float(np.clip(pitch_corr_raw, -self.pitch_corr_clip, self.pitch_corr_clip))
+            alpha = self.pitch_corr_filter_alpha
+            pitch_corr = (1.0 - alpha) * self._pitch_corr_filtered + alpha * pitch_corr_raw
+            self._pitch_corr_filtered = pitch_corr
+        else:
+            pitch_corr = 0.0
+            self._pitch_corr_filtered = 0.0
+        lateral_lean = self.lateral_lean_gain * float(cmd[1])
+
+        target = np.zeros(16, dtype=np.float32)
+        for leg_idx in range(4):
+            side = 1.0 if leg_idx in (0, 2) else -1.0
+            fore_aft = self.pitch_front_sign if leg_idx in (0, 1) else -self.pitch_front_sign
+            target[leg_idx * 3 + 0] = float(
+                np.clip(side * roll_corr + lateral_lean, -self.hip_abduction_clip, self.hip_abduction_clip)
+            )
+            target[leg_idx * 3 + 1] = float(
+                np.clip(hip_base + fore_aft * pitch_corr, self.hip_pitch_clip[0], self.hip_pitch_clip[1])
+            )
+            target[leg_idx * 3 + 2] = float(np.clip(knee_base, self.knee_clip[0], self.knee_clip[1]))
+        target[12:] = 0.0
+
+        stable = (
+            abs(np.degrees(roll)) <= self.stable_roll_deg
+            and abs(np.degrees(pitch)) <= self.stable_pitch_deg
+            and max(abs(np.degrees(roll_rate)), abs(np.degrees(pitch_rate))) <= self.stable_gyro_deg_s
+        )
+        self._stable_time = self._stable_time + self.control_dt if stable else 0.0
+        self._last_debug = StandBalanceDebug(
+            roll=roll,
+            pitch=pitch,
+            roll_rate=roll_rate,
+            pitch_rate=pitch_rate,
+            hip_base=hip_base,
+            knee_base=knee_base,
+            roll_corr=roll_corr,
+            pitch_corr=pitch_corr,
+            pitch_compensation_enabled=self.pitch_compensation_enabled,
+            target=target.tolist(),
+            stable=stable,
+        )
+        return target
+
+    def is_stable(self) -> bool:
+        return self._stable_time >= self.enter_hold_s