feat: 添加双臂 MuJoCo 运动学模型

xr_rm_mujoco/xr_rm_mujoco/__init__.py

@@ -0,0 +1 @@
+"""双 RM75 MuJoCo 运动学显示包。"""

xr_rm_mujoco/xr_rm_mujoco/dual_arm_simulator.py

@@ -0,0 +1,86 @@
+"""使用现有双 RM75 URDF 的 MuJoCo 运动学状态映射。"""
+
+from __future__ import annotations
+
+import math
+from pathlib import Path
+
+import mujoco
+
+
+ARM_JOINT_NAMES = {
+    "left": tuple(f"scissor_joint_{index}" for index in range(1, 8)),
+    "right": tuple(f"omnipic_joint_{index}" for index in range(1, 8)),
+}
+
+
+class DualArmKinematicModel:
+    """加载双臂 URDF，并按关节名更新 MuJoCo qpos。"""
+
+    def __init__(self, urdf_path: str) -> None:
+        path = Path(urdf_path).expanduser().resolve()
+        if not path.is_file():
+            raise FileNotFoundError(f"dual RM75 URDF not found: {path}")
+
+        self.model = mujoco.MjModel.from_xml_path(str(path))
+        self.data = mujoco.MjData(self.model)
+        self._qpos_addresses: dict[str, dict[str, int]] = {}
+        self._received_arms: set[str] = set()
+
+        for arm, names in ARM_JOINT_NAMES.items():
+            addresses = {}
+            for name in names:
+                joint_id = mujoco.mj_name2id(
+                    self.model,
+                    mujoco.mjtObj.mjOBJ_JOINT,
+                    name,
+                )
+                if joint_id < 0:
+                    raise RuntimeError(f"MuJoCo joint not found: {name}")
+                if self.model.jnt_type[joint_id] != mujoco.mjtJoint.mjJNT_HINGE:
+                    raise RuntimeError(f"MuJoCo joint must be hinge: {name}")
+                addresses[name] = int(self.model.jnt_qposadr[joint_id])
+            self._qpos_addresses[arm] = addresses
+
+    @property
+    def ready(self) -> bool:
+        return self._received_arms == set(ARM_JOINT_NAMES)
+
+    def apply_arm_state(
+        self,
+        arm: str,
+        names: list[str] | tuple[str, ...],
+        positions: list[float] | tuple[float, ...],
+    ) -> None:
+        if arm not in ARM_JOINT_NAMES:
+            raise ValueError("arm must be left or right")
+        if len(names) != len(positions):
+            raise ValueError("joint names and positions must have the same length")
+        if len(set(names)) != len(names):
+            raise ValueError("joint names must be unique")
+
+        expected = set(ARM_JOINT_NAMES[arm])
+        if set(names) != expected:
+            raise ValueError(f"joint names must match expected {arm} joints")
+
+        values = [float(value) for value in positions]
+        if not all(math.isfinite(value) for value in values):
+            raise ValueError("joint positions must be finite")
+
+        by_name = dict(zip(names, values))
+        updates = [
+            (self._qpos_addresses[arm][name], by_name[name])
+            for name in ARM_JOINT_NAMES[arm]
+        ]
+        for address, value in updates:
+            self.data.qpos[address] = value
+        mujoco.mj_forward(self.model, self.data)
+        self._received_arms.add(arm)
+
+    def joint_positions(self, arm: str) -> list[float]:
+        if arm not in ARM_JOINT_NAMES:
+            raise ValueError("arm must be left or right")
+        return [
+            float(self.data.qpos[self._qpos_addresses[arm][name]])
+            for name in ARM_JOINT_NAMES[arm]
+        ]