unitree-g1-mujoco / tests /test_models.py
feat: configurable yaw jitter of world objects (WORLD_RANDOMIZE_YAW)
d0e4913 11 days ago
32.3 kB
"""Real MuJoCo regression tests; DDS transport is isolated with a test double.
Run: MUJOCO_GL=egl python -m unittest discover -s tests -v
The separate smoke_live.py exercises the actual Gym/DDS/ZMQ dependencies.
"""

import contextlib
import importlib
import io
import os
from pathlib import Path
import sys
from types import ModuleType, SimpleNamespace
import unittest
import xml.etree.ElementTree as ET
from unittest.mock import Mock, patch

os.environ.setdefault("MUJOCO_GL", "egl")
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

import mujoco
import numpy as np
import yaml

from sim.image_publish_utils import ImagePublishProcess
from sim.model_config import DEFAULT_CAMERAS, select_end_effector

# The physics tests do not require a DDS installation or open DDS sockets.
channel = ModuleType("unitree_sdk2py.core.channel")
channel.ChannelFactoryInitialize = Mock()
bridge_module = ModuleType("sim.unitree_sdk2py_bridge")
bridge_module.ElasticBand = Mock()
bridge_module.UnitreeSdk2Bridge = Mock()
with patch.dict(sys.modules, {
    "unitree_sdk2py.core.channel": channel,
    "sim.unitree_sdk2py_bridge": bridge_module,
}):
    DefaultEnv = importlib.import_module("sim.base_sim").DefaultEnv

from sim.unitree_sdk2py_bridge import JOINT_SLOTS, MODE_MACHINE  # noqa: E402


def motor(q=0.0, kp=0.0, kd=0.0, tau=0.0):
    return SimpleNamespace(q=q, dq=0.0, kp=kp, kd=kd, tau=tau)


def load_config():
    return yaml.safe_load((ROOT / "config.yaml").read_text())


def make_physics(
    mode="dex1", render=False, body="29dof", head_mount="fixed", head_sensor="d435i", world="", randomize=0.0, world_grasp=None,
    randomize_yaw=None,
):
    config = select_end_effector(load_config(), mode, body, head_mount, head_sensor, world, world_grasp)
    config["ENABLE_ELASTIC_BAND"] = False
    config["WORLD_RANDOMIZE"] = randomize
    if randomize_yaw is not None:
        config["WORLD_RANDOMIZE_YAW"] = randomize_yaw
    cameras = {name: {"height": 480, "width": 640} for name in config["CAMERAS"]}
    env = DefaultEnv(config, camera_configs=cameras, offscreen=render)
    count = config["NUM_HAND_MOTORS"]
    num_body_motor = config["NUM_MOTORS"]
    env.set_unitree_bridge(SimpleNamespace(
        num_body_motor=num_body_motor, num_hand_motor=count, use_sensor=False, joystick=None,
        joint_slots=JOINT_SLOTS[body],
        low_cmd=SimpleNamespace(motor_cmd=[motor() for _ in range(35)]),
        left_hand_cmd=SimpleNamespace(motor_cmd=[motor() for _ in range(count)]),
        right_hand_cmd=SimpleNamespace(motor_cmd=[motor() for _ in range(count)]),
        PublishLowState=Mock(),
    ))
    return env


