diff --git a/docs/source/groot.mdx b/docs/source/groot.mdx index 729a6465643..84d2702bd01 100644 --- a/docs/source/groot.mdx +++ b/docs/source/groot.mdx @@ -103,7 +103,7 @@ Once you have trained your model using your parameters you can run inference in ```bash lerobot-record \ - --robot.type=bi_so100_follower \ + --robot.type=bi_so_follower \ --robot.left_arm_port=/dev/ttyACM1 \ --robot.right_arm_port=/dev/ttyACM0 \ --robot.id=bimanual_follower \ diff --git a/examples/so100_to_so100_EE/record.py b/examples/so100_to_so100_EE/record.py index db24f4b93ab..eead7a9a828 100644 --- a/examples/so100_to_so100_EE/record.py +++ b/examples/so100_to_so100_EE/record.py @@ -34,8 +34,7 @@ InverseKinematicsEEToJoints, ) from lerobot.scripts.lerobot_record import record_loop -from lerobot.teleoperators.so_leader import SO100LeaderConfig -from lerobot.teleoperators.so_leader.so100_leader import SO100Leader +from lerobot.teleoperators.so_leader import SO100Leader, SO100LeaderConfig from lerobot.utils.control_utils import init_keyboard_listener from lerobot.utils.utils import log_say from lerobot.utils.visualization_utils import init_rerun diff --git a/examples/so100_to_so100_EE/teleoperate.py b/examples/so100_to_so100_EE/teleoperate.py index d520a6eaf6e..71d2899de9d 100644 --- a/examples/so100_to_so100_EE/teleoperate.py +++ b/examples/so100_to_so100_EE/teleoperate.py @@ -29,8 +29,7 @@ ForwardKinematicsJointsToEE, InverseKinematicsEEToJoints, ) -from lerobot.teleoperators.so_leader import SO100LeaderConfig -from lerobot.teleoperators.so_leader.so100_leader import SO100Leader +from lerobot.teleoperators.so_leader import SO100Leader, SO100LeaderConfig from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.visualization_utils import init_rerun, log_rerun_data diff --git a/src/lerobot/async_inference/constants.py b/src/lerobot/async_inference/constants.py index f8b6d7bb303..081db050428 100644 --- a/src/lerobot/async_inference/constants.py +++ b/src/lerobot/async_inference/constants.py @@ -26,4 +26,4 @@ SUPPORTED_POLICIES = ["act", "smolvla", "diffusion", "tdmpc", "vqbet", "pi0", "pi05"] # TODO: Add all other robots -SUPPORTED_ROBOTS = ["so100_follower", "so101_follower", "bi_so100_follower", "omx_follower"] +SUPPORTED_ROBOTS = ["so100_follower", "so101_follower", "bi_so_follower", "omx_follower"] diff --git a/src/lerobot/async_inference/robot_client.py b/src/lerobot/async_inference/robot_client.py index c3668d40bd4..e4e3e0abc84 100644 --- a/src/lerobot/async_inference/robot_client.py +++ b/src/lerobot/async_inference/robot_client.py @@ -51,7 +51,6 @@ from lerobot.robots import ( # noqa: F401 Robot, RobotConfig, - bi_so100_follower, koch_follower, make_robot_from_config, omx_follower, diff --git a/src/lerobot/robots/bi_so100_follower/config_bi_so100_follower.py b/src/lerobot/robots/bi_so100_follower/config_bi_so100_follower.py deleted file mode 100644 index 5806d74159f..00000000000 --- a/src/lerobot/robots/bi_so100_follower/config_bi_so100_follower.py +++ /dev/null @@ -1,39 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2025 The HuggingFace Inc. team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from dataclasses import dataclass, field - -from lerobot.cameras import CameraConfig - -from ..config import RobotConfig - - -@RobotConfig.register_subclass("bi_so100_follower") -@dataclass -class BiSO100FollowerConfig(RobotConfig): - left_arm_port: str - right_arm_port: str - - # Optional - left_arm_disable_torque_on_disconnect: bool = True - left_arm_max_relative_target: float | dict[str, float] | None = None - left_arm_use_degrees: bool = False - right_arm_disable_torque_on_disconnect: bool = True - right_arm_max_relative_target: float | dict[str, float] | None = None - right_arm_use_degrees: bool = False - - # cameras (shared between both arms) - cameras: dict[str, CameraConfig] = field(default_factory=dict) diff --git a/src/lerobot/robots/so_follower/__init__.py b/src/lerobot/robots/so_follower/__init__.py index 82755250cd7..999e7cacc37 100644 --- a/src/lerobot/robots/so_follower/__init__.py +++ b/src/lerobot/robots/so_follower/__init__.py @@ -14,10 +14,9 @@ # See the License for the specific language governing permissions and # limitations under the License. - +from .bi_so_follower.bi_so_follower import BiSOFollower +from .bi_so_follower.config_bi_so_follower import BiSOFollowerConfig from .so100_follower.config_so100_follower import SO100FollowerConfig from .so100_follower.so100_follower import SO100Follower from .so101_follower.config_so101_follower import SO101FollowerConfig from .so101_follower.so101_follower import SO101Follower -from .so_follower_base import SOFollowerBase -from .so_follower_config_base import SOFollowerConfigBase diff --git a/src/lerobot/robots/bi_so100_follower/bi_so100_follower.py b/src/lerobot/robots/so_follower/bi_so_follower/bi_so_follower.py similarity index 52% rename from src/lerobot/robots/bi_so100_follower/bi_so100_follower.py rename to src/lerobot/robots/so_follower/bi_so_follower/bi_so_follower.py index 87a7edcc5de..6d2cbd76357 100644 --- a/src/lerobot/robots/bi_so100_follower/bi_so100_follower.py +++ b/src/lerobot/robots/so_follower/bi_so_follower/bi_so_follower.py @@ -15,66 +15,73 @@ # limitations under the License. import logging -import time from functools import cached_property from typing import Any -from lerobot.cameras.utils import make_cameras_from_configs -from lerobot.robots.so_follower import SO100Follower, SO100FollowerConfig - -from ..robot import Robot -from .config_bi_so100_follower import BiSO100FollowerConfig +from ...robot import Robot +from ..so_follower_base import SOFollowerBase +from ..so_follower_config_base import SOFollowerRobotConfigBase +from .config_bi_so_follower import BiSOFollowerConfig logger = logging.getLogger(__name__) -class BiSO100Follower(Robot): +class BiSOFollower(Robot): """ - [Bimanual SO-100 Follower Arms](https://github.com/TheRobotStudio/SO-ARM100) designed by TheRobotStudio - This bimanual robot can also be easily adapted to use SO-101 follower arms, just replace the SO100Follower class with SO101Follower and SO100FollowerConfig with SO101FollowerConfig. + [Bimanual SO Follower Arms](https://github.com/TheRobotStudio/SO-ARM100) designed by TheRobotStudio """ - config_class = BiSO100FollowerConfig - name = "bi_so100_follower" + config_class = BiSOFollowerConfig + name = "bi_so_follower" - def __init__(self, config: BiSO100FollowerConfig): + def __init__(self, config: BiSOFollowerConfig): super().__init__(config) self.config = config - left_arm_config = SO100FollowerConfig( + left_arm_config = SOFollowerRobotConfigBase( id=f"{config.id}_left" if config.id else None, calibration_dir=config.calibration_dir, - port=config.left_arm_port, - disable_torque_on_disconnect=config.left_arm_disable_torque_on_disconnect, - max_relative_target=config.left_arm_max_relative_target, - use_degrees=config.left_arm_use_degrees, - cameras={}, + port=config.left_arm_config.port, + disable_torque_on_disconnect=config.left_arm_config.disable_torque_on_disconnect, + max_relative_target=config.left_arm_config.max_relative_target, + use_degrees=config.left_arm_config.use_degrees, + cameras=config.left_arm_config.cameras, ) - right_arm_config = SO100FollowerConfig( + right_arm_config = SOFollowerRobotConfigBase( id=f"{config.id}_right" if config.id else None, calibration_dir=config.calibration_dir, - port=config.right_arm_port, - disable_torque_on_disconnect=config.right_arm_disable_torque_on_disconnect, - max_relative_target=config.right_arm_max_relative_target, - use_degrees=config.right_arm_use_degrees, - cameras={}, + port=config.right_arm_config.port, + disable_torque_on_disconnect=config.right_arm_config.disable_torque_on_disconnect, + max_relative_target=config.right_arm_config.max_relative_target, + use_degrees=config.right_arm_config.use_degrees, + cameras=config.right_arm_config.cameras, ) - self.left_arm = SO100Follower(left_arm_config) - self.right_arm = SO100Follower(right_arm_config) - self.cameras = make_cameras_from_configs(config.cameras) + self.left_arm = SOFollowerBase(left_arm_config) + self.right_arm = SOFollowerBase(right_arm_config) + + # Only for compatibility with other parts of the codebase that expect a `robot.cameras` attribute + self.cameras = {**self.left_arm.cameras, **self.right_arm.cameras} @property def _motors_ft(self) -> dict[str, type]: - return {f"left_{motor}.pos": float for motor in self.left_arm.bus.motors} | { - f"right_{motor}.pos": float for motor in self.right_arm.bus.motors + left_arm_motors_ft = self.left_arm._motors_ft + right_arm_motors_ft = self.right_arm._motors_ft + + return { + **{f"left_{k}": v for k, v in left_arm_motors_ft.items()}, + **{f"right_{k}": v for k, v in right_arm_motors_ft.items()}, } @property def _cameras_ft(self) -> dict[str, tuple]: + left_arm_cameras_ft = self.left_arm._cameras_ft + right_arm_cameras_ft = self.right_arm._cameras_ft + return { - cam: (self.config.cameras[cam].height, self.config.cameras[cam].width, 3) for cam in self.cameras + **{f"left_{k}": v for k, v in left_arm_cameras_ft.items()}, + **{f"right_{k}": v for k, v in right_arm_cameras_ft.items()}, } @cached_property @@ -87,19 +94,12 @@ def action_features(self) -> dict[str, type]: @property def is_connected(self) -> bool: - return ( - self.left_arm.bus.is_connected - and self.right_arm.bus.is_connected - and all(cam.is_connected for cam in self.cameras.values()) - ) + return self.left_arm.is_connected and self.right_arm.is_connected def connect(self, calibrate: bool = True) -> None: self.left_arm.connect(calibrate) self.right_arm.connect(calibrate) - for cam in self.cameras.values(): - cam.connect() - @property def is_calibrated(self) -> bool: return self.left_arm.is_calibrated and self.right_arm.is_calibrated @@ -127,12 +127,6 @@ def get_observation(self) -> dict[str, Any]: right_obs = self.right_arm.get_observation() obs_dict.update({f"right_{key}": value for key, value in right_obs.items()}) - for cam_key, cam in self.cameras.items(): - start = time.perf_counter() - obs_dict[cam_key] = cam.async_read() - dt_ms = (time.perf_counter() - start) * 1e3 - logger.debug(f"{self} read {cam_key}: {dt_ms:.1f}ms") - return obs_dict def send_action(self, action: dict[str, Any]) -> dict[str, Any]: @@ -145,18 +139,15 @@ def send_action(self, action: dict[str, Any]) -> dict[str, Any]: key.removeprefix("right_"): value for key, value in action.items() if key.startswith("right_") } - send_action_left = self.left_arm.send_action(left_action) - send_action_right = self.right_arm.send_action(right_action) + sent_action_left = self.left_arm.send_action(left_action) + sent_action_right = self.right_arm.send_action(right_action) # Add prefixes back - prefixed_send_action_left = {f"left_{key}": value for key, value in send_action_left.items()} - prefixed_send_action_right = {f"right_{key}": value for key, value in send_action_right.items()} + prefixed_sent_action_left = {f"left_{key}": value for key, value in sent_action_left.items()} + prefixed_sent_action_right = {f"right_{key}": value for key, value in sent_action_right.items()} - return {**prefixed_send_action_left, **prefixed_send_action_right} + return {**prefixed_sent_action_left, **prefixed_sent_action_right} def disconnect(self): self.left_arm.disconnect() self.right_arm.disconnect() - - for cam in self.cameras.values(): - cam.disconnect() diff --git a/src/lerobot/teleoperators/bi_so100_leader/config_bi_so100_leader.py b/src/lerobot/robots/so_follower/bi_so_follower/config_bi_so_follower.py similarity index 67% rename from src/lerobot/teleoperators/bi_so100_leader/config_bi_so100_leader.py rename to src/lerobot/robots/so_follower/bi_so_follower/config_bi_so_follower.py index 117e099131e..4b2ecc293fb 100644 --- a/src/lerobot/teleoperators/bi_so100_leader/config_bi_so100_leader.py +++ b/src/lerobot/robots/so_follower/bi_so_follower/config_bi_so_follower.py @@ -16,11 +16,14 @@ from dataclasses import dataclass -from ..config import TeleoperatorConfig +from ...config import RobotConfig +from ..so_follower_config_base import SOFollowerConfigBase -@TeleoperatorConfig.register_subclass("bi_so100_leader") +@RobotConfig.register_subclass("bi_so_follower") @dataclass -class BiSO100LeaderConfig(TeleoperatorConfig): - left_arm_port: str - right_arm_port: str +class BiSOFollowerConfig(RobotConfig): + """Configuration class for Bi SO Follower robots.""" + + left_arm_config: SOFollowerConfigBase + right_arm_config: SOFollowerConfigBase diff --git a/src/lerobot/robots/so_follower/so100_follower/config_so100_follower.py b/src/lerobot/robots/so_follower/so100_follower/config_so100_follower.py index 3e8dd7e9ff9..f0ffc634ffc 100644 --- a/src/lerobot/robots/so_follower/so100_follower/config_so100_follower.py +++ b/src/lerobot/robots/so_follower/so100_follower/config_so100_follower.py @@ -17,10 +17,10 @@ from dataclasses import dataclass from ...config import RobotConfig -from ..so_follower_config_base import SOFollowerConfigBase +from ..so_follower_config_base import SOFollowerRobotConfigBase @RobotConfig.register_subclass("so100_follower") @dataclass -class SO100FollowerConfig(SOFollowerConfigBase): +class SO100FollowerConfig(SOFollowerRobotConfigBase): pass diff --git a/src/lerobot/robots/so_follower/so101_follower/config_so101_follower.py b/src/lerobot/robots/so_follower/so101_follower/config_so101_follower.py index 950e8c839cb..b2b9d64835d 100644 --- a/src/lerobot/robots/so_follower/so101_follower/config_so101_follower.py +++ b/src/lerobot/robots/so_follower/so101_follower/config_so101_follower.py @@ -17,10 +17,10 @@ from dataclasses import dataclass from ...config import RobotConfig -from ..so_follower_config_base import SOFollowerConfigBase +from ..so_follower_config_base import SOFollowerRobotConfigBase @RobotConfig.register_subclass("so101_follower") @dataclass -class SO101FollowerConfig(SOFollowerConfigBase): +class SO101FollowerConfig(SOFollowerRobotConfigBase): pass diff --git a/src/lerobot/robots/so_follower/so_follower_base.py b/src/lerobot/robots/so_follower/so_follower_base.py index 5e6e3288821..3caaf8c0138 100644 --- a/src/lerobot/robots/so_follower/so_follower_base.py +++ b/src/lerobot/robots/so_follower/so_follower_base.py @@ -29,7 +29,7 @@ from ..robot import Robot from ..utils import ensure_safe_goal_position -from .so_follower_config_base import SOFollowerConfigBase +from .so_follower_config_base import SOFollowerRobotConfigBase logger = logging.getLogger(__name__) @@ -40,9 +40,10 @@ class SOFollowerBase(Robot): Designed to be subclassed with a per-hardware-model `config_class` and `name`. """ - # `config_class` and `name` should be set by subclasses + config_class = SOFollowerRobotConfigBase + name = "so_follower" - def __init__(self, config: SOFollowerConfigBase): + def __init__(self, config: SOFollowerRobotConfigBase): super().__init__(config) self.config = config # choose normalization mode depending on config if available diff --git a/src/lerobot/robots/so_follower/so_follower_config_base.py b/src/lerobot/robots/so_follower/so_follower_config_base.py index 6e42df2787f..a57a703a381 100644 --- a/src/lerobot/robots/so_follower/so_follower_config_base.py +++ b/src/lerobot/robots/so_follower/so_follower_config_base.py @@ -22,7 +22,7 @@ @dataclass -class SOFollowerConfigBase(RobotConfig): +class SOFollowerConfigBase: """Base configuration class for SO Follower robots.""" # Port to connect to the arm @@ -40,3 +40,8 @@ class SOFollowerConfigBase(RobotConfig): # Set to `True` for backward compatibility with previous policies/dataset use_degrees: bool = False + + +@dataclass +class SOFollowerRobotConfigBase(RobotConfig, SOFollowerConfigBase): + pass diff --git a/src/lerobot/robots/utils.py b/src/lerobot/robots/utils.py index ad6cc3da1e3..9fa4b6f8517 100644 --- a/src/lerobot/robots/utils.py +++ b/src/lerobot/robots/utils.py @@ -52,10 +52,10 @@ def make_robot_from_config(config: RobotConfig) -> Robot: from .hope_jr import HopeJrArm return HopeJrArm(config) - elif config.type == "bi_so100_follower": - from .bi_so100_follower import BiSO100Follower + elif config.type == "bi_so_follower": + from .so_follower import BiSOFollower - return BiSO100Follower(config) + return BiSOFollower(config) elif config.type == "reachy2": from .reachy2 import Reachy2Robot diff --git a/src/lerobot/scripts/lerobot_record.py b/src/lerobot/scripts/lerobot_record.py index a81a5d54edd..213acb48deb 100644 --- a/src/lerobot/scripts/lerobot_record.py +++ b/src/lerobot/scripts/lerobot_record.py @@ -40,21 +40,23 @@ Example recording with bimanual so100: ```shell lerobot-record \ - --robot.type=bi_so100_follower \ - --robot.left_arm_port=/dev/tty.usbmodem5A460851411 \ - --robot.right_arm_port=/dev/tty.usbmodem5A460812391 \ + --robot.type=bi_so_follower \ + --robot.left_arm_config.port=/dev/tty.usbmodem5A460822851 \ + --robot.right_arm_config.port=/dev/tty.usbmodem5A460814411 \ --robot.id=bimanual_follower \ - --robot.cameras='{ - left: {"type": "opencv", "index_or_path": 0, "width": 640, "height": 480, "fps": 30}, - top: {"type": "opencv", "index_or_path": 1, "width": 640, "height": 480, "fps": 30}, - right: {"type": "opencv", "index_or_path": 2, "width": 640, "height": 480, "fps": 30} + --robot.left_arm_config.cameras='{ + wrist: {"type": "opencv", "index_or_path": 1, "width": 640, "height": 480, "fps": 30}, + top: {"type": "opencv", "index_or_path": 3, "width": 640, "height": 480, "fps": 30}, + }' --robot.right_arm_config.cameras='{ + wrist: {"type": "opencv", "index_or_path": 2, "width": 640, "height": 480, "fps": 30}, + front: {"type": "opencv", "index_or_path": 4, "width": 640, "height": 480, "fps": 30}, }' \ - --teleop.type=bi_so100_leader \ - --teleop.left_arm_port=/dev/tty.usbmodem5A460828611 \ - --teleop.right_arm_port=/dev/tty.usbmodem5A460826981 \ + --teleop.type=bi_so_leader \ + --teleop.left_arm_config.port=/dev/tty.usbmodem5A460852721 \ + --teleop.right_arm_config.port=/dev/tty.usbmodem5A460819811 \ --teleop.id=bimanual_leader \ --display_data=true \ - --dataset.repo_id=${HF_USER}/bimanual-so100-handover-cube \ + --dataset.repo_id=${HF_USER}/bimanual-so-handover-cube \ --dataset.num_episodes=25 \ --dataset.single_task="Grab and handover the red cube to the other arm" ``` @@ -94,7 +96,6 @@ from lerobot.robots import ( # noqa: F401 Robot, RobotConfig, - bi_so100_follower, earthrover_mini_plus, hope_jr, koch_follower, @@ -105,7 +106,6 @@ from lerobot.teleoperators import ( # noqa: F401 Teleoperator, TeleoperatorConfig, - bi_so100_leader, homunculus, koch_leader, make_teleoperator_from_config, diff --git a/src/lerobot/scripts/lerobot_replay.py b/src/lerobot/scripts/lerobot_replay.py index af7c6336513..7808e2b88fb 100644 --- a/src/lerobot/scripts/lerobot_replay.py +++ b/src/lerobot/scripts/lerobot_replay.py @@ -29,7 +29,7 @@ Example replay with bimanual so100: ```shell lerobot-replay \ - --robot.type=bi_so100_follower \ + --robot.type=bi_so_follower \ --robot.left_arm_port=/dev/tty.usbmodem5A460851411 \ --robot.right_arm_port=/dev/tty.usbmodem5A460812391 \ --robot.id=bimanual_follower \ @@ -53,7 +53,6 @@ from lerobot.robots import ( # noqa: F401 Robot, RobotConfig, - bi_so100_follower, earthrover_mini_plus, hope_jr, koch_follower, diff --git a/src/lerobot/scripts/lerobot_teleoperate.py b/src/lerobot/scripts/lerobot_teleoperate.py index 2e0724574bf..835f1c69285 100644 --- a/src/lerobot/scripts/lerobot_teleoperate.py +++ b/src/lerobot/scripts/lerobot_teleoperate.py @@ -33,18 +33,18 @@ ```shell lerobot-teleoperate \ - --robot.type=bi_so100_follower \ - --robot.left_arm_port=/dev/tty.usbmodem5A460851411 \ - --robot.right_arm_port=/dev/tty.usbmodem5A460812391 \ + --robot.type=bi_so_follower \ + --robot.left_arm_config.port=/dev/tty.usbmodem5A460822851 \ + --robot.right_arm_config.port=/dev/tty.usbmodem5A460814411 \ --robot.id=bimanual_follower \ - --robot.cameras='{ - left: {"type": "opencv", "index_or_path": 0, "width": 1920, "height": 1080, "fps": 30}, - top: {"type": "opencv", "index_or_path": 1, "width": 1920, "height": 1080, "fps": 30}, - right: {"type": "opencv", "index_or_path": 2, "width": 1920, "height": 1080, "fps": 30} + --robot.left_arm_config.cameras='{ + wrist: {"type": "opencv", "index_or_path": 1, "width": 640, "height": 480, "fps": 30}, + }' --robot.right_arm_config.cameras='{ + wrist: {"type": "opencv", "index_or_path": 2, "width": 640, "height": 480, "fps": 30}, }' \ - --teleop.type=bi_so100_leader \ - --teleop.left_arm_port=/dev/tty.usbmodem5A460828611 \ - --teleop.right_arm_port=/dev/tty.usbmodem5A460826981 \ + --teleop.type=bi_so_leader \ + --teleop.left_arm_config.port=/dev/tty.usbmodem5A460852721 \ + --teleop.right_arm_config.port=/dev/tty.usbmodem5A460819811 \ --teleop.id=bimanual_leader \ --display_data=true ``` @@ -70,7 +70,6 @@ from lerobot.robots import ( # noqa: F401 Robot, RobotConfig, - bi_so100_follower, earthrover_mini_plus, hope_jr, koch_follower, @@ -81,7 +80,6 @@ from lerobot.teleoperators import ( # noqa: F401 Teleoperator, TeleoperatorConfig, - bi_so100_leader, gamepad, homunculus, keyboard, diff --git a/src/lerobot/teleoperators/bi_so100_leader/__init__.py b/src/lerobot/teleoperators/bi_so100_leader/__init__.py deleted file mode 100644 index 34313a61e6b..00000000000 --- a/src/lerobot/teleoperators/bi_so100_leader/__init__.py +++ /dev/null @@ -1,18 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2025 The HuggingFace Inc. team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from .bi_so100_leader import BiSO100Leader -from .config_bi_so100_leader import BiSO100LeaderConfig diff --git a/src/lerobot/teleoperators/so_leader/__init__.py b/src/lerobot/teleoperators/so_leader/__init__.py index a1017f3b9aa..80ac2236da6 100644 --- a/src/lerobot/teleoperators/so_leader/__init__.py +++ b/src/lerobot/teleoperators/so_leader/__init__.py @@ -14,9 +14,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +from .bi_so_leader.bi_so_leader import BiSOLeader +from .bi_so_leader.config_bi_so_leader import BiSOLeaderConfig from .so100_leader.config_so100_leader import SO100LeaderConfig from .so100_leader.so100_leader import SO100Leader from .so101_leader.config_so101_leader import SO101LeaderConfig from .so101_leader.so101_leader import SO101Leader -from .so_leader_base import SOLeaderBase -from .so_leader_config_base import SOLeaderConfigBase diff --git a/src/lerobot/teleoperators/bi_so100_leader/bi_so100_leader.py b/src/lerobot/teleoperators/so_leader/bi_so_leader/bi_so_leader.py similarity index 61% rename from src/lerobot/teleoperators/bi_so100_leader/bi_so100_leader.py rename to src/lerobot/teleoperators/so_leader/bi_so_leader/bi_so_leader.py index 93f66eb2e91..992dfb67e11 100644 --- a/src/lerobot/teleoperators/bi_so100_leader/bi_so100_leader.py +++ b/src/lerobot/teleoperators/so_leader/bi_so_leader/bi_so_leader.py @@ -17,46 +17,49 @@ import logging from functools import cached_property -from lerobot.teleoperators.so_leader import SO100Leader, SO100LeaderConfig - -from ..teleoperator import Teleoperator -from .config_bi_so100_leader import BiSO100LeaderConfig +from ...teleoperator import Teleoperator +from ..so_leader_base import SOLeaderBase +from ..so_leader_config_base import SOLeaderTeleopConfigBase +from .config_bi_so_leader import BiSOLeaderConfig logger = logging.getLogger(__name__) -class BiSO100Leader(Teleoperator): +class BiSOLeader(Teleoperator): """ - [Bimanual SO-100 Leader Arms](https://github.com/TheRobotStudio/SO-ARM100) designed by TheRobotStudio - This bimanual leader arm can also be easily adapted to use SO-101 leader arms, just replace the SO100Leader class with SO101Leader and SO100LeaderConfig with SO101LeaderConfig. + [Bimanual SO Leader Arms](https://github.com/TheRobotStudio/SO-ARM100) designed by TheRobotStudio """ - config_class = BiSO100LeaderConfig - name = "bi_so100_leader" + config_class = BiSOLeaderConfig + name = "bi_so_leader" - def __init__(self, config: BiSO100LeaderConfig): + def __init__(self, config: BiSOLeaderConfig): super().__init__(config) self.config = config - left_arm_config = SO100LeaderConfig( + left_arm_config = SOLeaderTeleopConfigBase( id=f"{config.id}_left" if config.id else None, calibration_dir=config.calibration_dir, - port=config.left_arm_port, + port=config.left_arm_config.port, ) - right_arm_config = SO100LeaderConfig( + right_arm_config = SOLeaderTeleopConfigBase( id=f"{config.id}_right" if config.id else None, calibration_dir=config.calibration_dir, - port=config.right_arm_port, + port=config.right_arm_config.port, ) - self.left_arm = SO100Leader(left_arm_config) - self.right_arm = SO100Leader(right_arm_config) + self.left_arm = SOLeaderBase(left_arm_config) + self.right_arm = SOLeaderBase(right_arm_config) @cached_property def action_features(self) -> dict[str, type]: - return {f"left_{motor}.pos": float for motor in self.left_arm.bus.motors} | { - f"right_{motor}.pos": float for motor in self.right_arm.bus.motors + left_arm_features = self.left_arm.action_features + right_arm_features = self.right_arm.action_features + + return { + **{f"left_{k}": v for k, v in left_arm_features.items()}, + **{f"right_{k}": v for k, v in right_arm_features.items()}, } @cached_property @@ -101,19 +104,8 @@ def get_action(self) -> dict[str, float]: return action_dict def send_feedback(self, feedback: dict[str, float]) -> None: - # Remove "left_" prefix - left_feedback = { - key.removeprefix("left_"): value for key, value in feedback.items() if key.startswith("left_") - } - # Remove "right_" prefix - right_feedback = { - key.removeprefix("right_"): value for key, value in feedback.items() if key.startswith("right_") - } - - if left_feedback: - self.left_arm.send_feedback(left_feedback) - if right_feedback: - self.right_arm.send_feedback(right_feedback) + # TODO: Implement force feedback + raise NotImplementedError def disconnect(self) -> None: self.left_arm.disconnect() diff --git a/src/lerobot/robots/bi_so100_follower/__init__.py b/src/lerobot/teleoperators/so_leader/bi_so_leader/config_bi_so_leader.py similarity index 62% rename from src/lerobot/robots/bi_so100_follower/__init__.py rename to src/lerobot/teleoperators/so_leader/bi_so_leader/config_bi_so_leader.py index 90f56516b69..c2e84a4fb8b 100644 --- a/src/lerobot/robots/bi_so100_follower/__init__.py +++ b/src/lerobot/teleoperators/so_leader/bi_so_leader/config_bi_so_leader.py @@ -14,5 +14,16 @@ # See the License for the specific language governing permissions and # limitations under the License. -from .bi_so100_follower import BiSO100Follower -from .config_bi_so100_follower import BiSO100FollowerConfig +from dataclasses import dataclass + +from ...config import TeleoperatorConfig +from ..so_leader_config_base import SOLeaderConfigBase + + +@TeleoperatorConfig.register_subclass("bi_so_leader") +@dataclass +class BiSOLeaderConfig(TeleoperatorConfig): + """Configuration class for Bi SO Leader teleoperators.""" + + left_arm_config: SOLeaderConfigBase + right_arm_config: SOLeaderConfigBase diff --git a/src/lerobot/teleoperators/so_leader/so100_leader/config_so100_leader.py b/src/lerobot/teleoperators/so_leader/so100_leader/config_so100_leader.py index 092eb0cfc03..4503bc6a79e 100644 --- a/src/lerobot/teleoperators/so_leader/so100_leader/config_so100_leader.py +++ b/src/lerobot/teleoperators/so_leader/so100_leader/config_so100_leader.py @@ -17,10 +17,10 @@ from dataclasses import dataclass from ...config import TeleoperatorConfig -from ..so_leader_config_base import SOLeaderConfigBase +from ..so_leader_config_base import SOLeaderTeleopConfigBase @TeleoperatorConfig.register_subclass("so100_leader") @dataclass -class SO100LeaderConfig(SOLeaderConfigBase): +class SO100LeaderConfig(SOLeaderTeleopConfigBase): pass diff --git a/src/lerobot/teleoperators/so_leader/so101_leader/config_so101_leader.py b/src/lerobot/teleoperators/so_leader/so101_leader/config_so101_leader.py index 1a6bb0df475..1e110d73e85 100644 --- a/src/lerobot/teleoperators/so_leader/so101_leader/config_so101_leader.py +++ b/src/lerobot/teleoperators/so_leader/so101_leader/config_so101_leader.py @@ -17,10 +17,10 @@ from dataclasses import dataclass from ...config import TeleoperatorConfig -from ..so_leader_config_base import SOLeaderConfigBase +from ..so_leader_config_base import SOLeaderTeleopConfigBase @TeleoperatorConfig.register_subclass("so101_leader") @dataclass -class SO101LeaderConfig(SOLeaderConfigBase): +class SO101LeaderConfig(SOLeaderTeleopConfigBase): pass diff --git a/src/lerobot/teleoperators/so_leader/so_leader_base.py b/src/lerobot/teleoperators/so_leader/so_leader_base.py index 6cdaca2e760..01936d8ac05 100644 --- a/src/lerobot/teleoperators/so_leader/so_leader_base.py +++ b/src/lerobot/teleoperators/so_leader/so_leader_base.py @@ -25,7 +25,7 @@ from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError from ..teleoperator import Teleoperator -from .so_leader_config_base import SOLeaderConfigBase +from .so_leader_config_base import SOLeaderTeleopConfigBase logger = logging.getLogger(__name__) @@ -33,7 +33,10 @@ class SOLeaderBase(Teleoperator): """Generic SO leader base for SO-100/101/10X teleoperators.""" - def __init__(self, config: SOLeaderConfigBase): + config_class = SOLeaderTeleopConfigBase + name = "so_leader" + + def __init__(self, config: SOLeaderTeleopConfigBase): super().__init__(config) self.config = config norm_mode_body = MotorNormMode.DEGREES if config.use_degrees else MotorNormMode.RANGE_M100_100 diff --git a/src/lerobot/teleoperators/so_leader/so_leader_config_base.py b/src/lerobot/teleoperators/so_leader/so_leader_config_base.py index b2f2dfcfa77..db58c18d5f5 100644 --- a/src/lerobot/teleoperators/so_leader/so_leader_config_base.py +++ b/src/lerobot/teleoperators/so_leader/so_leader_config_base.py @@ -20,7 +20,7 @@ @dataclass -class SOLeaderConfigBase(TeleoperatorConfig): +class SOLeaderConfigBase: """Base configuration class for SO Leader teleoperators.""" # Port to connect to the arm @@ -28,3 +28,8 @@ class SOLeaderConfigBase(TeleoperatorConfig): # Whether to use degrees for angles use_degrees: bool = False + + +@dataclass +class SOLeaderTeleopConfigBase(TeleoperatorConfig, SOLeaderConfigBase): + pass diff --git a/src/lerobot/teleoperators/utils.py b/src/lerobot/teleoperators/utils.py index 74e43ec9509..5baae36e2f2 100644 --- a/src/lerobot/teleoperators/utils.py +++ b/src/lerobot/teleoperators/utils.py @@ -73,10 +73,10 @@ def make_teleoperator_from_config(config: TeleoperatorConfig) -> Teleoperator: from .homunculus import HomunculusArm return HomunculusArm(config) - elif config.type == "bi_so100_leader": - from .bi_so100_leader import BiSO100Leader + elif config.type == "bi_so_leader": + from .so_leader import BiSOLeader - return BiSO100Leader(config) + return BiSOLeader(config) elif config.type == "reachy2_teleoperator": from .reachy2_teleoperator import Reachy2Teleoperator