mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-24 10:16:09 +00:00
Add Unitree Go2 robot + spatial-memory navigation base
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 <noreply@anthropic.com>
This commit is contained in:
@@ -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
|
||||
Reference in New Issue
Block a user