class ModelTests(unittest.TestCase):
    def test_all_variants_keep_body_interface_and_mapped_efforts(self):
        expected = mujoco.MjModel.from_xml_path(str(ROOT / "assets/g1_29dof.xml"))
        body_names = [expected.joint(i).name for i in range(1, expected.njnt)
                      if "_hand_" not in expected.joint(i).name]
        for mode, count in (("dex1", 2), ("dex3", 7), ("rubber_hand", 0), ("none", 0)):
            with self.subTest(mode=mode):
                env = make_physics(mode)
                model, data = env.mj_model, env.mj_data
                self.assertEqual(model.nu, 29 + 2 * count)
                self.assertEqual([model.joint(i).name for i in env.body_joint_index], body_names)
                self.assertEqual(len(env.left_hand_index), count)
                self.assertEqual(len(env.right_hand_index), count)
                body_efforts = np.arange(1, 30) / 10
                for command, effort in zip(env.unitree_bridge.low_cmd.motor_cmd, body_efforts):
                    command.tau = effort
                if count:
                    env.unitree_bridge.left_hand_cmd.motor_cmd[0].tau = 0.35
                    env.unitree_bridge.right_hand_cmd.motor_cmd[-1].tau = -0.45
                env.sim_step()
                np.testing.assert_allclose(data.ctrl[env.body_actuator_index], body_efforts)
                obs = env.prepare_obs()
                self.assertEqual(obs["body_q"].shape, (29,))
                if count:
                    self.assertAlmostEqual(data.ctrl[env.left_hand_actuator_index[0]], 0.35)
                    self.assertAlmostEqual(data.ctrl[env.right_hand_actuator_index[-1]], -0.45)
                    self.assertEqual(obs["left_hand_q"].shape, (count,))
                else:
                    self.assertNotIn("left_hand_q", obs)
                for i, name in enumerate(body_names):
                    joint = model.joint(name).id
                    self.assertEqual(obs["body_q"][i], data.qpos[model.jnt_qposadr[joint]])
                    actuator = np.flatnonzero(model.actuator_trnid[:, 0] == joint)[0]
                    self.assertEqual(obs["body_tau_est"][i], data.actuator_force[actuator])

    def test_grippers_close_with_external_pd_and_reset_open(self):
        env = make_physics()
        model, data = env.mj_model, env.mj_data
        model.opt.gravity[:] = 0
        data.qpos[2] = 2.0  # Clear the floor while testing the fingers.
        mujoco.mj_forward(model, data)
        for command in env.unitree_bridge.low_cmd.motor_cmd:
            command.kp, command.kd = 100.0, 5.0
        for side in ("left", "right"):
            for command in getattr(env.unitree_bridge, side + "_hand_cmd").motor_cmd:
                command.q, command.kp, command.kd = -0.023, 1000.0, 25.0
        finger_actuators = np.r_[env.left_hand_actuator_index, env.right_hand_actuator_index]
        self.assertTrue(np.all(model.actuator_biastype[finger_actuators] == mujoco.mjtBias.mjBIAS_NONE))
        max_force = 0.0
        for _ in range(250):
            env.sim_step()
            max_force = max(max_force, float(np.max(np.abs(data.actuator_force[finger_actuators]))))
        positions = np.r_[data.qpos[env.left_hand_qpos_index], data.qpos[env.right_hand_qpos_index]]
        self.assertTrue(np.all(positions < -0.021), positions)
        self.assertLessEqual(max_force, 20.00001)
        self.assertTrue(np.isfinite(data.qpos).all())
        self.assertFalse(np.any(data.warning.number))
        env.reset()
        np.testing.assert_allclose(data.qpos[env.left_hand_qpos_index], 0.0245)
        np.testing.assert_allclose(data.qpos[env.right_hand_qpos_index], 0.0245)

    def test_wrist_cameras_move_only_with_their_own_wrist(self):
        for mode in ("dex1", "dex3"):
            with self.subTest(mode=mode):
                env = make_physics(mode)
                model, data = env.mj_model, env.mj_data
                before = {name: data.cam_xmat[model.camera(name).id].copy() for name in DEFAULT_CAMERAS}
                joint = model.joint("left_wrist_roll_joint").id
                data.qpos[model.jnt_qposadr[joint]] = 0.5
                mujoco.mj_forward(model, data)
                self.assertFalse(np.allclose(before["left_wrist_cam"], data.cam_xmat[model.camera("left_wrist_cam").id]))
                for name in ("head_camera", "right_wrist_cam"):
                    np.testing.assert_allclose(before[name], data.cam_xmat[model.camera(name).id])

    def test_all_three_cameras_render_and_enter_publisher_buffers(self):
        for mode in ("dex1", "dex3"):
            with self.subTest(mode=mode):
                env = make_physics(mode, render=True)
                publisher = ImagePublishProcess(env.camera_configs, image_dt=1 / 30)
                env.image_publish_process = publisher
                try:
                    frames = env.update_render_caches()
                    self.assertEqual(set(frames), {name + "_image" for name in DEFAULT_CAMERAS})
                    self.assertTrue(publisher.data_ready_event.is_set())
                    for name in DEFAULT_CAMERAS:
                        frame = frames[name + "_image"]
                        self.assertEqual(frame.shape, (480, 640, 3))
                        self.assertEqual(frame.dtype, np.uint8)
                        self.assertGreater(float(frame.std()), 2.0)
                        received = np.ndarray(frame.shape, dtype=np.uint8,
                                              buffer=publisher.shared_memory_blocks[name].buf).copy()
                        np.testing.assert_array_equal(received, frame)
                finally:
                    publisher.stop()
                    for renderer in env.renderers.values():
                        renderer.close()

    def test_dex1_fingers_are_slide_joints(self):
        model = make_physics("dex1").mj_model
        for side in ("left", "right"):
            for finger in (1, 2):
                joint = model.joint(f"{side}_dex1_finger_joint_{finger}")
                self.assertEqual(joint.type, mujoco.mjtJoint.mjJNT_SLIDE)


