SO-101 Support

mani_skill/agents/robots/__init__.py

@@ -15,6 +15,7 @@
 from .koch import *
 from .panda import *
 from .so100 import *
+from .so101 import *
 from .stompy import Stompy
 from .trifingerpro import TriFingerPro
 from .unitree_g1 import *

mani_skill/agents/robots/lerobot/manipulator.py

@@ -13,10 +13,10 @@
 from mani_skill.utils.structs.types import Array
 
 try:
-    from lerobot.common.cameras.camera import Camera
-    from lerobot.common.motors.motors_bus import MotorNormMode
-    from lerobot.common.robots.robot import Robot
-    from lerobot.common.utils.robot_utils import busy_wait
+    from cameras.camera import Camera
+    from lerobot.motors.motors_bus import MotorNormMode
+    from lerobot.robots.robot import Robot
+    from lerobot.utils.robot_utils import busy_wait
 except ImportError:
     pass
 
@@ -41,7 +41,7 @@ def __init__(self, robot: Robot, use_cached_qpos: bool = True, **kwargs):
         self._cached_qpos = None
         self._motor_keys: List[str] = None
 
-        if self.real_robot.name == "so100_follower":
+        if self.real_robot.name == "so100_follower" or self.real_robot.name == "so101_follower":
             self.real_robot.bus.motors["gripper"].norm_mode = MotorNormMode.DEGREES
 
     def start(self):
@@ -58,6 +58,8 @@ def set_target_qpos(self, qpos: Array):
         # NOTE (stao): It seems the calibration from LeRobot has some offsets in some joints. We fix reading them here to match the expected behavior
         if self.real_robot.name == "so100_follower":
             qpos["elbow_flex.pos"] = qpos["elbow_flex.pos"] + 6.8
+        elif self.real_robot.name == "so101_follower":
+            pass
         self.real_robot.send_action(qpos)
 
     def reset(self, qpos: Array):
@@ -113,6 +115,8 @@ def get_qpos(self):
         # NOTE (stao): It seems the calibration from LeRobot has some offsets in some joints. We fix reading them here to match the expected behavior
         if self.real_robot.name == "so100_follower":
             qpos_deg["elbow_flex"] = qpos_deg["elbow_flex"] - 6.8
+        elif self.real_robot.name == "so101_follower":
+            pass
         if self._motor_keys is None:
             self._motor_keys = list(qpos_deg.keys())
         qpos_deg = common.flatten_state_dict(qpos_deg)

mani_skill/agents/robots/so101/__init__.py

@@ -0,0 +1 @@
+from .so_101 import SO101 

mani_skill/agents/robots/so101/so_101.py

