Add real robot devices and scripts to control real robot (#288)

Co-authored-by: Simon Alibert <alibert.sim@gmail.com>

tests/conftest.py

@@ -13,8 +13,25 @@
 # 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 pytest
+
 from .utils import DEVICE
 
 
 def pytest_collection_finish():
     print(f"\nTesting with {DEVICE=}")
+
+
+@pytest.fixture(scope="session")
+def is_koch_available():
+    try:
+        from lerobot.common.robot_devices.robots.factory import make_robot
+
+        robot = make_robot("koch")
+        robot.connect()
+        del robot
+        return True
+    except Exception as e:
+        print("An alexander koch robot is not available.")
+        print(e)
+        return False

tests/test_cameras.py

@@ -0,0 +1,125 @@
+import numpy as np
+import pytest
+
+from lerobot.common.robot_devices.cameras.opencv import OpenCVCamera, save_images_from_cameras
+from lerobot.common.robot_devices.utils import RobotDeviceAlreadyConnectedError, RobotDeviceNotConnectedError
+from tests.utils import require_koch
+
+CAMERA_INDEX = 2
+# Maximum absolute difference between two consecutive images recored by a camera.
+# This value differs with respect to the camera.
+MAX_PIXEL_DIFFERENCE = 25
+
+
+def compute_max_pixel_difference(first_image, second_image):
+    return np.abs(first_image.astype(float) - second_image.astype(float)).max()
+
+
+@require_koch
+def test_camera(request):
+    """Test assumes that `camera.read()` returns the same image when called multiple times in a row.
+    So the environment should not change (you shouldnt be in front of the camera) and the camera should not be moving.
+
+    Warning: The tests worked for a macbookpro camera, but I am getting assertion error (`np.allclose(color_image, async_color_image)`)
+    for my iphone camera and my LG monitor camera.
+    """
+    # TODO(rcadene): measure fps in nightly?
+    # TODO(rcadene): test logs
+    # TODO(rcadene): add compatibility with other camera APIs
+
+    # Test instantiating
+    camera = OpenCVCamera(CAMERA_INDEX)
+
+    # Test reading, async reading, disconnecting before connecting raises an error
+    with pytest.raises(RobotDeviceNotConnectedError):
+        camera.read()
+    with pytest.raises(RobotDeviceNotConnectedError):
+        camera.async_read()
+    with pytest.raises(RobotDeviceNotConnectedError):
+        camera.disconnect()
+
+    # Test deleting the object without connecting first
+    del camera
+
+    # Test connecting
+    camera = OpenCVCamera(CAMERA_INDEX)
+    camera.connect()
+    assert camera.is_connected
+    assert camera.fps is not None
+    assert camera.width is not None
+    assert camera.height is not None
+
+    # Test connecting twice raises an error
+    with pytest.raises(RobotDeviceAlreadyConnectedError):
+        camera.connect()
+
+    # Test reading from the camera
+    color_image = camera.read()
+    assert isinstance(color_image, np.ndarray)
+    assert color_image.ndim == 3
+    h, w, c = color_image.shape
+    assert c == 3
+    assert w > h
+
+    # Test read and async_read outputs similar images
+    # ...warming up as the first frames can be black
+    for _ in range(30):
+        camera.read()
+    color_image = camera.read()
+    async_color_image = camera.async_read()
+    print(
+        "max_pixel_difference between read() and async_read()",
+        compute_max_pixel_difference(color_image, async_color_image),
+    )
+    assert np.allclose(color_image, async_color_image, rtol=1e-5, atol=MAX_PIXEL_DIFFERENCE)
+
+    # Test disconnecting
+    camera.disconnect()
+    assert camera.camera is None
+    assert camera.thread is None
+
+    # Test disconnecting with `__del__`
+    camera = OpenCVCamera(CAMERA_INDEX)
+    camera.connect()
+    del camera
+
+    # Test acquiring a bgr image
+    camera = OpenCVCamera(CAMERA_INDEX, color_mode="bgr")
+    camera.connect()
+    assert camera.color_mode == "bgr"
+    bgr_color_image = camera.read()
+    assert np.allclose(color_image, bgr_color_image[:, :, [2, 1, 0]], rtol=1e-5, atol=MAX_PIXEL_DIFFERENCE)
+    del camera
+
+    # TODO(rcadene): Add a test for a camera that doesnt support fps=60 and raises an OSError
+    # TODO(rcadene): Add a test for a camera that supports fps=60
+
+    # Test fps=10 raises an OSError
+    camera = OpenCVCamera(CAMERA_INDEX, fps=10)
+    with pytest.raises(OSError):
+        camera.connect()
+    del camera
+
+    # Test width and height can be set
+    camera = OpenCVCamera(CAMERA_INDEX, fps=30, width=1280, height=720)
+    camera.connect()
+    assert camera.fps == 30
+    assert camera.width == 1280
+    assert camera.height == 720
+    color_image = camera.read()
+    h, w, c = color_image.shape
+    assert h == 720
+    assert w == 1280
+    assert c == 3
+    del camera
+
+    # Test not supported width and height raise an error
+    camera = OpenCVCamera(CAMERA_INDEX, fps=30, width=0, height=0)
+    with pytest.raises(OSError):
+        camera.connect()
+    del camera
+
+
+@require_koch
+def test_save_images_from_cameras(tmpdir, request):
+    save_images_from_cameras(tmpdir, record_time_s=1)

tests/test_control_robot.py

@@ -0,0 +1,48 @@
+from pathlib import Path
+
+from lerobot.common.policies.factory import make_policy
+from lerobot.common.robot_devices.robots.factory import make_robot
+from lerobot.common.utils.utils import init_hydra_config
+from lerobot.scripts.control_robot import record_dataset, replay_episode, run_policy, teleoperate
+from tests.utils import DEFAULT_CONFIG_PATH, DEVICE, require_koch
+
+
+@require_koch
+def test_teleoperate(request):
+    robot = make_robot("koch")
+    teleoperate(robot, teleop_time_s=1)
+    teleoperate(robot, fps=30, teleop_time_s=1)
+    teleoperate(robot, fps=60, teleop_time_s=1)
+    del robot
+
+
+@require_koch
+def test_record_dataset_and_replay_episode_and_run_policy(tmpdir, request):
+    robot_name = "koch"
+    env_name = "koch_real"
+    policy_name = "act_koch_real"
+
+    root = Path(tmpdir)
+    repo_id = "lerobot/debug"
+
+    robot = make_robot(robot_name)
+    dataset = record_dataset(
+        robot, fps=30, root=root, repo_id=repo_id, warmup_time_s=1, episode_time_s=1, num_episodes=2
+    )
+
+    replay_episode(robot, episode=0, fps=30, root=root, repo_id=repo_id)
+
+    cfg = init_hydra_config(
+        DEFAULT_CONFIG_PATH,
+        overrides=[
+            f"env={env_name}",
+            f"policy={policy_name}",
+            f"device={DEVICE}",
+        ],
+    )
+
+    policy = make_policy(hydra_cfg=cfg, dataset_stats=dataset.stats)
+
+    run_policy(robot, policy, cfg, run_time_s=1)
+
+    del robot

tests/test_motors.py

@@ -0,0 +1,92 @@
+import time
+
+import numpy as np
+import pytest
+
+from lerobot.common.robot_devices.utils import RobotDeviceAlreadyConnectedError, RobotDeviceNotConnectedError
+from tests.utils import require_koch
+
+
+@require_koch
+def test_motors_bus(request):
+    # TODO(rcadene): measure fps in nightly?
+    # TODO(rcadene): test logs
+    # TODO(rcadene): test calibration
+    # TODO(rcadene): add compatibility with other motors bus
+    from lerobot.common.robot_devices.motors.dynamixel import DynamixelMotorsBus
+
+    # Test instantiating a common motors structure.
+    # Here the one from Alexander Koch follower arm.
+    port = "/dev/tty.usbmodem575E0032081"
+    motors = {
+        # name: (index, model)
+        "shoulder_pan": (1, "xl430-w250"),
+        "shoulder_lift": (2, "xl430-w250"),
+        "elbow_flex": (3, "xl330-m288"),
+        "wrist_flex": (4, "xl330-m288"),
+        "wrist_roll": (5, "xl330-m288"),
+        "gripper": (6, "xl330-m288"),
+    }
+    motors_bus = DynamixelMotorsBus(port, motors)
+
+    # Test reading and writting before connecting raises an error
+    with pytest.raises(RobotDeviceNotConnectedError):
+        motors_bus.read("Torque_Enable")
+    with pytest.raises(RobotDeviceNotConnectedError):
+        motors_bus.write("Torque_Enable", 1)
+    with pytest.raises(RobotDeviceNotConnectedError):
+        motors_bus.disconnect()
+
+    # Test deleting the object without connecting first
+    del motors_bus
+
+    # Test connecting
+    motors_bus = DynamixelMotorsBus(port, motors)
+    motors_bus.connect()
+
+    # Test connecting twice raises an error
+    with pytest.raises(RobotDeviceAlreadyConnectedError):
+        motors_bus.connect()
+
+    # Test disabling torque and reading torque on all motors
+    motors_bus.write("Torque_Enable", 0)
+    values = motors_bus.read("Torque_Enable")
+    assert isinstance(values, np.ndarray)
+    assert len(values) == len(motors)
+    assert (values == 0).all()
+
+    # Test writing torque on a specific motor
+    motors_bus.write("Torque_Enable", 1, "gripper")
+
+    # Test reading torque from this specific motor. It is now 1
+    values = motors_bus.read("Torque_Enable", "gripper")
+    assert len(values) == 1
+    assert values[0] == 1
+
+    # Test reading torque from all motors. It is 1 for the specific motor,
+    # and 0 on the others.
+    values = motors_bus.read("Torque_Enable")
+    gripper_index = motors_bus.motor_names.index("gripper")
+    assert values[gripper_index] == 1
+    assert values.sum() == 1  # gripper is the only motor to have torque 1
+
+    # Test writing torque on all motors and it is 1 for all.
+    motors_bus.write("Torque_Enable", 1)
+    values = motors_bus.read("Torque_Enable")
+    assert (values == 1).all()
+
+    # Test ordering the motors to move slightly (+1 value among 4096) and this move
+    # can be executed and seen by the motor position sensor
+    values = motors_bus.read("Present_Position")
+    motors_bus.write("Goal_Position", values + 1)
+    # Give time for the motors to move to the goal position
+    time.sleep(1)
+    new_values = motors_bus.read("Present_Position")
+    assert (new_values == values).all()
+
+
+@require_koch
+def test_find_port(request):
+    from lerobot.common.robot_devices.motors.dynamixel import find_port
+
+    find_port()

tests/test_robots.py

@@ -0,0 +1,128 @@
+import pickle
+from pathlib import Path
+
+import pytest
+import torch
+
+from lerobot.common.robot_devices.robots.factory import make_robot
+from lerobot.common.robot_devices.utils import RobotDeviceAlreadyConnectedError, RobotDeviceNotConnectedError
+from tests.utils import require_koch
+
+
+@require_koch
+def test_robot(tmpdir, request):
+    # TODO(rcadene): measure fps in nightly?
+    # TODO(rcadene): test logs
+    # TODO(rcadene): add compatibility with other robots
+    from lerobot.common.robot_devices.robots.koch import KochRobot
+
+    # Save calibration preset
+    calibration = {
+        "follower_main": {
+            "shoulder_pan": (-2048, False),
+            "shoulder_lift": (2048, True),
+            "elbow_flex": (-1024, False),
+            "wrist_flex": (2048, True),
+            "wrist_roll": (2048, True),
+            "gripper": (2048, True),
+        },
+        "leader_main": {
+            "shoulder_pan": (-2048, False),
+            "shoulder_lift": (1024, True),
+            "elbow_flex": (2048, True),
+            "wrist_flex": (-2048, False),
+            "wrist_roll": (2048, True),
+            "gripper": (2048, True),
+        },
+    }
+    tmpdir = Path(tmpdir)
+    calibration_path = tmpdir / "calibration.pkl"
+    calibration_path.parent.mkdir(parents=True, exist_ok=True)
+    with open(calibration_path, "wb") as f:
+        pickle.dump(calibration, f)
+
+    # Test connecting without devices raises an error
+    robot = KochRobot()
+    with pytest.raises(ValueError):
+        robot.connect()
+    del robot
+
+    # Test using robot before connecting raises an error
+    robot = KochRobot()
+    with pytest.raises(RobotDeviceNotConnectedError):
+        robot.teleop_step()
+    with pytest.raises(RobotDeviceNotConnectedError):
+        robot.teleop_step(record_data=True)
+    with pytest.raises(RobotDeviceNotConnectedError):
+        robot.capture_observation()
+    with pytest.raises(RobotDeviceNotConnectedError):
+        robot.send_action(None)
+    with pytest.raises(RobotDeviceNotConnectedError):
+        robot.disconnect()
+
+    # Test deleting the object without connecting first
+    del robot
+
+    # Test connecting
+    robot = make_robot("koch")
+    # TODO(rcadene): proper monkey patch
+    robot.calibration_path = calibration_path
+    robot.connect()  # run the manual calibration precedure
+    assert robot.is_connected
+
+    # Test connecting twice raises an error
+    with pytest.raises(RobotDeviceAlreadyConnectedError):
+        robot.connect()
+
+    # Test disconnecting with `__del__`
+    del robot
+
+    # Test teleop can run
+    robot = make_robot("koch")
+    robot.calibration_path = calibration_path
+    robot.connect()
+    robot.teleop_step()
+
+    # Test data recorded during teleop are well formated
+    observation, action = robot.teleop_step(record_data=True)
+    # State
+    assert "observation.state" in observation
+    assert isinstance(observation["observation.state"], torch.Tensor)
+    assert observation["observation.state"].ndim == 1
+    dim_state = sum(len(robot.follower_arms[name].motors) for name in robot.follower_arms)
+    assert observation["observation.state"].shape[0] == dim_state
+    # Cameras
+    for name in robot.cameras:
+        assert f"observation.images.{name}" in observation
+        assert isinstance(observation[f"observation.images.{name}"], torch.Tensor)
+        assert observation[f"observation.images.{name}"].ndim == 3
+    # Action
+    assert "action" in action
+    assert isinstance(action["action"], torch.Tensor)
+    assert action["action"].ndim == 1
+    dim_action = sum(len(robot.follower_arms[name].motors) for name in robot.follower_arms)
+    assert action["action"].shape[0] == dim_action
+    # TODO(rcadene): test if observation and action data are returned as expected
+
+    # Test capture_observation can run and observation returned are the same (since the arm didnt move)
+    captured_observation = robot.capture_observation()
+    assert set(captured_observation.keys()) == set(observation.keys())
+    for name in captured_observation:
+        if "image" in name:
+            # TODO(rcadene): skipping image for now as it's challenging to assess equality between two consecutive frames
+            continue
+        assert torch.allclose(captured_observation[name], observation[name], atol=1)
+
+    # Test send_action can run
+    robot.send_action(action["action"])
+
+    # Test disconnecting
+    robot.disconnect()
+    assert not robot.is_connected
+    for name in robot.follower_arms:
+        assert not robot.follower_arms[name].is_connected
+    for name in robot.leader_arms:
+        assert not robot.leader_arms[name].is_connected
+    for name in robot.cameras:
+        assert not robot.cameras[name].is_connected
+    del robot

tests/test_visualize_dataset.py

@@ -20,21 +20,6 @@ import pytest
 from lerobot.scripts.visualize_dataset import visualize_dataset
 
 
-@pytest.mark.parametrize(
-    "repo_id",
-    ["lerobot/pusht"],
-)
-def test_visualize_dataset(tmpdir, repo_id):
-    rrd_path = visualize_dataset(
-        repo_id,
-        episode_index=0,
-        batch_size=32,
-        save=True,
-        output_dir=tmpdir,
-    )
-    assert rrd_path.exists()
-
-
 @pytest.mark.parametrize(
     "repo_id",
     ["lerobot/pusht"],

tests/utils.py

@@ -147,3 +147,22 @@ def require_package(package_name):
         return wrapper
 
     return decorator
+
+
+def require_koch(func):
+    """
+    Decorator that skips the test if an alexander koch robot is not available
+    """
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        # Access the pytest request context to get the is_koch_available fixture
+        request = kwargs.get("request")
+        if request is None:
+            raise ValueError("The 'request' fixture must be passed to the test function as a parameter.")
+
+        if not request.getfixturevalue("is_koch_available"):
+            pytest.skip("An alexander koch robot is not available.")
+        return func(*args, **kwargs)
+
+    return wrapper