class Body23DofTests(unittest.TestCase):
    def test_model_loads_with_23_actuators_in_slot_order(self):
        env = make_physics(mode="rubber_hand", body="23dof")
        model = env.mj_model
        self.assertEqual(model.nu, 23)
        self.assertEqual(len(env.body_joint_index), 23)
        # Actuator order must match JOINT_SLOTS["23dof"]'s relative order: legs, waist_yaw,
        # then each arm without wrist_pitch/wrist_yaw -- the 29-slot SDK layout minus
        # lerobot's G1_23_INVALID_SDK_SLOTS (unitree_g1/g1_utils.py).
        expected_suffixes = [
            "hip_pitch", "hip_roll", "hip_yaw", "knee", "ankle_pitch", "ankle_roll",
        ]
        actuator_names = [model.actuator(i).name for i in range(model.nu)]
        for side, offset in (("left", 0), ("right", 6)):
            for i, suffix in enumerate(expected_suffixes):
                self.assertEqual(actuator_names[offset + i], f"{side}_{suffix}")
        self.assertEqual(actuator_names[12], "waist_yaw")
        arm_suffixes = ["shoulder_pitch", "shoulder_roll", "shoulder_yaw", "elbow", "wrist_roll"]
        for side, offset in (("left", 13), ("right", 18)):
            for i, suffix in enumerate(arm_suffixes):
                self.assertEqual(actuator_names[offset + i], f"{side}_{suffix}")

    def test_23dof_commands_route_through_sim_step_via_joint_slots(self):
        # Body torques are computed from unitree_bridge.low_cmd, which is keyed by DDS slot
        # (sparse for 23dof), while data.ctrl is keyed by MuJoCo actuator order (dense 0..22).
        # Write torques only at the real DDS slots -- as the live bridge's LowCmdHandler would
        # receive them -- and check compute_body_torques (via sim_step) routes them correctly.
        env = make_physics(mode="rubber_hand", body="23dof")
        for slot, effort in zip(JOINT_SLOTS["23dof"], np.arange(1, 24) / 10, strict=True):
            env.unitree_bridge.low_cmd.motor_cmd[slot].tau = effort
        env.sim_step()
        np.testing.assert_allclose(
            env.mj_data.ctrl[env.body_actuator_index], np.arange(1, 24) / 10
        )

    def test_23dof_drives_dex_hands_over_dds(self):
        for mode, count in (("dex1", 2), ("dex3", 7)):
            with self.subTest(mode=mode):
                env = make_physics(mode=mode, body="23dof")
                self.assertEqual(len(env.body_joint_index), 23)
                self.assertEqual(len(env.left_hand_index), count)
                self.assertEqual(len(env.right_hand_index), count)
                env.unitree_bridge.left_hand_cmd.motor_cmd[0].tau = 0.35
                env.sim_step()
                self.assertAlmostEqual(env.mj_data.ctrl[env.left_hand_actuator_index[0]], 0.35)

    def test_mode_machine_is_4_for_23dof_and_2_for_29dof(self):
        self.assertEqual(MODE_MACHINE["23dof"], 4)
        self.assertEqual(MODE_MACHINE["29dof"], 2)

    def test_bridge_publishes_zero_at_absent_slots_and_sets_mode_machine(self):
        # Exercise the real UnitreeSdk2Bridge.PublishLowState mapping (not the SimpleNamespace
        # test double `make_physics` uses elsewhere), against a mocked low_state DDS message
        # (the real SDK, unlike the ChannelFactory used by other tests, is a light import).
        bridge_mod = importlib.import_module("sim.unitree_sdk2py_bridge")
        low_state = Mock()
        low_state.motor_state = [SimpleNamespace(q=None, dq=None, ddq=None, tau_est=None) for _ in range(35)]
        low_state.imu_state = SimpleNamespace(quaternion=np.zeros(4), gyroscope=np.zeros(3), accelerometer=np.zeros(3))
        bridge = Mock()
        bridge.joint_slots = JOINT_SLOTS["23dof"]
        bridge.mode_machine = MODE_MACHINE["23dof"]
        bridge.low_state = low_state
        bridge.torso_imu_state = SimpleNamespace(quaternion=np.zeros(4), gyroscope=np.zeros(3))
        bridge.num_body_motor = 23
        bridge.num_hand_motor = 0
        bridge.use_sensor = False
        bridge.have_frame_sensor_ = False
        bridge.joystick = None

        obs = {
            "body_q": np.arange(23, dtype=np.float64),
            "body_dq": np.arange(23, dtype=np.float64) * 0.1,
            "body_ddq": np.zeros(23),
            "body_tau_est": np.zeros(23),
            "floating_base_pose": np.array([0, 0, 0, 1, 0, 0, 0], dtype=np.float64),
            "floating_base_vel": np.zeros(6),
            "floating_base_acc": np.zeros(6),
            "secondary_imu_quat": np.array([1, 0, 0, 0], dtype=np.float64),
            "secondary_imu_vel": np.zeros(6),
            "time": 0.0,
        }
        bridge_mod.UnitreeSdk2Bridge.PublishLowState(bridge, obs)

        absent_slots = set(range(29)) - set(JOINT_SLOTS["23dof"])
        self.assertEqual(absent_slots, {13, 14, 20, 21, 27, 28})
        for slot in absent_slots:
            self.assertEqual(low_state.motor_state[slot].q, None)
        for i, slot in enumerate(JOINT_SLOTS["23dof"]):
            self.assertEqual(low_state.motor_state[slot].q, float(i))
            self.assertEqual(low_state.motor_state[slot].dq, i * 0.1)
        self.assertEqual(low_state.mode_machine, 4)


