diff --git a/lerobot/scripts/control_robot.py b/lerobot/scripts/control_robot.py index a6506a3f..a23c21c6 100644 --- a/lerobot/scripts/control_robot.py +++ b/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: diff --git a/poetry.lock b/poetry.lock index 40bf29ef..d8165371 100644 --- a/poetry.lock +++ b/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"}, ] diff --git a/tests/test_control_robot.py b/tests/test_control_robot.py index fe73d9ff..14fe9158 100644 --- a/tests/test_control_robot.py +++ b/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 diff --git a/tests/utils.py b/tests/utils.py index 726d04a0..aeafd9d1 100644 --- a/tests/utils.py +++ b/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"