feat(robots): consolidates bi SO setups

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 \

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

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
 

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

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,

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

src/lerobot/robots/bi_so100_follower/bi_so100_follower.py → 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()

src/lerobot/teleoperators/bi_so100_leader/config_bi_so100_leader.py → 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

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

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

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

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

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