AH_SCENE = select_end_effector(load_config(), "amazing_hand", "23dof", "pan_tilt", "d455")["ROBOT_SCENE"]


class AmazingHandD455Tests(unittest.TestCase):
    """The 23dof + AmazingHand + D455 pan/tilt head variant (end_effectors/amazing_hand.xml, heads/pan_tilt.xml, sensors/d455.xml)."""

    HAND_JOINT_NAMES = tuple(
        [f"left_hand_motor{i}_joint" for i in range(11, 19)] + [f"right_hand_motor{i}_joint" for i in range(1, 9)]
    )

    def test_model_loads_with_41_actuators_in_body_head_hand_order(self):
        model = mujoco.MjModel.from_xml_path(AH_SCENE)
        self.assertEqual(model.nu, 41)
        body_names = [
            f"{side}_{suffix}"
            for side, offset in (("left", 0), ("right", 6))
            for suffix in ["hip_pitch", "hip_roll", "hip_yaw", "knee", "ankle_pitch", "ankle_roll"]
        ]
        arm_names = [
            f"{side}_{suffix}"
            for side in ("left", "right")
            for suffix in ["shoulder_pitch", "shoulder_roll", "shoulder_yaw", "elbow", "wrist_roll"]
        ]
        expected = (
            body_names[:12] + ["waist_yaw"] + arm_names[:5] + arm_names[5:]
            + ["head_yaw_joint", "head_pitch_joint"]
            + list(self.HAND_JOINT_NAMES)
        )
        actuator_names = [model.actuator(i).name for i in range(model.nu)]
        self.assertEqual(actuator_names, expected)

    def test_joint_ranges(self):
        model = mujoco.MjModel.from_xml_path(AH_SCENE)
        np.testing.assert_allclose(model.joint("head_yaw_joint").range, [-0.7, 0.7])
        np.testing.assert_allclose(model.joint("head_pitch_joint").range, [-1.5708, 0.8])
        for name in self.HAND_JOINT_NAMES:
            np.testing.assert_allclose(
                model.joint(name).range, [-1.5707963267948966, 1.5707963267948966]
            )

    def test_body_actuator_list_is_23_in_slot_order_excluding_head_and_hands(self):
        env = make_physics(mode="amazing_hand", body="23dof", head_mount="pan_tilt", head_sensor="d455")
        model = env.mj_model
        self.assertEqual(model.nu, 41)
        self.assertEqual(len(env.body_joint_index), 23)
        expected_suffixes = ["hip_pitch", "hip_roll", "hip_yaw", "knee", "ankle_pitch", "ankle_roll"]
        body_names = [model.joint(i).name for i in env.body_joint_index]
        for side, offset in (("left", 0), ("right", 6)):
            for i, suffix in enumerate(expected_suffixes):
                self.assertEqual(body_names[offset + i], f"{side}_{suffix}_joint")
        self.assertEqual(body_names[12], "waist_yaw_joint")
        arm_suffixes = ["shoulder_pitch", "shoulder_roll", "shoulder_yaw", "elbow", "wrist_roll"]
        for side, offset in (("left", 13), ("right", 18)):
            for i, suffix in enumerate(arm_suffixes):
                self.assertEqual(body_names[offset + i], f"{side}_{suffix}_joint")
        # Head (head_yaw_joint/head_pitch_joint) and hand actuators never reach DDS: neither the body
        # scan nor the DDS actuator index picks them up, only the 23 body joints/actuators.
        self.assertEqual(len(env.left_hand_index), 0)
        self.assertEqual(len(env.right_hand_index), 0)
        self.assertEqual(list(env.dds_actuator_index), list(range(23)))

    def test_head_camera_on_the_head_sensor_renders_nonuniform_image(self):
        env = make_physics(mode="amazing_hand", body="23dof", render=True, head_mount="pan_tilt", head_sensor="d455")
        model, data = env.mj_model, env.mj_data
        cam_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_CAMERA, "head_camera")
        self.assertGreaterEqual(cam_id, 0)
        self.assertEqual(model.body(model.cam_bodyid[cam_id]).name, "head_sensor")
        self.assertEqual(model.body(model.body_parentid[model.cam_bodyid[cam_id]]).name, "head_sensor_mount")
        try:
            frames = env.update_render_caches()
            frame = frames["head_camera_image"]
            self.assertEqual(frame.shape, (480, 640, 3))
            self.assertGreater(float(frame.std()), 2.0)
        finally:
            for renderer in env.renderers.values():
                renderer.close()

    def test_select_end_effector_takes_head_and_hand_independently(self):
        config = select_end_effector(load_config(), "amazing_hand", "23dof", "fixed", "d455")
        self.assertEqual(config["NUM_MOTORS"], 23)
        self.assertEqual((config["HEAD_MOUNT"], config["HEAD_SENSOR"]), ("fixed", "d455"))
        from_yaml = select_end_effector(load_config())
        self.assertEqual((from_yaml["HEAD_MOUNT"], from_yaml["HEAD_SENSOR"]), ("pan_tilt", "d455"))
        defaults = select_end_effector({})
        self.assertEqual((defaults["END_EFFECTOR"], defaults["HEAD_MOUNT"], defaults["HEAD_SENSOR"]), ("rubber_hand", "fixed", "d435i"))
        with self.assertRaises(ValueError):
            select_end_effector(load_config(), head_mount="d455_pan_tilt")
        for old_alias in ("dummy", "grippers", "hands"):
            with self.assertRaises(ValueError):
                select_end_effector(load_config(), old_alias, "23dof")
        with self.assertRaises(ValueError):
            select_end_effector(load_config(), "gloves", "23dof")


