Merge pull request #5 from pollen-robotics/separate-parts-in-config

Separate parts in config

src/lerobot/robots/reachy2/configuration_reachy2.py

@@ -31,6 +31,10 @@ class Reachy2RobotConfig(RobotConfig):
     ip_address: str | None = "localhost"
     use_external_commands: bool = False
     with_mobile_base: bool = True
+    with_l_arm: bool = True
+    with_r_arm: bool = True
+    with_neck: bool = True
+    with_antennas: bool = True
 
     mock: bool = False
 

src/lerobot/robots/reachy2/robot_reachy2.py

@@ -26,10 +26,18 @@ from ..robot import Robot
 from .configuration_reachy2 import Reachy2RobotConfig
 
 # {lerobot_keys: reachy2_sdk_keys}
-REACHY2_JOINTS = {
+REACHY2_NECK_JOINTS = {
     "neck_yaw.pos": "head.neck.yaw",
     "neck_pitch.pos": "head.neck.pitch",
     "neck_roll.pos": "head.neck.roll",
+}
+
+REACHY2_ANTENNAS_JOINTS = {
+    "l_antenna.pos": "head.l_antenna",
+    "r_antenna.pos": "head.r_antenna",
+}
+
+REACHY2_R_ARM_JOINTS = {
     "r_shoulder_pitch.pos": "r_arm.shoulder.pitch",
     "r_shoulder_roll.pos": "r_arm.shoulder.roll",
     "r_elbow_yaw.pos": "r_arm.elbow.yaw",
@@ -38,6 +46,9 @@ REACHY2_JOINTS = {
     "r_wrist_pitch.pos": "r_arm.wrist.pitch",
     "r_wrist_yaw.pos": "r_arm.wrist.yaw",
     "r_gripper.pos": "r_arm.gripper",
+}
+
+REACHY2_L_ARM_JOINTS = {
     "l_shoulder_pitch.pos": "l_arm.shoulder.pitch",
     "l_shoulder_roll.pos": "l_arm.shoulder.roll",
     "l_elbow_yaw.pos": "l_arm.elbow.yaw",
@@ -46,8 +57,6 @@ REACHY2_JOINTS = {
     "l_wrist_pitch.pos": "l_arm.wrist.pitch",
     "l_wrist_yaw.pos": "l_arm.wrist.yaw",
     "l_gripper.pos": "l_arm.gripper",
-    "l_antenna.pos": "head.l_antenna",
-    "r_antenna.pos": "head.r_antenna",
 }
 
 REACHY2_VEL = {
@@ -77,6 +86,9 @@ class Reachy2Robot(Robot):
 
         self.logs = {}
 
+        self.joints_dict: dict[str, str] = {}
+        self.generate_joints_dict()
+
     @property
     def observation_features(self) -> dict:
         return {**self.motors_features, **self.camera_features}
@@ -95,14 +107,14 @@ class Reachy2Robot(Robot):
     def motors_features(self) -> dict:
         if self.config.with_mobile_base:
             return {**dict.fromkeys(
-                REACHY2_JOINTS.keys(),
+                        self.joints_dict.keys(),
-                float,
+                        float,
-            ), **dict.fromkeys(
+                    ), **dict.fromkeys(
-                REACHY2_VEL.keys(),
+                        REACHY2_VEL.keys(),
-                float,
+                        float,
-            )}
+                    )}
         else:
-            return dict.fromkeys(REACHY2_JOINTS.keys(), float)
+            return dict.fromkeys(self.joints_dict.keys(), float)
 
     @property
     def is_connected(self) -> bool:
@@ -130,8 +142,18 @@ class Reachy2Robot(Robot):
     def calibrate(self) -> None:
         pass
 
+    def generate_joints_dict(self) -> dict[str, str]:
+        if self.config.with_neck:
+            self.joints_dict.update(REACHY2_NECK_JOINTS)
+        if self.config.with_l_arm:
+            self.joints_dict.update(REACHY2_L_ARM_JOINTS)
+        if self.config.with_r_arm:
+            self.joints_dict.update(REACHY2_R_ARM_JOINTS)
+        if self.config.with_antennas:
+            self.joints_dict.update(REACHY2_ANTENNAS_JOINTS)
+
     def _get_state(self) -> dict:
-        pos_dict = {k: self.reachy.joints[v].present_position for k, v in REACHY2_JOINTS.items()}
+        pos_dict = {k: self.reachy.joints[v].present_position for k, v in self.joints_dict.items()}
         if not self.config.with_mobile_base:
             return pos_dict
         vel_dict = {k: self.reachy.mobile_base.odometry[v] for k, v in REACHY2_VEL.items()}
@@ -159,13 +181,13 @@ class Reachy2Robot(Robot):
 
         vel = {}
         for key, val in action.items():
-            if key not in REACHY2_JOINTS:
+            if key not in self.joints_dict:
                 if key not in REACHY2_VEL:
                     raise KeyError(f"Key '{key}' is not a valid motor key in Reachy 2.")
                 else:
                     vel[REACHY2_VEL[key]] = val
             else:
-                self.reachy.joints[REACHY2_JOINTS[key]].goal_position = val
+                self.reachy.joints[self.joints_dict[key]].goal_position = val
 
         if self.config.with_mobile_base:
             self.reachy.mobile_base.set_goal_speed(vel["vx"], vel["vy"], vel["vtheta"])

src/lerobot/robots/reachy2/test_replay.py

@@ -44,6 +44,8 @@ l_arm_goal = [action["l_shoulder_pitch.pos"],
 
 robot.reachy.head.goto(neck_goal)
 robot.reachy.r_arm.goto(r_arm_goal)
+robot.reachy.r_arm.gripper.goto(100.0, percentage=True)
+robot.reachy.l_arm.gripper.goto(100.0, percentage=True)
 robot.reachy.l_arm.goto(l_arm_goal, wait=True)
 
 for idx in range(dataset.num_frames):

src/lerobot/robots/reachy2/test_robot_client.py

@@ -0,0 +1,47 @@
+import threading
+from lerobot.scripts.server.configs import RobotClientConfig
+from lerobot.robots.reachy2 import Reachy2Robot, Reachy2RobotConfig
+
+from lerobot.scripts.server.robot_client import RobotClient
+from lerobot.scripts.server.helpers import visualize_action_queue_size
+
+
+robot_config = Reachy2RobotConfig(
+    ip_address="192.168.0.199",
+    id="reachy2-pvt02",
+    with_mobile_base=False,
+    with_l_arm=False,
+    with_neck=False,
+    with_antennas=False,
+)
+
+# 3. Create client configuration
+client_cfg = RobotClientConfig(
+    robot=robot_config,
+    server_address="localhost:8080",
+    policy_device="cuda",
+    policy_type="act",
+    pretrained_name_or_path="CarolinePascal/pick_and_place_bottle",
+    chunk_size_threshold=0.5,
+    actions_per_chunk=20,  # make sure this is less than the max actions of the policy
+)
+
+# 4. Create and start client
+client = RobotClient(client_cfg)
+
+# 5. Specify the task
+task = "Don't do anything, stay still"
+
+if client.start():
+    # Start action receiver thread
+    action_receiver_thread = threading.Thread(target=client.receive_actions, daemon=True)
+    action_receiver_thread.start()
+
+    try:
+        # Run the control loop
+        client.control_loop(task)
+    except KeyboardInterrupt:
+        client.stop()
+        action_receiver_thread.join()
+        # (Optionally) plot the action queue size
+        visualize_action_queue_size(client.action_queue_size)

src/lerobot/teleoperators/reachy2_fake_teleoperator/config_reachy2_fake_teleoperator.py

@@ -25,3 +25,7 @@ class Reachy2FakeTeleoperatorConfig(TeleoperatorConfig):
     # Port to connect to the arm
     ip_address: str | None = "localhost"
     with_mobile_base: bool = True
+    with_l_arm: bool = True
+    with_r_arm: bool = True
+    with_neck: bool = True
+    with_antennas: bool = True

src/lerobot/teleoperators/reachy2_fake_teleoperator/reachy2_fake_teleoperator.py

@@ -32,10 +32,18 @@ from .config_reachy2_fake_teleoperator import Reachy2FakeTeleoperatorConfig
 logger = logging.getLogger(__name__)
 
 # {lerobot_keys: reachy2_sdk_keys}
-REACHY2_JOINTS = {
+REACHY2_NECK_JOINTS = {
     "neck_yaw.pos": "head.neck.yaw",
     "neck_pitch.pos": "head.neck.pitch",
     "neck_roll.pos": "head.neck.roll",
+}
+
+REACHY2_ANTENNAS_JOINTS = {
+    "l_antenna.pos": "head.l_antenna",
+    "r_antenna.pos": "head.r_antenna",
+}
+
+REACHY2_R_ARM_JOINTS = {
     "r_shoulder_pitch.pos": "r_arm.shoulder.pitch",
     "r_shoulder_roll.pos": "r_arm.shoulder.roll",
     "r_elbow_yaw.pos": "r_arm.elbow.yaw",
@@ -44,6 +52,9 @@ REACHY2_JOINTS = {
     "r_wrist_pitch.pos": "r_arm.wrist.pitch",
     "r_wrist_yaw.pos": "r_arm.wrist.yaw",
     "r_gripper.pos": "r_arm.gripper",
+}
+
+REACHY2_L_ARM_JOINTS = {
     "l_shoulder_pitch.pos": "l_arm.shoulder.pitch",
     "l_shoulder_roll.pos": "l_arm.shoulder.roll",
     "l_elbow_yaw.pos": "l_arm.elbow.yaw",
@@ -52,14 +63,12 @@ REACHY2_JOINTS = {
     "l_wrist_pitch.pos": "l_arm.wrist.pitch",
     "l_wrist_yaw.pos": "l_arm.wrist.yaw",
     "l_gripper.pos": "l_arm.gripper",
-    "l_antenna.pos": "head.l_antenna",
-    "r_antenna.pos": "head.r_antenna",
 }
 
 REACHY2_VEL = {
-    "mobile_base.vx": "x",
+    "mobile_base.vx": "vx",
-    "mobile_base.vy": "y",
+    "mobile_base.vy": "vy",
-    "mobile_base.vtheta": "theta",
+    "mobile_base.vtheta": "vtheta",
 }
 
 
@@ -76,18 +85,31 @@ class Reachy2FakeTeleoperator(Teleoperator):
         self.config = config
         self.reachy: None | ReachySDK = None
 
+        self.joints_dict: dict[str, str] = {}
+        self.generate_joints_dict()
+
+    def generate_joints_dict(self) -> dict[str, str]:
+        if self.config.with_neck:
+            self.joints_dict.update(REACHY2_NECK_JOINTS)
+        if self.config.with_l_arm:
+            self.joints_dict.update(REACHY2_L_ARM_JOINTS)
+        if self.config.with_r_arm:
+            self.joints_dict.update(REACHY2_R_ARM_JOINTS)
+        if self.config.with_antennas:
+            self.joints_dict.update(REACHY2_ANTENNAS_JOINTS)
+
     @property
     def action_features(self) -> dict[str, type]:
         if self.config.with_mobile_base:
             return {**dict.fromkeys(
-                REACHY2_JOINTS.keys(),
+                        self.joints_dict.keys(),
-                float,
+                        float,
-            ), **dict.fromkeys(
+                    ), **dict.fromkeys(
-                REACHY2_VEL.keys(),
+                        REACHY2_VEL.keys(),
-                float,
+                        float,
-            )}
+                    )}
         else:
-            return dict.fromkeys(REACHY2_JOINTS.keys(), float)
+            return dict.fromkeys(self.joints_dict.keys(), float)
 
     @property
     def feedback_features(self) -> dict[str, type]: