yam-common 0.1.0__tar.gz

yam_common-0.1.0/PKG-INFO

@@ -0,0 +1,22 @@
+Metadata-Version: 2.4
+Name: yam-common
+Version: 0.1.0
+Summary: Shared utilities for YAM arm LeRobot plugins
+Author: Praveen Selvaraj
+License: MIT
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: numpy
+Requires-Dist: mujoco
+Requires-Dist: lerobot>=0.4
+
+# yam-common
+
+Shared utilities for the YAM arm LeRobot plugins.
+
+Includes the motor chain controller, MuJoCo gravity compensation, and gripper
+utilities used by both the follower robot and leader teleoperator packages.
+

yam_common-0.1.0/README.md

@@ -0,0 +1,7 @@
+# yam-common
+
+Shared utilities for the YAM arm LeRobot plugins.
+
+Includes the motor chain controller, MuJoCo gravity compensation, and gripper
+utilities used by both the follower robot and leader teleoperator packages.
+

yam_common-0.1.0/pyproject.toml

@@ -0,0 +1,29 @@
+[build-system]
+requires = ["setuptools>=61", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "yam-common"
+version = "0.1.0"
+description = "Shared utilities for YAM arm LeRobot plugins"
+readme = "README.md"
+authors = [
+  { name = "Praveen Selvaraj" },
+]
+license = { text = "MIT" }
+requires-python = ">=3.10"
+classifiers = [
+  "Programming Language :: Python :: 3",
+  "License :: OSI Approved :: MIT License",
+  "Operating System :: OS Independent",
+]
+dependencies = [
+  "numpy",
+  "mujoco",
+  "lerobot>=0.4",
+]
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["yam_common*"]
+

yam_common-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

yam_common-0.1.0/yam_common/__init__.py

@@ -0,0 +1,13 @@
+from .mujoco_kdl import MuJoCoKDL, get_yam_mujoco_kdl
+from .motor_chain_robot import MotorChainRobot
+from .utils import GripperForceLimiter, GripperType, JointMapper
+
+__all__ = [
+    "MuJoCoKDL",
+    "get_yam_mujoco_kdl",
+    "MotorChainRobot",
+    "GripperForceLimiter",
+    "GripperType",
+    "JointMapper",
+]
+

yam_common-0.1.0/yam_common/motor_chain_robot.py

@@ -0,0 +1,404 @@
+"""i2rt-authentic MotorChainRobot port for YAM."""
+
+from __future__ import annotations
+
+import copy
+import logging
+import os
+import threading
+import time
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Union
+
+import numpy as np
+
+from lerobot.motors.dm.dm_driver import MotorChain, MotorInfo
+from yam_common.mujoco_kdl import MuJoCoKDL
+from yam_common.utils import GripperForceLimiter, GripperType, JointMapper
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class JointStates:
+    names: List[str]
+    pos: np.ndarray
+    vel: np.ndarray
+    eff: np.ndarray
+    temp_mos: np.ndarray
+    temp_rotor: np.ndarray
+
+    def asdict(self) -> Dict[str, Any]:
+        return {
+            "names": self.names,
+            "pos": self.pos.flatten().tolist(),
+            "vel": self.vel.flatten().tolist(),
+            "eff": self.eff.flatten().tolist(),
+        }
+
+
+@dataclass
+class JointCommands:
+    torques: np.ndarray
+    pos: np.ndarray
+    vel: np.ndarray
+    kp: np.ndarray
+    kd: np.ndarray
+    indices: Optional[List[int]] = None
+
+    @classmethod
+    def init_all_zero(cls, n_joints: int) -> "JointCommands":
+        return cls(
+            torques=np.zeros(n_joints),
+            pos=np.zeros(n_joints),
+            vel=np.zeros(n_joints),
+            kp=np.zeros(n_joints),
+            kd=np.zeros(n_joints),
+        )
+
+
+class MotorChainRobot:
+    """i2rt-authentic motor-chain robot controller."""
+
+    def __init__(
+        self,
+        motor_chain: MotorChain,
+        xml_path: Optional[str] = None,
+        use_gravity_comp: bool = True,
+        gravity: Optional[np.ndarray] = None,
+        gravity_comp_factor: float = 1.0,
+        gripper_index: Optional[int] = None,
+        kp: Union[float, List[float]] = 10.0,
+        kd: Union[float, List[float]] = 1.0,
+        joint_limits: Optional[np.ndarray] = None,
+        gripper_limits: Optional[np.ndarray] = None,
+        limit_gripper_force: float = -1,
+        clip_motor_torque: float = np.inf,
+        gripper_type: GripperType = GripperType.CRANK_4310,
+        temp_record_flag: bool = False,
+        enable_gripper_calibration: bool = False,
+        zero_gravity_mode: bool = True,
+        test_torque: float = 0.5,
+        test_duration: float = 2.0,
+        position_threshold: float = 0.01,
+        check_interval: float = 0.05,
+    ) -> None:
+        self.temp_record_flag = temp_record_flag
+        if gripper_index is not None:
+            assert gripper_index == len(motor_chain) - 1, (
+                "Gripper index should be the last one, but got {gripper_index}"
+            )
+            if gripper_limits is None and enable_gripper_calibration:
+                from lerobot.robots.yam_follower.utils import detect_gripper_limits
+
+                logger.info("Auto-detecting gripper limits...")
+                detected_limits = detect_gripper_limits(
+                    motor_chain=motor_chain,
+                    gripper_index=gripper_index,
+                    test_torque=test_torque,
+                    max_duration=test_duration,
+                    position_threshold=position_threshold,
+                    check_interval=check_interval,
+                )
+                gripper_limits = np.array(detected_limits)
+                logger.info(f"Gripper limits auto-detected: {gripper_limits}")
+            elif gripper_limits is None:
+                raise ValueError(
+                    f"{self}: Gripper limits are required if gripper index is provided and auto-calibration is disabled."
+                )
+            else:
+                logger.info(f"Using provided gripper limits: {gripper_limits}")
+
+        self._last_gripper_command_qpos = 1
+        assert clip_motor_torque >= 0.0
+        self._clip_motor_torque = clip_motor_torque
+        self.motor_chain = motor_chain
+        self.use_gravity_comp = use_gravity_comp
+        self.gravity_comp_factor = gravity_comp_factor
+
+        self._gripper_index = gripper_index
+        self.remapper = JointMapper({}, len(motor_chain))
+        self._gripper_limits = gripper_limits
+
+        if self._gripper_index is not None:
+            self._gripper_force_limiter = GripperForceLimiter(
+                max_force=limit_gripper_force, gripper_type=gripper_type, kp=kp[gripper_index]
+            )
+            self._limit_gripper_force = limit_gripper_force
+            self.remapper = JointMapper(
+                index_range_map={gripper_index: gripper_limits},
+                total_dofs=len(motor_chain),
+            )
+
+        self._kp = (
+            np.array([kp] * len(motor_chain)) if isinstance(kp, float) else np.array(kp)
+        )
+        self._kd = (
+            np.array([kd] * len(motor_chain)) if isinstance(kd, float) else np.array(kd)
+        )
+
+        self._joint_limits: Optional[np.ndarray] = None
+        if xml_path is not None:
+            self.xml_path = os.path.expanduser(xml_path)
+            self.kdl = MuJoCoKDL(self.xml_path)
+            if gravity is not None:
+                self.kdl.set_gravity(gravity)
+            self._joint_limits = self.kdl.joint_limits
+        else:
+            assert use_gravity_comp is False, "Gravity compensation requires a valid XML path."
+
+        if joint_limits is not None:
+            joint_limits = np.array(joint_limits)
+            assert np.all(joint_limits[:, 0] < joint_limits[:, 1]), (
+                "Lower joint limits must be smaller than upper limits"
+            )
+            self._joint_limits = joint_limits
+
+        self._command_lock = threading.Lock()
+        self._state_lock = threading.Lock()
+        self._joint_state: Optional[JointStates] = None
+        while self._joint_state is None:
+            time.sleep(0.05)
+            self._joint_state = self._motor_state_to_joint_state(self.motor_chain.read_states())
+        self._commands = JointCommands.init_all_zero(len(motor_chain))
+        self._check_current_qpos_in_joint_limits()
+
+        self._stop_event = threading.Event()
+        self._server_thread = threading.Thread(target=self.start_server, name="robot_server")
+        self._server_thread.start()
+
+        if not zero_gravity_mode:
+            self.command_joint_pos(self._joint_state.pos)
+
+    def __repr__(self) -> str:
+        return f"MotorChainRobot(motor_chain={self.motor_chain})"
+
+    def _check_current_qpos_in_joint_limits(self, buffer_rad: float = 0.1) -> None:
+        if self._joint_state is None or self._joint_limits is None:
+            raise RuntimeError(f"{self}: Joint limits:{self._joint_limits} or joint state:{self._joint_state} are not set.")
+
+        current_pos = self._joint_state.pos
+        if self._gripper_index is not None:
+            arm_pos = current_pos[: self._gripper_index]
+            arm_limits = self._joint_limits
+        else:
+            arm_pos = current_pos
+            arm_limits = self._joint_limits
+
+        lower_limits = arm_limits[:, 0] - buffer_rad
+        upper_limits = arm_limits[:, 1] + buffer_rad
+        lower_violations = arm_pos < lower_limits
+        upper_violations = arm_pos > upper_limits
+
+        if np.any(lower_violations) or np.any(upper_violations):
+            violation_details = []
+            for i, (pos, lower, upper) in enumerate(zip(arm_pos, lower_limits, upper_limits)):
+                if pos < lower:
+                    violation_details.append(f"Joint {i}: {pos:.4f} < {lower:.4f} (lower limit)")
+                elif pos > upper:
+                    violation_details.append(f"Joint {i}: {pos:.4f} > {upper:.4f} (upper limit)")
+            violation_msg = "; ".join(violation_details)
+            self.motor_chain.running = False
+            raise RuntimeError(
+                f"{self}: Joint limit violation detected: {violation_msg}, the root reason should be zero position "
+                "offset. possible solution: 1. move the arm to zero position and power cycle the robot. "
+                "2. Recalibrate the motor zero position."
+            )
+
+    def start_server(self) -> None:
+        last_time = time.time()
+        iteration_count = 0
+        try:
+            self.update()
+            logging.info("initializing, ....")
+            while not self._stop_event.is_set():
+                current_time = time.time()
+                elapsed_time = current_time - last_time
+                self.update()
+                if not self.motor_chain.running:
+                    raise RuntimeError(
+                        f"{self}: motor_chain_robot's motor chain is not running, exiting the robot server"
+                    )
+                time.sleep(0.004)
+                iteration_count += 1
+                if elapsed_time >= 10.0:
+                    control_frequency = iteration_count / elapsed_time
+                    logging.info(f"{self}: Grav Comp Control Frequency: {control_frequency:.2f} Hz")
+                    if control_frequency < 100:
+                        logging.warning(
+                            f"{self}: Gravity compensation control loop is slow, current frequency: {control_frequency:.2f} Hz"
+                        )
+                    last_time = current_time
+                    iteration_count = 0
+        except Exception as exc:
+            logging.error(f"{self}: robot server error, entering zero-torque mode: {exc}")
+            try:
+                self.zero_torque_mode()
+            except Exception:
+                pass
+            self._stop_event.set()
+            raise
+
+    def update(self) -> None:
+        with self._command_lock:
+            joint_commands = copy.deepcopy(self._commands)
+        with self._state_lock:
+            g = self._compute_gravity_compensation(self._joint_state)
+            motor_torques = joint_commands.torques + g * self.gravity_comp_factor
+            motor_torques = np.clip(motor_torques, -self._clip_motor_torque, self._clip_motor_torque)
+
+            if self._gripper_index is not None:
+                if self._limit_gripper_force > 0 and self._joint_state is not None:
+                    gripper_state = {
+                        "target_qpos": joint_commands.pos[self._gripper_index],
+                        "current_qpos": self.remapper.to_robot_joint_pos_space(self._joint_state.pos)[
+                            self._gripper_index
+                        ],
+                        "current_qvel": self._joint_state.vel[self._gripper_index],
+                        "current_eff": self._joint_state.eff[self._gripper_index],
+                        "current_normalized_qpos": self._joint_state.pos[self._gripper_index],
+                        "target_normalized_qpos": self.remapper.to_command_joint_pos_space(joint_commands.pos)[
+                            self._gripper_index
+                        ],
+                        "last_command_qpos": self._last_gripper_command_qpos,
+                    }
+                    joint_commands.pos[self._gripper_index] = self._gripper_force_limiter.update(gripper_state)
+
+                joint_commands.pos[self._gripper_index] = np.clip(
+                    joint_commands.pos[self._gripper_index],
+                    min(self._gripper_limits),
+                    max(self._gripper_limits),
+                )
+                self._last_gripper_command_qpos = joint_commands.pos[self._gripper_index]
+
+            if not self.motor_chain.start_thread_flag:
+                self.motor_chain.set_commands(
+                    motor_torques,
+                    pos=joint_commands.pos,
+                    vel=joint_commands.vel,
+                    kp=joint_commands.kp,
+                    kd=joint_commands.kd,
+                )
+                self.motor_chain.start_thread()
+                self.motor_chain.start_thread_flag = True
+
+            motor_state = self.motor_chain.set_commands(
+                motor_torques,
+                pos=joint_commands.pos,
+                vel=joint_commands.vel,
+                kp=joint_commands.kp,
+                kd=joint_commands.kd,
+            )
+            self._joint_state = self._motor_state_to_joint_state(motor_state)
+            self._check_current_qpos_in_joint_limits()
+
+    def _motor_state_to_joint_state(self, motor_state: List[MotorInfo]) -> JointStates:
+        names = [str(i) for i in range(len(motor_state))]
+        pos = np.array([motor.pos for motor in motor_state])
+        pos = self.remapper.to_command_joint_pos_space(pos)
+        vel = np.array([motor.vel for motor in motor_state])
+        vel = self.remapper.to_command_joint_vel_space(vel)
+        eff = np.array([motor.eff for motor in motor_state])
+        temp_mos = np.array([motor.temp_mos for motor in motor_state])
+        temp_rotor = np.array([motor.temp_rotor for motor in motor_state])
+        return JointStates(
+            names=names,
+            pos=pos,
+            vel=vel,
+            eff=eff,
+            temp_mos=temp_mos,
+            temp_rotor=temp_rotor,
+        )
+
+    def _compute_gravity_compensation(self, joint_state: Optional[JointStates]) -> np.ndarray:
+        if joint_state is None or not self.use_gravity_comp:
+            return np.zeros(len(self.motor_chain))
+        q = joint_state.pos[: self._gripper_index] if self._gripper_index is not None else joint_state.pos
+        t = self.kdl.compute_inverse_dynamics(q, np.zeros(q.shape), np.zeros(q.shape))
+        if np.max(np.abs(t)) > 20.0:
+            print([f"{s:.2f}" for s in t])
+            raise RuntimeError(f"{self}: too large torques")
+        if self._gripper_index is None:
+            return t
+        return np.append(t, 0.0)
+
+    def num_dofs(self) -> int:
+        return len(self.motor_chain)
+
+    def get_joint_pos(self) -> np.ndarray:
+        with self._state_lock:
+            return self._joint_state.pos
+
+    def _clip_robot_joint_pos_command(self, pos: np.ndarray) -> np.ndarray:
+        if self._joint_limits is not None:
+            if self._gripper_index is not None:
+                pos[: self._gripper_index] = np.clip(
+                    pos[: self._gripper_index],
+                    self._joint_limits[:, 0],
+                    self._joint_limits[:, 1],
+                )
+            else:
+                pos = np.clip(pos, self._joint_limits[:, 0], self._joint_limits[:, 1])
+        return pos
+
+    def command_joint_pos(self, joint_pos: np.ndarray) -> None:
+        pos = self._clip_robot_joint_pos_command(joint_pos)
+        with self._command_lock:
+            self._commands = JointCommands.init_all_zero(len(self.motor_chain))
+            self._commands.pos = self.remapper.to_robot_joint_pos_space(pos)
+            self._commands.kp = self._kp
+            self._commands.kd = self._kd
+
+    def command_joint_state(self, joint_state: Dict[str, np.ndarray]) -> None:
+        pos = self._clip_robot_joint_pos_command(joint_state["pos"])
+        vel = joint_state["vel"]
+        self._commands = JointCommands.init_all_zero(len(self.motor_chain))
+        kp = joint_state.get("kp", self._kp)
+        kd = joint_state.get("kd", self._kd)
+        with self._command_lock:
+            self._commands.pos = self.remapper.to_robot_joint_pos_space(pos)
+            self._commands.vel = self.remapper.to_robot_joint_vel_space(vel)
+            self._commands.kp = kp
+            self._commands.kd = kd
+
+    def zero_torque_mode(self) -> None:
+        logging.info(f"Entering zero_torque_mode for {self}")
+        with self._command_lock:
+            self._commands = JointCommands.init_all_zero(len(self.motor_chain))
+            self._kp = np.zeros(len(self.motor_chain))
+            self._kd = np.zeros(len(self.motor_chain))
+
+    def get_observations(self) -> Dict[str, np.ndarray]:
+        with self._state_lock:
+            if self._gripper_index is None:
+                result = {
+                    "joint_pos": self._joint_state.pos,
+                    "joint_vel": self._joint_state.vel,
+                    "joint_eff": self._joint_state.eff,
+                }
+            else:
+                result = {
+                    "joint_pos": self._joint_state.pos[: self._gripper_index],
+                    "gripper_pos": np.array([self._joint_state.pos[self._gripper_index]]),
+                    "joint_vel": self._joint_state.vel,
+                    "joint_eff": self._joint_state.eff,
+                }
+            if self.temp_record_flag:
+                result["temp_mos"] = self._joint_state.temp_mos
+                result["temp_rotor"] = self._joint_state.temp_rotor
+            return result
+
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+        self.close()
+
+    def close(self) -> None:
+        self._stop_event.set()
+        self._server_thread.join()
+        self.motor_chain.close()
+        print("Robot closed with all torques set to zero.")
+
+    def update_kp_kd(self, kp: np.ndarray, kd: np.ndarray) -> None:
+        assert kp.shape == self._kp.shape == kd.shape
+        self._kp = kp
+        self._kd = kd
+

yam_common-0.1.0/yam_common/mujoco_kdl.py

@@ -0,0 +1,206 @@
+"""
+MuJoCo-based gravity compensation for YAM robot.
+
+Faithfully ported from i2rt/utils/mujoco_utils.py.
+
+This module uses MuJoCo's inverse dynamics to compute the torques needed
+to compensate for gravity at any given joint configuration.
+"""
+
+import os
+from typing import Optional
+
+import numpy as np
+
+try:
+    import mujoco
+except ImportError:
+    mujoco = None
+
+
+class MuJoCoKDL:
+    """
+    A class for computing inverse dynamics (gravity compensation) using MuJoCo.
+
+    Source: i2rt/utils/mujoco_utils.py:MuJoCoKDL
+
+    This class loads a MuJoCo XML model and uses mujoco.mj_inverse() to compute
+    the torques needed to maintain a given configuration against gravity.
+
+    For gravity compensation specifically:
+    - Set qpos to current joint positions
+    - Set qvel to zero (no velocity)
+    - Set qacc to zero (no acceleration)
+    - Call mj_inverse() to get the torques needed to maintain this configuration
+
+    Usage:
+        kdl = MuJoCoKDL("/path/to/yam.xml")
+        gravity_torques = kdl.compute_gravity_compensation(current_joint_positions)
+    """
+
+    def __init__(self, xml_path: str, gravity: Optional[np.ndarray] = None):
+        """
+        Initialize MuJoCoKDL with a robot model.
+
+        Args:
+            xml_path: Path to MuJoCo XML model file
+            gravity: Optional gravity vector (default: [0, 0, -9.81])
+        """
+        if mujoco is None:
+            raise ImportError(
+                "mujoco is required for MuJoCoKDL gravity compensation. "
+                "Install with: pip install mujoco"
+            )
+
+        # Load model - expand user path as in i2rt
+        self.xml_path = os.path.expanduser(xml_path)
+        self.model = mujoco.MjModel.from_xml_path(self.xml_path)
+        self.data = mujoco.MjData(self.model)
+
+        # Set gravity (default: Earth gravity pointing down)
+        if gravity is None:
+            gravity = np.array([0.0, 0.0, -9.81])
+        self.set_gravity(gravity)
+
+        # Disable all collisions - we only care about dynamics, not contacts
+        # Source: i2rt/utils/mujoco_utils.py lines 16-17
+        self.model.geom_contype[:] = 0
+        self.model.geom_conaffinity[:] = 0
+
+        # Disable all joint limits - we handle limits separately
+        # Source: i2rt/utils/mujoco_utils.py line 19
+        self.model.jnt_limited[:] = 0
+
+    @property
+    def num_joints(self) -> int:
+        """Number of joints in the model."""
+        return self.model.nq
+
+    @property
+    def joint_limits(self) -> np.ndarray:
+        """
+        Get joint limits from model.
+
+        Returns:
+            Array of shape (num_joints, 2) with [min, max] for each joint
+        """
+        return self.model.jnt_range.copy()
+
+    def set_gravity(self, gravity: np.ndarray) -> None:
+        """
+        Set the gravity vector for the robot.
+
+        Source: i2rt/utils/mujoco_utils.py:MuJoCoKDL.set_gravity()
+
+        Args:
+            gravity: Gravity vector as a 3D NumPy array (e.g., [0, 0, -9.81])
+        """
+        assert gravity.shape == (3,), f"Gravity must be 3D vector, got shape {gravity.shape}"
+        self.model.opt.gravity = gravity
+
+    def compute_inverse_dynamics(
+        self,
+        q: np.ndarray,
+        qdot: np.ndarray,
+        qdotdot: np.ndarray,
+    ) -> np.ndarray:
+        """
+        Compute inverse dynamics to get required joint torques.
+
+        Source: i2rt/utils/mujoco_utils.py:MuJoCoKDL.compute_inverse_dynamics()
+
+        This computes the torques needed to achieve the given acceleration
+        from the given position and velocity, accounting for gravity,
+        Coriolis forces, etc.
+
+        Args:
+            q: Joint positions (radians)
+            qdot: Joint velocities (rad/s)
+            qdotdot: Joint accelerations (rad/s^2)
+
+        Returns:
+            Joint torques (Nm) needed to achieve the given motion
+        """
+        assert len(q) == len(qdot) == len(qdotdot), (
+            f"Input dimensions must match: q={len(q)}, qdot={len(qdot)}, qdotdot={len(qdotdot)}"
+        )
+
+        length = len(q)
+
+        # Set state
+        self.data.qpos[:length] = q
+        self.data.qvel[:length] = qdot
+        self.data.qacc[:length] = qdotdot
+
+        # Compute inverse dynamics
+        mujoco.mj_inverse(self.model, self.data)
+
+        # Return the computed torques
+        return self.data.qfrc_inverse[:length].copy()
+
+    def compute_gravity_compensation(self, q: np.ndarray) -> np.ndarray:
+        """
+        Compute gravity compensation torques for a given joint configuration.
+
+        This is a convenience method that calls compute_inverse_dynamics with
+        zero velocity and zero acceleration, which gives the torques needed
+        to hold the arm stationary against gravity.
+
+        Args:
+            q: Joint positions (radians)
+
+        Returns:
+            Gravity compensation torques (Nm)
+        """
+        zeros = np.zeros_like(q)
+        return self.compute_inverse_dynamics(q, zeros, zeros)
+
+
+def get_yam_mujoco_kdl(gripper_type: str = "crank_4310") -> MuJoCoKDL:
+    """
+    Get MuJoCoKDL instance for YAM robot with appropriate XML model.
+
+    Args:
+        gripper_type: Type of gripper ("crank_4310", "linear_3507", "linear_4310",
+                      "yam_teaching_handle", "no_gripper")
+
+    Returns:
+        MuJoCoKDL instance configured for YAM robot
+    """
+    # Map gripper type to XML file
+    # Source: i2rt/robots/utils.py:GripperType.get_xml_path()
+    xml_paths = {
+        "crank_4310": "yam.xml",
+        "linear_3507": "yam_lw_gripper.xml",
+        "linear_4310": "yam_4310_linear.xml",
+        "yam_teaching_handle": "yam_teaching_handle.xml",
+        "no_gripper": "yam_no_gripper.xml",
+    }
+
+    if gripper_type not in xml_paths:
+        raise ValueError(f"Unknown gripper type: {gripper_type}. Valid types: {list(xml_paths.keys())}")
+
+    # Try to find the XML file in several locations
+    xml_filename = xml_paths[gripper_type]
+
+    # Search paths (in order of preference):
+    # 1. i2rt installation (if available)
+    # 2. Local lerobot models directory
+    search_paths = [
+        # i2rt installation
+        os.path.expanduser("~/Desktop/code/i2rt/i2rt/robot_models/yam"),
+        # Relative to this file
+        os.path.join(os.path.dirname(__file__), "robot_models"),
+        # Package data
+        os.path.join(os.path.dirname(__file__), "..", "..", "data", "robot_models", "yam"),
+    ]
+
+    for base_path in search_paths:
+        xml_path = os.path.join(base_path, xml_filename)
+        if os.path.exists(xml_path):
+            return MuJoCoKDL(xml_path)
+
+    raise FileNotFoundError(
+        f"Could not find MuJoCo XML file '{xml_filename}' for YAM robot. "
+        f"Searched paths: {search_paths}"
+    )