(refactor) refactor gym_manipulator.py to use the new kinematics from placo

Author: michel-aractingi

lerobot/common/model/kinematics.py

@@ -1,123 +1,117 @@
 import numpy as np
 
 
-class PlacoRobotKinematics:
+class RobotKinematics:
     """Robot kinematics using placo library for forward and inverse kinematics."""
-    
+
     def __init__(self, urdf_path: str, ee_frame_name: str = "gripper_tip", joint_names: list[str] = None):
         """
         Initialize placo-based kinematics solver.
-        
+
         Args:
             urdf_path: Path to the robot URDF file
             ee_frame_name: Name of the end-effector frame in the URDF
             joint_names: List of joint names to control (if None, will use default naming)
         """
-        try:
-            import placo
-        except ImportError:
-            raise ImportError(
-                "placo is required for PlacoRobotKinematics. "
-                "Please install it with: pip install placo"
-            )
-        
+        import placo
+
         self.robot = placo.RobotWrapper(urdf_path)
         self.solver = placo.KinematicsSolver(self.robot)
         self.solver.mask_fbase(True)  # Fix the base
-        
+
         self.ee_frame_name = ee_frame_name
-        
+
         # Set joint names
         if joint_names is None:
             # Default joint names for SO-ARM100
             self.joint_names = ["1", "2", "3", "4", "5"]
         else:
             self.joint_names = joint_names
-        
+
         # Initialize frame task for IK
         self.tip_starting_pose = np.eye(4)
         self.tip_frame = self.solver.add_frame_task(self.ee_frame_name, self.tip_starting_pose)
         self.tip_frame.configure(self.ee_frame_name, "soft", 1.0, 1.0)
-    
+
     def forward_kinematics(self, robot_pos_deg, frame=None):
         """
         Compute forward kinematics for given joint configuration.
-        
+
         Args:
             robot_pos_deg: Joint positions in degrees (numpy array)
             frame: Target frame name (if None, uses ee_frame_name)
-        
+
         Returns:
             4x4 transformation matrix of the end-effector pose
         """
         if frame is None:
             frame = self.ee_frame_name
-        
+
         # Convert degrees to radians
-        robot_pos_rad = np.deg2rad(robot_pos_deg[:len(self.joint_names)])
-        
+        robot_pos_rad = np.deg2rad(robot_pos_deg[: len(self.joint_names)])
+
         # Update joint positions in placo robot
         for i, joint_name in enumerate(self.joint_names):
             self.robot.set_joint(joint_name, robot_pos_rad[i])
-        
+
         # Update kinematics
         self.robot.update_kinematics()
-        
+
         # Get the transformation matrix
         return self.robot.get_T_world_frame(frame)
-    
+
     def ik(self, current_joint_pos, desired_ee_pose, position_only=True, frame=None):
         """
         Compute inverse kinematics using placo solver.
-        
+
         Args:
             current_joint_pos: Current joint positions in degrees (used as initial guess)
             desired_ee_pose: Target end-effector pose as a 4x4 transformation matrix
             position_only: If True, only match position (not orientation)
             frame: Target frame name (if None, uses ee_frame_name)
-        
+
         Returns:
             Joint positions in degrees that achieve the desired end-effector pose
         """
         if frame is None:
             frame = self.ee_frame_name
-        
+
         # Convert current joint positions to radians for initial guess
-        current_joint_rad = np.deg2rad(current_joint_pos[:len(self.joint_names)])
-        
+        current_joint_rad = np.deg2rad(current_joint_pos[: len(self.joint_names)])
+
         # Set current joint positions as initial guess
         for i, joint_name in enumerate(self.joint_names):
             self.robot.set_joint(joint_name, current_joint_rad[i])
-        
+
         # Update the target pose for the frame task
         self.tip_frame.T_world_frame = desired_ee_pose
-        
+
         # Configure the task based on position_only flag
         if position_only:
             # Only constrain position, not orientation
             self.tip_frame.configure(self.ee_frame_name, "soft", 1.0, 0.0)
         else:
             # Constrain both position and orientation
             self.tip_frame.configure(self.ee_frame_name, "soft", 1.0, 1.0)
-        
+
         # Solve IK
         self.solver.solve(True)
         self.robot.update_kinematics()
