Implement dual-arm teleoperation with RM75 QP controller and MuJoCo backend

ik_qp/tests/test_mujoco_stage2.py

@@ -0,0 +1,157 @@
+import os
+from pathlib import Path
+
+import mujoco
+import numpy as np
+import pytest
+
+from rm75_ik import DualArmAssembly, pose_errors
+from rm75_ik.mujoco_backend import CONTROLLED_HOME_Q_RAD, DualArmMuJoCo
+from rm75_ik.mujoco_model import build_normalized_dual_mjcf
+from rm75_ik.mujoco_trajectories import (
+    cartesian_demo_targets,
+    joint_demo_trajectory,
+    se3_target_trajectory,
+    solve_pose_trajectory,
+)
+from rm75_ik import RM75IkSolver, RM75Kinematics, teleop_joint_limits
+
+
+def test_normalized_model_has_standard_dual_arm_structure():
+    xml, assets = build_normalized_dual_mjcf()
+    model = mujoco.MjModel.from_xml_string(xml, assets)
+
+    assert model.nq == model.nv == model.njnt == 14
+    assert model.nu == 0
+    assert model.nmesh == 27
+    assert len([key for key in assets if key.endswith(".obj")]) == 19
+    assert [model.joint(index).name for index in range(14)] == [
+        f"{arm}_joint_{joint}"
+        for arm in ("left", "right")
+        for joint in range(1, 8)
+    ]
+
+
+@pytest.mark.parametrize("arm", ["left", "right"])
+def test_mujoco_flange_matches_stage1_dual_assembly(arm):
+    scene = DualArmMuJoCo(controlled_arm=arm)
+    assembly = DualArmAssembly.from_source_urdf()
+    q = np.deg2rad([20, -10, 30, 40, -20, 15, 80])
+
+    scene.set_arm_configuration(arm, q)
+    errors = pose_errors(scene.get_flange_pose(arm), assembly.forward(arm, q))
+
+    assert errors[0] < 1e-9
+    assert errors[1] < 1e-9
+
+
+def test_invalid_configuration_does_not_change_state():
+    scene = DualArmMuJoCo()
+    before = scene.get_arm_configuration("left")
+
+    with pytest.raises(ValueError):
+        scene.set_arm_configuration("left", np.full(7, np.nan))
+
+    np.testing.assert_array_equal(scene.get_arm_configuration("left"), before)
+
+
+def test_playback_moves_only_controlled_arm():
+    scene = DualArmMuJoCo(controlled_arm="left")
+    inactive_before = scene.get_arm_configuration("right")
+    trajectory = joint_demo_trajectory(CONTROLLED_HOME_Q_RAD, 20)
+
+    result = scene.play_trajectory(trajectory)
+
+    assert result.samples == 20
+    np.testing.assert_array_equal(
+        scene.get_arm_configuration("right"), inactive_before
+    )
+    np.testing.assert_allclose(
+        scene.get_arm_configuration("left"), trajectory[-1], atol=0.0
+    )
+
+
+def test_se3_target_trajectory_preserves_endpoints():
+    kinematics = RM75Kinematics(limits=teleop_joint_limits())
+    start = kinematics.forward(CONTROLLED_HOME_Q_RAD)
+    target = start.copy()
+    target.translation = target.translation + np.array([0.04, 0.01, 0.0])
+
+    trajectory = se3_target_trajectory(start, target, 25)
+
+    assert pose_errors(trajectory[0], start) == pytest.approx((0.0, 0.0), abs=1e-12)
+    assert pose_errors(trajectory[-1], target) == pytest.approx((0.0, 0.0), abs=1e-12)
+
+
+def test_manual_drag_dynamics_moves_only_controlled_arm():
+    scene = DualArmMuJoCo(controlled_arm="left")
+    controlled_before = scene.get_arm_configuration("left")
+    inactive_before = scene.get_arm_configuration("right")
+    scene.configure_manual_drag()
+    body_id = mujoco.mj_name2id(
+        scene.model, mujoco.mjtObj.mjOBJ_BODY, "left_link_7"
+    )
+    scene.data.xfrc_applied[body_id, :3] = [0.0, 10.0, 0.0]
+
+    for _ in range(50):
+        scene.step_manual_drag()
+
+    assert np.max(np.abs(scene.get_arm_configuration("left") - controlled_before)) > 1e-5
+    np.testing.assert_array_equal(
+        scene.get_arm_configuration("right"), inactive_before
+    )
+
+
+@pytest.mark.parametrize("kind", ["line", "arc", "orientation", "combined"])
+def test_demo_cartesian_trajectories_are_solvable(kind):
+    kinematics = RM75Kinematics(limits=teleop_joint_limits())
+    solver = RM75IkSolver(kinematics)
+    targets = cartesian_demo_targets(
+        kind, kinematics.forward(CONTROLLED_HOME_Q_RAD), 30
+    )
+
+    solutions = solve_pose_trajectory(solver, targets, CONTROLLED_HOME_Q_RAD)
+
+    assert solutions.shape == (30, 7)
+    assert np.all(np.isfinite(solutions))
+
+
+@pytest.mark.skipif(
+    os.environ.get("MUJOCO_GL") != "egl",
+    reason="EGL rendering test requires MUJOCO_GL=egl",
+)
+def test_offscreen_render_contains_both_arms():
+    scene = DualArmMuJoCo()
+    rgb = scene.render(640, 360)
+    segmentation = scene.render(640, 360, segmentation=True)
+    pairs = np.unique(segmentation.reshape(-1, 2), axis=0)
+    names = {
+        mujoco.mj_id2name(scene.model, mujoco.mjtObj.mjOBJ_GEOM, int(object_id))
+        for object_id, object_type in pairs
+        if object_type == mujoco.mjtObj.mjOBJ_GEOM and object_id >= 0
+    }
+
+    assert rgb.std() > 5.0
+    assert any(name.startswith("left_") for name in names if name)
+    assert any(name.startswith("right_") for name in names if name)
+
+
+@pytest.fixture(scope="module")
+def realman_sdk_root():
+    value = os.environ.get("REALMAN_SDK_ROOT")
+    if not value:
+        pytest.skip("REALMAN_SDK_ROOT is not set")
+    return Path(value)
+
+
+def test_quick_stage2_validation_passes(realman_sdk_root):
+    from rm75_ik.realman_reference import RealManFkReference
+    from rm75_ik.stage2_validation import Stage2Settings, Stage2Validator
+
+    validator = Stage2Validator(
+        RealManFkReference(realman_sdk_root), settings=Stage2Settings.quick()
+    )
+    summary = validator.run()
+
+    assert summary["passed"] is True
+    assert summary["failure_count"] == 0

