121 lines
4.8 KiB
Python
121 lines
4.8 KiB
Python
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
|
|
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.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.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._last_debug = StandBalanceDebug(0.0, 0.0, 0.0, 0.0, 0.9, -1.8, 0.0, 0.0, False)
|
|
|
|
@property
|
|
def last_debug(self) -> StandBalanceDebug:
|
|
return self._last_debug
|
|
|
|
def reset(self) -> None:
|
|
self._stable_time = 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
|
|
pitch_corr = -self.kp_pitch * pitch - self.kd_pitch_rate * pitch_rate
|
|
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
|
|
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 + 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,
|
|
stable=stable,
|
|
)
|
|
return target
|
|
|
|
def is_stable(self) -> bool:
|
|
return self._stable_time >= self.enter_hold_s
|