import copy import importlib.metadata import json import math import os import subprocess import sys import unittest from pathlib import Path from unittest.mock import patch import numpy as np import torch from lerobot.robots.robot import Robot from lerobot.robots.utils import make_robot_from_config from lerobot.utils.errors import DeviceNotConnectedError from lerobot_robot_mujoco import LeKiwiSim, LeKiwiSimConfig from lerobot_robot_mujoco.lekiwi_sim import ARM, BASE, FEATURES, PROFILE, from_sim, to_sim from mujoco_control_bridge import RobotError from mujoco_control_bridge.protocol import values def descriptor(): channels = [] for i, key in enumerate(FEATURES): gripper = key == "arm_gripper.pos" base = key in BASE limit = PROFILE["baseAngularLimit"] if key == "theta.vel" else PROFILE["baseLinearLimit"] channels.append( { "id": key, "unit": "ratio" if gripper else "rad/s" if key == "theta.vel" else "m/s" if base else "rad", "mode": "opening" if gripper else "velocity" if base else "position", "min": 0 if gripper else -limit if base else PROFILE["arm"][i]["min"], "max": 1 if gripper else limit if base else PROFILE["arm"][i]["max"], } ) return { "protocolVersion": 1, "profileId": "lekiwi-v1", "profileVersion": 1, "modelFingerprint": "a" * 64, "frame": "x-forward-y-left-z-up", "capabilities": {"reset": True, "lockstep": False, "cameras": False, "training": False}, "actionChannels": channels, "observationChannels": copy.deepcopy(channels), } class FakeClient: def __init__(self, *_): self.is_connected = False self.sent = [] def describe(self): return descriptor() def connect(self): self.is_connected = True return self.describe() def get_observation(self): return {"values": {**dict.fromkeys(FEATURES, 0.1), "arm_gripper.pos": 0.25}} def send_action(self, value): self.sent.append(value) return {"values": values(value, self.describe()["actionChannels"], True)} def disconnect(self): self.is_connected = False class LeRobotPluginTests(unittest.TestCase): def test_actual_upstream_factory_and_no_hardware_calibration(self): self.assertEqual(importlib.metadata.version("lerobot"), "0.6.1") with patch.object(Path, "mkdir", side_effect=AssertionError("hardware calibration I/O")): robot = make_robot_from_config(LeKiwiSimConfig()) self.assertIsInstance(robot, Robot) self.assertIsInstance(robot, LeKiwiSim) self.assertEqual(tuple(robot.action_features), FEATURES) self.assertEqual(robot.observation_features, robot.action_features) self.assertTrue(robot.is_calibrated) self.assertFalse(robot.is_connected) self.assertEqual(robot.cameras, {}) self.assertFalse(torch.cuda.is_initialized()) with self.assertRaises(DeviceNotConnectedError): robot.get_observation() def test_fresh_process_plugin_discovery(self): result = subprocess.run( [ sys.executable, "-c", """ import sys from lerobot.utils.import_utils import register_third_party_plugins from lerobot.robots.config import RobotConfig from lerobot.robots.utils import make_robot_from_config assert 'lerobot_robot_mujoco' not in sys.modules register_third_party_plugins() config = RobotConfig.get_choice_class('lekiwi_sim')() robot = make_robot_from_config(config) assert robot.name == 'lekiwi_sim' assert len(robot.action_features) == 9 assert not any(m in sys.modules for m in ('serial','zmq','pyrealsense2','scservo_sdk')) """, ], env={**os.environ, "PYTHONPATH": "", "PYTHONNOUSERSITE": "1"}, capture_output=True, text=True, ) self.assertEqual(result.returncode, 0, result.stderr) def test_profile_resource_matches_canonical(self): root = Path(__file__).resolve().parents[3] self.assertEqual(PROFILE, json.loads((root / "robot_profiles/lekiwi-v1.json").read_text())) def test_bidirectional_units_and_upstream_scalar_values(self): for key in FEATURES: for value in (-30.0, 0.0, 30.0): self.assertAlmostEqual(from_sim(key, to_sim(key, value)), value) self.assertAlmostEqual(to_sim(ARM[0], 90), math.pi / 2) self.assertEqual(to_sim("arm_gripper.pos", 50), 0.5) self.assertAlmostEqual(to_sim("theta.vel", 30), math.pi / 6) self.assertAlmostEqual(to_sim(ARM[0], np.float32(10)), math.radians(10)) self.assertAlmostEqual(to_sim(ARM[0], torch.tensor(10.0)), math.radians(10)) for bad in ( True, np.bool_(True), "2", [1], np.array([1]), torch.tensor([1.0]), float("nan"), float("inf"), ): with self.subTest(value=repr(bad)), self.assertRaises(RobotError): to_sim(ARM[0], bad) def test_partial_action_holds_confirmed_arm_and_stops_omitted_base(self): with patch("lerobot_robot_mujoco.lekiwi_sim.SimRobotClient", FakeClient): robot = LeKiwiSim(LeKiwiSimConfig()) robot.connect() first = robot.send_action({ARM[0]: 180, "arm_gripper.pos": 150, "x.vel": 0.1}) self.assertAlmostEqual(first[ARM[0]], math.degrees(PROFILE["arm"][0]["max"])) self.assertEqual(first["arm_gripper.pos"], 100) self.assertEqual(first["x.vel"], 0.1) second = robot.send_action({ARM[1]: 5}) self.assertEqual(second[ARM[0]], first[ARM[0]]) self.assertEqual(second["arm_gripper.pos"], 100) self.assertEqual(second["x.vel"], 0) self.assertEqual(second["theta.vel"], 0) self.assertEqual(robot.get_observation()["arm_gripper.pos"], 25) self.assertNotEqual(robot.get_observation()[ARM[0]], first[ARM[0]]) with self.assertRaises(RobotError): robot.send_action({"unknown": 1}) robot.stop_base() robot.disconnect() robot.disconnect() def test_reject_unsupported_configuration_and_descriptor(self): for kwargs in ( {"use_degrees": False}, {"cameras": {"camera": object()}}, {"calibration_dir": Path("/tmp/never-touch")}, {"endpoint": "http://token@localhost"}, ): with self.assertRaises(ValueError): LeKiwiSimConfig(**kwargs) config = LeKiwiSimConfig(token="secret-not-real") self.assertNotIn("secret-not-real", repr(config)) robot = LeKiwiSim(config) for field in ("profileId", "profileVersion", "unit", "mode"): desc = descriptor() if field == "profileVersion": desc[field] = 2 elif field == "profileId": desc[field] = "other" else: desc["actionChannels"][0][field] = "wrong" with self.assertRaises(RobotError): robot._validate_descriptor(desc) def test_keyboard_helpers_match_upstream_key_semantics(self): robot = LeKiwiSim(LeKiwiSimConfig()) self.assertEqual( robot._from_keyboard_to_base_action(["w", "a", "z"]), {"x.vel": 0.1, "y.vel": 0.1, "theta.vel": 30}, ) self.assertEqual(robot._from_keyboard_to_base_action(["w", "s", "r"])["x.vel"], 0) self.assertEqual(robot._from_keyboard_to_base_action(["d"])["y.vel"], -0.2) if __name__ == "__main__": unittest.main()