GoWelcome/gowelcome/robot/go2_robot.py

441 lines
17 KiB
Python

"""Concrete :class:`RobotInterface` backed by the official Unitree Go2 SDK.
This module talks to a physical Unitree Go2 over CycloneDDS using the
``unitree_sdk2py`` package. All heavy/optional dependencies (``unitree_sdk2py``,
``cv2``) are imported *lazily* inside :class:`Go2Robot.__init__`, so importing
this module never fails on an off-robot machine that lacks the SDK -- only
*instantiating* :class:`Go2Robot` requires it.
Velocity convention (matches ``SportClient.Move``):
vx forward (+) / backward (-) m/s, body frame
vy left (+) / right (-) m/s, body frame
vyaw turn-left/CCW (+) / right/CW (-) rad/s
Driving is routed through one of two paths:
* ``ObstaclesAvoidClient.Move`` -- the firmware LiDAR avoidance layer filters
and hard-stops these commands (the *safe* path), used when
``cfg.safety.use_lidar_avoidance`` is enabled.
* ``SportClient.Move`` -- the raw locomotion path, used otherwise.
A background daemon thread continuously pulls + decodes camera frames so that
:meth:`get_frame` is non-blocking and decoupled from the control rate.
"""
from __future__ import annotations
import logging
import threading
import time
from typing import TYPE_CHECKING, Callable, Dict, Optional, Tuple
import numpy as np
from gowelcome.robot.interface import GESTURES, RobotInterface
if TYPE_CHECKING: # pragma: no cover - type hints only
from config import GoWelcomeConfig
logger = logging.getLogger(__name__)
class Go2Robot(RobotInterface):
"""Drive a physical Unitree Go2 via the official ``unitree_sdk2py`` SDK.
Args:
cfg: The fully-populated :class:`~config.GoWelcomeConfig`.
Raises:
ImportError: If ``unitree_sdk2py`` (or ``cv2``) is not installed.
"""
def __init__(self, cfg: "GoWelcomeConfig") -> None:
# --- Lazy heavy imports (kept out of module scope on purpose) --------
try:
from unitree_sdk2py.core.channel import ChannelFactoryInitialize
from unitree_sdk2py.go2.obstacles_avoid.obstacles_avoid_client import (
ObstaclesAvoidClient,
)
from unitree_sdk2py.go2.sport.sport_client import SportClient
from unitree_sdk2py.go2.video.video_client import VideoClient
except ImportError as exc: # pragma: no cover - off-robot path
raise ImportError(
"Go2Robot requires the official Unitree SDK. Install it with: "
"pip install unitree_sdk2py (or build from "
"github.com/unitreerobotics/unitree_sdk2_python)."
) from exc
try:
import cv2 # noqa: F401 (imported for an early, clear failure)
except ImportError as exc: # pragma: no cover - off-robot path
raise ImportError(
"Go2Robot requires OpenCV to decode camera frames. Install it "
"with: pip install opencv-python"
) from exc
self.cfg = cfg
self._cv2 = cv2 # cache the module for the camera thread
# --- DDS channel: initialise EXACTLY ONCE before any client ---------
ChannelFactoryInitialize(cfg.network.domain_id, cfg.network.interface)
# --- Sport (locomotion / posture / gestures) ------------------------
self._sport = SportClient()
self._sport.SetTimeout(10.0)
self._sport.Init()
# --- Video (front camera) -------------------------------------------
self._video = VideoClient()
self._video.SetTimeout(3.0)
self._video.Init()
# --- Obstacle avoidance (firmware LiDAR hard-stop layer) ------------
self._oa = ObstaclesAvoidClient()
self._oa.SetTimeout(3.0)
self._oa.Init()
# --- Vui (Go2 on-board volume only) ---------------------------------
self._vui = None
try:
from unitree_sdk2py.go2.vui.vui_client import VuiClient
vui = VuiClient()
vui.SetTimeout(3.0)
vui.Init()
self._vui = vui
except Exception as exc: # pragma: no cover - firmware-dependent
logger.warning("VuiClient unavailable (volume control disabled): %s", exc)
# --- Audio backend (greeting playback) ------------------------------
from gowelcome.robot.audio import build_audio_backend
self.audio = build_audio_backend(cfg)
self._greet_volume_set = False
# --- Gesture name -> SportClient method map -------------------------
self._gestures: Dict[str, Callable[[], int]] = {
"hello": self._sport.Hello,
"heart": self._sport.Heart,
"stretch": self._sport.Stretch,
"dance1": self._sport.Dance1,
"dance2": self._sport.Dance2,
"scrape": self._sport.Scrape,
"content": self._sport.Content,
"sit": self._sport.Sit,
"rise_sit": self._sport.RiseSit,
}
# --- Camera frame buffer (written by the camera thread) -------------
self._frame: Optional[np.ndarray] = None
self._frame_size: Tuple[int, int] = (0, 0)
self._frame_lock = threading.Lock()
# --- Avoidance routing state ----------------------------------------
self._avoidance_on = False
if cfg.safety.use_lidar_avoidance:
self.set_avoidance(True)
# --- Shutdown guard (idempotent) ------------------------------------
self._shutdown = False
self._shutdown_lock = threading.Lock()
# --- Start the camera thread last, once everything is wired up ------
self._cam_stop = threading.Event()
self._cam_thread = threading.Thread(
target=self._camera_loop, name="Go2CameraThread", daemon=True
)
self._cam_thread.start()
logger.info(
"Go2Robot ready (iface=%s domain=%d avoidance=%s dry_run=%s)",
cfg.network.interface,
cfg.network.domain_id,
self._avoidance_on,
cfg.dry_run,
)
# ------------------------------------------------------------------ #
# Camera thread #
# ------------------------------------------------------------------ #
def _camera_loop(self) -> None:
"""Continuously pull + decode camera frames into ``self._frame``.
Runs on a daemon thread; failures are logged and retried so a transient
decode error never kills the stream. Throttled to ~camera target_fps
when frames are unavailable to avoid a busy-spin.
"""
cv2 = self._cv2
min_period = 1.0 / max(1.0, self.cfg.camera.target_fps)
while not self._cam_stop.is_set():
start = time.monotonic()
try:
code, data = self._video.GetImageSample()
except Exception as exc: # pragma: no cover - hardware path
logger.debug("GetImageSample raised: %s", exc)
self._cam_stop.wait(min_period)
continue
if code != 0 or not data:
logger.debug("GetImageSample non-zero code=%s", code)
self._cam_stop.wait(min_period)
continue
try:
buf = np.frombuffer(bytes(data), dtype=np.uint8)
image = cv2.imdecode(buf, cv2.IMREAD_COLOR)
except Exception as exc: # pragma: no cover - hardware path
logger.debug("Frame decode failed: %s", exc)
self._cam_stop.wait(min_period)
continue
if image is None:
self._cam_stop.wait(min_period)
continue
h, w = image.shape[:2]
with self._frame_lock:
self._frame = image
self._frame_size = (int(w), int(h))
# Gentle pace so we don't outrun the camera / burn CPU.
elapsed = time.monotonic() - start
if elapsed < min_period:
self._cam_stop.wait(min_period - elapsed)
# ------------------------------------------------------------------ #
# Perception input #
# ------------------------------------------------------------------ #
def get_frame(self) -> "Optional[np.ndarray]":
"""Return the latest BGR camera frame, or ``None`` if none yet."""
with self._frame_lock:
return self._frame
def frame_size(self) -> "tuple[int, int]":
"""Current ``(width, height)``; ``(0, 0)`` until the first frame."""
with self._frame_lock:
return self._frame_size
# ------------------------------------------------------------------ #
# Locomotion #
# ------------------------------------------------------------------ #
def drive(self, vx: float, vy: float, vyaw: float) -> None:
"""Command a body-frame velocity for one tick.
The :class:`~config.SafetyConfig` hard caps are applied first, then the
command is routed through the active path (avoidance ``Move`` if
enabled, else sport ``Move``). When ``cfg.dry_run`` is set the intended
command is logged but only a zero velocity is actually sent.
"""
safety = self.cfg.safety
cvx = _clamp(vx, -safety.max_vx, safety.max_vx)
cvy = _clamp(vy, -safety.max_vy, safety.max_vy)
cvyaw = _clamp(vyaw, -safety.max_vyaw, safety.max_vyaw)
if self.cfg.dry_run:
logger.info(
"[dry_run] drive intent vx=%.3f vy=%.3f vyaw=%.3f (sending 0,0,0)",
cvx,
cvy,
cvyaw,
)
self._send_move(0.0, 0.0, 0.0)
return
self._send_move(cvx, cvy, cvyaw)
def _send_move(self, vx: float, vy: float, vyaw: float) -> None:
"""Dispatch a velocity through the currently-active drive path."""
if self._avoidance_on:
self._call("oa.Move", lambda: self._oa.Move(vx, vy, vyaw))
else:
self._call("sport.Move", lambda: self._sport.Move(vx, vy, vyaw))
def stop(self) -> None:
"""Command zero velocity through the active path (stay standing)."""
self._send_move(0.0, 0.0, 0.0)
def set_avoidance(self, on: bool) -> None:
"""Enable/disable the on-robot LiDAR hard-stop avoidance layer.
When enabling: poll ``SwitchSet/SwitchGet`` until the firmware reports
the switch as on, then claim remote control via
``UseRemoteCommandFromApi(True)``. When disabling: release remote
control and turn the switch off. ``drive`` routes through ``oa.Move``
only while enabled.
"""
if on:
ok = self._enable_avoidance()
self._avoidance_on = ok
if not ok:
logger.warning(
"Avoidance could not be enabled; falling back to sport.Move path"
)
else:
self._call(
"oa.UseRemoteCommandFromApi(False)",
lambda: self._oa.UseRemoteCommandFromApi(False),
)
self._call("oa.SwitchSet(False)", lambda: self._oa.SwitchSet(False))
self._avoidance_on = False
logger.info("LiDAR avoidance disabled")
def _enable_avoidance(self) -> bool:
"""Turn the firmware avoidance switch on and claim API control.
Returns:
``True`` if the switch is confirmed on and API control was claimed.
"""
deadline = time.monotonic() + 5.0
enabled = False
while time.monotonic() < deadline:
try:
self._oa.SwitchSet(True)
code, enable = self._oa.SwitchGet()
except Exception as exc: # pragma: no cover - hardware path
logger.debug("Avoidance SwitchGet/Set raised: %s", exc)
time.sleep(0.1)
continue
if code == 0 and enable:
enabled = True
break
time.sleep(0.1)
if not enabled:
return False
rc = self._call(
"oa.UseRemoteCommandFromApi(True)",
lambda: self._oa.UseRemoteCommandFromApi(True),
)
if rc != 0:
return False
logger.info("LiDAR avoidance enabled (API remote control claimed)")
return True
# ------------------------------------------------------------------ #
# Posture / expression #
# ------------------------------------------------------------------ #
def balance_stand(self) -> None:
"""Enter balanced standing (ready-to-move)."""
self._call("sport.BalanceStand", self._sport.BalanceStand)
def stand_up(self) -> None:
"""Stiff stand up."""
self._call("sport.StandUp", self._sport.StandUp)
def damp(self) -> None:
"""Emergency soft-stop: limp the motors (safe failure posture)."""
self._call("sport.Damp", self._sport.Damp)
def gesture(self, name: str) -> None:
"""Perform an expressive gesture (see :data:`GESTURES`).
Unknown names are logged and ignored (no-op).
"""
fn = self._gestures.get(name)
if fn is None:
logger.warning(
"Unknown gesture %r (known: %s) -- ignoring",
name,
", ".join(GESTURES),
)
return
self._call(f"sport.{name}", fn)
# ------------------------------------------------------------------ #
# Greeting payload #
# ------------------------------------------------------------------ #
def play_greeting(self) -> None:
"""Play the configured greeting clip (non-blocking).
Best-effort sets the Go2 on-board volume once via ``VuiClient`` (the
gesture sequence itself is driven by the state machine calling
:meth:`gesture`).
"""
if self._vui is not None and not self._greet_volume_set:
try:
self._vui.SetVolume(int(self.cfg.greet.audio_volume))
except Exception as exc: # pragma: no cover - firmware-dependent
logger.debug("VuiClient.SetVolume failed: %s", exc)
finally:
self._greet_volume_set = True
try:
self.audio.play(self.cfg.greet.wav_path, blocking=False)
except Exception as exc:
logger.warning("Greeting playback failed: %s", exc)
# ------------------------------------------------------------------ #
# Lifecycle #
# ------------------------------------------------------------------ #
def shutdown(self) -> None:
"""Stop motion, release the avoidance API, close camera + audio.
Idempotent and safe to call from a signal handler / ``finally`` block.
Every step is wrapped so one failure does not block the rest.
"""
with self._shutdown_lock:
if self._shutdown:
return
self._shutdown = True
# 1. Stop the camera thread.
try:
self._cam_stop.set()
if self._cam_thread.is_alive():
self._cam_thread.join(timeout=2.0)
except Exception as exc:
logger.warning("Camera thread shutdown failed: %s", exc)
# 2. Release the firmware avoidance path (zero velocity + drop API).
if self._avoidance_on:
try:
self._oa.Move(0.0, 0.0, 0.0)
except Exception as exc:
logger.warning("oa.Move(0,0,0) on shutdown failed: %s", exc)
try:
self._oa.UseRemoteCommandFromApi(False)
except Exception as exc:
logger.warning("oa.UseRemoteCommandFromApi(False) failed: %s", exc)
self._avoidance_on = False
# 3. Stop locomotion.
try:
self._sport.StopMove()
except Exception as exc:
logger.warning("sport.StopMove() on shutdown failed: %s", exc)
# 4. Close audio.
try:
self.audio.close()
except Exception as exc:
logger.warning("audio.close() on shutdown failed: %s", exc)
logger.info("Go2Robot shutdown complete")
# ------------------------------------------------------------------ #
# Helpers #
# ------------------------------------------------------------------ #
def _call(self, what: str, fn: Callable[[], int]) -> int:
"""Invoke an SDK call, logging (never raising) on a non-zero/error.
Args:
what: Human-readable label for log messages.
fn: A zero-arg callable returning an SDK int code (0 == ok).
Returns:
The SDK return code, or ``-1`` if the call raised.
"""
try:
code = fn()
except Exception as exc: # pragma: no cover - hardware path
logger.error("%s raised: %s", what, exc)
return -1
if code is not None and code != 0:
logger.warning("%s returned non-zero code=%s", what, code)
return code if code is not None else 0
def _clamp(value: float, lo: float, hi: float) -> float:
"""Clamp ``value`` into the inclusive ``[lo, hi]`` range."""
return lo if value < lo else hi if value > hi else value