Files
lerobot/tests/robots/test_rebot_b601_follower.py
T
Pepijn 462a8aea90 feat(robots): natively integrate Seeed Studio reBot B601-DM arm
Add first-class LeRobot support for the Seeed Studio reBot arm, replacing
the out-of-tree `lerobot-robot-seeed-b601` / `lerobot-teleoperator-rebot-arm-102`
plugin packages.

New devices:
- robot `rebot_b601_follower` — single-arm B601-DM follower (6-DOF + gripper,
  Damiao CAN motors via `motorbridge`)
- robot `bi_rebot_b601_follower` — bimanual follower composing two single arms
- teleoperator `rebot_102_leader` — single-arm StarArm102 / reBot Arm 102 leader
  (FashionStar UART servos via `motorbridge-smart-servo`)
- teleoperator `bi_rebot_102_leader` — bimanual leader composing two single arms

The bimanual variants reuse the single-arm classes and namespace each arm's
observation/action keys with `left_` / `right_` prefixes, so a bimanual
StarArm102 leader can teleoperate a bimanual reBot B601 follower.

Optional SDK imports are guarded; a `rebot` extra installs `motorbridge` and
`motorbridge-smart-servo`.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-18 11:33:46 +02:00

117 lines
4.2 KiB
Python

#!/usr/bin/env python
# Copyright 2026 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 math
from unittest.mock import MagicMock, patch
import pytest
from lerobot.robots.bi_rebot_b601_follower import BiRebotB601Follower, BiRebotB601FollowerConfig
from lerobot.robots.rebot_b601_follower import (
RebotB601Follower,
RebotB601FollowerConfig,
RebotB601FollowerRobotConfig,
)
_MODULE = "lerobot.robots.rebot_b601_follower.rebot_b601_follower"
def _make_motor_mock(position_rad: float = 0.0) -> MagicMock:
motor = MagicMock(name="MotorMock")
state = MagicMock()
state.pos = position_rad
motor.get_state.return_value = state
return motor
def _make_bus_mock() -> MagicMock:
bus = MagicMock(name="MotorBridgeControllerMock")
# add_damiao_motor returns a fresh motor mock; position encodes the call order.
bus._motor_count = 0
def _add_motor(_send_id, _recv_id, _model):
bus._motor_count += 1
return _make_motor_mock(position_rad=math.radians(bus._motor_count))
bus.add_damiao_motor.side_effect = _add_motor
return bus
@pytest.fixture
def follower():
bus_mock = _make_bus_mock()
with (
patch(f"{_MODULE}.require_package", lambda *a, **kw: None),
patch(f"{_MODULE}.MotorBridgeController") as controller_cls,
patch(f"{_MODULE}.MotorBridgeMode", MagicMock()),
):
controller_cls.from_dm_serial.return_value = bus_mock
cfg = RebotB601FollowerRobotConfig(port="/dev/null")
robot = RebotB601Follower(cfg)
robot.connect(calibrate=False)
yield robot
if robot.is_connected:
robot.disconnect()
def test_features_match_joints():
with patch(f"{_MODULE}.require_package", lambda *a, **kw: None):
robot = RebotB601Follower(RebotB601FollowerRobotConfig(port="/dev/null"))
expected = {f"{m}.pos" for m in robot.motor_names}
assert set(robot.action_features) == expected
assert set(robot.observation_features) == expected
assert "gripper.pos" in expected
def test_connect_disconnect(follower):
assert follower.is_connected
follower.disconnect()
assert not follower.is_connected
def test_get_observation_converts_to_degrees(follower):
obs = follower.get_observation()
assert set(obs) == {f"{m}.pos" for m in follower.motor_names}
# The bus mock seeds each motor's position with its 1-indexed creation order (radians).
for idx, motor in enumerate(follower.motor_names, 1):
assert obs[f"{motor}.pos"] == pytest.approx(math.degrees(math.radians(idx)))
def test_send_action_clips_to_joint_limits(follower):
# shoulder_pan limit is (-145, 145); request beyond the upper bound.
returned = follower.send_action({"shoulder_pan.pos": 999.0})
assert returned["shoulder_pan.pos"] == 145.0
follower.motors["shoulder_pan"].send_pos_vel.assert_called_once()
def test_send_action_routes_gripper_to_force_pos(follower):
follower.send_action({"gripper.pos": -10.0})
follower.motors["gripper"].send_force_pos.assert_called_once()
follower.motors["gripper"].send_pos_vel.assert_not_called()
def test_bimanual_prefixes_features():
with patch(f"{_MODULE}.require_package", lambda *a, **kw: None):
cfg = BiRebotB601FollowerConfig(
left_arm_config=RebotB601FollowerConfig(port="/dev/null0"),
right_arm_config=RebotB601FollowerConfig(port="/dev/null1"),
)
robot = BiRebotB601Follower(cfg)
assert any(k.startswith("left_") for k in robot.action_features)
assert any(k.startswith("right_") for k in robot.action_features)
assert "left_gripper.pos" in robot.action_features
assert "right_gripper.pos" in robot.action_features