ik_qp/tests/test_stage3_control.py

@@ -0,0 +1,195 @@
+import os
+from dataclasses import replace
+from pathlib import Path
+from types import SimpleNamespace
+
+import numpy as np
+import pinocchio as pin
+import pytest
+
+from rm75_ik import pose_errors
+from rm75_ik.mujoco_robot import MujocoRobot
+from rm75_ik.teleop_config import load_dual_arm_profiles
+from rm75_ik.teleop_control import (
+    ControllerSample,
+    DualArmQpTeleopController,
+    RelativePoseMapper,
+    SafetyState,
+)
+
+
+ROOT = Path(__file__).resolve().parents[2]
+
+
+@pytest.fixture(scope="module")
+def profiles():
+    return load_dual_arm_profiles(
+        ROOT / "xr_rm_bringup/config/dual_arm_rm75.yaml",
+        ROOT / "xr_rm_bringup/config/peripherals_rm75.yaml",
+    )
+
+
+def sample(arm, grip, position=(0.0, 0.0, 0.0), quaternion=(0, 0, 0, 1)):
+    return ControllerSample(
+        hand=arm,
+        grip=grip,
+        trigger=0.0,
+        position_m=np.asarray(position, dtype=float),
+        quaternion_xyzw=np.asarray(quaternion, dtype=float),
+    )
+
+
+def test_profiles_use_expected_tools_and_mapping(profiles):
+    assert profiles["left"].active_tool_name == "minisci"
+    assert profiles["right"].active_tool_name == "omnipic"
+    np.testing.assert_array_equal(
+        profiles["left"].xr_to_robot,
+        [[0, -1, 0], [0, 0, 1], [-1, 0, 0]],
+    )
+    np.testing.assert_array_equal(
+        profiles["right"].xr_to_robot,
+        [[0, 1, 0], [0, 0, 1], [1, 0, 0]],
+    )
+
+
+@pytest.mark.parametrize(
+    ("arm", "expected_delta"),
+    [("left", [-0.01, 0.0, 0.0]), ("right", [0.01, 0.0, 0.0])],
+)
+def test_relative_position_mapping_matches_ros_configuration(
+    profiles, arm, expected_delta
+):
+    profile = replace(
+        profiles[arm],
+        scale=1.0,
+        deadband_m=0.0,
+        target_filter_alpha=1.0,
+        target_filter_alpha_fast=1.0,
+        max_linear_speed_m_s=10.0,
+        orientation_filter_alpha=1.0,
+        max_orientation_speed_rad_s=10.0,
+    )
+    mapper = RelativePoseMapper(profile)
+    start = profile.initial_tcp
+    mapper.map(sample(arm, True), start, 0.1)
+
+    mapped = mapper.map(sample(arm, True, [0.0, 0.01, 0.0]), start, 0.1)
+
+    np.testing.assert_allclose(
+        mapped.target_tcp.translation - start.translation,
+        expected_delta,
+        atol=1e-12,
+    )
+
+
+def test_mujoco_initializes_both_configured_tool_tcp_poses(profiles):
+    robot = MujocoRobot(profiles)
+    try:
+        for arm in ("left", "right"):
+            position_error, orientation_error = pose_errors(
+                robot.get_tcp_pose(arm), profiles[arm].initial_tcp
+            )
+            assert position_error < 1e-3
+            assert orientation_error < np.deg2rad(0.1)
+            diagnostic = robot.initial_pose_diagnostics[arm]
+            assert diagnostic.configured_joint_position_error_m > 0.05
+    finally:
+        robot.close()
+
+
+def test_mujoco_dual_command_is_atomic(profiles):
+    robot = MujocoRobot(profiles)
+    try:
+        before = robot.read_joint_positions().positions_rad
+        with pytest.raises(ValueError):
+            robot.command_joint_positions(
+                {
+                    "left": before["left"] + 0.01,
+                    "right": np.full(7, np.nan),
+                }
+            )
+        after = robot.read_joint_positions().positions_rad
+        np.testing.assert_array_equal(after["left"], before["left"])
+        np.testing.assert_array_equal(after["right"], before["right"])
+    finally:
+        robot.close()
+
+
+def test_dual_controller_has_no_grip_jump_and_moves_both_arms(profiles):
+    robot = MujocoRobot(profiles)
+    controller = DualArmQpTeleopController(robot, profiles)
+    try:
+        initial = robot.read_joint_positions().positions_rad
+        for arm in ("left", "right"):
+            controller.update_sample(sample(arm, True), 0.0)
+        engaged = controller.step(0.0)
+        after_engage = robot.read_joint_positions().positions_rad
+        assert engaged.state is SafetyState.ACTIVE
+        for arm in ("left", "right"):
+            np.testing.assert_allclose(after_engage[arm], initial[arm], atol=1e-10)
+
+        for arm in ("left", "right"):
+            controller.update_sample(sample(arm, True, [0, 0.005, 0]), 0.01)
+        moved_result = controller.step(0.01)
+        moved = robot.read_joint_positions().positions_rad
+        assert moved_result.state is SafetyState.ACTIVE
+        for arm in ("left", "right"):
+            assert np.max(np.abs(moved[arm] - initial[arm])) > 1e-5
+    finally:
+        controller.close()
+
+
+def test_active_arm_timeout_faults_both_and_requires_release(profiles):
+    robot = MujocoRobot(profiles)
+    controller = DualArmQpTeleopController(robot, profiles)
+    try:
+        for arm in ("left", "right"):
+            controller.update_sample(sample(arm, True), 0.0)
+        assert controller.step(0.0).state is SafetyState.ACTIVE
+
+        fault = controller.step(1.0)
+        assert fault.state is SafetyState.FAULT
+        assert "timed out" in fault.reason
+
+        for arm in ("left", "right"):
+            controller.update_sample(sample(arm, False), 1.01)
+        cleared = controller.step(1.01)
+        assert cleared.state is SafetyState.IDLE
+    finally:
+        controller.close()
+
+
+def test_message_compatible_adapter_rejects_wrong_hand():
+    message = SimpleNamespace(
+        hand="right",
+        grip=True,
+        trigger=0.0,
+        pose=SimpleNamespace(
+            position=SimpleNamespace(x=0.0, y=0.0, z=0.0),
+            orientation=SimpleNamespace(x=0.0, y=0.0, z=0.0, w=1.0),
+        ),
+    )
+
+    with pytest.raises(ValueError, match="does not match"):
+        ControllerSample.from_message(message, expected_arm="left")
+
+
+@pytest.mark.skipif(
+    not os.environ.get("REALMAN_SDK_ROOT"),
+    reason="REALMAN_SDK_ROOT is not set",
+)
+def test_quick_stage3_validation_passes():
+    from rm75_ik.realman_reference import RealManFkReference
+    from rm75_ik.stage3_validation import Stage3Settings, Stage3Validator
+
+    validator = Stage3Validator(
+        RealManFkReference(Path(os.environ["REALMAN_SDK_ROOT"])),
+        ROOT / "xr_rm_bringup/config/dual_arm_rm75.yaml",
+        ROOT / "xr_rm_bringup/config/peripherals_rm75.yaml",
+        Stage3Settings.quick(),
+    )
+
+    summary = validator.run()
+
+    assert summary["passed"] is True
+    assert summary["failure_count"] == 0