-        
+
         # Extract joint positions
         joint_positions_rad = []
         for joint_name in self.joint_names:
             joint = self.robot.get_joint(joint_name)
             joint_positions_rad.append(joint)
-        
+
         # Convert back to degrees
         joint_positions_deg = np.rad2deg(joint_positions_rad)
-        
+
         # Preserve gripper position if present in current_joint_pos
         if len(current_joint_pos) > len(self.joint_names):
             result = np.zeros_like(current_joint_pos)
-            result[:len(self.joint_names)] = joint_positions_deg
-            result[len(self.joint_names):] = current_joint_pos[len(self.joint_names):]
+            result[: len(self.joint_names)] = joint_positions_deg
+            result[len(self.joint_names) :] = current_joint_pos[len(self.joint_names) :]
             return result
         else:
-            return joint_positions_deg
+            return joint_positions_deg

lerobot/common/robots/so100_follower/config_so100_follower.py

@@ -46,10 +46,10 @@ class SO100FollowerEndEffectorConfig(SO100FollowerConfig):
 
     # Path to URDF file for kinematics
     urdf_path: str | None = None
-    
+
     # End-effector frame name in the URDF
     ee_frame_name: str = "gripperframe"
-    
+
     # Joint names for kinematics (if None, uses default naming)
     joint_names: list[str] | None = None
 

lerobot/common/robots/so100_follower/so100_follower_end_effector.py

@@ -22,7 +22,7 @@
 
 from lerobot.common.cameras import make_cameras_from_configs
 from lerobot.common.errors import DeviceNotConnectedError
-from lerobot.common.model.kinematics_utils import PlacoRobotKinematics
+from lerobot.common.model.kinematics import RobotKinematics
 from lerobot.common.motors import Motor, MotorNormMode
 from lerobot.common.motors.feetech import FeetechMotorsBus
 
@@ -68,11 +68,11 @@ def __init__(self, config: SO100FollowerEndEffectorConfig):
                 "urdf_path must be provided in the configuration for end-effector control. "
                 "Please set urdf_path in your SO100FollowerEndEffectorConfig."
             )