@@ -0,0 +1,137 @@
+import copy
+
+import numpy as np
+import sapien
+import sapien.render
+import torch
+from transforms3d.euler import euler2quat
+
+from mani_skill import PACKAGE_ASSET_DIR
+from mani_skill.agents.base_agent import BaseAgent, Keyframe
+from mani_skill.agents.controllers import *
+from mani_skill.agents.registration import register_agent
+from mani_skill.utils import common
+from mani_skill.utils.structs.actor import Actor
+from mani_skill.utils.structs.pose import Pose
+
+
+@register_agent()
+class SO101(BaseAgent):
+    uid = "so101"
+    urdf_path = f"{PACKAGE_ASSET_DIR}/robots/so101/so101.urdf"
+    urdf_config = dict(
+        _materials=dict(
+            gripper=dict(static_friction=2.5, dynamic_friction=2.5, restitution=0.0)
+        ),
+        link=dict(
+            gripper_link=dict(material="gripper", patch_radius=0.1, min_patch_radius=0.1),
+            moving_jaw_so101_v1_link=dict(material="gripper", patch_radius=0.1, min_patch_radius=0.1),
+            finger1_tip=dict(material="gripper", patch_radius=0.1, min_patch_radius=0.1),
+            finger2_tip=dict(material="gripper", patch_radius=0.1, min_patch_radius=0.1),
+        ),
+    )
+
+    keyframes = dict(
+        rest=Keyframe(
+            qpos=np.array([0, -1.5708, 1.5708, 0.66, 0.0, -10 * np.pi/180]),  # Fully closed gripper
+            pose=sapien.Pose(q=list(euler2quat(0, 0, np.pi / 2))),
+        ),
+        zero=Keyframe(
+            qpos=np.array([0, 0, 0, 0, 0, 0]),
+            pose=sapien.Pose(q=list(euler2quat(0, 0, np.pi / 2))),
+        ),
+        extended=Keyframe(
+            qpos=np.array([0, -0.7854, 0.7854, 0, 0, 100 * np.pi/180]),  # Fully open gripper
+            pose=sapien.Pose(q=list(euler2quat(0, 0, np.pi / 2))),
+        ),
+    )
+
+    arm_joint_names = [
+        "shoulder_pan",
+        "shoulder_lift", 
+        "elbow_flex",
+        "wrist_flex",
+        "wrist_roll",
+    ]
+    gripper_joint_names = [
+        "gripper",
+    ]
+
+    @property
+    def _controller_configs(self):
+        pd_joint_pos = PDJointPosControllerConfig(
+            [joint.name for joint in self.robot.active_joints],
+            lower=None,
+            upper=None,
+            stiffness=1e3,
+            damping=1e2,
+            force_limit=100,
+            normalize_action=False,
+        )
+
+        # Improved delta position control for SO101
+        pd_joint_delta_pos = PDJointPosControllerConfig(
+            [joint.name for joint in self.robot.active_joints],
+            [-0.05, -0.05, -0.05, -0.05, -0.05, -0.2],  # Match gripper joint limits
+            [0.05, 0.05, 0.05, 0.05, 0.05, 0.2],
+            stiffness=[1e3] * 6,
+            damping=[1e2] * 6,
+            force_limit=100,
+            use_delta=True,
+            use_target=False,
+        )
+
+        pd_joint_target_delta_pos = copy.deepcopy(pd_joint_delta_pos)
+        pd_joint_target_delta_pos.use_target = True
+
+        controller_configs = dict(
+            pd_joint_delta_pos=pd_joint_delta_pos,
+            pd_joint_pos=pd_joint_pos,
+            pd_joint_target_delta_pos=pd_joint_target_delta_pos,
+        )
+        return deepcopy_dict(controller_configs)
+
+    def _after_loading_articulation(self):
+        super()._after_loading_articulation()
+        self.finger1_link = self.robot.links_map["gripper_link"]
+        self.finger2_link = self.robot.links_map["moving_jaw_so101_v1_link"]
+        self.finger1_tip = self.robot.links_map["finger1_tip"]
+        self.finger2_tip = self.robot.links_map["finger2_tip"]
+
+    @property
+    def tcp_pos(self):
+        # computes the tool center point as the mid point between the the fixed and moving jaw's tips
+        return (self.finger1_tip.pose.p + self.finger2_tip.pose.p) / 2
+
+    @property
+    def tcp_pose(self):
+        return Pose.create_from_pq(self.tcp_pos, self.finger1_link.pose.q)
+
+    def is_grasping(self, object: Actor, min_force=0.5, max_angle=110):
+        """Check if the robot is grasping an object (more lenient parameters)"""
+        l_contact_forces = self.scene.get_pairwise_contact_forces(
+            self.finger1_link, object
+        )
+        r_contact_forces = self.scene.get_pairwise_contact_forces(
+            self.finger2_link, object
+        )
+        lforce = torch.linalg.norm(l_contact_forces, axis=1)
+        rforce = torch.linalg.norm(r_contact_forces, axis=1)
+
+        # direction to open the gripper
+        ldirection = self.finger1_link.pose.to_transformation_matrix()[..., :3, 1]
+        rdirection = -self.finger2_link.pose.to_transformation_matrix()[..., :3, 1]
+        langle = common.compute_angle_between(ldirection, l_contact_forces)
+        rangle = common.compute_angle_between(rdirection, r_contact_forces)
+        lflag = torch.logical_and(
+            lforce >= min_force, torch.rad2deg(langle) <= max_angle
+        )
+        rflag = torch.logical_and(
+            rforce >= min_force, torch.rad2deg(rangle) <= max_angle
+        )
+        return torch.logical_and(lflag, rflag)
+
+    def is_static(self, threshold=0.15):
+        """Check if the robot is static (improved for SO101)"""
+        qvel = self.robot.get_qvel()[:, :-1]  # exclude the gripper joint
+        return torch.max(torch.abs(qvel), 1)[0] <= threshold 

mani_skill/agents/robots/so101/so_101_real.py

@@ -0,0 +1,8 @@
+from mani_skill.agents.robots.lerobot.manipulator import LeRobotRealAgent
+
+class SO101RealAgent(LeRobotRealAgent):
+    """
+    SO101RealAgent is a class for controlling the SO101 real robot via the LeRobot system.
+    It inherits from LeRobotRealAgent and uses the same functionality with the SO101 robot configuration.
+    """
+    pass

mani_skill/assets/robots/so101/meshes/Fixed_part_1.obj

@@ -0,0 +1,39 @@
+# Blender 4.4.3
+# www.blender.org
+mtllib Fixed_part_1.mtl
+o Cube
+v -0.035590 -0.021408 0.044782
+v -0.035590 0.024292 0.044782
+v -0.035590 -0.021408 0.000782
+v -0.035590 0.024292 0.000782
+v 0.029610 -0.021408 0.044782
+v 0.029610 0.024292 0.044782
+v 0.029610 -0.021408 0.000782
+v 0.029610 0.024292 0.000782
+vn -1.0000 -0.0000 -0.0000
+vn -0.0000 -0.0000 -1.0000
+vn 1.0000 -0.0000 -0.0000
+vn -0.0000 -0.0000 1.0000
+vn -0.0000 -1.0000 -0.0000
+vn -0.0000 1.0000 -0.0000
+vt 0.375000 0.000000
+vt 0.625000 0.000000
+vt 0.625000 0.250000
+vt 0.375000 0.250000
+vt 0.625000 0.500000
+vt 0.375000 0.500000
+vt 0.625000 0.750000
+vt 0.375000 0.750000
+vt 0.625000 1.000000
+vt 0.375000 1.000000
+vt 0.125000 0.500000
+vt 0.125000 0.750000
+vt 0.875000 0.500000
+vt 0.875000 0.750000
+s 0
+f 1/1/1 2/2/1 4/3/1 3/4/1
+f 3/4/2 4/3/2 8/5/2 7/6/2
+f 7/6/3 8/5/3 6/7/3 5/8/3
+f 5/8/4 6/7/4 2/9/4 1/10/4
+f 3/11/5 7/6/5 5/8/5 1/12/5
+f 8/5/6 4/13/6 2/14/6 6/7/6