Files
lerobot/tests/test_control_robot.py
2024-09-26 16:28:08 +02:00

166 lines
4.4 KiB
Python

"""
Tests for physical robots and their mocked versions.
If the physical robots are not connected to the computer, or not working,
the test will be skipped.
Example of running a specific test:
```bash
pytest -sx tests/test_control_robot.py::test_teleoperate
```
Example of running test on real robots connected to the computer:
```bash
pytest -sx 'tests/test_control_robot.py::test_teleoperate[koch-False]'
pytest -sx 'tests/test_control_robot.py::test_teleoperate[koch_bimanual-False]'
pytest -sx 'tests/test_control_robot.py::test_teleoperate[aloha-False]'
```
Example of running test on a mocked version of robots:
```bash
pytest -sx 'tests/test_control_robot.py::test_teleoperate[koch-True]'
pytest -sx 'tests/test_control_robot.py::test_teleoperate[koch_bimanual-True]'
pytest -sx 'tests/test_control_robot.py::test_teleoperate[aloha-True]'
```
"""
from pathlib import Path
import pytest
from lerobot import available_robots
from lerobot.common.policies.factory import make_policy
from lerobot.common.utils.utils import init_hydra_config
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, require_mock_robot, require_robot
def _test_teleoperate(robot_type):
robot = make_robot(robot_type)
teleoperate(robot, teleop_time_s=1)
teleoperate(robot, fps=30, teleop_time_s=1)
teleoperate(robot, fps=60, teleop_time_s=1)
del robot
def _test_calibrate(robot_type):
robot = make_robot(robot_type)
calibrate(robot, arms=get_available_arms(robot))
del robot
def _test_record_without_cameras(tmpdir, robot_type):
root = Path(tmpdir)
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,
run_compute_stats=False,
push_to_hub=False,
video=False,
)
def _test_record_and_replay_and_policy(tmpdir, robot_type):
env_name = "koch_real"
policy_name = "act_koch_real"
root = Path(tmpdir)
repo_id = "lerobot/debug"
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,
push_to_hub=False,
# TODO(rcadene, aliberts): test video=True
video=False,
)
replay(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)
record(
robot,
policy,
cfg,
warmup_time_s=1,
episode_time_s=1,
num_episodes=2,
run_compute_stats=False,
push_to_hub=False,
video=False,
)
del robot
@pytest.mark.parametrize("robot_type", available_robots)
@require_mock_robot
def test_teleoperate_mock(monkeypatch, robot_type):
_test_teleoperate(robot_type)
@pytest.mark.parametrize("robot_type", available_robots)
@require_robot
def test_teleoperate(request, robot_type):
_test_teleoperate(robot_type)
@pytest.mark.parametrize("robot_type", available_robots)
@require_mock_robot
def test_calibrate_mock(monkeypatch, robot_type):
_test_calibrate(robot_type)
@pytest.mark.parametrize("robot_type", available_robots)
@require_robot
def test_calibrate(request, robot_type):
_test_calibrate(robot_type)
@pytest.mark.parametrize("robot_type", available_robots)
@require_mock_robot
def test_record_without_cameras_mock(tmpdir, monkeypatch, robot_type):
_test_record_without_cameras(tmpdir, robot_type)
@pytest.mark.parametrize("robot_type", available_robots)
@require_robot
def test_record_without_cameras(tmpdir, request, robot_type):
_test_record_without_cameras(tmpdir, robot_type)
@pytest.mark.parametrize("robot_type", available_robots)
@require_mock_robot
def test_record_and_replay_and_policy_mock(tmpdir, monkeypatch, robot_type):
_test_record_and_replay_and_policy(tmpdir, robot_type)
@pytest.mark.parametrize("robot_type", available_robots)
@require_robot
def test_record_and_replay_and_policy(tmpdir, request, robot_type):
_test_record_and_replay_and_policy(tmpdir, robot_type)