class CompositionTests(unittest.TestCase):
    """Every end effector on either body, with or without the pan/tilt head (sim/mjcf/compose.py)."""

    HANDS = {"rubber_hand": 0, "none": 0, "dex1": 2, "dex3": 7, "amazing_hand": 0}

    def test_every_combination_loads_with_its_dds_hands_and_cameras(self):
        for body, num_body in (("29dof", 29), ("23dof", 23)):
            for mode, count in self.HANDS.items():
                for head in ("fixed", "pan_tilt"):
                    with self.subTest(body=body, mode=mode, head=head):
                        env = make_physics(mode=mode, body=body, head_mount=head)
                        model = env.mj_model
                        self.assertEqual(len(env.body_joint_index), num_body)
                        self.assertEqual(len(env.left_hand_index), count)
                        self.assertEqual(len(env.right_hand_index), count)
                        head_joint = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "head_yaw_joint")
                        self.assertEqual(head_joint >= 0, head == "pan_tilt")
                        ah_joint = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "left_hand_motor11_joint")
                        self.assertEqual(ah_joint >= 0, mode == "amazing_hand")
                        rubber = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_MESH, "left_rubber_hand")
                        self.assertEqual(rubber >= 0, mode == "rubber_hand")
                        env.sim_step()
                        self.assertTrue(np.isfinite(env.mj_data.qpos).all())

    def test_hands_mount_at_the_same_place_on_both_bodies(self):
        # The 23dof mount sits where the 29dof wrist-yaw frame is with wrist pitch/yaw at zero.
        offsets = {}
        for body, wrist in (("29dof", "left_wrist_roll_link"), ("23dof", "left_wrist_roll_link")):
            model = mujoco.MjModel.from_xml_path(select_end_effector(load_config(), "dex3", body, "fixed")["ROBOT_SCENE"])
            data = mujoco.MjData(model)
            mujoco.mj_forward(model, data)
            rotation = data.xmat[model.body(wrist).id].reshape(3, 3)
            offsets[body] = rotation.T @ (data.xpos[model.body("left_ee_mount").id] - data.xpos[model.body(wrist).id])
        np.testing.assert_allclose(offsets["29dof"], offsets["23dof"], atol=1e-9)
        np.testing.assert_allclose(offsets["23dof"], [0.084, 0, 0], atol=1e-9)

    def test_each_hand_adds_the_same_mass_to_either_body(self):
        def mass(body, mode):
            model = mujoco.MjModel.from_xml_path(select_end_effector(load_config(), mode, body, "fixed")["ROBOT_SCENE"])
            return model.body_subtreemass[1]

        for mode in ("rubber_hand", "dex1", "dex3", "amazing_hand"):
            with self.subTest(mode=mode):
                # to the gram: the 29dof rubber hand body carries a 0.1 g placeholder inertial
                self.assertAlmostEqual(
                    mass("29dof", mode) - mass("29dof", "none"), mass("23dof", mode) - mass("23dof", "none"), places=3
                )
        self.assertGreater(mass("23dof", "none"), 0)

    def test_29dof_models_keep_the_unitree_reference_masses(self):
        # Total masses of Unitree's full models these parts were extracted from (this repo @ 1801640):
        # g1_29dof_rubber_hand.xml, the Dex1-1 gripper model and g1_29dof_dex3.xml (palm lumped in).
        for mode, reference in (("rubber_hand", 35.1126), ("dex1", 35.5031), ("dex3", 36.1661)):
            with self.subTest(mode=mode):
                model = mujoco.MjModel.from_xml_path(select_end_effector(load_config(), mode, "29dof", "fixed")["ROBOT_SCENE"])
                self.assertAlmostEqual(model.body_subtreemass[1], reference, places=3)

    def test_every_head_mount_and_sensor_follow_the_convention(self):
        from sim.mjcf.compose import HEAD_MOUNTS, HEAD_SENSORS

        for mount in HEAD_MOUNTS:
            with self.subTest(mount=mount):
                files = [f"g1_{body}.xml" for body in ("29dof", "23dof")] if mount == "fixed" else [f"heads/{mount}.xml"]
                for name in files:
                    root = ET.parse(ROOT / "assets" / name).getroot()
                    self.assertIsNotNone(root.find(".//body[@name='head_sensor_mount']"))
                    self.assertIsNone(root.find(".//camera[@name='head_camera']"))
        for sensor in HEAD_SENSORS:
            with self.subTest(sensor=sensor):
                root = ET.parse(ROOT / f"assets/sensors/{sensor}.xml").getroot()
                self.assertIsNotNone(root.find("worldbody/body[@name='head_sensor']"))
                self.assertIsNotNone(root.find(".//camera[@name='head_camera']"))

    def test_head_sensors_carry_their_intrinsics(self):
        # d435i: datasheet, centred. d455: the lab unit's factory calibration (Fx Fy PPX PPY at color 640x480).
        for sensor, fx, fy, ppx, ppy in (("d435i", 623.01, 623.01, 320, 240), ("d455", 385.66, 385.24, 327.14, 245.97)):
            with self.subTest(sensor=sensor):
                env = make_physics(mode="rubber_hand", body="29dof", head_sensor=sensor)
                model = env.mj_model
                camera = model.camera("head_camera").id
                np.testing.assert_array_equal(model.cam_resolution[camera], [640, 480])
                pitch = model.cam_sensorsize[camera] / model.cam_resolution[camera]
                np.testing.assert_allclose(model.cam_intrinsic[camera][:2] / pitch, [fx, fy], rtol=1e-6)
                # MuJoCo's principal offset has the opposite sign to the image offset of the principal point.
                np.testing.assert_allclose(model.cam_intrinsic[camera][2:] / pitch, [320 - ppx, 240 - ppy], atol=1e-6)

    def test_any_sensor_on_any_mount_looks_forward(self):
        for mount in ("fixed", "pan_tilt"):
            for sensor in ("d435i", "d455"):
                with self.subTest(mount=mount, sensor=sensor):
                    env = make_physics(mode="rubber_hand", body="29dof", head_mount=mount, head_sensor=sensor)
                    model, data = env.mj_model, env.mj_data
                    camera = model.camera("head_camera").id
                    self.assertEqual(model.body(model.cam_bodyid[camera]).name, "head_sensor")
                    view = -data.cam_xmat[camera].reshape(3, 3)[:, 2]  # camera looks along its -z
                    self.assertGreater(view[0], 0.5)  # forward (+x), from any mount

    def test_every_parts_file_mounts_both_sides(self):
        for path in sorted((ROOT / "assets/end_effectors").glob("*.xml")):
            with self.subTest(parts=path.name):
                root = ET.parse(path).getroot()
                for side in ("left", "right"):
                    self.assertIsNotNone(root.find(f"worldbody/body[@name='{side}_ee_mount']"))
                    self.assertIsNotNone(root.find(f"worldbody/body[@name='{side}_ee_mount']/site[@name='{side}_tcp']"))

    def test_head_camera_rides_the_pan_tilt_head_on_29dof(self):
        env = make_physics(mode="rubber_hand", body="29dof", head_mount="pan_tilt", head_sensor="d455")
        model = env.mj_model
        self.assertEqual(model.body(model.cam_bodyid[model.camera("head_camera").id]).name, "head_sensor")


