Mujoco_WASM/integrations/lerobot/tests/test_plugin.py

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()