Divide joinits into several dicts in teleoperator

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]: