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 b1e48387c..afc9e6946 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 cca089e89..6fdf796b0 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 @@ 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.vy": "y", - "mobile_base.vtheta": "theta", + "mobile_base.vx": "vx", + "mobile_base.vy": "vy", + "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(), - 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]: