Add end effector action space to hil-serl (#861)

Co-authored-by: Adil Zouitine <adilzouitinegm@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

lerobot/scripts/server/find_joint_limits.py

@@ -7,25 +7,26 @@ import numpy as np
 from lerobot.common.robot_devices.control_utils import is_headless
 from lerobot.common.robot_devices.robots.factory import make_robot
 from lerobot.common.utils.utils import init_hydra_config
+from lerobot.scripts.server.kinematics import RobotKinematics
 
 
 def find_joint_bounds(
     robot,
-    control_time_s=20,
+    control_time_s=30,
     display_cameras=False,
 ):
-    # TODO(rcadene): Add option to record logs
     if not robot.is_connected:
         robot.connect()
 
-        control_time_s = float("inf")
-
-    timestamp = 0
     start_episode_t = time.perf_counter()
     pos_list = []
-    while timestamp < control_time_s:
+    while True:
         observation, action = robot.teleop_step(record_data=True)
 
+        # Wait for 5 seconds to stabilize the robot initial position
+        if time.perf_counter() - start_episode_t < 5:
+            continue
+
         pos_list.append(robot.follower_arms["main"].read("Present_Position"))
 
         if display_cameras and not is_headless():
@@ -36,8 +37,7 @@ def find_joint_bounds(
                 )
             cv2.waitKey(1)
 
-        timestamp = time.perf_counter() - start_episode_t
-        if timestamp > 60:
+        if time.perf_counter() - start_episode_t > control_time_s:
             max = np.max(np.stack(pos_list), 0)
             min = np.min(np.stack(pos_list), 0)
             print(f"Max angle position per joint {max}")
@@ -45,6 +45,43 @@ def find_joint_bounds(
             break
 
 
+def find_ee_bounds(
+    robot,
+    control_time_s=30,
+    display_cameras=False,
+):
+    if not robot.is_connected:
+        robot.connect()
+
+    start_episode_t = time.perf_counter()
+    ee_list = []
+    while True:
+        observation, action = robot.teleop_step(record_data=True)
+
+        # Wait for 5 seconds to stabilize the robot initial position
+        if time.perf_counter() - start_episode_t < 5:
+            continue
+
+        joint_positions = robot.follower_arms["main"].read("Present_Position")
+        print(f"Joint positions: {joint_positions}")
+        ee_list.append(RobotKinematics.fk_gripper_tip(joint_positions)[:3, 3])
+
+        if display_cameras and not is_headless():
+            image_keys = [key for key in observation if "image" in key]
+            for key in image_keys:
+                cv2.imshow(
+                    key, cv2.cvtColor(observation[key].numpy(), cv2.COLOR_RGB2BGR)
+                )
+            cv2.waitKey(1)
+
+        if time.perf_counter() - start_episode_t > control_time_s:
+            max = np.max(np.stack(ee_list), 0)
+            min = np.min(np.stack(ee_list), 0)
+            print(f"Max ee position {max}")
+            print(f"Min ee position {min}")
+            break
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument(
@@ -59,14 +96,26 @@ if __name__ == "__main__":
         nargs="*",
         help="Any key=value arguments to override config values (use dots for.nested=overrides)",
     )
+    parser.add_argument(
+        "--mode",
+        type=str,
+        default="joint",
+        choices=["joint", "ee"],
+        help="Mode to run the script in. Can be 'joint' or 'ee'.",
+    )
     parser.add_argument(
         "--control-time-s",
-        type=float,
-        default=20,
-        help="Maximum episode length in seconds",
+        type=int,
+        default=30,
+        help="Time step to use for control.",
     )
     args = parser.parse_args()
     robot_cfg = init_hydra_config(args.robot_path, args.robot_overrides)
 
     robot = make_robot(robot_cfg)
-    find_joint_bounds(robot, control_time_s=args.control_time_s)
+    if args.mode == "joint":
+        find_joint_bounds(robot, args.control_time_s)
+    elif args.mode == "ee":
+        find_ee_bounds(robot, args.control_time_s)
+    if robot.is_connected:
+        robot.disconnect()