From ff94e6385b15148861cdd5506dd7e13a8f813013 Mon Sep 17 00:00:00 2001 From: Pepijn Date: Mon, 20 Jul 2026 17:45:41 +0200 Subject: [PATCH] Add Unitree Go2 robot + spatial-memory navigation base MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit unitree_go2 robot: high-level sport-mode body-velocity control over DDS (unitree_sdk2py) straight from the host — no companion computer, unlike the G1's low-level ZMQ bridge. Actions x.vel/y.vel/theta.vel; observation is planar odometry (rt/sportmodestate) + the built-in front camera via VideoClient. SDK imported only in connect() so configs/features/tests work without it; 20 tests via mocks. unitree_go2 pyproject extra. lerobot.navigation package: BaseController protocol + StubBaseController + SafeBaseController (velocity clamp, occupancy gate, keyframe watchdog, e-stop latch) + RobotBaseController wrapping any Robot on the standard REP-103 mobile-base contract, carrying the world<->body velocity and odometry<->world pose frame math. Robot-agnostic: the SDK lives only in the robot class. 29 tests, SDK/torch-free. First step of consolidating the dyna360 DynaMem navigation stack into lerobot; dyna360 is a source to copy from, not a runtime dependency. Co-Authored-By: Claude Opus 4.8 --- pyproject.toml | 5 + src/lerobot/navigation/README.md | 60 +++ src/lerobot/navigation/__init__.py | 42 ++ src/lerobot/navigation/base_controller.py | 389 ++++++++++++++++++ src/lerobot/robots/unitree_go2/__init__.py | 20 + .../robots/unitree_go2/config_unitree_go2.py | 57 +++ src/lerobot/robots/unitree_go2/unitree_go2.py | 260 ++++++++++++ tests/navigation/test_base_controller.py | 320 ++++++++++++++ tests/robots/test_unitree_go2.py | 202 +++++++++ 9 files changed, 1355 insertions(+) create mode 100644 src/lerobot/navigation/README.md create mode 100644 src/lerobot/navigation/__init__.py create mode 100644 src/lerobot/navigation/base_controller.py create mode 100644 src/lerobot/robots/unitree_go2/__init__.py create mode 100644 src/lerobot/robots/unitree_go2/config_unitree_go2.py create mode 100644 src/lerobot/robots/unitree_go2/unitree_go2.py create mode 100644 tests/navigation/test_base_controller.py create mode 100644 tests/robots/test_unitree_go2.py diff --git a/pyproject.toml b/pyproject.toml index d07080232..72c89bebb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -187,6 +187,11 @@ unitree_g1 = [ "lerobot[matplotlib-dep]", "lerobot[pygame-dep]", ] +# Go2 talks plain DDS from the host — no bridge server, no extra deps beyond +# the SDK itself (cyclonedds-based, hence Linux-only). +unitree_go2 = [ + "unitree_sdk2py>=1.0.1; sys_platform == 'linux'", +] # reachy2-sdk caps grpcio<=1.73.1 and protobuf<=6.32.0; quarantined here so downstream users aren't held back. reachy2-sdk is unlikely to release new versions. reachy2 = [ "reachy2_sdk>=1.0.15,<1.1.0", diff --git a/src/lerobot/navigation/README.md b/src/lerobot/navigation/README.md new file mode 100644 index 000000000..782e35e82 --- /dev/null +++ b/src/lerobot/navigation/README.md @@ -0,0 +1,60 @@ +# `lerobot.navigation` — spatial-memory navigation + +Online spatio-semantic mapping (DynaMem-style), A* planning, obstacle +avoidance and open-vocabulary goto/explore for LeRobot mobile bases. +Ported from the dyna360 research stack; the physical robot layer lives in +`lerobot.robots` (e.g. [`unitree_go2`](../robots/unitree_go2)). + +## Idea + +Drive any LeRobot `Robot` on the standard REP-103 mobile-base contract — +body-velocity actions `x.vel`/`y.vel`/`theta.vel` and planar odometry +`x.pos`/`y.pos`/`theta.pos` — from a spatial memory that is built and +updated online from the robot's camera. With no prompt the base explores +autonomously; given a text prompt it queries the map and navigates to the +matching object, or explores to find it if it isn't there (or has moved). + +## Architecture + +The navigation layer talks to hardware only through LeRobot's own `Robot` +interface, so it is robot-agnostic and carries no SDK dependency. + +``` +BaseController (protocol) world-frame move()/pose() seam +├── StubBaseController kinematic integrator (sim, tests) +├── RobotBaseController wraps any Robot; world<->body + +│ odometry<->world frame math +└── SafeBaseController velocity clamp, occupancy gate, + keyframe watchdog, e-stop latch +``` + +World frame is OpenCV (x right, y down, z forward); the base moves in the +XZ plane. `RobotBaseController.feed_observation(obs)` updates pose from +the observation the navigation loop already fetches (closed-loop +odometry), avoiding an extra camera read; absent odometry it integrates +open-loop so sim matches hardware. + +## Status (branch `feat/unitree-go2`) + +Implemented: +- `base_controller.py` — the controller seam above. SDK/torch-free; + tests in `tests/navigation/test_base_controller.py`. + +Planned (porting from dyna360, everything lands here — dyna360 is a +source to copy from, not a runtime dependency): +- `geometry.py` — LingBot-Map streaming reconstruction runner + (`{points, local_points, conf, camera_poses}`), scale-anchored to + odometry. (Pi3X is not carried over.) +- `voxel_map.py` — 5 cm `SegmentVoxelMap`, no point-cloud retention; + one SigLIP2 feature per segment. +- `occupancy.py` — occupancy grid + A* + frontiers; `value_map.py`. +- `features.py` (SigLIP2), `segmenter.py` (SAM2) — run offboard. +- `agent.py`, `skills.py`, `dog_cli.py` — the interactive explore/query + REPL (the deliverable). + +## Target platform + +Unitree Go2 EDU, no companion computer: the workstation (single RTX 5090) +talks DDS straight to the dog; geometry is monocular LingBot-Map from the +built-in front camera, scale-anchored to sport-mode odometry; the map is +5 cm voxels. See [`robots/unitree_go2`](../robots/unitree_go2). diff --git a/src/lerobot/navigation/__init__.py b/src/lerobot/navigation/__init__.py new file mode 100644 index 000000000..8b1b8c19c --- /dev/null +++ b/src/lerobot/navigation/__init__.py @@ -0,0 +1,42 @@ +#!/usr/bin/env python + +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Spatial-memory navigation for LeRobot mobile bases. + +Online spatio-semantic mapping (DynaMem-style), A* planning, obstacle +avoidance and open-vocabulary goto/explore, driving any LeRobot ``Robot`` +that exposes body-velocity actions and planar odometry. Ported from the +dyna360 research stack; the physical robot layer lives in +``lerobot.robots`` (e.g. ``unitree_go2``). +""" + +from .base_controller import ( + BaseController, + RobotBaseController, + SafeBaseController, + StubBaseController, + odometry_to_world_pose, + world_velocity_to_body, +) + +__all__ = [ + "BaseController", + "RobotBaseController", + "SafeBaseController", + "StubBaseController", + "odometry_to_world_pose", + "world_velocity_to_body", +] diff --git a/src/lerobot/navigation/base_controller.py b/src/lerobot/navigation/base_controller.py new file mode 100644 index 000000000..f270c8cb6 --- /dev/null +++ b/src/lerobot/navigation/base_controller.py @@ -0,0 +1,389 @@ +#!/usr/bin/env python + +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Base controller for spatial-memory navigation. + +The navigation/skills layer commands motion in a single **world frame** +(OpenCV convention: x right, y down, z forward — the base lives in the XZ +plane, y is gravity) and reads back an SE(3) pose. :class:`BaseController` +is that seam. Three implementations: + + - :class:`StubBaseController` — kinematic integrator, no hardware; sim + + unit tests. + - :class:`RobotBaseController` — drives any LeRobot :class:`Robot` whose + action space is body-frame velocities ``x.vel`` (forward, m/s), + ``y.vel`` (left, m/s), ``theta.vel`` (CCW yaw, rad/s) and whose + observation carries planar odometry ``x.pos``/``y.pos``/``theta.pos`` + (REP-103: x forward, y left, yaw CCW). The Unitree Go2 satisfies this + out of the box; so would a LeKiwi base. + - :class:`SafeBaseController` — wraps any of the above with velocity + clamping, an optional occupancy gate, a keyframe watchdog and an + e-stop latch. + +All frame conversions between the world frame and a robot's body/odometry +frame live in :func:`world_velocity_to_body` and +:func:`odometry_to_world_pose`; nothing else needs to know the mapping. +""" + +from __future__ import annotations + +import logging +import math +import time +from abc import abstractmethod +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Protocol, runtime_checkable + +import numpy as np + +if TYPE_CHECKING: + from lerobot.robots import Robot + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------- # +# BaseController protocol +# --------------------------------------------------------------------- # + + +@runtime_checkable +class BaseController(Protocol): + """Mobile-base interface used by the navigation/skills layer. + + Velocities are in **world** frame XZ (m/s); ``yaw_rate`` is rad/s + about the world's −Y axis (turning around the up vector). ``pose`` + is 4×4 SE(3) camera-to-world (OpenCV). + """ + + @abstractmethod + def move(self, vx: float, vz: float, yaw_rate: float = 0.0, dt: float = 0.05) -> None: ... + + @abstractmethod + def stop(self) -> None: ... + + @abstractmethod + def pose(self) -> np.ndarray: ... + + @abstractmethod + def position(self) -> tuple[float, float, float]: ... + + +# --------------------------------------------------------------------- # +# Frame math (pure functions) +# --------------------------------------------------------------------- # + + +def world_velocity_to_body( + vx_world: float, + vz_world: float, + yaw_rate_rad_s: float, + heading_rad: float, +) -> tuple[float, float, float]: + """World-frame velocity → body-frame ``(x.vel, y.vel, theta.vel)``. + + Returns ``(vx_forward, vy_left, vyaw)`` in m/s, m/s, rad/s — the + action a REP-103 base expects. At heading ``h`` the body axes in the + world XZ plane are forward = (sin h, cos h), left = (−cos h, sin h) + (left = up × forward, up = −y). The navigation world's positive yaw + is clockwise about the up vector; a REP-103 base's ``theta.vel`` is + counter-clockwise, hence the sign flip. + """ + s, c = math.sin(heading_rad), math.cos(heading_rad) + vx_fwd = vx_world * s + vz_world * c + vy_left = -vx_world * c + vz_world * s + return vx_fwd, vy_left, -yaw_rate_rad_s + + +def odometry_to_world_pose( + x_fwd: float, + y_left: float, + yaw: float, + origin: tuple[float, float, float], +) -> tuple[np.ndarray, float]: + """Planar odometry ``(x_fwd, y_left, yaw)`` → world pose + heading. + + ``origin`` is the ``(x_fwd, y_left, yaw)`` sample captured when the + controller first saw odometry, so the run starts at identity + regardless of where the robot's odometry origin sits. The result is + the OpenCV world convention, planarized: height Y is 0 and only yaw + survives of the orientation — pitch/roll gait wobble is the camera's + concern, not the base's. + + Odometry frame is REP-103 (x forward, y left, yaw CCW about z-up). + Mapping to OpenCV world: ``x_world = −y_odom``, ``z_world = x_odom``, + ``heading = −yaw``. + """ + ox, oy, oyaw = origin + dx, dy = x_fwd - ox, y_left - oy + c0, s0 = math.cos(-oyaw), math.sin(-oyaw) + x_rel = c0 * dx - s0 * dy + y_rel = s0 * dx + c0 * dy + yaw_rel = yaw - oyaw + + x_world, z_world = -y_rel, x_rel + heading = -yaw_rel + + ch, sh = math.cos(heading), math.sin(heading) + pose = np.eye(4, dtype=np.float64) + pose[0, 0], pose[0, 2] = ch, sh + pose[2, 0], pose[2, 2] = -sh, ch + pose[0, 3], pose[2, 3] = x_world, z_world + return pose, heading + + +def _heading_pose(x: float, z: float, heading: float) -> np.ndarray: + """Build a planar world pose from position + heading.""" + c, s = math.cos(heading), math.sin(heading) + pose = np.eye(4, dtype=np.float64) + pose[0, 0], pose[0, 2] = c, s + pose[2, 0], pose[2, 2] = -s, c + pose[0, 3], pose[2, 3] = x, z + return pose + + +# --------------------------------------------------------------------- # +# Stub controller (kinematic, no hardware) +# --------------------------------------------------------------------- # + + +@dataclass +class StubBaseController: + """Kinematic stub: integrates each ``move()`` into pose exactly. + + No latency, slip or dynamics — for sim and skill-layer unit tests. + """ + + initial_pose: np.ndarray | None = None + max_lin_speed: float = 1.0 + max_yaw_rate: float = 1.0 + + def __post_init__(self) -> None: + self._pose = ( + np.asarray(self.initial_pose, dtype=np.float64).copy() + if self.initial_pose is not None + else np.eye(4, dtype=np.float64) + ) + if self._pose.shape != (4, 4): + raise ValueError(f"initial_pose must be (4, 4); got {self._pose.shape}") + self._heading = 0.0 + self._stopped = False + + def move(self, vx: float, vz: float, yaw_rate: float = 0.0, dt: float = 0.05) -> None: + vx = float(np.clip(vx, -self.max_lin_speed, self.max_lin_speed)) + vz = float(np.clip(vz, -self.max_lin_speed, self.max_lin_speed)) + yaw_rate = float(np.clip(yaw_rate, -self.max_yaw_rate, self.max_yaw_rate)) + if dt <= 0: + return + self._pose[0, 3] += vx * dt + self._pose[2, 3] += vz * dt + if yaw_rate != 0.0: + self._heading += yaw_rate * dt + self._pose = _heading_pose(self._pose[0, 3], self._pose[2, 3], self._heading) + self._stopped = False + + def stop(self) -> None: + self._stopped = True + + def pose(self) -> np.ndarray: + return self._pose.copy() + + def position(self) -> tuple[float, float, float]: + p = self._pose[:3, 3] + return float(p[0]), float(p[1]), float(p[2]) + + @property + def is_stopped(self) -> bool: + return self._stopped + + +# --------------------------------------------------------------------- # +# Robot-backed controller +# --------------------------------------------------------------------- # + + +@dataclass(frozen=True) +class RobotBaseControllerConfig: + """Behaviour knobs for :class:`RobotBaseController`.""" + + max_lin_speed: float = 0.6 + """Hard cap on per-axis world linear velocity (m/s).""" + + max_yaw_rate: float = 1.2 + """Hard cap on yaw rate (rad/s).""" + + pose_from_odometry: bool = True + """Report pose from the robot's odometry (closed-loop). When False, + integrate pose open-loop from commanded velocities.""" + + +class RobotBaseController(BaseController): + """:class:`BaseController` over any LeRobot :class:`Robot`. + + The robot must accept body-velocity actions ``x.vel`` (forward), + ``y.vel`` (left), ``theta.vel`` (CCW yaw) and — for closed-loop pose + — report odometry ``x.pos``/``y.pos``/``theta.pos`` in its + observation. This is the standard REP-103 mobile-base contract, which + ``UnitreeGo2`` implements. + + Pose is refreshed from observations the navigation loop already + fetches: call :meth:`feed_observation` each keyframe rather than + having the controller poll the robot (which would trigger an extra + camera read). Absent any fed observation, pose falls back to + open-loop integration so sim/dry-run behaves like the stub. + """ + + def __init__(self, robot: Robot, cfg: RobotBaseControllerConfig | None = None) -> None: + self.robot = robot + self.cfg = cfg or RobotBaseControllerConfig() + self._pose = np.eye(4, dtype=np.float64) + self._heading = 0.0 + self._stopped = False + self._odom_origin: tuple[float, float, float] | None = None + self._have_odom = False + + # ----- odometry feed -------------------------------------------------- + + def feed_observation(self, obs: dict) -> None: + """Update pose from an observation the nav loop already fetched.""" + if not self.cfg.pose_from_odometry: + return + if not {"x.pos", "y.pos", "theta.pos"} <= obs.keys(): + return + sample = (float(obs["x.pos"]), float(obs["y.pos"]), float(obs["theta.pos"])) + if self._odom_origin is None: + self._odom_origin = sample + self._pose, self._heading = odometry_to_world_pose(*sample, self._odom_origin) + self._have_odom = True + + # ----- BaseController API -------------------------------------------- + + def move(self, vx: float, vz: float, yaw_rate: float = 0.0, dt: float = 0.05) -> None: + vx = float(np.clip(vx, -self.cfg.max_lin_speed, self.cfg.max_lin_speed)) + vz = float(np.clip(vz, -self.cfg.max_lin_speed, self.cfg.max_lin_speed)) + yaw_rate = float(np.clip(yaw_rate, -self.cfg.max_yaw_rate, self.cfg.max_yaw_rate)) + if dt <= 0: + return + + vx_fwd, vy_left, vyaw = world_velocity_to_body(vx, vz, yaw_rate, self._heading) + self.robot.send_action({"x.vel": vx_fwd, "y.vel": vy_left, "theta.vel": vyaw}) + + # Open-loop pose only when we have no odometry to trust. + if not (self.cfg.pose_from_odometry and self._have_odom): + self._pose[0, 3] += vx * dt + self._pose[2, 3] += vz * dt + if yaw_rate != 0.0: + self._heading += yaw_rate * dt + self._pose = _heading_pose(self._pose[0, 3], self._pose[2, 3], self._heading) + self._stopped = False + + def stop(self) -> None: + self._stopped = True + try: + self.robot.send_action({"x.vel": 0.0, "y.vel": 0.0, "theta.vel": 0.0}) + except Exception: + logger.exception("stop(): failed to send zero-velocity action") + + def pose(self) -> np.ndarray: + return self._pose.copy() + + def position(self) -> tuple[float, float, float]: + p = self._pose[:3, 3] + return float(p[0]), float(p[1]), float(p[2]) + + @property + def is_stopped(self) -> bool: + return self._stopped + + +# --------------------------------------------------------------------- # +# Safety wrapper +# --------------------------------------------------------------------- # + + +@dataclass +class SafeBaseController(BaseController): + """Wrap any :class:`BaseController` with safety layers: + + - **velocity clamp** on every ``move()``; + - **occupancy gate**: when ``occupancy_provider`` is set, predict + the next position and refuse (latch e-stop) if it lands in an + obstacle cell. The provider returns an object exposing + ``world_to_cell(x, z) -> (iz, ix)`` and an ``is_obstacle(iz, ix) + -> bool`` predicate; ``None`` means "no map yet, allow"; + - **watchdog**: if no keyframe has been fed in + ``watchdog_timeout_s`` (caller ticks :meth:`feed_watchdog` per + map update), ``move()`` latches stop until :meth:`reset_watchdog`. + """ + + inner: BaseController + max_lin_speed: float = 0.6 + max_yaw_rate: float = 1.2 + occupancy_provider: object = None # callable[[], grid | None] when set + watchdog_timeout_s: float = 2.0 + e_stop_latched: bool = False + _last_keyframe_walltime: float = field(default_factory=time.monotonic, init=False) + + def feed_watchdog(self) -> None: + self._last_keyframe_walltime = time.monotonic() + + def reset_watchdog(self) -> None: + self.e_stop_latched = False + self._last_keyframe_walltime = time.monotonic() + + def latch_estop(self, reason: str = "external") -> None: + logger.warning("SafeBaseController e-stop latched: %s", reason) + self.e_stop_latched = True + self.inner.stop() + + def move(self, vx: float, vz: float, yaw_rate: float = 0.0, dt: float = 0.05) -> None: + if self.e_stop_latched: + return + if (time.monotonic() - self._last_keyframe_walltime) > self.watchdog_timeout_s: + self.latch_estop(f"watchdog: no keyframe in last {self.watchdog_timeout_s:.2f}s") + return + + vx = float(np.clip(vx, -self.max_lin_speed, self.max_lin_speed)) + vz = float(np.clip(vz, -self.max_lin_speed, self.max_lin_speed)) + yaw_rate = float(np.clip(yaw_rate, -self.max_yaw_rate, self.max_yaw_rate)) + + if self.occupancy_provider is not None: + try: + grid = self.occupancy_provider() + except Exception: + logger.exception("occupancy_provider raised; refusing move") + return + if grid is not None and self._would_enter_obstacle(grid, vx, vz, dt): + self.latch_estop("about to enter obstacle cell") + return + + self.inner.move(vx, vz, yaw_rate, dt) + + def stop(self) -> None: + self.inner.stop() + + def pose(self) -> np.ndarray: + return self.inner.pose() + + def position(self) -> tuple[float, float, float]: + return self.inner.position() + + def _would_enter_obstacle(self, grid, vx: float, vz: float, dt: float) -> bool: + pos = self.inner.position() + next_x = pos[0] + vx * dt + next_z = pos[2] + vz * dt + iz, ix = grid.world_to_cell(next_x, next_z) + return bool(grid.is_obstacle(iz, ix)) diff --git a/src/lerobot/robots/unitree_go2/__init__.py b/src/lerobot/robots/unitree_go2/__init__.py new file mode 100644 index 000000000..f4f396845 --- /dev/null +++ b/src/lerobot/robots/unitree_go2/__init__.py @@ -0,0 +1,20 @@ +#!/usr/bin/env python + +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .config_unitree_go2 import UnitreeGo2Config +from .unitree_go2 import UnitreeGo2 + +__all__ = ["UnitreeGo2", "UnitreeGo2Config"] diff --git a/src/lerobot/robots/unitree_go2/config_unitree_go2.py b/src/lerobot/robots/unitree_go2/config_unitree_go2.py new file mode 100644 index 000000000..8624b88ac --- /dev/null +++ b/src/lerobot/robots/unitree_go2/config_unitree_go2.py @@ -0,0 +1,57 @@ +#!/usr/bin/env python + +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass, field + +from lerobot.cameras import CameraConfig + +from ..config import RobotConfig + + +@RobotConfig.register_subclass("unitree_go2") +@dataclass +class UnitreeGo2Config(RobotConfig): + """Configuration for the Unitree Go2 quadruped (EDU). + + The host machine talks DDS directly to the dog over Ethernet/WiFi via + ``unitree_sdk2py`` — no onboard companion computer is required. Actions + are high-level sport-mode body velocities; observations are sport-mode + odometry plus the dog's built-in front camera. + """ + + # Network interface on the host that is wired/bridged to the Go2 + # (the dog lives on 192.168.123.x when connected over Ethernet). + network_interface: str = "eth0" + + # DDS domain id (0 for a stock Go2). + domain_id: int = 0 + + # Safety clamps applied in send_action() before commands reach the dog. + # The Go2 accepts far more (vx up to ~3.7 m/s) — keep indoor-sane defaults. + max_x_vel: float = 1.0 # m/s, body forward + max_y_vel: float = 0.5 # m/s, body left + max_theta_vel: float = 1.5 # rad/s, CCW about z-up + + # Built-in front camera, served through the SDK VideoClient. + use_front_camera: bool = True + front_camera_width: int = 1280 + front_camera_height: int = 720 + + # Send BalanceStand once on connect so the dog is ready to walk. + stand_on_connect: bool = True + + # Additional external cameras (standard LeRobot camera configs). + cameras: dict[str, CameraConfig] = field(default_factory=dict) diff --git a/src/lerobot/robots/unitree_go2/unitree_go2.py b/src/lerobot/robots/unitree_go2/unitree_go2.py new file mode 100644 index 000000000..9a68d0e20 --- /dev/null +++ b/src/lerobot/robots/unitree_go2/unitree_go2.py @@ -0,0 +1,260 @@ +#!/usr/bin/env python + +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unitree Go2 quadruped (EDU) — high-level sport-mode integration. + +Unlike :class:`~lerobot.robots.unitree_g1.UnitreeG1` (low-level joint +control at 250 Hz through an on-robot ZMQ bridge), the Go2 is driven with +sport-mode **body velocity commands** at tens of Hz, which work fine over +plain DDS from any Linux host on the dog's network — no bridge server, no +onboard companion computer. + +Setup: + 1. Connect the host to the Go2 via Ethernet (dog is on 192.168.123.x) + or put both on the same WiFi network. + 2. ``pip install unitree_sdk2py`` (Linux only — rides on cyclonedds). + 3. Find your interface name (``ip link``), then e.g.:: + + lerobot-teleoperate \ + --robot.type=unitree_go2 \ + --robot.network_interface=enp2s0 \ + --teleop.type=gamepad + +Actions are body-frame velocities ``x.vel`` (forward, m/s), ``y.vel`` +(left, m/s), ``theta.vel`` (CCW yaw, rad/s) — the exact arguments of the +SDK's ``SportClient.Move``. Observations are planar sport-mode odometry +(``*.pos`` pose + ``*.vel`` body velocities) and the built-in front +camera, plus any extra configured cameras. +""" + +from __future__ import annotations + +import logging +import threading +from functools import cached_property +from typing import Any + +import cv2 +import numpy as np + +from lerobot.cameras import make_cameras_from_configs +from lerobot.types import RobotAction, RobotObservation +from lerobot.utils.import_utils import require_package + +from ..robot import Robot +from .config_unitree_go2 import UnitreeGo2Config + +logger = logging.getLogger(__name__) + +# DDS topic names follow Unitree SDK naming conventions +SPORT_MODE_STATE_TOPIC = "rt/sportmodestate" + + +class UnitreeGo2(Robot): + """LeRobot interface to a Unitree Go2 over unitree_sdk2py sport mode.""" + + config_class = UnitreeGo2Config + name = "unitree_go2" + + def __init__(self, config: UnitreeGo2Config): + super().__init__(config) + self.config = config + + self._cameras = make_cameras_from_configs(config.cameras) + + # SDK handles — populated in connect(); the SDK import lives there + # too so that configs, features and tests work on SDK-less hosts. + self._sport = None + self._video = None + self._state_subscriber = None + + self._state_lock = threading.Lock() + self._latest_state = None # last SportModeState_ message + self._connected = False + + # ------------------------------------------------------------------ # + # Features + # ------------------------------------------------------------------ # + + @cached_property + def _odom_ft(self) -> dict[str, type]: + return { + "x.pos": float, + "y.pos": float, + "theta.pos": float, + "x.vel": float, + "y.vel": float, + "theta.vel": float, + } + + @property + def _cameras_ft(self) -> dict[str, tuple]: + ft: dict[str, tuple] = {} + if self.config.use_front_camera: + ft["front"] = (self.config.front_camera_height, self.config.front_camera_width, 3) + for name, cam in self._cameras.items(): + ft[name] = (cam.height, cam.width, 3) + return ft + + @property + def observation_features(self) -> dict: + return {**self._odom_ft, **self._cameras_ft} + + @property + def action_features(self) -> dict: + return {"x.vel": float, "y.vel": float, "theta.vel": float} + + # ------------------------------------------------------------------ # + # Lifecycle + # ------------------------------------------------------------------ # + + @property + def is_connected(self) -> bool: + return self._connected + + def connect(self, calibrate: bool = True) -> None: + if self._connected: + return + require_package("unitree-sdk2py", extra="unitree_go2", import_name="unitree_sdk2py") + + from unitree_sdk2py.core.channel import ChannelFactoryInitialize, ChannelSubscriber + from unitree_sdk2py.go2.sport.sport_client import SportClient + from unitree_sdk2py.idl.unitree_go.msg.dds_ import SportModeState_ + + ChannelFactoryInitialize(self.config.domain_id, self.config.network_interface) + + sport = SportClient() + sport.SetTimeout(5.0) + sport.Init() + self._sport = sport + + subscriber = ChannelSubscriber(SPORT_MODE_STATE_TOPIC, SportModeState_) + subscriber.Init(self._on_sport_state, 10) + self._state_subscriber = subscriber + + if self.config.use_front_camera: + from unitree_sdk2py.go2.video.video_client import VideoClient + + video = VideoClient() + video.SetTimeout(3.0) + video.Init() + self._video = video + + for cam in self._cameras.values(): + cam.connect() + + if self.config.stand_on_connect: + self._sport.BalanceStand() + + self._connected = True + self.configure() + logger.info( + "%s connected (iface=%s, domain=%d)", + self, + self.config.network_interface, + self.config.domain_id, + ) + + def disconnect(self) -> None: + if self._sport is not None: + try: + self._sport.StopMove() + except Exception: + logger.exception("StopMove on disconnect failed") + for cam in self._cameras.values(): + try: + cam.disconnect() + except Exception: + logger.exception("camera disconnect failed") + self._sport = None + self._video = None + self._state_subscriber = None + self._connected = False + + # Sport mode needs no calibration. + @property + def is_calibrated(self) -> bool: + return True + + def calibrate(self) -> None: + pass + + def configure(self) -> None: + pass + + # ------------------------------------------------------------------ # + # I/O + # ------------------------------------------------------------------ # + + def get_observation(self) -> RobotObservation: + if not self._connected: + raise ConnectionError(f"{self} is not connected.") + + obs: dict[str, Any] = dict.fromkeys(self._odom_ft, 0.0) + with self._state_lock: + state = self._latest_state + if state is not None: + obs["x.pos"] = float(state.position[0]) + obs["y.pos"] = float(state.position[1]) + obs["theta.pos"] = float(state.imu_state.rpy[2]) + obs["x.vel"] = float(state.velocity[0]) + obs["y.vel"] = float(state.velocity[1]) + obs["theta.vel"] = float(state.yaw_speed) + + if self.config.use_front_camera: + obs["front"] = self._read_front_camera() + + for name, cam in self._cameras.items(): + obs[name] = cam.async_read() + + return obs + + def send_action(self, action: RobotAction) -> RobotAction: + if not self._connected: + raise ConnectionError(f"{self} is not connected.") + + vx = float(np.clip(action.get("x.vel", 0.0), -self.config.max_x_vel, self.config.max_x_vel)) + vy = float(np.clip(action.get("y.vel", 0.0), -self.config.max_y_vel, self.config.max_y_vel)) + vyaw = float( + np.clip(action.get("theta.vel", 0.0), -self.config.max_theta_vel, self.config.max_theta_vel) + ) + + self._sport.Move(vx, vy, vyaw) + return {"x.vel": vx, "y.vel": vy, "theta.vel": vyaw} + + # ------------------------------------------------------------------ # + # Internals + # ------------------------------------------------------------------ # + + def _on_sport_state(self, msg) -> None: + with self._state_lock: + self._latest_state = msg + + def _read_front_camera(self) -> np.ndarray: + """Fetch one frame from the built-in front camera (RGB, HxWx3).""" + h, w = self.config.front_camera_height, self.config.front_camera_width + code, data = self._video.GetImageSample() + if code != 0 or data is None: + logger.warning("front camera GetImageSample failed (code=%s)", code) + return np.zeros((h, w, 3), dtype=np.uint8) + frame = cv2.imdecode(np.frombuffer(bytes(data), dtype=np.uint8), cv2.IMREAD_COLOR) + if frame is None: + logger.warning("front camera frame failed to decode") + return np.zeros((h, w, 3), dtype=np.uint8) + frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + if frame.shape[:2] != (h, w): + frame = cv2.resize(frame, (w, h), interpolation=cv2.INTER_AREA) + return frame diff --git a/tests/navigation/test_base_controller.py b/tests/navigation/test_base_controller.py new file mode 100644 index 000000000..5c9bb0376 --- /dev/null +++ b/tests/navigation/test_base_controller.py @@ -0,0 +1,320 @@ +#!/usr/bin/env python + +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the navigation base controller. + +Hardware-free and SDK-free: the frame math is pure, the stub is +kinematic, and the robot-backed controller is exercised through a fake +Robot that records actions and serves canned odometry. +""" + +from __future__ import annotations + +import math +import time + +import numpy as np +import pytest + +from lerobot.navigation.base_controller import ( + BaseController, + RobotBaseController, + RobotBaseControllerConfig, + SafeBaseController, + StubBaseController, + odometry_to_world_pose, + world_velocity_to_body, +) + +# ----- world_velocity_to_body --------------------------------------------- + + +def test_forward_maps_to_body_x(): + """heading=0, world +z (forward) → (vx>0, 0, 0).""" + vx_f, vy_l, vyaw = world_velocity_to_body(0.0, 0.3, 0.0, heading_rad=0.0) + assert vx_f == pytest.approx(0.3) + assert vy_l == pytest.approx(0.0) + assert vyaw == pytest.approx(0.0) + + +def test_world_right_maps_to_negative_left(): + """heading=0, world +x is the robot's RIGHT → negative y.vel.""" + vx_f, vy_l, _ = world_velocity_to_body(0.3, 0.0, 0.0, heading_rad=0.0) + assert vx_f == pytest.approx(0.0) + assert vy_l == pytest.approx(-0.3) + + +def test_world_x_is_forward_after_quarter_turn(): + vx_f, vy_l, _ = world_velocity_to_body(0.3, 0.0, 0.0, heading_rad=math.pi / 2) + assert vx_f == pytest.approx(0.3) + assert vy_l == pytest.approx(0.0, abs=1e-9) + + +def test_yaw_rate_sign_flips(): + _, _, vyaw = world_velocity_to_body(0.0, 0.0, 0.5, heading_rad=0.0) + assert vyaw == pytest.approx(-0.5) + + +def test_velocity_magnitude_preserved_under_rotation(): + vx_f, vy_l, _ = world_velocity_to_body(0.3, 0.4, 0.0, heading_rad=1.234) + assert math.hypot(vx_f, vy_l) == pytest.approx(0.5) + + +# ----- odometry_to_world_pose --------------------------------------------- + + +def test_odometry_at_origin_is_identity(): + pose, heading = odometry_to_world_pose(1.0, 2.0, 0.3, origin=(1.0, 2.0, 0.3)) + np.testing.assert_allclose(pose, np.eye(4), atol=1e-12) + assert heading == pytest.approx(0.0) + + +def test_odometry_forward_maps_to_world_z(): + pose, heading = odometry_to_world_pose(1.0, 0.0, 0.0, origin=(0.0, 0.0, 0.0)) + assert pose[0, 3] == pytest.approx(0.0) + assert pose[2, 3] == pytest.approx(1.0) + assert heading == pytest.approx(0.0) + + +def test_odometry_left_maps_to_world_negative_x(): + pose, _ = odometry_to_world_pose(0.0, 1.0, 0.0, origin=(0.0, 0.0, 0.0)) + assert pose[0, 3] == pytest.approx(-1.0) + assert pose[2, 3] == pytest.approx(0.0) + + +def test_odometry_yaw_sign_flip(): + pose, heading = odometry_to_world_pose(0.0, 0.0, 0.5, origin=(0.0, 0.0, 0.0)) + assert heading == pytest.approx(-0.5) + fwd = pose[:3, 2] + np.testing.assert_allclose(fwd, [math.sin(-0.5), 0.0, math.cos(-0.5)], atol=1e-12) + + +def test_odometry_origin_yaw_is_derotated(): + """Motion along the boot-time heading is always world +z, whatever + direction the robot faced when odometry started.""" + origin = (0.0, 0.0, math.pi / 2) + pose, heading = odometry_to_world_pose(0.0, 1.0, math.pi / 2, origin=origin) + assert pose[0, 3] == pytest.approx(0.0, abs=1e-12) + assert pose[2, 3] == pytest.approx(1.0) + assert heading == pytest.approx(0.0) + + +# ----- StubBaseController -------------------------------------------------- + + +def test_stub_is_basecontroller(): + assert isinstance(StubBaseController(), BaseController) + + +def test_stub_integrates_forward(): + c = StubBaseController() + c.move(0.0, 0.2, dt=1.0) + assert c.position()[2] == pytest.approx(0.2) + + +def test_stub_clamps_velocity(): + c = StubBaseController(max_lin_speed=0.1) + c.move(5.0, 0.0, dt=1.0) + assert c.position()[0] == pytest.approx(0.1) + + +# ----- RobotBaseController ------------------------------------------------- + + +class FakeRobot: + """Minimal Robot stand-in: records actions, serves canned odometry.""" + + def __init__(self) -> None: + self.actions: list[dict] = [] + self.obs: dict = {} + + def send_action(self, action: dict) -> dict: + self.actions.append(action) + return action + + def get_observation(self) -> dict: + return self.obs + + +def _robot_controller(**cfg_kwargs) -> tuple[RobotBaseController, FakeRobot]: + robot = FakeRobot() + cfg = RobotBaseControllerConfig(**cfg_kwargs) + return RobotBaseController(robot, cfg), robot + + +def _odom(x=0.0, y=0.0, yaw=0.0) -> dict: + return {"x.pos": x, "y.pos": y, "theta.pos": yaw} + + +def test_robot_controller_is_basecontroller(): + ctl, _ = _robot_controller() + assert isinstance(ctl, BaseController) + + +def test_forward_command_reaches_send_action(): + ctl, robot = _robot_controller() + ctl.feed_observation(_odom()) # heading 0 + ctl.move(vx=0.0, vz=0.3, dt=0.05) + assert robot.actions[-1] == { + "x.vel": pytest.approx(0.3), + "y.vel": pytest.approx(0.0), + "theta.vel": pytest.approx(0.0), + } + + +def test_command_uses_odometry_heading(): + """After the robot turns to heading +π/2, a world +x command comes out + as pure body-forward. First sample fixes the origin.""" + ctl, robot = _robot_controller() + ctl.feed_observation(_odom()) # origin, heading 0 + ctl.feed_observation(_odom(yaw=-math.pi / 2)) # turned; heading +π/2 + ctl.move(vx=0.3, vz=0.0, dt=0.05) + assert robot.actions[-1]["x.vel"] == pytest.approx(0.3) + assert robot.actions[-1]["y.vel"] == pytest.approx(0.0, abs=1e-9) + + +def test_command_is_clamped_before_send(): + ctl, robot = _robot_controller(max_lin_speed=0.1) + ctl.feed_observation(_odom()) + ctl.move(vx=0.0, vz=9.0, dt=0.05) + assert robot.actions[-1]["x.vel"] == pytest.approx(0.1) + + +def test_pose_comes_from_odometry_not_integration(): + ctl, _ = _robot_controller() + ctl.feed_observation(_odom()) + ctl.move(0.0, 0.3, dt=1.0) # would integrate 0.3 m open-loop + ctl.feed_observation(_odom(x=0.05)) # ...but odometry says 5 cm forward + assert ctl.position()[2] == pytest.approx(0.05) + + +def test_origin_is_first_odometry_sample(): + ctl, _ = _robot_controller() + ctl.feed_observation(_odom(x=3.0, y=-1.0, yaw=0.7)) + np.testing.assert_allclose(ctl.pose(), np.eye(4), atol=1e-12) + + +def test_open_loop_fallback_without_odometry(): + """No odometry fed → integrate open-loop like the stub.""" + ctl, _ = _robot_controller() + ctl.move(0.0, 0.2, dt=1.0) + assert ctl.position()[2] == pytest.approx(0.2) + + +def test_stop_sends_zero_velocity(): + ctl, robot = _robot_controller() + ctl.stop() + assert robot.actions[-1] == {"x.vel": 0.0, "y.vel": 0.0, "theta.vel": 0.0} + assert ctl.is_stopped + + +def test_robot_controller_matches_stub_open_loop(): + """Open-loop pose integration matches StubBaseController for the same + command sequence — sim runs must transfer to the real base.""" + ctl, _ = _robot_controller(max_lin_speed=1.0) + stub = StubBaseController() + for vx, vz, yaw in [(0.2, 0.0, 0.0), (0.0, 0.3, 0.5), (0.1, 0.1, -0.2)]: + ctl.move(vx, vz, yaw, dt=0.5) + stub.move(vx, vz, yaw, dt=0.5) + np.testing.assert_allclose(ctl.pose(), stub.pose(), atol=1e-9) + + +# ----- SafeBaseController -------------------------------------------------- + + +class FakeGrid: + """Occupancy stand-in with a single obstacle cell.""" + + def __init__(self, obstacle_cell=(5, 6), cell_size=0.1, origin_x=-0.5, origin_z=-0.5): + self.obstacle_cell = obstacle_cell + self.cell_size = cell_size + self.origin_x = origin_x + self.origin_z = origin_z + + def world_to_cell(self, x: float, z: float) -> tuple[int, int]: + ix = int((x - self.origin_x) / self.cell_size) + iz = int((z - self.origin_z) / self.cell_size) + return iz, ix + + def is_obstacle(self, iz: int, ix: int) -> bool: + return (iz, ix) == self.obstacle_cell + + +def test_safe_passes_normal_moves(): + inner = StubBaseController() + safe = SafeBaseController(inner=inner) + safe.feed_watchdog() + safe.move(0.0, 0.1, dt=1.0) + assert inner.position()[2] == pytest.approx(0.1) + + +def test_safe_clamps_speed(): + inner = StubBaseController(max_lin_speed=100.0) + safe = SafeBaseController(inner=inner, max_lin_speed=0.5) + safe.feed_watchdog() + safe.move(10.0, 0.0, dt=1.0) + assert inner.position()[0] == pytest.approx(0.5) + + +def test_safe_watchdog_latches_on_stale_keyframes(): + inner = StubBaseController() + safe = SafeBaseController(inner=inner, watchdog_timeout_s=0.05) + safe.feed_watchdog() + time.sleep(0.1) + safe.move(0.0, 0.1, dt=1.0) + assert safe.e_stop_latched + safe.move(0.0, 10.0, dt=1.0) # refused + assert inner.position()[2] == pytest.approx(0.0, abs=1e-6) + + +def test_safe_reset_watchdog_re_enables_motion(): + inner = StubBaseController() + safe = SafeBaseController(inner=inner, watchdog_timeout_s=0.05) + safe.feed_watchdog() + time.sleep(0.1) + safe.move(0.0, 0.1) + assert safe.e_stop_latched + safe.reset_watchdog() + safe.move(0.0, 0.1, dt=1.0) + assert inner.position()[2] == pytest.approx(0.1) + + +def test_safe_refuses_move_into_obstacle(): + inner = StubBaseController() + grid = FakeGrid(obstacle_cell=(5, 6)) + safe = SafeBaseController(inner=inner, occupancy_provider=lambda: grid) + safe.feed_watchdog() + # +0.15 m in x from origin lands mid-column ix=6 (origin_x=-0.5, cell=0.1). + safe.move(vx=0.15, vz=0.0, dt=1.0) + assert safe.e_stop_latched + assert inner.position()[0] == pytest.approx(0.0, abs=1e-6) + + +def test_safe_allows_move_into_free_cell(): + inner = StubBaseController() + grid = FakeGrid(obstacle_cell=(99, 99)) + safe = SafeBaseController(inner=inner, occupancy_provider=lambda: grid) + safe.feed_watchdog() + safe.move(vx=0.1, vz=0.0, dt=1.0) + assert inner.position()[0] == pytest.approx(0.1) + + +def test_safe_allows_when_no_map_yet(): + inner = StubBaseController() + safe = SafeBaseController(inner=inner, occupancy_provider=lambda: None) + safe.feed_watchdog() + safe.move(vx=0.1, vz=0.0, dt=1.0) + assert inner.position()[0] == pytest.approx(0.1) diff --git a/tests/robots/test_unitree_go2.py b/tests/robots/test_unitree_go2.py new file mode 100644 index 000000000..49994d68c --- /dev/null +++ b/tests/robots/test_unitree_go2.py @@ -0,0 +1,202 @@ +#!/usr/bin/env python + +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the Unitree Go2 robot. + +The SDK is only imported inside ``UnitreeGo2.connect()``, so everything +here runs without unitree_sdk2py installed: the sport client, state +subscriber and video client are replaced with mocks. +""" + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import cv2 +import numpy as np +import pytest + +from lerobot.robots.unitree_go2 import UnitreeGo2, UnitreeGo2Config + +# --------------------------------------------------------------------------- +# Config (no SDK needed) +# --------------------------------------------------------------------------- + + +class TestUnitreeGo2Config: + def test_registered_type_name(self): + assert UnitreeGo2Config().type == "unitree_go2" + + def test_default_config(self): + cfg = UnitreeGo2Config() + assert cfg.domain_id == 0 + assert cfg.use_front_camera is True + assert cfg.stand_on_connect is True + assert cfg.cameras == {} + + def test_safety_clamps_are_positive(self): + cfg = UnitreeGo2Config() + assert cfg.max_x_vel > 0 + assert cfg.max_y_vel > 0 + assert cfg.max_theta_vel > 0 + + +# --------------------------------------------------------------------------- +# Features (no SDK needed) +# --------------------------------------------------------------------------- + + +def _make_robot(**cfg_kwargs) -> UnitreeGo2: + cfg = UnitreeGo2Config(id="test_go2", **cfg_kwargs) + return UnitreeGo2(cfg) + + +class TestFeatures: + def test_action_features(self): + robot = _make_robot() + assert robot.action_features == {"x.vel": float, "y.vel": float, "theta.vel": float} + + def test_observation_features_with_front_camera(self): + robot = _make_robot() + ft = robot.observation_features + assert ft["front"] == (720, 1280, 3) + for key in ("x.pos", "y.pos", "theta.pos", "x.vel", "y.vel", "theta.vel"): + assert ft[key] is float + + def test_observation_features_without_front_camera(self): + robot = _make_robot(use_front_camera=False) + assert "front" not in robot.observation_features + + def test_features_available_before_connect(self): + robot = _make_robot() + assert not robot.is_connected + assert robot.observation_features + assert robot.action_features + + def test_is_calibrated_always_true(self): + assert _make_robot().is_calibrated is True + + +# --------------------------------------------------------------------------- +# I/O with mocked SDK handles +# --------------------------------------------------------------------------- + + +def _connected_robot(**cfg_kwargs) -> UnitreeGo2: + """A robot with mocked SDK handles, as if connect() had run.""" + robot = _make_robot(**cfg_kwargs) + robot._sport = MagicMock() + robot._video = MagicMock() + robot._connected = True + return robot + + +def _fake_state(x=0.0, y=0.0, yaw=0.0, vx=0.0, vy=0.0, yaw_speed=0.0): + return SimpleNamespace( + position=[x, y, 0.0], + velocity=[vx, vy, 0.0], + yaw_speed=yaw_speed, + imu_state=SimpleNamespace(rpy=[0.0, 0.0, yaw]), + ) + + +class TestSendAction: + def test_action_reaches_sport_move(self): + robot = _connected_robot() + sent = robot.send_action({"x.vel": 0.3, "y.vel": -0.1, "theta.vel": 0.5}) + robot._sport.Move.assert_called_once_with(0.3, -0.1, 0.5) + assert sent == {"x.vel": 0.3, "y.vel": -0.1, "theta.vel": 0.5} + + def test_action_is_clamped(self): + robot = _connected_robot(max_x_vel=0.5, max_y_vel=0.2, max_theta_vel=1.0) + sent = robot.send_action({"x.vel": 9.0, "y.vel": -9.0, "theta.vel": -9.0}) + robot._sport.Move.assert_called_once_with(0.5, -0.2, -1.0) + assert sent == {"x.vel": 0.5, "y.vel": -0.2, "theta.vel": -1.0} + + def test_missing_keys_default_to_zero(self): + robot = _connected_robot() + sent = robot.send_action({}) + robot._sport.Move.assert_called_once_with(0.0, 0.0, 0.0) + assert sent == {"x.vel": 0.0, "y.vel": 0.0, "theta.vel": 0.0} + + def test_raises_when_not_connected(self): + robot = _make_robot() + with pytest.raises(ConnectionError): + robot.send_action({"x.vel": 0.1}) + + +class TestGetObservation: + def test_odometry_fields(self): + robot = _connected_robot(use_front_camera=False) + robot._latest_state = _fake_state(x=1.0, y=2.0, yaw=0.3, vx=0.1, vy=-0.05, yaw_speed=0.2) + obs = robot.get_observation() + assert obs["x.pos"] == pytest.approx(1.0) + assert obs["y.pos"] == pytest.approx(2.0) + assert obs["theta.pos"] == pytest.approx(0.3) + assert obs["x.vel"] == pytest.approx(0.1) + assert obs["y.vel"] == pytest.approx(-0.05) + assert obs["theta.vel"] == pytest.approx(0.2) + + def test_odometry_zero_before_first_state(self): + robot = _connected_robot(use_front_camera=False) + obs = robot.get_observation() + assert all(obs[k] == 0.0 for k in robot._odom_ft) + + def test_front_camera_decodes_to_configured_shape(self): + robot = _connected_robot(front_camera_width=64, front_camera_height=48) + raw = np.full((48, 64, 3), 128, dtype=np.uint8) + ok, jpeg = cv2.imencode(".jpg", raw) + assert ok + robot._video.GetImageSample.return_value = (0, jpeg.tobytes()) + obs = robot.get_observation() + assert obs["front"].shape == (48, 64, 3) + assert obs["front"].dtype == np.uint8 + + def test_front_camera_resizes_native_frames(self): + robot = _connected_robot(front_camera_width=64, front_camera_height=48) + native = np.zeros((720, 1280, 3), dtype=np.uint8) + ok, jpeg = cv2.imencode(".jpg", native) + assert ok + robot._video.GetImageSample.return_value = (0, jpeg.tobytes()) + assert robot.get_observation()["front"].shape == (48, 64, 3) + + def test_front_camera_failure_returns_black_frame(self): + robot = _connected_robot(front_camera_width=64, front_camera_height=48) + robot._video.GetImageSample.return_value = (1, None) + frame = robot.get_observation()["front"] + assert frame.shape == (48, 64, 3) + assert frame.sum() == 0 + + def test_observation_matches_features(self): + robot = _connected_robot(front_camera_width=64, front_camera_height=48) + raw = np.zeros((48, 64, 3), dtype=np.uint8) + _, jpeg = cv2.imencode(".jpg", raw) + robot._video.GetImageSample.return_value = (0, jpeg.tobytes()) + obs = robot.get_observation() + assert set(obs.keys()) == set(robot.observation_features.keys()) + + def test_raises_when_not_connected(self): + robot = _make_robot() + with pytest.raises(ConnectionError): + robot.get_observation() + + +class TestDisconnect: + def test_disconnect_stops_motion(self): + robot = _connected_robot() + sport = robot._sport + robot.disconnect() + sport.StopMove.assert_called_once() + assert not robot.is_connected