Fix unit tests

lerobot/scripts/control_robot.py

@@ -242,7 +242,8 @@ def is_headless():
 ########################################################################################
 
 
-def calibrate(robot: Robot, arms: list[str] | None):
+def get_available_arms(robot):
+    # TODO(rcadene): moves this function in manipulator class?
     available_arms = []
     for name in robot.follower_arms:
         arm_id = get_arm_id(name, "follower")
@@ -250,9 +251,12 @@ def calibrate(robot: Robot, arms: list[str] | None):
     for name in robot.leader_arms:
         arm_id = get_arm_id(name, "leader")
         available_arms.append(arm_id)
+    return available_arms
 
+
+def calibrate(robot: Robot, arms: list[str] | None):
+    available_arms = get_available_arms(robot)
     unknown_arms = [arm_id for arm_id in arms if arm_id not in available_arms]
-
     available_arms_str = " ".join(available_arms)
     unknown_arms_str = " ".join(unknown_arms)
 
@@ -445,6 +449,7 @@ def record(
     # Using `with` to exist smoothly if an execption is raised.
     futures = []
     num_image_writers = num_image_writers_per_camera * len(robot.cameras)
+    num_image_writers = max(num_image_writers, 1)
     with concurrent.futures.ThreadPoolExecutor(max_workers=num_image_writers) as executor:
         # Start recording all episodes
         while episode_index < num_episodes:

poetry.lock

@@ -1,4 +1,4 @@
-# This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand.
+# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand.
 
 [[package]]
 name = "absl-py"
@@ -2406,7 +2406,6 @@ description = "Nvidia JIT LTO Library"
 optional = false
 python-versions = ">=3"
 files = [
-    {file = "nvidia_nvjitlink_cu12-12.5.82-py3-none-manylinux2014_aarch64.whl", hash = "sha256:98103729cc5226e13ca319a10bbf9433bbbd44ef64fe72f45f067cacc14b8d27"},
     {file = "nvidia_nvjitlink_cu12-12.5.82-py3-none-manylinux2014_x86_64.whl", hash = "sha256:f9b37bc5c8cf7509665cb6ada5aaa0ce65618f2332b7d3e78e9790511f111212"},
     {file = "nvidia_nvjitlink_cu12-12.5.82-py3-none-win_amd64.whl", hash = "sha256:e782564d705ff0bf61ac3e1bf730166da66dd2fe9012f111ede5fc49b64ae697"},
 ]

tests/test_control_robot.py

@@ -29,7 +29,7 @@ import pytest
 
 from lerobot.common.policies.factory import make_policy
 from lerobot.common.utils.utils import init_hydra_config
-from lerobot.scripts.control_robot import calibrate, record, replay, teleoperate
+from lerobot.scripts.control_robot import calibrate, get_available_arms, record, replay, teleoperate
 from tests.test_robots import make_robot
 from tests.utils import DEFAULT_CONFIG_PATH, DEVICE, TEST_ROBOT_TYPES, require_robot
 
@@ -48,7 +48,7 @@ def test_teleoperate(request, robot_type, mock):
 @require_robot
 def test_calibrate(request, robot_type, mock):
     robot = make_robot(robot_type)
-    calibrate(robot)
+    calibrate(robot, arms=get_available_arms(robot))
     del robot
 
 
@@ -59,7 +59,17 @@ def test_record_without_cameras(tmpdir, request, robot_type, mock):
     repo_id = "lerobot/debug"
 
     robot = make_robot(robot_type, overrides=["~cameras"])
-    record(robot, fps=30, root=root, repo_id=repo_id, warmup_time_s=1, episode_time_s=1, num_episodes=2)
+    record(
+        robot,
+        fps=30,
+        root=root,
+        repo_id=repo_id,
+        warmup_time_s=1,
+        episode_time_s=1,
+        num_episodes=2,
+        run_compute_stats=False,
+        push_to_hub=False,
+    )
 
 
 @pytest.mark.parametrize("robot_type, mock", TEST_ROBOT_TYPES)
@@ -73,7 +83,14 @@ def test_record_and_replay_and_policy(tmpdir, request, robot_type, mock):
 
     robot = make_robot(robot_type)
     dataset = record(
-        robot, fps=30, root=root, repo_id=repo_id, warmup_time_s=1, episode_time_s=1, num_episodes=2
+        robot,
+        fps=30,
+        root=root,
+        repo_id=repo_id,
+        warmup_time_s=1,
+        episode_time_s=1,
+        num_episodes=2,
+        push_to_hub=False,
     )
 
     replay(robot, episode=0, fps=30, root=root, repo_id=repo_id)
@@ -89,6 +106,6 @@ def test_record_and_replay_and_policy(tmpdir, request, robot_type, mock):
 
     policy = make_policy(hydra_cfg=cfg, dataset_stats=dataset.stats)
 
-    record(robot, policy, cfg, run_time_s=1)
+    record(robot, policy, cfg, warmup_time_s=1, episode_time_s=1, run_compute_stats=False, push_to_hub=False)
 
     del robot

tests/utils.py

@@ -204,7 +204,7 @@ def require_robot(func):
         # Run test with a monkeypatched version of the robot devices.
         if mock:
             # TODO(rcadene): redesign mocking to not have this hardcoded logic
-            if robot_type == "koch":
+            if robot_type in ["koch", "koch_bimanual"]:
                 camera_type = "opencv"
             elif robot_type == "aloha":
                 camera_type = "intelrealsense"