217 lines
7.7 KiB
Python
217 lines
7.7 KiB
Python
import time
|
|
from typing import Callable, Dict, Tuple
|
|
|
|
import numpy as np
|
|
|
|
from interface.imu_client import IMUClient
|
|
from interface.motor_driver import HardwareIO
|
|
from tools.math_utils import LowPassFilter, MahonyFilter, get_gravity_orientation
|
|
|
|
|
|
def _quat_yaw_wxyz(quat) -> float:
|
|
w, x, y, z = [float(v) for v in quat]
|
|
return float(np.arctan2(2.0 * (w * z + x * y), 1.0 - 2.0 * (y * y + z * z)))
|
|
|
|
|
|
def _wrap_pi(angle: float) -> float:
|
|
return float((angle + np.pi) % (2.0 * np.pi) - np.pi)
|
|
|
|
|
|
class OdomTracker:
|
|
def __init__(self, jump_distance_m: float = 0.5, jump_yaw_rad: float = 0.8):
|
|
self.jump_distance_m = float(jump_distance_m)
|
|
self.jump_yaw_rad = float(jump_yaw_rad)
|
|
self.origin_pos = None
|
|
self.origin_yaw = 0.0
|
|
self.last_local_pos = None
|
|
self.last_local_yaw = 0.0
|
|
|
|
def reset(self):
|
|
self.origin_pos = None
|
|
self.origin_yaw = 0.0
|
|
self.last_local_pos = None
|
|
self.last_local_yaw = 0.0
|
|
|
|
def update(self, odom):
|
|
if odom is None:
|
|
return None
|
|
pos = np.asarray(odom.get("pos", [0.0, 0.0, 0.0]), dtype=np.float32)
|
|
yaw = _quat_yaw_wxyz(odom.get("quat_wxyz", [1.0, 0.0, 0.0, 0.0]))
|
|
if self.origin_pos is None:
|
|
self.origin_pos = pos.copy()
|
|
self.origin_yaw = yaw
|
|
local_pos = pos - self.origin_pos
|
|
local_yaw = _wrap_pi(yaw - self.origin_yaw)
|
|
jump_detected = False
|
|
jump_distance = 0.0
|
|
jump_yaw = 0.0
|
|
if self.last_local_pos is not None:
|
|
jump_distance = float(np.linalg.norm(local_pos[:2] - self.last_local_pos[:2]))
|
|
jump_yaw = abs(_wrap_pi(local_yaw - self.last_local_yaw))
|
|
jump_detected = jump_distance > self.jump_distance_m or jump_yaw > self.jump_yaw_rad
|
|
self.last_local_pos = local_pos.copy()
|
|
self.last_local_yaw = local_yaw
|
|
tracked = dict(odom)
|
|
tracked.update(
|
|
{
|
|
"local_pos": local_pos.tolist(),
|
|
"local_yaw": local_yaw,
|
|
"jump_detected": bool(jump_detected),
|
|
"jump_distance_m": jump_distance,
|
|
"jump_yaw_rad": jump_yaw,
|
|
}
|
|
)
|
|
return tracked
|
|
|
|
|
|
class RealIO:
|
|
def __init__(
|
|
self,
|
|
driver_factory: Callable[[str, str, bool], Tuple[object, object]],
|
|
motor_model: str,
|
|
can1_port: str,
|
|
can2_port: str,
|
|
imu_lib_path: str,
|
|
control_dt: float = 0.02,
|
|
motor_dt: float = 0.005,
|
|
kp_leg: float = 80.0,
|
|
kd_leg: float = 2.5,
|
|
hold_kp_leg: float | None = None,
|
|
hold_kd_leg: float | None = None,
|
|
kd_wheel: float = 2.0,
|
|
debug: bool = False,
|
|
dry_run: bool = False,
|
|
):
|
|
self.control_dt = control_dt
|
|
self.motor_dt = motor_dt
|
|
self.kp_leg = kp_leg
|
|
self.kd_leg = kd_leg
|
|
self.hold_kp_leg = kp_leg if hold_kp_leg is None else float(hold_kp_leg)
|
|
self.hold_kd_leg = kd_leg if hold_kd_leg is None else float(hold_kd_leg)
|
|
self.kd_wheel = kd_wheel
|
|
|
|
print("[RealIO] 初始化电机驱动...")
|
|
self.hw = HardwareIO(driver_factory, motor_model, can1_port, can2_port, debug)
|
|
print("[RealIO] 初始化 IMU...")
|
|
self.imu = IMUClient(lib_path=imu_lib_path, dry_run=dry_run)
|
|
|
|
# 使用 motor_dt 初始化滤波器，因为它们都在 200Hz 电机控制循环中更新
|
|
self.imu_filter = MahonyFilter(kp=2.0, ki=0.0, dt=motor_dt)
|
|
self.quat_wxyz = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32)
|
|
|
|
self.lpf_legs = LowPassFilter(cutoff_freq=5.0, dt=motor_dt, dim=12)
|
|
self.lpf_wheels = LowPassFilter(cutoff_freq=15.0, dt=motor_dt, dim=4)
|
|
|
|
self._last_imu_age_ms = -1.0
|
|
self._last_imu_fresh = False
|
|
self.odom_tracker = OdomTracker()
|
|
self._last_read_time = None
|
|
|
|
def connect(self, imu_timeout_ms: int = 8000):
|
|
self.hw.connect()
|
|
self.imu.start(timeout_ms=imu_timeout_ms)
|
|
if self.imu.initial_gravity is not None:
|
|
self.imu_filter.reset_with_accel(self.imu.initial_gravity)
|
|
self.quat_wxyz = self.imu_filter.q.copy()
|
|
self.odom_tracker.reset()
|
|
self._last_read_time = None
|
|
|
|
def disconnect(self):
|
|
try:
|
|
self.hw.disable_all()
|
|
finally:
|
|
self.imu.stop()
|
|
self.hw.disconnect()
|
|
|
|
def enable_motors(self):
|
|
self.hw.enable_all()
|
|
|
|
def disable_motors(self):
|
|
self.hw.disable_all()
|
|
|
|
def damping_brake(self):
|
|
self.hw.damping_brake(self.kd_leg, self.kd_wheel)
|
|
|
|
def wait_feedback_ready(self, max_attempts: int = 20, poll_interval: float = 0.05):
|
|
return self.hw.wait_feedback_ready(max_attempts=max_attempts, poll_interval=poll_interval)
|
|
|
|
def read_measured_pose(self) -> np.ndarray:
|
|
return self.hw.read_measured_pose()
|
|
|
|
def read_state(self) -> Dict[str, object]:
|
|
joint_pos, joint_vel, joint_torque, motor_diag = self.hw.read_state()
|
|
gyro, accel, age_ms, fresh = self.imu.get_latest()
|
|
odom = self.odom_tracker.update(self.imu.get_latest_odom())
|
|
self._last_imu_age_ms = age_ms
|
|
self._last_imu_fresh = fresh
|
|
|
|
# 动态测量 dt，以适应 POLL (5Hz) 与 RUNTIME (200Hz) 的不同频率切换
|
|
t_now = time.perf_counter()
|
|
if self._last_read_time is not None:
|
|
dt = t_now - self._last_read_time
|
|
if dt <= 0.0 or dt > 0.5:
|
|
dt = self.motor_dt
|
|
else:
|
|
dt = self.motor_dt
|
|
self._last_read_time = t_now
|
|
|
|
self.quat_wxyz = self.imu_filter.update(accel, gyro, dt=dt)
|
|
projected_gravity = get_gravity_orientation(self.quat_wxyz)
|
|
|
|
return {
|
|
"joint_pos": joint_pos,
|
|
"joint_vel": joint_vel,
|
|
"joint_torque": joint_torque,
|
|
"imu_gyro": gyro,
|
|
"imu_accel": accel,
|
|
"quat_wxyz": self.quat_wxyz.copy(),
|
|
"projected_gravity": projected_gravity,
|
|
"imu_age_ms": age_ms,
|
|
"imu_fresh": fresh,
|
|
"odom": odom,
|
|
"motor_stale": motor_diag,
|
|
}
|
|
|
|
def get_obs_policy(
|
|
self,
|
|
state: Dict[str, object],
|
|
command: np.ndarray,
|
|
default_dof_pos: np.ndarray,
|
|
last_actions_raw: np.ndarray,
|
|
) -> np.ndarray:
|
|
gyro = state["imu_gyro"]
|
|
joint_pos = state["joint_pos"]
|
|
joint_vel = state["joint_vel"]
|
|
projected_gravity = state["projected_gravity"]
|
|
|
|
base_ang_vel = (gyro * 0.25).astype(np.float32)
|
|
joint_pos_rel = (joint_pos[:12] - default_dof_pos[:12]).astype(np.float32)
|
|
joint_vel_leg = (joint_vel[:12] * 0.05).astype(np.float32)
|
|
wheel_vel = (joint_vel[12:] * 0.05).astype(np.float32)
|
|
|
|
return np.concatenate(
|
|
[
|
|
base_ang_vel,
|
|
projected_gravity,
|
|
command.astype(np.float32),
|
|
joint_pos_rel,
|
|
joint_vel_leg,
|
|
wheel_vel,
|
|
last_actions_raw,
|
|
]
|
|
).astype(np.float32)
|
|
|
|
def send_actions(self, scaled_actions: np.ndarray, default_dof_pos: np.ndarray):
|
|
act = (scaled_actions + default_dof_pos).astype(np.float32)
|
|
act = np.clip(act, -100.0, 100.0)
|
|
act[:12] = self.lpf_legs.filter(act[:12])
|
|
act[12:] = self.lpf_wheels.filter(act[12:])
|
|
self.hw.send_control(act, self.kp_leg, self.kd_leg, self.kd_wheel)
|
|
return act
|
|
|
|
def hold_pose(self, sim_target_pose: np.ndarray, kp_scale: float = 1.0):
|
|
target = np.clip(sim_target_pose.astype(np.float32), -100.0, 100.0)
|
|
kp_scale = float(np.clip(kp_scale, 0.0, 1.0))
|
|
self.hw.send_control(target, self.hold_kp_leg * kp_scale, self.hold_kd_leg, self.kd_wheel)
|
|
return target
|