class HubCacheLayoutTests(unittest.TestCase):
    """LeRobot loads this repo from the Hugging Face cache, where every file is a symlink into blobs/."""

    def test_composes_through_symlinked_files(self):
        import importlib.util
        import tempfile

        with tempfile.TemporaryDirectory() as tmp:
            snapshot, blobs = Path(tmp) / "snapshots" / "rev", Path(tmp) / "blobs"
            blobs.mkdir(parents=True)
            for i, path in enumerate(p for p in ROOT.rglob("*") if p.is_file() and ".git" not in p.parts):
                blob = blobs / f"{i:05d}"  # blob names carry no extension, like the Hub cache's hashes
                blob.symlink_to(path.resolve())
                link = snapshot / path.relative_to(ROOT)
                link.parent.mkdir(parents=True, exist_ok=True)
                link.symlink_to(os.path.relpath(blob, link.parent))
            spec = importlib.util.spec_from_file_location("hub_compose", snapshot / "sim/mjcf/compose.py")
            compose = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(compose)
            self.assertEqual(compose.ASSETS, snapshot / "assets")
            scene = compose.build_scene("23dof", "amazing_hand", "pan_tilt", "d455", True)
            model = mujoco.MjModel.from_xml_path(str(scene))
            self.assertGreater(model.nmesh, 0)


