diff --git a/src/lerobot/robots/reachy2/configuration_reachy2.py b/src/lerobot/robots/reachy2/configuration_reachy2.py index d79e0a0f87..350a02eaab 100644 --- a/src/lerobot/robots/reachy2/configuration_reachy2.py +++ b/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 diff --git a/src/lerobot/robots/reachy2/robot_reachy2.py b/src/lerobot/robots/reachy2/robot_reachy2.py index 491b5528e3..7e9575e44a 100644 --- a/src/lerobot/robots/reachy2/robot_reachy2.py +++ b/src/lerobot/robots/reachy2/robot_reachy2.py @@ -26,10 +26,18 @@ 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 @@ "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 @@ "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 @@ def __init__(self, config: Reachy2RobotConfig): 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 @@ def camera_features(self) -> dict[str, dict]: def motors_features(self) -> dict: if self.config.with_mobile_base: return {**dict.fromkeys( - REACHY2_JOINTS.keys(), - float, - ), **dict.fromkeys( - REACHY2_VEL.keys(), - float, - )} + self.joints_dict.keys(), + float, + ), **dict.fromkeys( + REACHY2_VEL.keys(), + 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 @@ def is_calibrated(self) -> bool: 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 @@ def send_action(self, action: dict[str, Any]) -> dict[str, Any]: 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"]) diff --git a/src/lerobot/robots/reachy2/test_replay.py b/src/lerobot/robots/reachy2/test_replay.py index c8d1951fed..6bcb3249bc 100644 --- a/src/lerobot/robots/reachy2/test_replay.py +++ b/src/lerobot/robots/reachy2/test_replay.py @@ -44,6 +44,8 @@ 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): diff --git a/src/lerobot/robots/reachy2/test_robot_client.py b/src/lerobot/robots/reachy2/test_robot_client.py new file mode 100644 index 0000000000..41f04fe447 --- /dev/null +++ b/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) diff --git a/src/lerobot/teleoperators/reachy2_fake_teleoperator/config_reachy2_fake_teleoperator.py b/src/lerobot/teleoperators/reachy2_fake_teleoperator/config_reachy2_fake_teleoperator.py index b1e48387c1..afc9e69461 100644 --- a/src/lerobot/teleoperators/reachy2_fake_teleoperator/config_reachy2_fake_teleoperator.py +++ b/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 diff --git a/src/lerobot/teleoperators/reachy2_fake_teleoperator/reachy2_fake_teleoperator.py b/src/lerobot/teleoperators/reachy2_fake_teleoperator/reachy2_fake_teleoperator.py index cca089e89c..6fdf796b09 100644 --- a/src/lerobot/teleoperators/reachy2_fake_teleoperator/reachy2_fake_teleoperator.py +++ b/src/lerobot/teleoperators/reachy2_fake_teleoperator/reachy2_fake_teleoperator.py @@ -32,10 +32,18 @@ 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 @@ "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 @@ "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.vy": "y", - "mobile_base.vtheta": "theta", + "mobile_base.vx": "vx", + "mobile_base.vy": "vy", + "mobile_base.vtheta": "vtheta", } @@ -76,18 +85,31 @@ def __init__(self, config: Reachy2FakeTeleoperatorConfig): 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(), - float, - ), **dict.fromkeys( - REACHY2_VEL.keys(), - float, - )} + self.joints_dict.keys(), + float, + ), **dict.fromkeys( + REACHY2_VEL.keys(), + 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]: