hangX/lerobot
1
0
Fork 0
forked from tangger/lerobot
Code Pull Requests Activity

Control aloha robot natively (#316)

Browse Source 
Co-authored-by: Simon Alibert <75076266+aliberts@users.noreply.github.com>
This commit is contained in:
Remi
2024-09-04 19:28:05 +02:00
committed by GitHub
parent 27ba2951d1
commit 429a463aff
32 changed files with 898 additions and 390 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
tests/conftest.py
View File

@@ -13,28 +13,31 @@
# 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 traceback
import pytest
from lerobot.common.utils.utils import init_hydra_config
from .utils import DEVICE, KOCH_ROBOT_CONFIG_PATH
from .utils import DEVICE, ROBOT_CONFIG_PATH_TEMPLATE
def pytest_collection_finish():
print(f"\nTesting with {DEVICE=}")
@pytest.fixture(scope="session")
def is_koch_available():
@pytest.fixture
def is_robot_available(robot_type):
try:
from lerobot.common.robot_devices.robots.factory import make_robot
robot_cfg = init_hydra_config(KOCH_ROBOT_CONFIG_PATH)
config_path = ROBOT_CONFIG_PATH_TEMPLATE.format(robot=robot_type)
robot_cfg = init_hydra_config(config_path)
robot = make_robot(robot_cfg)
robot.connect()
del robot
return True
except Exception as e:
print("A koch robot is not available.")
print(e)
except Exception:
traceback.print_exc()
print(f"\nA {robot_type} robot is not available.")
return False

22
tests/test_cameras.py
View File

@@ -1,9 +1,19 @@
"""
Tests meant to be used locally and launched manually.
Example usage:
```bash
pytest -sx tests/test_cameras.py::test_camera
```
"""
import numpy as np
import pytest
from lerobot import available_robots
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
from tests.utils import require_robot
CAMERA_INDEX = 2
# Maximum absolute difference between two consecutive images recored by a camera.
@@ -15,8 +25,9 @@ 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):
@pytest.mark.parametrize("robot_type", available_robots)
@require_robot
def test_camera(request, robot_type):
"""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.
@@ -120,6 +131,7 @@ def test_camera(request):
del camera
@require_koch
def test_save_images_from_cameras(tmpdir, request):
@pytest.mark.parametrize("robot_type", available_robots)
@require_robot
def test_save_images_from_cameras(tmpdir, request, robot_type):
save_images_from_cameras(tmpdir, record_time_s=1)

42
tests/test_control_robot.py
View File

@@ -1,53 +1,53 @@
from pathlib import Path
import pytest
from lerobot import available_robots
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 calibrate, record, replay, teleoperate
from tests.utils import DEFAULT_CONFIG_PATH, DEVICE, KOCH_ROBOT_CONFIG_PATH, require_koch
from tests.test_robots import make_robot
from tests.utils import DEFAULT_CONFIG_PATH, DEVICE, require_robot
def make_robot_(overrides=None):
robot_cfg = init_hydra_config(KOCH_ROBOT_CONFIG_PATH, overrides)
robot = make_robot(robot_cfg)
return robot
@require_koch
# `require_koch` uses `request` to access `is_koch_available` fixture
def test_teleoperate(request):
robot = make_robot_()
@pytest.mark.parametrize("robot_type", available_robots)
@require_robot
def test_teleoperate(request, 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
@require_koch
def test_calibrate(request):
robot = make_robot_()
@pytest.mark.parametrize("robot_type", available_robots)
@require_robot
def test_calibrate(request, robot_type):
robot = make_robot(robot_type)
calibrate(robot)
del robot
@require_koch
def test_record_without_cameras(tmpdir, request):
@pytest.mark.parametrize("robot_type", available_robots)
@require_robot
def test_record_without_cameras(tmpdir, request, robot_type):
root = Path(tmpdir)
repo_id = "lerobot/debug"
robot = make_robot_(overrides=["~cameras"])
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)
@require_koch
def test_record_and_replay_and_policy(tmpdir, request):
@pytest.mark.parametrize("robot_type", available_robots)
@require_robot
def test_record_and_replay_and_policy(tmpdir, request, robot_type):
env_name = "koch_real"
policy_name = "act_koch_real"
root = Path(tmpdir)
repo_id = "lerobot/debug"
robot = make_robot_()
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
)

52
tests/test_motors.py
View File

@@ -1,3 +1,13 @@
"""
Tests meant to be used locally and launched manually.
Example usage:
```bash
pytest -sx tests/test_motors.py::test_find_port
pytest -sx tests/test_motors.py::test_motors_bus
```
"""
# TODO(rcadene): measure fps in nightly?
# TODO(rcadene): test logs
# TODO(rcadene): test calibration
@@ -5,34 +15,41 @@
import time
import hydra
import numpy as np
import pytest
from lerobot import available_robots
from lerobot.common.robot_devices.motors.utils import MotorsBus
from lerobot.common.robot_devices.robots.factory import make_robot
from lerobot.common.robot_devices.utils import RobotDeviceAlreadyConnectedError, RobotDeviceNotConnectedError
from lerobot.common.utils.utils import init_hydra_config
from tests.utils import KOCH_ROBOT_CONFIG_PATH, require_koch
from tests.utils import ROBOT_CONFIG_PATH_TEMPLATE, require_robot
def make_motors_bus():
robot_cfg = init_hydra_config(KOCH_ROBOT_CONFIG_PATH)
# Instantiating a common motors structure.
# Here the one from Alexander Koch follower arm.
motors_bus = hydra.utils.instantiate(robot_cfg.leader_arms.main)
def make_motors_bus(robot_type: str) -> MotorsBus:
# Instantiate a robot and return one of its leader arms
config_path = ROBOT_CONFIG_PATH_TEMPLATE.format(robot=robot_type)
robot_cfg = init_hydra_config(config_path)
robot = make_robot(robot_cfg)
first_bus_name = list(robot.leader_arms.keys())[0]
motors_bus = robot.leader_arms[first_bus_name]
return motors_bus
@require_koch
def test_find_port(request):
@pytest.mark.parametrize("robot_type", available_robots)
@require_robot
def test_find_port(request, robot_type):
from lerobot.common.robot_devices.motors.dynamixel import find_port
find_port()
@require_koch
def test_configure_motors_all_ids_1(request):
@pytest.mark.parametrize("robot_type", available_robots)
@require_robot
def test_configure_motors_all_ids_1(request, robot_type):
input("Are you sure you want to re-configure the motors? Press enter to continue...")
# This test expect the configuration was already correct.
motors_bus = make_motors_bus()
motors_bus = make_motors_bus(robot_type)
motors_bus.connect()
motors_bus.write("Baud_Rate", [0] * len(motors_bus.motors))
motors_bus.set_bus_baudrate(9_600)
@@ -40,15 +57,16 @@ def test_configure_motors_all_ids_1(request):
del motors_bus
# Test configure
motors_bus = make_motors_bus()
motors_bus = make_motors_bus(robot_type)
motors_bus.connect()
assert motors_bus.are_motors_configured()
del motors_bus
@require_koch
def test_motors_bus(request):
motors_bus = make_motors_bus()
@pytest.mark.parametrize("robot_type", available_robots)
@require_robot
def test_motors_bus(request, robot_type):
motors_bus = make_motors_bus(robot_type)
# Test reading and writting before connecting raises an error
with pytest.raises(RobotDeviceNotConnectedError):
@@ -62,7 +80,7 @@ def test_motors_bus(request):
del motors_bus
# Test connecting
motors_bus = make_motors_bus()
motors_bus = make_motors_bus(robot_type)
motors_bus.connect()
# Test connecting twice raises an error

66
tests/test_robots.py
View File

@@ -1,54 +1,52 @@
import pickle
"""
Tests meant to be used locally and launched manually.
Example usage:
```bash
pytest -sx tests/test_robots.py::test_robot
```
"""
from pathlib import Path
import pytest
import torch
from lerobot.common.robot_devices.robots.factory import make_robot
from lerobot import available_robots
from lerobot.common.robot_devices.robots.factory import make_robot as make_robot_from_cfg
from lerobot.common.robot_devices.robots.utils import Robot
from lerobot.common.robot_devices.utils import RobotDeviceAlreadyConnectedError, RobotDeviceNotConnectedError
from tests.utils import require_koch
from lerobot.common.utils.utils import init_hydra_config
from tests.utils import ROBOT_CONFIG_PATH_TEMPLATE, require_robot
@require_koch
def test_robot(tmpdir, request):
def make_robot(robot_type: str, overrides: list[str] | None = None) -> Robot:
config_path = ROBOT_CONFIG_PATH_TEMPLATE.format(robot=robot_type)
robot_cfg = init_hydra_config(config_path, overrides)
robot = make_robot_from_cfg(robot_cfg)
return robot
@pytest.mark.parametrize("robot_type", available_robots)
@require_robot
def test_robot(tmpdir, request, robot_type):
# 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
from lerobot.common.robot_devices.robots.manipulator import ManipulatorRobot
# 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)
calibration_dir = tmpdir / robot_type
# Test connecting without devices raises an error
robot = KochRobot()
robot = ManipulatorRobot()
with pytest.raises(ValueError):
robot.connect()
del robot
# Test using robot before connecting raises an error
robot = KochRobot()
robot = ManipulatorRobot()
with pytest.raises(RobotDeviceNotConnectedError):
robot.teleop_step()
with pytest.raises(RobotDeviceNotConnectedError):
@@ -64,9 +62,7 @@ def test_robot(tmpdir, request):
del robot
# Test connecting
robot = make_robot("koch")
# TODO(rcadene): proper monkey patch
robot.calibration_path = calibration_path
robot = make_robot(robot_type, overrides=[f"calibration_dir={calibration_dir}"])
robot.connect() # run the manual calibration precedure
assert robot.is_connected
@@ -78,8 +74,8 @@ def test_robot(tmpdir, request):
del robot
# Test teleop can run
robot = make_robot("koch")
robot.calibration_path = calibration_path
robot = make_robot(robot_type, overrides=[f"calibration_dir={calibration_dir}"])
robot.calibration_dir = calibration_dir
robot.connect()
robot.teleop_step()

31
tests/utils.py
View File

@@ -21,11 +21,12 @@ import torch
from lerobot.common.utils.import_utils import is_package_available
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
# Pass this as the first argument to init_hydra_config.
DEFAULT_CONFIG_PATH = "lerobot/configs/default.yaml"
KOCH_ROBOT_CONFIG_PATH = "lerobot/configs/robot/koch.yaml"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
ROBOT_CONFIG_PATH_TEMPLATE = "lerobot/configs/robot/{robot}.yaml"
def require_x86_64_kernel(func):
@@ -150,21 +151,35 @@ def require_package(package_name):
return decorator
def require_koch(func):
def require_robot(func):
"""
Decorator that skips the test if an alexander koch robot is not available
Decorator that skips the test if a robot is not available
The decorated function must have two arguments `request` and `robot_type`.
Example of usage:
```python
@pytest.mark.parametrize(
"robot_type", ["koch", "aloha"]
)
@require_robot
def test_require_robot(request, robot_type):
pass
```
"""
@wraps(func)
def wrapper(*args, **kwargs):
# Access the pytest request context to get the is_koch_available fixture
# Access the pytest request context to get the is_robot_available fixture
request = kwargs.get("request")
robot_type = kwargs.get("robot_type")
if request is None:
raise ValueError("The 'request' fixture must be passed to the test function as a parameter.")
# The function `is_koch_available` is defined in `tests/conftest.py`
if not request.getfixturevalue("is_koch_available"):
pytest.skip("An alexander koch robot is not available.")
# The function `is_robot_available` is defined in `tests/conftest.py`
if not request.getfixturevalue("is_robot_available"):
pytest.skip(f"A {robot_type} robot is not available.")
return func(*args, **kwargs)
return wrapper