class ToolsTests(unittest.TestCase):
    """The generators in tools/ still reproduce the committed parts files."""

    def test_generated_parts_files_are_up_to_date(self):
        import tempfile

        sys.path.insert(0, str(ROOT / "tools"))
        import build_amazing_hand
        import build_dex1_model

        for tool, name in ((build_dex1_model, "dex1.xml"), (build_amazing_hand, "amazing_hand.xml")):
            with self.subTest(parts=name), tempfile.TemporaryDirectory() as out:
                with contextlib.redirect_stdout(io.StringIO()):
                    tool.build(out_dir=out)
                self.assertEqual((Path(out) / name).read_text(), (ROOT / "assets/end_effectors" / name).read_text())

    def test_soft_pad_variant_has_a_limited_slide_per_distal_body(self):
        import tempfile

        sys.path.insert(0, str(ROOT / "tools"))
        import build_amazing_hand

        with tempfile.TemporaryDirectory() as out:
            with contextlib.redirect_stdout(io.StringIO()):
                build_amazing_hand.build(out_dir=out, pads=True)
            text = (Path(out) / "amazing_hand.xml").read_text()
        self.assertEqual(text.count("_pad_joint"), 8)
        self.assertNotIn("_pad_joint", (ROOT / "assets/end_effectors/amazing_hand.xml").read_text())
        parts = ET.fromstring(text.split("-->", 1)[1])
        pads = [j for j in parts.iter("joint") if j.get("name", "").endswith("_pad_joint")]
        self.assertTrue(all(j.get("type") == "slide" and j.get("limited") == "true" for j in pads))
        axes = {j.get("name").split("_hand_")[0] + ("_thumb" if j.get("name").endswith("_4_pad_joint") else ""): j.get("axis") for j in pads}
        self.assertEqual(axes["right"], axes["left"])
        self.assertEqual(axes["right_thumb"], "0 -0.97 -0.23")
        self.assertEqual(axes["left_thumb"], "0 -0.97 0.23")

