From 0938a1d816c70d689ab6df3298a71572798c15a6 Mon Sep 17 00:00:00 2001 From: Pepijn <138571049+pkooij@users.noreply.github.com> Date: Wed, 16 Jul 2025 17:50:36 +0200 Subject: [PATCH] Feat/add bimanual so100 robot (#1509) --- src/lerobot/record.py | 24 +++ src/lerobot/replay.py | 15 +- .../robots/bi_so100_follower/__init__.py | 18 ++ .../bi_so100_follower/bi_so100_follower.py | 163 ++++++++++++++++++ .../config_bi_so100_follower.py | 39 +++++ .../so100_follower/config_so100_follower.py | 4 +- src/lerobot/robots/utils.py | 4 + src/lerobot/teleoperate.py | 23 +++ .../teleoperators/bi_so100_leader/__init__.py | 18 ++ .../bi_so100_leader/bi_so100_leader.py | 121 +++++++++++++ .../bi_so100_leader/config_bi_so100_leader.py | 26 +++ src/lerobot/teleoperators/utils.py | 4 + 12 files changed, 457 insertions(+), 2 deletions(-) create mode 100644 src/lerobot/robots/bi_so100_follower/__init__.py create mode 100644 src/lerobot/robots/bi_so100_follower/bi_so100_follower.py create mode 100644 src/lerobot/robots/bi_so100_follower/config_bi_so100_follower.py create mode 100644 src/lerobot/teleoperators/bi_so100_leader/__init__.py create mode 100644 src/lerobot/teleoperators/bi_so100_leader/bi_so100_leader.py create mode 100644 src/lerobot/teleoperators/bi_so100_leader/config_bi_so100_leader.py diff --git a/src/lerobot/record.py b/src/lerobot/record.py index 9fc0dc7ed..c8184d40b 100644 --- a/src/lerobot/record.py +++ b/src/lerobot/record.py @@ -33,6 +33,28 @@ python -m lerobot.record \ # <- Policy optional if you want to record with a policy \ # --policy.path=${HF_USER}/my_policy \ ``` + +Example recording with bimanual so100: +```shell +python -m lerobot.record \ + --robot.type=bi_so100_follower \ + --robot.left_arm_port=/dev/tty.usbmodem5A460851411 \ + --robot.right_arm_port=/dev/tty.usbmodem5A460812391 \ + --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} + }' \ + --teleop.type=bi_so100_leader \ + --teleop.left_arm_port=/dev/tty.usbmodem5A460828611 \ + --teleop.right_arm_port=/dev/tty.usbmodem5A460826981 \ + --teleop.id=bimanual_leader \ + --display_data=true \ + --dataset.repo_id=${HF_USER}/bimanual-so100-handover-cube \ + --dataset.num_episodes=25 \ + --dataset.single_task="Grab and handover the red cube to the other arm" +``` """ import logging @@ -57,6 +79,7 @@ from lerobot.policies.pretrained import PreTrainedPolicy from lerobot.robots import ( # noqa: F401 Robot, RobotConfig, + bi_so100_follower, hope_jr, koch_follower, make_robot_from_config, @@ -66,6 +89,7 @@ from lerobot.robots import ( # noqa: F401 from lerobot.teleoperators import ( # noqa: F401 Teleoperator, TeleoperatorConfig, + bi_so100_leader, homunculus, koch_leader, make_teleoperator_from_config, diff --git a/src/lerobot/replay.py b/src/lerobot/replay.py index c51c55cee..afe54d90f 100644 --- a/src/lerobot/replay.py +++ b/src/lerobot/replay.py @@ -15,7 +15,7 @@ """ Replays the actions of an episode from a dataset on a robot. -Example: +Examples: ```shell python -m lerobot.replay \ @@ -25,6 +25,18 @@ python -m lerobot.replay \ --dataset.repo_id=aliberts/record-test \ --dataset.episode=2 ``` + +Example replay with bimanual so100: +```shell +python -m lerobot.replay \ + --robot.type=bi_so100_follower \ + --robot.left_arm_port=/dev/tty.usbmodem5A460851411 \ + --robot.right_arm_port=/dev/tty.usbmodem5A460812391 \ + --robot.id=bimanual_follower \ + --dataset.repo_id=${HF_USER}/bimanual-so100-handover-cube \ + --dataset.episode=0 +``` + """ import logging @@ -39,6 +51,7 @@ from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.robots import ( # noqa: F401 Robot, RobotConfig, + bi_so100_follower, hope_jr, koch_follower, make_robot_from_config, diff --git a/src/lerobot/robots/bi_so100_follower/__init__.py b/src/lerobot/robots/bi_so100_follower/__init__.py new file mode 100644 index 000000000..90f56516b --- /dev/null +++ b/src/lerobot/robots/bi_so100_follower/__init__.py @@ -0,0 +1,18 @@ +#!/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_follower import BiSO100Follower +from .config_bi_so100_follower import BiSO100FollowerConfig diff --git a/src/lerobot/robots/bi_so100_follower/bi_so100_follower.py b/src/lerobot/robots/bi_so100_follower/bi_so100_follower.py new file mode 100644 index 000000000..7992b79fd --- /dev/null +++ b/src/lerobot/robots/bi_so100_follower/bi_so100_follower.py @@ -0,0 +1,163 @@ +#!/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. + +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.so100_follower import SO100Follower +from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig + +from ..robot import Robot +from .config_bi_so100_follower import BiSO100FollowerConfig + +logger = logging.getLogger(__name__) + + +class BiSO100Follower(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. + """ + + config_class = BiSO100FollowerConfig + name = "bi_so100_follower" + + def __init__(self, config: BiSO100FollowerConfig): + super().__init__(config) + self.config = config + + left_arm_config = SO100FollowerConfig( + 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={}, + ) + + right_arm_config = SO100FollowerConfig( + 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={}, + ) + + self.left_arm = SO100Follower(left_arm_config) + self.right_arm = SO100Follower(right_arm_config) + self.cameras = make_cameras_from_configs(config.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 + } + + @property + def _cameras_ft(self) -> dict[str, tuple]: + return { + cam: (self.config.cameras[cam].height, self.config.cameras[cam].width, 3) for cam in self.cameras + } + + @cached_property + def observation_features(self) -> dict[str, type | tuple]: + return {**self._motors_ft, **self._cameras_ft} + + @cached_property + def action_features(self) -> dict[str, type]: + return self._motors_ft + + @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()) + ) + + 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 + + def calibrate(self) -> None: + self.left_arm.calibrate() + self.right_arm.calibrate() + + def configure(self) -> None: + self.left_arm.configure() + self.right_arm.configure() + + def setup_motors(self) -> None: + self.left_arm.setup_motors() + self.right_arm.setup_motors() + + def get_observation(self) -> dict[str, Any]: + obs_dict = {} + + # Add "left_" prefix + left_obs = self.left_arm.get_observation() + obs_dict.update({f"left_{key}": value for key, value in left_obs.items()}) + + # Add "right_" prefix + 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]: + # Remove "left_" prefix + left_action = { + key.removeprefix("left_"): value for key, value in action.items() if key.startswith("left_") + } + # Remove "right_" prefix + right_action = { + 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) + + # 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()} + + return {**prefixed_send_action_left, **prefixed_send_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/robots/bi_so100_follower/config_bi_so100_follower.py b/src/lerobot/robots/bi_so100_follower/config_bi_so100_follower.py new file mode 100644 index 000000000..00643b85f --- /dev/null +++ b/src/lerobot/robots/bi_so100_follower/config_bi_so100_follower.py @@ -0,0 +1,39 @@ +#!/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: int | None = None + left_arm_use_degrees: bool = False + right_arm_disable_torque_on_disconnect: bool = True + right_arm_max_relative_target: int | 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/so100_follower/config_so100_follower.py b/src/lerobot/robots/so100_follower/config_so100_follower.py index 7cd23d340..ea8b9f1c2 100644 --- a/src/lerobot/robots/so100_follower/config_so100_follower.py +++ b/src/lerobot/robots/so100_follower/config_so100_follower.py @@ -1,4 +1,6 @@ -# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +#!/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. diff --git a/src/lerobot/robots/utils.py b/src/lerobot/robots/utils.py index 911d40465..7486ee499 100644 --- a/src/lerobot/robots/utils.py +++ b/src/lerobot/robots/utils.py @@ -57,6 +57,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 + + return BiSO100Follower(config) elif config.type == "mock_robot": from tests.mocks.mock_robot import MockRobot diff --git a/src/lerobot/teleoperate.py b/src/lerobot/teleoperate.py index 168f898c4..9836f1393 100644 --- a/src/lerobot/teleoperate.py +++ b/src/lerobot/teleoperate.py @@ -28,6 +28,27 @@ python -m lerobot.teleoperate \ --teleop.id=blue \ --display_data=true ``` + +Example teleoperation with bimanual so100: + +```shell +python -m lerobot.teleoperate \ + --robot.type=bi_so100_follower \ + --robot.left_arm_port=/dev/tty.usbmodem5A460851411 \ + --robot.right_arm_port=/dev/tty.usbmodem5A460812391 \ + --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} + }' \ + --teleop.type=bi_so100_leader \ + --teleop.left_arm_port=/dev/tty.usbmodem5A460828611 \ + --teleop.right_arm_port=/dev/tty.usbmodem5A460826981 \ + --teleop.id=bimanual_leader \ + --display_data=true +``` + """ import logging @@ -43,6 +64,7 @@ from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraCon from lerobot.robots import ( # noqa: F401 Robot, RobotConfig, + bi_so100_follower, hope_jr, koch_follower, make_robot_from_config, @@ -52,6 +74,7 @@ from lerobot.robots import ( # noqa: F401 from lerobot.teleoperators import ( # noqa: F401 Teleoperator, TeleoperatorConfig, + bi_so100_leader, gamepad, homunculus, koch_leader, diff --git a/src/lerobot/teleoperators/bi_so100_leader/__init__.py b/src/lerobot/teleoperators/bi_so100_leader/__init__.py new file mode 100644 index 000000000..34313a61e --- /dev/null +++ b/src/lerobot/teleoperators/bi_so100_leader/__init__.py @@ -0,0 +1,18 @@ +#!/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/bi_so100_leader/bi_so100_leader.py b/src/lerobot/teleoperators/bi_so100_leader/bi_so100_leader.py new file mode 100644 index 000000000..769669655 --- /dev/null +++ b/src/lerobot/teleoperators/bi_so100_leader/bi_so100_leader.py @@ -0,0 +1,121 @@ +#!/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. + +import logging +from functools import cached_property + +from lerobot.teleoperators.so100_leader.config_so100_leader import SO100LeaderConfig +from lerobot.teleoperators.so100_leader.so100_leader import SO100Leader + +from ..teleoperator import Teleoperator +from .config_bi_so100_leader import BiSO100LeaderConfig + +logger = logging.getLogger(__name__) + + +class BiSO100Leader(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. + """ + + config_class = BiSO100LeaderConfig + name = "bi_so100_leader" + + def __init__(self, config: BiSO100LeaderConfig): + super().__init__(config) + self.config = config + + left_arm_config = SO100LeaderConfig( + id=f"{config.id}_left" if config.id else None, + calibration_dir=config.calibration_dir, + port=config.left_arm_port, + ) + + right_arm_config = SO100LeaderConfig( + id=f"{config.id}_right" if config.id else None, + calibration_dir=config.calibration_dir, + port=config.right_arm_port, + ) + + self.left_arm = SO100Leader(left_arm_config) + self.right_arm = SO100Leader(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 + } + + @cached_property + def feedback_features(self) -> dict[str, type]: + return {} + + @property + def is_connected(self) -> bool: + 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) + + @property + def is_calibrated(self) -> bool: + return self.left_arm.is_calibrated and self.right_arm.is_calibrated + + def calibrate(self) -> None: + self.left_arm.calibrate() + self.right_arm.calibrate() + + def configure(self) -> None: + self.left_arm.configure() + self.right_arm.configure() + + def setup_motors(self) -> None: + self.left_arm.setup_motors() + self.right_arm.setup_motors() + + def get_action(self) -> dict[str, float]: + action_dict = {} + + # Add "left_" prefix + left_action = self.left_arm.get_action() + action_dict.update({f"left_{key}": value for key, value in left_action.items()}) + + # Add "right_" prefix + right_action = self.right_arm.get_action() + action_dict.update({f"right_{key}": value for key, value in right_action.items()}) + + 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) + + def disconnect(self) -> None: + self.left_arm.disconnect() + self.right_arm.disconnect() diff --git a/src/lerobot/teleoperators/bi_so100_leader/config_bi_so100_leader.py b/src/lerobot/teleoperators/bi_so100_leader/config_bi_so100_leader.py new file mode 100644 index 000000000..117e09913 --- /dev/null +++ b/src/lerobot/teleoperators/bi_so100_leader/config_bi_so100_leader.py @@ -0,0 +1,26 @@ +#!/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 + +from ..config import TeleoperatorConfig + + +@TeleoperatorConfig.register_subclass("bi_so100_leader") +@dataclass +class BiSO100LeaderConfig(TeleoperatorConfig): + left_arm_port: str + right_arm_port: str diff --git a/src/lerobot/teleoperators/utils.py b/src/lerobot/teleoperators/utils.py index 8a667fd41..344a95d72 100644 --- a/src/lerobot/teleoperators/utils.py +++ b/src/lerobot/teleoperators/utils.py @@ -61,5 +61,9 @@ 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 + + return BiSO100Leader(config) else: raise ValueError(config.type)