-        
-        self.kinematics = PlacoRobotKinematics(
+
+        self.kinematics = RobotKinematics(
             urdf_path=self.config.urdf_path,
             ee_frame_name=self.config.ee_frame_name,
-            joint_names=self.config.joint_names
+            joint_names=self.config.joint_names,
         )
 
         # Store the bounds for end-effector position
@@ -135,7 +135,7 @@ def send_action(self, action: dict[str, Any]) -> dict[str, Any]:
 
         # Calculate current end-effector position using forward kinematics
         if self.current_ee_pos is None:
-            self.current_ee_pos = self.kinematics.forward_kinematics(self.current_joint_pos, frame=self.config.ee_frame_name)
+            self.current_ee_pos = self.kinematics.forward_kinematics(self.current_joint_pos)
 
         # Set desired end-effector position by adding delta
         desired_ee_pos = np.eye(4)
@@ -152,7 +152,7 @@ def send_action(self, action: dict[str, Any]) -> dict[str, Any]:
 
         # Compute inverse kinematics to get joint positions
         target_joint_values_in_degrees = self.kinematics.ik(
-            self.current_joint_pos, desired_ee_pos, position_only=True, frame=self.config.ee_frame_name
+            self.current_joint_pos, desired_ee_pos, position_only=True
         )
 
         target_joint_values_in_degrees = np.clip(target_joint_values_in_degrees, -180.0, 180.0)

lerobot/scripts/rl/gym_manipulator.py

@@ -264,7 +264,7 @@ def __init__(
     def _get_observation(self) -> np.ndarray:
         """Helper to convert a dictionary from bus.sync_read to an ordered numpy array."""
         obs_dict = self.robot.get_observation()
-        joint_positions = np.array([obs_dict[name] for name in self._joint_names], dtype=np.float32)
+        joint_positions = np.array([obs_dict[name] for name in self._joint_names])
 
         images = {key: obs_dict[key] for key in self._image_keys}
         return {"agent_pos": joint_positions, "pixels": images}
@@ -1090,13 +1090,11 @@ def __init__(self, env, ee_pose_limits):
             dtype=np.float32,
         )
 
-        # Initialize kinematics instance for the appropriate robot type
-        robot_type = getattr(env.unwrapped.robot.config, "robot_type", "so101")
-        if "so100" in robot_type or "so101" in robot_type:
-            # Note to be compatible with the rest of the codebase,
-            # we are using the new calibration method for so101 and so100
-            robot_type = "so_new_calibration"
-        self.kinematics = RobotKinematics(robot_type)
+        self.kinematics = RobotKinematics(
+            urdf_path=env.unwrapped.robot.config.urdf_path,
+            ee_frame_name=env.unwrapped.robot.config.ee_frame_name,
+            joint_names=env.unwrapped.robot.config.joint_names,
+        )
 
     def observation(self, observation):
         """
@@ -1110,7 +1108,7 @@ def observation(self, observation):
         """
         current_joint_pos = self.unwrapped._get_observation()["agent_pos"]
 
-        current_ee_pos = self.kinematics.forward_kinematics(current_joint_pos, frame="gripper_tip")[:3, 3]
+        current_ee_pos = self.kinematics.forward_kinematics(current_joint_pos)[:3, 3]
         observation["agent_pos"] = np.concatenate([observation["agent_pos"], current_ee_pos], -1)
         return observation
 
@@ -1157,12 +1155,11 @@ def __init__(
         self.event_lock = Lock()  # Thread-safe access to events
 
         # Initialize robot control
-        robot_type = getattr(env.unwrapped.robot.config, "robot_type", "so101")
-        if "so100" in robot_type or "so101" in robot_type:
-            # Note to be compatible with the rest of the codebase,
-            # we are using the new calibration method for so101 and so100
-            robot_type = "so_new_calibration"
-        self.kinematics = RobotKinematics(robot_type)
+        self.kinematics = RobotKinematics(
+            urdf_path=env.unwrapped.robot.config.urdf_path,
+            ee_frame_name=env.unwrapped.robot.config.ee_frame_name,
+            joint_names=env.unwrapped.robot.config.joint_names,
+        )
         self.leader_torque_enabled = True
         self.prev_leader_gripper = None
 
@@ -1260,14 +1257,14 @@ def _handle_intervention(self, action):
         leader_pos_dict = self.robot_leader.bus.sync_read("Present_Position")
         follower_pos_dict = self.robot_follower.bus.sync_read("Present_Position")
 
-        leader_pos = np.array([leader_pos_dict[name] for name in leader_pos_dict], dtype=np.float32)
-        follower_pos = np.array([follower_pos_dict[name] for name in follower_pos_dict], dtype=np.float32)
+        leader_pos = np.array([leader_pos_dict[name] for name in leader_pos_dict])
+        follower_pos = np.array([follower_pos_dict[name] for name in follower_pos_dict])
 
         self.leader_tracking_error_queue.append(np.linalg.norm(follower_pos[:-1] - leader_pos[:-1]))
 
         # [:3, 3] Last column of the transformation matrix corresponds to the xyz translation
-        leader_ee = self.kinematics.forward_kinematics(leader_pos, frame="gripper_tip")[:3, 3]
-        follower_ee = self.kinematics.forward_kinematics(follower_pos, frame="gripper_tip")[:3, 3]
+        leader_ee = self.kinematics.forward_kinematics(leader_pos)[:3, 3]
+        follower_ee = self.kinematics.forward_kinematics(follower_pos)[:3, 3]
 
         action = np.clip(leader_ee - follower_ee, -self.end_effector_step_sizes, self.end_effector_step_sizes)
         # Normalize the action to the range [-1, 1]
@@ -1293,7 +1290,6 @@ def _handle_intervention(self, action):
                 gripper_action = 1
 
             action = np.append(action, gripper_action)
-            action = torch.from_numpy(action).float()
 
         return action
 
@@ -1342,9 +1338,12 @@ def step(self, action):
         # NOTE:
         obs, reward, terminated, truncated, info = self.env.step(action)
 
+        if isinstance(action, np.ndarray):
+            action = torch.from_numpy(action)
+
         # Add intervention info
         info["is_intervention"] = is_intervention
-        info["action_intervention"] = action if is_intervention else None
+        info["action_intervention"] = action
 
         self.prev_leader_gripper = np.clip(
             self.robot_leader.bus.sync_read("Present_Position")["gripper"],