if __name__ == "__main__":
    unittest.main()


class HubRepoIdTests(unittest.TestCase):
    def test_repo_id_from_cache_path(self):
        from pathlib import Path

        import env

        cache = Path("/x/hub/models--k-valentin--unitree-g1-mujoco/snapshots/abc123")
        self.assertEqual(env._hub_repo_id_from_cache(cache), ("k-valentin/unitree-g1-mujoco", "abc123"))
        self.assertEqual(env._hub_repo_id_from_cache(Path("/some/checkout")), (env._UPSTREAM_REPO_ID, None))

    def test_partial_xr_snapshot_is_hydrated(self):
        import tempfile

        import env

        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            self.assertTrue(env._needs_hydration(root))
            for name in ("sim/mjcf/compose.py", "assets/scene.xml", "assets/g1_29dof.xml"):  # the XR teleop's allow_patterns
                (root / name).parent.mkdir(parents=True, exist_ok=True)
                (root / name).touch()
            self.assertTrue(env._needs_hydration(root))
            for name in env._SNAPSHOT_SENTINELS:
                (root / name).parent.mkdir(parents=True, exist_ok=True)
                (root / name).touch()
            self.assertFalse(env._needs_hydration(root))
        self.assertFalse(env._needs_hydration(ROOT))

    def test_env_option_parsing(self):
        import env

        self.assertIs(env._parse_env_option("0"), False)
        self.assertIs(env._parse_env_option("true"), True)
        self.assertEqual(env._parse_env_option("5556"), 5556)
        self.assertEqual(env._parse_env_option("head_camera, left_wrist_cam"), ["head_camera", "left_wrist_cam"])
        self.assertEqual(env._parse_env_option("23dof"), "23dof")


class JoystickMappingTests(unittest.TestCase):
    class _FakeJoystick:
        def __init__(self, buttons=(), hat=(0, 0), numhats=0):
            self.buttons, self.hat, self.numhats = set(buttons), hat, numhats

        def get_button(self, i):
            return 1 if i in self.buttons else 0

        def get_numhats(self):
            return self.numhats

        def get_hat(self, _i):
            return self.hat

    def _bridge(self, **kw):
        from sim.unitree_sdk2py_bridge import UnitreeSdk2Bridge

        bridge = UnitreeSdk2Bridge.__new__(UnitreeSdk2Bridge)
        bridge.joystick = self._FakeJoystick(**kw)
        bridge.dpad_buttons = None
        return bridge

    def test_dpad_from_buttons(self):
        bridge = self._bridge(buttons={11, 14})
        bridge.dpad_buttons = {"up": 11, "down": 12, "left": 13, "right": 14}
        self.assertEqual(bridge._dpad_xy(), (1, 1))

    def test_dpad_from_hat_and_without_hat(self):
        self.assertEqual(self._bridge(hat=(-1, 0), numhats=1)._dpad_xy(), (-1, 0))
        self.assertEqual(self._bridge(numhats=0)._dpad_xy(), (0, 0))