Divide joinits into several dicts in teleoperator

This commit is contained in:
glannuzel
2025-08-08 14:22:49 +02:00
parent 8236f7b52b
commit cba73399f1
2 changed files with 39 additions and 13 deletions
@@ -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
@@ -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]: