mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-23 01:41:54 +00:00
Merge remote-tracking branch 'origin/main' into user/rcadene/2025_04_11_dataset_v3
This commit is contained in:
Remi Cadene
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:9dc9df05797dc0e7b92edc845caab2e4c37c3cfcabb4ee6339c67212b5baba3b
|
||||
size 38023
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:7e11af87616b83c1cdb30330e951b91e86b51c64a1326e1ba5b4a3fbcdec1a11
|
||||
size 55698
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:b8840fb643afe903191248703b1f95a57faf5812ecd9978ac502ee939646fdb2
|
||||
size 121115
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:f79d14daafb1c0cf2fec5d46ee8029a73fe357402fdd31a7cd4a4794d7319a7c
|
||||
size 260367
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:a8d6e64d6cb0e02c94ae125630ee758055bd2e695772c0463a30d63ddc6c5e17
|
||||
size 3520862
|
||||
@@ -1,101 +0,0 @@
|
||||
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# 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.
|
||||
from functools import cache
|
||||
|
||||
import numpy as np
|
||||
|
||||
CAP_V4L2 = 200
|
||||
CAP_DSHOW = 700
|
||||
CAP_AVFOUNDATION = 1200
|
||||
CAP_ANY = -1
|
||||
|
||||
CAP_PROP_FPS = 5
|
||||
CAP_PROP_FRAME_WIDTH = 3
|
||||
CAP_PROP_FRAME_HEIGHT = 4
|
||||
COLOR_RGB2BGR = 4
|
||||
COLOR_BGR2RGB = 4
|
||||
|
||||
ROTATE_90_COUNTERCLOCKWISE = 2
|
||||
ROTATE_90_CLOCKWISE = 0
|
||||
ROTATE_180 = 1
|
||||
|
||||
|
||||
@cache
|
||||
def _generate_image(width: int, height: int):
|
||||
return np.random.randint(0, 256, size=(height, width, 3), dtype=np.uint8)
|
||||
|
||||
|
||||
def cvtColor(color_image, color_conversion): # noqa: N802
|
||||
if color_conversion in [COLOR_RGB2BGR, COLOR_BGR2RGB]:
|
||||
return color_image[:, :, [2, 1, 0]]
|
||||
else:
|
||||
raise NotImplementedError(color_conversion)
|
||||
|
||||
|
||||
def rotate(color_image, rotation):
|
||||
if rotation is None:
|
||||
return color_image
|
||||
elif rotation == ROTATE_90_CLOCKWISE:
|
||||
return np.rot90(color_image, k=1)
|
||||
elif rotation == ROTATE_180:
|
||||
return np.rot90(color_image, k=2)
|
||||
elif rotation == ROTATE_90_COUNTERCLOCKWISE:
|
||||
return np.rot90(color_image, k=3)
|
||||
else:
|
||||
raise NotImplementedError(rotation)
|
||||
|
||||
|
||||
class VideoCapture:
|
||||
def __init__(self, *args, **kwargs):
|
||||
self._mock_dict = {
|
||||
CAP_PROP_FPS: 30,
|
||||
CAP_PROP_FRAME_WIDTH: 640,
|
||||
CAP_PROP_FRAME_HEIGHT: 480,
|
||||
}
|
||||
self._is_opened = True
|
||||
|
||||
def isOpened(self): # noqa: N802
|
||||
return self._is_opened
|
||||
|
||||
def set(self, propId: int, value: float) -> bool: # noqa: N803
|
||||
if not self._is_opened:
|
||||
raise RuntimeError("Camera is not opened")
|
||||
self._mock_dict[propId] = value
|
||||
return True
|
||||
|
||||
def get(self, propId: int) -> float: # noqa: N803
|
||||
if not self._is_opened:
|
||||
raise RuntimeError("Camera is not opened")
|
||||
value = self._mock_dict[propId]
|
||||
if value == 0:
|
||||
if propId == CAP_PROP_FRAME_HEIGHT:
|
||||
value = 480
|
||||
elif propId == CAP_PROP_FRAME_WIDTH:
|
||||
value = 640
|
||||
return value
|
||||
|
||||
def read(self):
|
||||
if not self._is_opened:
|
||||
raise RuntimeError("Camera is not opened")
|
||||
h = self.get(CAP_PROP_FRAME_HEIGHT)
|
||||
w = self.get(CAP_PROP_FRAME_WIDTH)
|
||||
ret = True
|
||||
return ret, _generate_image(width=w, height=h)
|
||||
|
||||
def release(self):
|
||||
self._is_opened = False
|
||||
|
||||
def __del__(self):
|
||||
if self._is_opened:
|
||||
self.release()
|
||||
@@ -1,148 +0,0 @@
|
||||
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# 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 enum
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
class stream(enum.Enum): # noqa: N801
|
||||
color = 0
|
||||
depth = 1
|
||||
|
||||
|
||||
class format(enum.Enum): # noqa: N801
|
||||
rgb8 = 0
|
||||
z16 = 1
|
||||
|
||||
|
||||
class config: # noqa: N801
|
||||
def enable_device(self, device_id: str):
|
||||
self.device_enabled = device_id
|
||||
|
||||
def enable_stream(self, stream_type: stream, width=None, height=None, color_format=None, fps=None):
|
||||
self.stream_type = stream_type
|
||||
# Overwrite default values when possible
|
||||
self.width = 848 if width is None else width
|
||||
self.height = 480 if height is None else height
|
||||
self.color_format = format.rgb8 if color_format is None else color_format
|
||||
self.fps = 30 if fps is None else fps
|
||||
|
||||
|
||||
class RSColorProfile:
|
||||
def __init__(self, config):
|
||||
self.config = config
|
||||
|
||||
def fps(self):
|
||||
return self.config.fps
|
||||
|
||||
def width(self):
|
||||
return self.config.width
|
||||
|
||||
def height(self):
|
||||
return self.config.height
|
||||
|
||||
|
||||
class RSColorStream:
|
||||
def __init__(self, config):
|
||||
self.config = config
|
||||
|
||||
def as_video_stream_profile(self):
|
||||
return RSColorProfile(self.config)
|
||||
|
||||
|
||||
class RSProfile:
|
||||
def __init__(self, config):
|
||||
self.config = config
|
||||
|
||||
def get_stream(self, color_format):
|
||||
del color_format # unused
|
||||
return RSColorStream(self.config)
|
||||
|
||||
|
||||
class pipeline: # noqa: N801
|
||||
def __init__(self):
|
||||
self.started = False
|
||||
self.config = None
|
||||
|
||||
def start(self, config):
|
||||
self.started = True
|
||||
self.config = config
|
||||
return RSProfile(self.config)
|
||||
|
||||
def stop(self):
|
||||
if not self.started:
|
||||
raise RuntimeError("You need to start the camera before stop.")
|
||||
self.started = False
|
||||
self.config = None
|
||||
|
||||
def wait_for_frames(self, timeout_ms=50000):
|
||||
del timeout_ms # unused
|
||||
return RSFrames(self.config)
|
||||
|
||||
|
||||
class RSFrames:
|
||||
def __init__(self, config):
|
||||
self.config = config
|
||||
|
||||
def get_color_frame(self):
|
||||
return RSColorFrame(self.config)
|
||||
|
||||
def get_depth_frame(self):
|
||||
return RSDepthFrame(self.config)
|
||||
|
||||
|
||||
class RSColorFrame:
|
||||
def __init__(self, config):
|
||||
self.config = config
|
||||
|
||||
def get_data(self):
|
||||
data = np.ones((self.config.height, self.config.width, 3), dtype=np.uint8)
|
||||
# Create a difference between rgb and bgr
|
||||
data[:, :, 0] = 2
|
||||
return data
|
||||
|
||||
|
||||
class RSDepthFrame:
|
||||
def __init__(self, config):
|
||||
self.config = config
|
||||
|
||||
def get_data(self):
|
||||
return np.ones((self.config.height, self.config.width), dtype=np.uint16)
|
||||
|
||||
|
||||
class RSDevice:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def get_info(self, camera_info) -> str:
|
||||
del camera_info # unused
|
||||
# return fake serial number
|
||||
return "123456789"
|
||||
|
||||
|
||||
class context: # noqa: N801
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def query_devices(self):
|
||||
return [RSDevice()]
|
||||
|
||||
|
||||
class camera_info: # noqa: N801
|
||||
# fake name
|
||||
name = "Intel RealSense D435I"
|
||||
|
||||
def __init__(self, serial_number):
|
||||
del serial_number
|
||||
pass
|
||||
@@ -1,252 +0,0 @@
|
||||
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# 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.
|
||||
"""
|
||||
Tests for physical cameras and their mocked versions.
|
||||
If the physical camera is not connected to the computer, or not working,
|
||||
the test will be skipped.
|
||||
|
||||
Example of running a specific test:
|
||||
```bash
|
||||
pytest -sx tests/test_cameras.py::test_camera
|
||||
```
|
||||
|
||||
Example of running test on a real camera connected to the computer:
|
||||
```bash
|
||||
pytest -sx 'tests/test_cameras.py::test_camera[opencv-False]'
|
||||
pytest -sx 'tests/test_cameras.py::test_camera[intelrealsense-False]'
|
||||
```
|
||||
|
||||
Example of running test on a mocked version of the camera:
|
||||
```bash
|
||||
pytest -sx 'tests/test_cameras.py::test_camera[opencv-True]'
|
||||
pytest -sx 'tests/test_cameras.py::test_camera[intelrealsense-True]'
|
||||
```
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from lerobot.common.robot_devices.utils import RobotDeviceAlreadyConnectedError, RobotDeviceNotConnectedError
|
||||
from tests.utils import TEST_CAMERA_TYPES, make_camera, require_camera
|
||||
|
||||
# Maximum absolute difference between two consecutive images recorded 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()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("camera_type, mock", TEST_CAMERA_TYPES)
|
||||
@require_camera
|
||||
def test_camera(request, camera_type, mock):
|
||||
"""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
|
||||
|
||||
if camera_type == "opencv" and not mock:
|
||||
pytest.skip("TODO(rcadene): fix test for opencv physical camera")
|
||||
|
||||
camera_kwargs = {"camera_type": camera_type, "mock": mock}
|
||||
|
||||
# Test instantiating
|
||||
camera = make_camera(**camera_kwargs)
|
||||
|
||||
# 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 = make_camera(**camera_kwargs)
|
||||
camera.connect()
|
||||
assert camera.is_connected
|
||||
assert camera.fps is not None
|
||||
assert camera.capture_width is not None
|
||||
assert camera.capture_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()
|
||||
error_msg = (
|
||||
"max_pixel_difference between read() and async_read()",
|
||||
compute_max_pixel_difference(color_image, async_color_image),
|
||||
)
|
||||
# TODO(rcadene): properly set `rtol`
|
||||
np.testing.assert_allclose(
|
||||
color_image, async_color_image, rtol=1e-5, atol=MAX_PIXEL_DIFFERENCE, err_msg=error_msg
|
||||
)
|
||||
|
||||
# Test disconnecting
|
||||
camera.disconnect()
|
||||
assert camera.camera is None
|
||||
assert camera.thread is None
|
||||
|
||||
# Test disconnecting with `__del__`
|
||||
camera = make_camera(**camera_kwargs)
|
||||
camera.connect()
|
||||
del camera
|
||||
|
||||
# Test acquiring a bgr image
|
||||
camera = make_camera(**camera_kwargs, color_mode="bgr")
|
||||
camera.connect()
|
||||
assert camera.color_mode == "bgr"
|
||||
bgr_color_image = camera.read()
|
||||
np.testing.assert_allclose(
|
||||
color_image, bgr_color_image[:, :, [2, 1, 0]], rtol=1e-5, atol=MAX_PIXEL_DIFFERENCE, err_msg=error_msg
|
||||
)
|
||||
del camera
|
||||
|
||||
# Test acquiring a rotated image
|
||||
camera = make_camera(**camera_kwargs)
|
||||
camera.connect()
|
||||
ori_color_image = camera.read()
|
||||
del camera
|
||||
|
||||
for rotation in [None, 90, 180, -90]:
|
||||
camera = make_camera(**camera_kwargs, rotation=rotation)
|
||||
camera.connect()
|
||||
|
||||
if mock:
|
||||
import tests.cameras.mock_cv2 as cv2
|
||||
else:
|
||||
import cv2
|
||||
|
||||
if rotation is None:
|
||||
manual_rot_img = ori_color_image
|
||||
assert camera.rotation is None
|
||||
elif rotation == 90:
|
||||
manual_rot_img = np.rot90(color_image, k=1)
|
||||
assert camera.rotation == cv2.ROTATE_90_CLOCKWISE
|
||||
elif rotation == 180:
|
||||
manual_rot_img = np.rot90(color_image, k=2)
|
||||
assert camera.rotation == cv2.ROTATE_180
|
||||
elif rotation == -90:
|
||||
manual_rot_img = np.rot90(color_image, k=3)
|
||||
assert camera.rotation == cv2.ROTATE_90_COUNTERCLOCKWISE
|
||||
|
||||
rot_color_image = camera.read()
|
||||
|
||||
np.testing.assert_allclose(
|
||||
rot_color_image, manual_rot_img, rtol=1e-5, atol=MAX_PIXEL_DIFFERENCE, err_msg=error_msg
|
||||
)
|
||||
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 width and height can be set
|
||||
camera = make_camera(**camera_kwargs, 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 = make_camera(**camera_kwargs, fps=30, width=0, height=0)
|
||||
with pytest.raises(OSError):
|
||||
camera.connect()
|
||||
del camera
|
||||
|
||||
|
||||
@pytest.mark.parametrize("camera_type, mock", TEST_CAMERA_TYPES)
|
||||
@require_camera
|
||||
def test_save_images_from_cameras(tmp_path, request, camera_type, mock):
|
||||
# TODO(rcadene): refactor
|
||||
if camera_type == "opencv":
|
||||
from lerobot.common.robot_devices.cameras.opencv import save_images_from_cameras
|
||||
elif camera_type == "intelrealsense":
|
||||
from lerobot.common.robot_devices.cameras.intelrealsense import save_images_from_cameras
|
||||
|
||||
# Small `record_time_s` to speedup unit tests
|
||||
save_images_from_cameras(tmp_path, record_time_s=0.02, mock=mock)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("camera_type, mock", TEST_CAMERA_TYPES)
|
||||
@require_camera
|
||||
def test_camera_rotation(request, camera_type, mock):
|
||||
config_kwargs = {"camera_type": camera_type, "mock": mock, "width": 640, "height": 480, "fps": 30}
|
||||
|
||||
# No rotation.
|
||||
camera = make_camera(**config_kwargs, rotation=None)
|
||||
camera.connect()
|
||||
assert camera.capture_width == 640
|
||||
assert camera.capture_height == 480
|
||||
assert camera.width == 640
|
||||
assert camera.height == 480
|
||||
no_rot_img = camera.read()
|
||||
h, w, c = no_rot_img.shape
|
||||
assert h == 480 and w == 640 and c == 3
|
||||
camera.disconnect()
|
||||
|
||||
# Rotation = 90 (clockwise).
|
||||
camera = make_camera(**config_kwargs, rotation=90)
|
||||
camera.connect()
|
||||
# With a 90° rotation, we expect the metadata dimensions to be swapped.
|
||||
assert camera.capture_width == 640
|
||||
assert camera.capture_height == 480
|
||||
assert camera.width == 480
|
||||
assert camera.height == 640
|
||||
import cv2
|
||||
|
||||
assert camera.rotation == cv2.ROTATE_90_CLOCKWISE
|
||||
rot_img = camera.read()
|
||||
h, w, c = rot_img.shape
|
||||
assert h == 640 and w == 480 and c == 3
|
||||
camera.disconnect()
|
||||
|
||||
# Rotation = 180.
|
||||
camera = make_camera(**config_kwargs, rotation=None)
|
||||
camera.connect()
|
||||
assert camera.capture_width == 640
|
||||
assert camera.capture_height == 480
|
||||
assert camera.width == 640
|
||||
assert camera.height == 480
|
||||
no_rot_img = camera.read()
|
||||
h, w, c = no_rot_img.shape
|
||||
assert h == 480 and w == 640 and c == 3
|
||||
camera.disconnect()
|
||||
@@ -0,0 +1,188 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# 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.
|
||||
|
||||
# Example of running a specific test:
|
||||
# ```bash
|
||||
# pytest tests/cameras/test_opencv.py::test_connect
|
||||
# ```
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from lerobot.common.cameras.configs import Cv2Rotation
|
||||
from lerobot.common.cameras.opencv import OpenCVCamera, OpenCVCameraConfig
|
||||
from lerobot.common.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
|
||||
|
||||
# NOTE(Steven): more tests + assertions?
|
||||
TEST_ARTIFACTS_DIR = Path(__file__).parent.parent / "artifacts" / "cameras"
|
||||
DEFAULT_PNG_FILE_PATH = TEST_ARTIFACTS_DIR / "image_160x120.png"
|
||||
TEST_IMAGE_SIZES = ["128x128", "160x120", "320x180", "480x270"]
|
||||
TEST_IMAGE_PATHS = [TEST_ARTIFACTS_DIR / f"image_{size}.png" for size in TEST_IMAGE_SIZES]
|
||||
|
||||
|
||||
def test_abc_implementation():
|
||||
"""Instantiation should raise an error if the class doesn't implement abstract methods/properties."""
|
||||
config = OpenCVCameraConfig(index_or_path=0)
|
||||
|
||||
_ = OpenCVCamera(config)
|
||||
|
||||
|
||||
def test_connect():
|
||||
config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH)
|
||||
camera = OpenCVCamera(config)
|
||||
|
||||
camera.connect(warmup=False)
|
||||
|
||||
assert camera.is_connected
|
||||
|
||||
|
||||
def test_connect_already_connected():
|
||||
config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH)
|
||||
camera = OpenCVCamera(config)
|
||||
camera.connect(warmup=False)
|
||||
|
||||
with pytest.raises(DeviceAlreadyConnectedError):
|
||||
camera.connect(warmup=False)
|
||||
|
||||
|
||||
def test_connect_invalid_camera_path():
|
||||
config = OpenCVCameraConfig(index_or_path="nonexistent/camera.png")
|
||||
camera = OpenCVCamera(config)
|
||||
|
||||
with pytest.raises(ConnectionError):
|
||||
camera.connect(warmup=False)
|
||||
|
||||
|
||||
def test_invalid_width_connect():
|
||||
config = OpenCVCameraConfig(
|
||||
index_or_path=DEFAULT_PNG_FILE_PATH,
|
||||
width=99999, # Invalid width to trigger error
|
||||
height=480,
|
||||
)
|
||||
camera = OpenCVCamera(config)
|
||||
|
||||
with pytest.raises(RuntimeError):
|
||||
camera.connect(warmup=False)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("index_or_path", TEST_IMAGE_PATHS, ids=TEST_IMAGE_SIZES)
|
||||
def test_read(index_or_path):
|
||||
config = OpenCVCameraConfig(index_or_path=index_or_path)
|
||||
camera = OpenCVCamera(config)
|
||||
camera.connect(warmup=False)
|
||||
|
||||
img = camera.read()
|
||||
|
||||
assert isinstance(img, np.ndarray)
|
||||
|
||||
|
||||
def test_read_before_connect():
|
||||
config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH)
|
||||
camera = OpenCVCamera(config)
|
||||
|
||||
with pytest.raises(DeviceNotConnectedError):
|
||||
_ = camera.read()
|
||||
|
||||
|
||||
def test_disconnect():
|
||||
config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH)
|
||||
camera = OpenCVCamera(config)
|
||||
camera.connect(warmup=False)
|
||||
|
||||
camera.disconnect()
|
||||
|
||||
assert not camera.is_connected
|
||||
|
||||
|
||||
def test_disconnect_before_connect():
|
||||
config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH)
|
||||
camera = OpenCVCamera(config)
|
||||
|
||||
with pytest.raises(DeviceNotConnectedError):
|
||||
_ = camera.disconnect()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("index_or_path", TEST_IMAGE_PATHS, ids=TEST_IMAGE_SIZES)
|
||||
def test_async_read(index_or_path):
|
||||
config = OpenCVCameraConfig(index_or_path=index_or_path)
|
||||
camera = OpenCVCamera(config)
|
||||
camera.connect(warmup=False)
|
||||
|
||||
try:
|
||||
img = camera.async_read()
|
||||
|
||||
assert camera.thread is not None
|
||||
assert camera.thread.is_alive()
|
||||
assert isinstance(img, np.ndarray)
|
||||
finally:
|
||||
if camera.is_connected:
|
||||
camera.disconnect() # To stop/join the thread. Otherwise get warnings when the test ends
|
||||
|
||||
|
||||
def test_async_read_timeout():
|
||||
config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH)
|
||||
camera = OpenCVCamera(config)
|
||||
camera.connect(warmup=False)
|
||||
|
||||
try:
|
||||
with pytest.raises(TimeoutError):
|
||||
camera.async_read(timeout_ms=0)
|
||||
finally:
|
||||
if camera.is_connected:
|
||||
camera.disconnect()
|
||||
|
||||
|
||||
def test_async_read_before_connect():
|
||||
config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH)
|
||||
camera = OpenCVCamera(config)
|
||||
|
||||
with pytest.raises(DeviceNotConnectedError):
|
||||
_ = camera.async_read()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("index_or_path", TEST_IMAGE_PATHS, ids=TEST_IMAGE_SIZES)
|
||||
@pytest.mark.parametrize(
|
||||
"rotation",
|
||||
[
|
||||
Cv2Rotation.NO_ROTATION,
|
||||
Cv2Rotation.ROTATE_90,
|
||||
Cv2Rotation.ROTATE_180,
|
||||
Cv2Rotation.ROTATE_270,
|
||||
],
|
||||
ids=["no_rot", "rot90", "rot180", "rot270"],
|
||||
)
|
||||
def test_rotation(rotation, index_or_path):
|
||||
filename = Path(index_or_path).name
|
||||
dimensions = filename.split("_")[-1].split(".")[0] # Assumes filenames format (_wxh.png)
|
||||
original_width, original_height = map(int, dimensions.split("x"))
|
||||
|
||||
config = OpenCVCameraConfig(index_or_path=index_or_path, rotation=rotation)
|
||||
camera = OpenCVCamera(config)
|
||||
camera.connect(warmup=False)
|
||||
|
||||
img = camera.read()
|
||||
assert isinstance(img, np.ndarray)
|
||||
|
||||
if rotation in (Cv2Rotation.ROTATE_90, Cv2Rotation.ROTATE_270):
|
||||
assert camera.width == original_height
|
||||
assert camera.height == original_width
|
||||
assert img.shape[:2] == (original_width, original_height)
|
||||
else:
|
||||
assert camera.width == original_width
|
||||
assert camera.height == original_height
|
||||
assert img.shape[:2] == (original_height, original_width)
|
||||
@@ -0,0 +1,204 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# 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.
|
||||
|
||||
# Example of running a specific test:
|
||||
# ```bash
|
||||
# pytest tests/cameras/test_opencv.py::test_connect
|
||||
# ```
|
||||
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from lerobot.common.cameras.configs import Cv2Rotation
|
||||
from lerobot.common.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
|
||||
|
||||
pytest.importorskip("pyrealsense2")
|
||||
|
||||
from lerobot.common.cameras.realsense import RealSenseCamera, RealSenseCameraConfig
|
||||
|
||||
TEST_ARTIFACTS_DIR = Path(__file__).parent.parent / "artifacts" / "cameras"
|
||||
BAG_FILE_PATH = TEST_ARTIFACTS_DIR / "test_rs.bag"
|
||||
|
||||
# NOTE(Steven): For some reason these tests take ~20sec in macOS but only ~2sec in Linux.
|
||||
|
||||
|
||||
def mock_rs_config_enable_device_from_file(rs_config_instance, _sn):
|
||||
return rs_config_instance.enable_device_from_file(str(BAG_FILE_PATH), repeat_playback=True)
|
||||
|
||||
|
||||
def mock_rs_config_enable_device_bad_file(rs_config_instance, _sn):
|
||||
return rs_config_instance.enable_device_from_file("non_existent_file.bag", repeat_playback=True)
|
||||
|
||||
|
||||
@pytest.fixture(name="patch_realsense", autouse=True)
|
||||
def fixture_patch_realsense():
|
||||
"""Automatically mock pyrealsense2.config.enable_device for all tests."""
|
||||
with patch(
|
||||
"pyrealsense2.config.enable_device", side_effect=mock_rs_config_enable_device_from_file
|
||||
) as mock:
|
||||
yield mock
|
||||
|
||||
|
||||
def test_abc_implementation():
|
||||
"""Instantiation should raise an error if the class doesn't implement abstract methods/properties."""
|
||||
config = RealSenseCameraConfig(serial_number_or_name="042")
|
||||
_ = RealSenseCamera(config)
|
||||
|
||||
|
||||
def test_connect():
|
||||
config = RealSenseCameraConfig(serial_number_or_name="042")
|
||||
camera = RealSenseCamera(config)
|
||||
|
||||
camera.connect(warmup=False)
|
||||
assert camera.is_connected
|
||||
|
||||
|
||||
def test_connect_already_connected():
|
||||
config = RealSenseCameraConfig(serial_number_or_name="042")
|
||||
camera = RealSenseCamera(config)
|
||||
camera.connect(warmup=False)
|
||||
|
||||
with pytest.raises(DeviceAlreadyConnectedError):
|
||||
camera.connect(warmup=False)
|
||||
|
||||
|
||||
def test_connect_invalid_camera_path(patch_realsense):
|
||||
patch_realsense.side_effect = mock_rs_config_enable_device_bad_file
|
||||
config = RealSenseCameraConfig(serial_number_or_name="042")
|
||||
camera = RealSenseCamera(config)
|
||||
|
||||
with pytest.raises(ConnectionError):
|
||||
camera.connect(warmup=False)
|
||||
|
||||
|
||||
def test_invalid_width_connect():
|
||||
config = RealSenseCameraConfig(serial_number_or_name="042", width=99999, height=480, fps=30)
|
||||
camera = RealSenseCamera(config)
|
||||
|
||||
with pytest.raises(ConnectionError):
|
||||
camera.connect(warmup=False)
|
||||
|
||||
|
||||
def test_read():
|
||||
config = RealSenseCameraConfig(serial_number_or_name="042", width=640, height=480, fps=30)
|
||||
camera = RealSenseCamera(config)
|
||||
camera.connect(warmup=False)
|
||||
|
||||
img = camera.read()
|
||||
assert isinstance(img, np.ndarray)
|
||||
|
||||
|
||||
def test_read_depth():
|
||||
config = RealSenseCameraConfig(serial_number_or_name="042", width=640, height=480, fps=30, use_depth=True)
|
||||
camera = RealSenseCamera(config)
|
||||
camera.connect(warmup=False)
|
||||
|
||||
img = camera.read_depth(timeout_ms=1000) # NOTE(Steven): Reading depth takes longer
|
||||
assert isinstance(img, np.ndarray)
|
||||
|
||||
|
||||
def test_read_before_connect():
|
||||
config = RealSenseCameraConfig(serial_number_or_name="042")
|
||||
camera = RealSenseCamera(config)
|
||||
|
||||
with pytest.raises(DeviceNotConnectedError):
|
||||
_ = camera.read()
|
||||
|
||||
|
||||
def test_disconnect():
|
||||
config = RealSenseCameraConfig(serial_number_or_name="042")
|
||||
camera = RealSenseCamera(config)
|
||||
camera.connect(warmup=False)
|
||||
|
||||
camera.disconnect()
|
||||
|
||||
assert not camera.is_connected
|
||||
|
||||
|
||||
def test_disconnect_before_connect():
|
||||
config = RealSenseCameraConfig(serial_number_or_name="042")
|
||||
camera = RealSenseCamera(config)
|
||||
|
||||
with pytest.raises(DeviceNotConnectedError):
|
||||
camera.disconnect()
|
||||
|
||||
|
||||
def test_async_read():
|
||||
config = RealSenseCameraConfig(serial_number_or_name="042", width=640, height=480, fps=30)
|
||||
camera = RealSenseCamera(config)
|
||||
camera.connect(warmup=False)
|
||||
|
||||
try:
|
||||
img = camera.async_read()
|
||||
|
||||
assert camera.thread is not None
|
||||
assert camera.thread.is_alive()
|
||||
assert isinstance(img, np.ndarray)
|
||||
finally:
|
||||
if camera.is_connected:
|
||||
camera.disconnect() # To stop/join the thread. Otherwise get warnings when the test ends
|
||||
|
||||
|
||||
def test_async_read_timeout():
|
||||
config = RealSenseCameraConfig(serial_number_or_name="042", width=640, height=480, fps=30)
|
||||
camera = RealSenseCamera(config)
|
||||
camera.connect(warmup=False)
|
||||
|
||||
try:
|
||||
with pytest.raises(TimeoutError):
|
||||
camera.async_read(timeout_ms=0)
|
||||
finally:
|
||||
if camera.is_connected:
|
||||
camera.disconnect()
|
||||
|
||||
|
||||
def test_async_read_before_connect():
|
||||
config = RealSenseCameraConfig(serial_number_or_name="042")
|
||||
camera = RealSenseCamera(config)
|
||||
|
||||
with pytest.raises(DeviceNotConnectedError):
|
||||
_ = camera.async_read()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"rotation",
|
||||
[
|
||||
Cv2Rotation.NO_ROTATION,
|
||||
Cv2Rotation.ROTATE_90,
|
||||
Cv2Rotation.ROTATE_180,
|
||||
Cv2Rotation.ROTATE_270,
|
||||
],
|
||||
ids=["no_rot", "rot90", "rot180", "rot270"],
|
||||
)
|
||||
def test_rotation(rotation):
|
||||
config = RealSenseCameraConfig(serial_number_or_name="042", rotation=rotation)
|
||||
camera = RealSenseCamera(config)
|
||||
camera.connect(warmup=False)
|
||||
|
||||
img = camera.read()
|
||||
assert isinstance(img, np.ndarray)
|
||||
|
||||
if rotation in (Cv2Rotation.ROTATE_90, Cv2Rotation.ROTATE_270):
|
||||
assert camera.width == 480
|
||||
assert camera.height == 640
|
||||
assert img.shape[:2] == (640, 480)
|
||||
else:
|
||||
assert camera.width == 640
|
||||
assert camera.height == 480
|
||||
assert img.shape[:2] == (480, 640)
|
||||
@@ -16,6 +16,7 @@
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from packaging import version
|
||||
from safetensors.torch import load_file
|
||||
from torchvision.transforms import v2
|
||||
from torchvision.transforms.v2 import functional as F # noqa: N812
|
||||
@@ -253,7 +254,14 @@ def test_backward_compatibility_single_transforms(
|
||||
|
||||
|
||||
@require_x86_64_kernel
|
||||
@pytest.mark.skipif(
|
||||
version.parse(torch.__version__) < version.parse("2.7.0"),
|
||||
reason="Test artifacts were generated with PyTorch >= 2.7.0 which has different multinomial behavior",
|
||||
)
|
||||
def test_backward_compatibility_default_config(img_tensor, default_transforms):
|
||||
# NOTE: PyTorch versions have different randomness, it might break this test.
|
||||
# See this PR: https://github.com/huggingface/lerobot/pull/1127.
|
||||
|
||||
cfg = ImageTransformsConfig(enable=True)
|
||||
default_tf = ImageTransforms(cfg)
|
||||
|
||||
|
||||
@@ -0,0 +1,580 @@
|
||||
import abc
|
||||
from typing import Callable
|
||||
|
||||
import dynamixel_sdk as dxl
|
||||
import serial
|
||||
from mock_serial.mock_serial import MockSerial
|
||||
|
||||
from lerobot.common.motors.dynamixel.dynamixel import _split_into_byte_chunks
|
||||
|
||||
from .mock_serial_patch import WaitableStub
|
||||
|
||||
# https://emanual.robotis.com/docs/en/dxl/crc/
|
||||
DXL_CRC_TABLE = [
|
||||
0x0000, 0x8005, 0x800F, 0x000A, 0x801B, 0x001E, 0x0014, 0x8011,
|
||||
0x8033, 0x0036, 0x003C, 0x8039, 0x0028, 0x802D, 0x8027, 0x0022,
|
||||
0x8063, 0x0066, 0x006C, 0x8069, 0x0078, 0x807D, 0x8077, 0x0072,
|
||||
0x0050, 0x8055, 0x805F, 0x005A, 0x804B, 0x004E, 0x0044, 0x8041,
|
||||
0x80C3, 0x00C6, 0x00CC, 0x80C9, 0x00D8, 0x80DD, 0x80D7, 0x00D2,
|
||||
0x00F0, 0x80F5, 0x80FF, 0x00FA, 0x80EB, 0x00EE, 0x00E4, 0x80E1,
|
||||
0x00A0, 0x80A5, 0x80AF, 0x00AA, 0x80BB, 0x00BE, 0x00B4, 0x80B1,
|
||||
0x8093, 0x0096, 0x009C, 0x8099, 0x0088, 0x808D, 0x8087, 0x0082,
|
||||
0x8183, 0x0186, 0x018C, 0x8189, 0x0198, 0x819D, 0x8197, 0x0192,
|
||||
0x01B0, 0x81B5, 0x81BF, 0x01BA, 0x81AB, 0x01AE, 0x01A4, 0x81A1,
|
||||
0x01E0, 0x81E5, 0x81EF, 0x01EA, 0x81FB, 0x01FE, 0x01F4, 0x81F1,
|
||||
0x81D3, 0x01D6, 0x01DC, 0x81D9, 0x01C8, 0x81CD, 0x81C7, 0x01C2,
|
||||
0x0140, 0x8145, 0x814F, 0x014A, 0x815B, 0x015E, 0x0154, 0x8151,
|
||||
0x8173, 0x0176, 0x017C, 0x8179, 0x0168, 0x816D, 0x8167, 0x0162,
|
||||
0x8123, 0x0126, 0x012C, 0x8129, 0x0138, 0x813D, 0x8137, 0x0132,
|
||||
0x0110, 0x8115, 0x811F, 0x011A, 0x810B, 0x010E, 0x0104, 0x8101,
|
||||
0x8303, 0x0306, 0x030C, 0x8309, 0x0318, 0x831D, 0x8317, 0x0312,
|
||||
0x0330, 0x8335, 0x833F, 0x033A, 0x832B, 0x032E, 0x0324, 0x8321,
|
||||
0x0360, 0x8365, 0x836F, 0x036A, 0x837B, 0x037E, 0x0374, 0x8371,
|
||||
0x8353, 0x0356, 0x035C, 0x8359, 0x0348, 0x834D, 0x8347, 0x0342,
|
||||
0x03C0, 0x83C5, 0x83CF, 0x03CA, 0x83DB, 0x03DE, 0x03D4, 0x83D1,
|
||||
0x83F3, 0x03F6, 0x03FC, 0x83F9, 0x03E8, 0x83ED, 0x83E7, 0x03E2,
|
||||
0x83A3, 0x03A6, 0x03AC, 0x83A9, 0x03B8, 0x83BD, 0x83B7, 0x03B2,
|
||||
0x0390, 0x8395, 0x839F, 0x039A, 0x838B, 0x038E, 0x0384, 0x8381,
|
||||
0x0280, 0x8285, 0x828F, 0x028A, 0x829B, 0x029E, 0x0294, 0x8291,
|
||||
0x82B3, 0x02B6, 0x02BC, 0x82B9, 0x02A8, 0x82AD, 0x82A7, 0x02A2,
|
||||
0x82E3, 0x02E6, 0x02EC, 0x82E9, 0x02F8, 0x82FD, 0x82F7, 0x02F2,
|
||||
0x02D0, 0x82D5, 0x82DF, 0x02DA, 0x82CB, 0x02CE, 0x02C4, 0x82C1,
|
||||
0x8243, 0x0246, 0x024C, 0x8249, 0x0258, 0x825D, 0x8257, 0x0252,
|
||||
0x0270, 0x8275, 0x827F, 0x027A, 0x826B, 0x026E, 0x0264, 0x8261,
|
||||
0x0220, 0x8225, 0x822F, 0x022A, 0x823B, 0x023E, 0x0234, 0x8231,
|
||||
0x8213, 0x0216, 0x021C, 0x8219, 0x0208, 0x820D, 0x8207, 0x0202
|
||||
] # fmt: skip
|
||||
|
||||
|
||||
class MockDynamixelPacketv2(abc.ABC):
|
||||
@classmethod
|
||||
def build(cls, dxl_id: int, params: list[int], length: int, *args, **kwargs) -> bytes:
|
||||
packet = cls._build(dxl_id, params, length, *args, **kwargs)
|
||||
packet = cls._add_stuffing(packet)
|
||||
packet = cls._add_crc(packet)
|
||||
return bytes(packet)
|
||||
|
||||
@abc.abstractclassmethod
|
||||
def _build(cls, dxl_id: int, params: list[int], length: int, *args, **kwargs) -> list[int]:
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def _add_stuffing(packet: list[int]) -> list[int]:
|
||||
"""
|
||||
Byte stuffing is a method of adding additional data to generated instruction packets to ensure that
|
||||
the packets are processed successfully. When the byte pattern "0xFF 0xFF 0xFD" appears in a packet,
|
||||
byte stuffing adds 0xFD to the end of the pattern to convert it to “0xFF 0xFF 0xFD 0xFD” to ensure
|
||||
that it is not interpreted as the header at the start of another packet.
|
||||
|
||||
Source: https://emanual.robotis.com/docs/en/dxl/protocol2/#transmission-process
|
||||
|
||||
Args:
|
||||
packet (list[int]): The raw packet without stuffing.
|
||||
|
||||
Returns:
|
||||
list[int]: The packet stuffed if it contained a "0xFF 0xFF 0xFD" byte sequence in its data bytes.
|
||||
"""
|
||||
packet_length_in = dxl.DXL_MAKEWORD(packet[dxl.PKT_LENGTH_L], packet[dxl.PKT_LENGTH_H])
|
||||
packet_length_out = packet_length_in
|
||||
|
||||
temp = [0] * dxl.TXPACKET_MAX_LEN
|
||||
|
||||
# FF FF FD XX ID LEN_L LEN_H
|
||||
temp[dxl.PKT_HEADER0 : dxl.PKT_HEADER0 + dxl.PKT_LENGTH_H + 1] = packet[
|
||||
dxl.PKT_HEADER0 : dxl.PKT_HEADER0 + dxl.PKT_LENGTH_H + 1
|
||||
]
|
||||
|
||||
index = dxl.PKT_INSTRUCTION
|
||||
|
||||
for i in range(0, packet_length_in - 2): # except CRC
|
||||
temp[index] = packet[i + dxl.PKT_INSTRUCTION]
|
||||
index = index + 1
|
||||
if (
|
||||
packet[i + dxl.PKT_INSTRUCTION] == 0xFD
|
||||
and packet[i + dxl.PKT_INSTRUCTION - 1] == 0xFF
|
||||
and packet[i + dxl.PKT_INSTRUCTION - 2] == 0xFF
|
||||
):
|
||||
# FF FF FD
|
||||
temp[index] = 0xFD
|
||||
index = index + 1
|
||||
packet_length_out = packet_length_out + 1
|
||||
|
||||
temp[index] = packet[dxl.PKT_INSTRUCTION + packet_length_in - 2]
|
||||
temp[index + 1] = packet[dxl.PKT_INSTRUCTION + packet_length_in - 1]
|
||||
index = index + 2
|
||||
|
||||
if packet_length_in != packet_length_out:
|
||||
packet = [0] * index
|
||||
|
||||
packet[0:index] = temp[0:index]
|
||||
|
||||
packet[dxl.PKT_LENGTH_L] = dxl.DXL_LOBYTE(packet_length_out)
|
||||
packet[dxl.PKT_LENGTH_H] = dxl.DXL_HIBYTE(packet_length_out)
|
||||
|
||||
return packet
|
||||
|
||||
@staticmethod
|
||||
def _add_crc(packet: list[int]) -> list[int]:
|
||||
"""Computes and add CRC to the packet.
|
||||
|
||||
https://emanual.robotis.com/docs/en/dxl/crc/
|
||||
https://en.wikipedia.org/wiki/Cyclic_redundancy_check
|
||||
|
||||
Args:
|
||||
packet (list[int]): The raw packet without CRC (but with placeholders for it).
|
||||
|
||||
Returns:
|
||||
list[int]: The raw packet with a valid CRC.
|
||||
"""
|
||||
crc = 0
|
||||
for j in range(len(packet) - 2):
|
||||
i = ((crc >> 8) ^ packet[j]) & 0xFF
|
||||
crc = ((crc << 8) ^ DXL_CRC_TABLE[i]) & 0xFFFF
|
||||
|
||||
packet[-2] = dxl.DXL_LOBYTE(crc)
|
||||
packet[-1] = dxl.DXL_HIBYTE(crc)
|
||||
|
||||
return packet
|
||||
|
||||
|
||||
class MockInstructionPacket(MockDynamixelPacketv2):
|
||||
"""
|
||||
Helper class to build valid Dynamixel Protocol 2.0 Instruction Packets.
|
||||
|
||||
Protocol 2.0 Instruction Packet structure
|
||||
https://emanual.robotis.com/docs/en/dxl/protocol2/#instruction-packet
|
||||
|
||||
| Header | Packet ID | Length | Instruction | Params | CRC |
|
||||
| ------------------- | --------- | ----------- | ----------- | ----------------- | ----------- |
|
||||
| 0xFF 0xFF 0xFD 0x00 | ID | Len_L Len_H | Instr | Param 1 … Param N | CRC_L CRC_H |
|
||||
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def _build(cls, dxl_id: int, params: list[int], length: int, instruction: int) -> list[int]:
|
||||
length = len(params) + 3
|
||||
return [
|
||||
0xFF, 0xFF, 0xFD, 0x00, # header
|
||||
dxl_id, # servo id
|
||||
dxl.DXL_LOBYTE(length), # length_l
|
||||
dxl.DXL_HIBYTE(length), # length_h
|
||||
instruction, # instruction type
|
||||
*params, # data bytes
|
||||
0x00, 0x00 # placeholder for CRC
|
||||
] # fmt: skip
|
||||
|
||||
@classmethod
|
||||
def ping(
|
||||
cls,
|
||||
dxl_id: int,
|
||||
) -> bytes:
|
||||
"""
|
||||
Builds a "Ping" broadcast instruction.
|
||||
https://emanual.robotis.com/docs/en/dxl/protocol2/#ping-0x01
|
||||
|
||||
No parameters required.
|
||||
"""
|
||||
return cls.build(dxl_id=dxl_id, params=[], length=3, instruction=dxl.INST_PING)
|
||||
|
||||
@classmethod
|
||||
def read(
|
||||
cls,
|
||||
dxl_id: int,
|
||||
start_address: int,
|
||||
data_length: int,
|
||||
) -> bytes:
|
||||
"""
|
||||
Builds a "Read" instruction.
|
||||
https://emanual.robotis.com/docs/en/dxl/protocol2/#read-0x02
|
||||
|
||||
The parameters for Read (Protocol 2.0) are:
|
||||
param[0] = start_address L
|
||||
param[1] = start_address H
|
||||
param[2] = data_length L
|
||||
param[3] = data_length H
|
||||
|
||||
And 'length' = data_length + 5, where:
|
||||
+1 is for instruction byte,
|
||||
+2 is for the length bytes,
|
||||
+2 is for the CRC at the end.
|
||||
"""
|
||||
params = [
|
||||
dxl.DXL_LOBYTE(start_address),
|
||||
dxl.DXL_HIBYTE(start_address),
|
||||
dxl.DXL_LOBYTE(data_length),
|
||||
dxl.DXL_HIBYTE(data_length),
|
||||
]
|
||||
length = len(params) + 3
|
||||
# length = data_length + 5
|
||||
return cls.build(dxl_id=dxl_id, params=params, length=length, instruction=dxl.INST_READ)
|
||||
|
||||
@classmethod
|
||||
def write(
|
||||
cls,
|
||||
dxl_id: int,
|
||||
value: int,
|
||||
start_address: int,
|
||||
data_length: int,
|
||||
) -> bytes:
|
||||
"""
|
||||
Builds a "Write" instruction.
|
||||
https://emanual.robotis.com/docs/en/dxl/protocol2/#write-0x03
|
||||
|
||||
The parameters for Write (Protocol 2.0) are:
|
||||
param[0] = start_address L
|
||||
param[1] = start_address H
|
||||
param[2] = 1st Byte
|
||||
param[3] = 2nd Byte
|
||||
...
|
||||
param[1+X] = X-th Byte
|
||||
|
||||
And 'length' = data_length + 5, where:
|
||||
+1 is for instruction byte,
|
||||
+2 is for the length bytes,
|
||||
+2 is for the CRC at the end.
|
||||
"""
|
||||
data = _split_into_byte_chunks(value, data_length)
|
||||
params = [
|
||||
dxl.DXL_LOBYTE(start_address),
|
||||
dxl.DXL_HIBYTE(start_address),
|
||||
*data,
|
||||
]
|
||||
length = data_length + 5
|
||||
return cls.build(dxl_id=dxl_id, params=params, length=length, instruction=dxl.INST_WRITE)
|
||||
|
||||
@classmethod
|
||||
def sync_read(
|
||||
cls,
|
||||
dxl_ids: list[int],
|
||||
start_address: int,
|
||||
data_length: int,
|
||||
) -> bytes:
|
||||
"""
|
||||
Builds a "Sync_Read" broadcast instruction.
|
||||
https://emanual.robotis.com/docs/en/dxl/protocol2/#sync-read-0x82
|
||||
|
||||
The parameters for Sync_Read (Protocol 2.0) are:
|
||||
param[0] = start_address L
|
||||
param[1] = start_address H
|
||||
param[2] = data_length L
|
||||
param[3] = data_length H
|
||||
param[4+] = motor IDs to read from
|
||||
|
||||
And 'length' = (number_of_params + 7), where:
|
||||
+1 is for instruction byte,
|
||||
+2 is for the address bytes,
|
||||
+2 is for the length bytes,
|
||||
+2 is for the CRC at the end.
|
||||
"""
|
||||
params = [
|
||||
dxl.DXL_LOBYTE(start_address),
|
||||
dxl.DXL_HIBYTE(start_address),
|
||||
dxl.DXL_LOBYTE(data_length),
|
||||
dxl.DXL_HIBYTE(data_length),
|
||||
*dxl_ids,
|
||||
]
|
||||
length = len(dxl_ids) + 7
|
||||
return cls.build(
|
||||
dxl_id=dxl.BROADCAST_ID, params=params, length=length, instruction=dxl.INST_SYNC_READ
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def sync_write(
|
||||
cls,
|
||||
ids_values: dict[int, int],
|
||||
start_address: int,
|
||||
data_length: int,
|
||||
) -> bytes:
|
||||
"""
|
||||
Builds a "Sync_Write" broadcast instruction.
|
||||
https://emanual.robotis.com/docs/en/dxl/protocol2/#sync-write-0x83
|
||||
|
||||
The parameters for Sync_Write (Protocol 2.0) are:
|
||||
param[0] = start_address L
|
||||
param[1] = start_address H
|
||||
param[2] = data_length L
|
||||
param[3] = data_length H
|
||||
param[5] = [1st motor] ID
|
||||
param[5+1] = [1st motor] 1st Byte
|
||||
param[5+2] = [1st motor] 2nd Byte
|
||||
...
|
||||
param[5+X] = [1st motor] X-th Byte
|
||||
param[6] = [2nd motor] ID
|
||||
param[6+1] = [2nd motor] 1st Byte
|
||||
param[6+2] = [2nd motor] 2nd Byte
|
||||
...
|
||||
param[6+X] = [2nd motor] X-th Byte
|
||||
|
||||
And 'length' = ((number_of_params * 1 + data_length) + 7), where:
|
||||
+1 is for instruction byte,
|
||||
+2 is for the address bytes,
|
||||
+2 is for the length bytes,
|
||||
+2 is for the CRC at the end.
|
||||
"""
|
||||
data = []
|
||||
for id_, value in ids_values.items():
|
||||
split_value = _split_into_byte_chunks(value, data_length)
|
||||
data += [id_, *split_value]
|
||||
params = [
|
||||
dxl.DXL_LOBYTE(start_address),
|
||||
dxl.DXL_HIBYTE(start_address),
|
||||
dxl.DXL_LOBYTE(data_length),
|
||||
dxl.DXL_HIBYTE(data_length),
|
||||
*data,
|
||||
]
|
||||
length = len(ids_values) * (1 + data_length) + 7
|
||||
return cls.build(
|
||||
dxl_id=dxl.BROADCAST_ID, params=params, length=length, instruction=dxl.INST_SYNC_WRITE
|
||||
)
|
||||
|
||||
|
||||
class MockStatusPacket(MockDynamixelPacketv2):
|
||||
"""
|
||||
Helper class to build valid Dynamixel Protocol 2.0 Status Packets.
|
||||
|
||||
Protocol 2.0 Status Packet structure
|
||||
https://emanual.robotis.com/docs/en/dxl/protocol2/#status-packet
|
||||
|
||||
| Header | Packet ID | Length | Instruction | Error | Params | CRC |
|
||||
| ------------------- | --------- | ----------- | ----------- | ----- | ----------------- | ----------- |
|
||||
| 0xFF 0xFF 0xFD 0x00 | ID | Len_L Len_H | 0x55 | Err | Param 1 … Param N | CRC_L CRC_H |
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def _build(cls, dxl_id: int, params: list[int], length: int, error: int = 0) -> list[int]:
|
||||
return [
|
||||
0xFF, 0xFF, 0xFD, 0x00, # header
|
||||
dxl_id, # servo id
|
||||
dxl.DXL_LOBYTE(length), # length_l
|
||||
dxl.DXL_HIBYTE(length), # length_h
|
||||
0x55, # instruction = 'status'
|
||||
error, # error
|
||||
*params, # data bytes
|
||||
0x00, 0x00 # placeholder for CRC
|
||||
] # fmt: skip
|
||||
|
||||
@classmethod
|
||||
def ping(cls, dxl_id: int, model_nb: int = 1190, firm_ver: int = 50, error: int = 0) -> bytes:
|
||||
"""
|
||||
Builds a 'Ping' status packet.
|
||||
https://emanual.robotis.com/docs/en/dxl/protocol2/#ping-0x01
|
||||
|
||||
Args:
|
||||
dxl_id (int): ID of the servo responding.
|
||||
model_nb (int, optional): Desired 'model number' to be returned in the packet. Defaults to 1190
|
||||
which corresponds to a XL330-M077-T.
|
||||
firm_ver (int, optional): Desired 'firmware version' to be returned in the packet.
|
||||
Defaults to 50.
|
||||
|
||||
Returns:
|
||||
bytes: The raw 'Ping' status packet ready to be sent through serial.
|
||||
"""
|
||||
params = [dxl.DXL_LOBYTE(model_nb), dxl.DXL_HIBYTE(model_nb), firm_ver]
|
||||
length = 7
|
||||
return cls.build(dxl_id, params=params, length=length, error=error)
|
||||
|
||||
@classmethod
|
||||
def read(cls, dxl_id: int, value: int, param_length: int, error: int = 0) -> bytes:
|
||||
"""
|
||||
Builds a 'Read' status packet (also works for 'Sync Read')
|
||||
https://emanual.robotis.com/docs/en/dxl/protocol2/#read-0x02
|
||||
https://emanual.robotis.com/docs/en/dxl/protocol2/#sync-read-0x82
|
||||
|
||||
Args:
|
||||
dxl_id (int): ID of the servo responding.
|
||||
value (int): Desired value to be returned in the packet.
|
||||
param_length (int): The address length as reported in the control table.
|
||||
|
||||
Returns:
|
||||
bytes: The raw 'Present_Position' status packet ready to be sent through serial.
|
||||
"""
|
||||
params = _split_into_byte_chunks(value, param_length)
|
||||
length = param_length + 4
|
||||
return cls.build(dxl_id, params=params, length=length, error=error)
|
||||
|
||||
|
||||
class MockPortHandler(dxl.PortHandler):
|
||||
"""
|
||||
This class overwrite the 'setupPort' method of the Dynamixel PortHandler because it can specify
|
||||
baudrates that are not supported with a serial port on MacOS.
|
||||
"""
|
||||
|
||||
def setupPort(self, cflag_baud): # noqa: N802
|
||||
if self.is_open:
|
||||
self.closePort()
|
||||
|
||||
self.ser = serial.Serial(
|
||||
port=self.port_name,
|
||||
# baudrate=self.baudrate, <- This will fail on MacOS
|
||||
# parity = serial.PARITY_ODD,
|
||||
# stopbits = serial.STOPBITS_TWO,
|
||||
bytesize=serial.EIGHTBITS,
|
||||
timeout=0,
|
||||
)
|
||||
self.is_open = True
|
||||
self.ser.reset_input_buffer()
|
||||
self.tx_time_per_byte = (1000.0 / self.baudrate) * 10.0
|
||||
|
||||
return True
|
||||
|
||||
|
||||
class MockMotors(MockSerial):
|
||||
"""
|
||||
This class will simulate physical motors by responding with valid status packets upon receiving some
|
||||
instruction packets. It is meant to test MotorsBus classes.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
@property
|
||||
def stubs(self) -> dict[str, WaitableStub]:
|
||||
return super().stubs
|
||||
|
||||
def stub(self, *, name=None, **kwargs):
|
||||
new_stub = WaitableStub(**kwargs)
|
||||
self._MockSerial__stubs[name or new_stub.receive_bytes] = new_stub
|
||||
return new_stub
|
||||
|
||||
def build_broadcast_ping_stub(
|
||||
self, ids_models: dict[int, list[int]] | None = None, num_invalid_try: int = 0
|
||||
) -> str:
|
||||
ping_request = MockInstructionPacket.ping(dxl.BROADCAST_ID)
|
||||
return_packets = b"".join(MockStatusPacket.ping(id_, model) for id_, model in ids_models.items())
|
||||
ping_response = self._build_send_fn(return_packets, num_invalid_try)
|
||||
|
||||
stub_name = "Ping_" + "_".join([str(id_) for id_ in ids_models])
|
||||
self.stub(
|
||||
name=stub_name,
|
||||
receive_bytes=ping_request,
|
||||
send_fn=ping_response,
|
||||
)
|
||||
return stub_name
|
||||
|
||||
def build_ping_stub(
|
||||
self, dxl_id: int, model_nb: int, firm_ver: int = 50, num_invalid_try: int = 0, error: int = 0
|
||||
) -> str:
|
||||
ping_request = MockInstructionPacket.ping(dxl_id)
|
||||
return_packet = MockStatusPacket.ping(dxl_id, model_nb, firm_ver, error)
|
||||
ping_response = self._build_send_fn(return_packet, num_invalid_try)
|
||||
stub_name = f"Ping_{dxl_id}"
|
||||
self.stub(
|
||||
name=stub_name,
|
||||
receive_bytes=ping_request,
|
||||
send_fn=ping_response,
|
||||
)
|
||||
return stub_name
|
||||
|
||||
def build_read_stub(
|
||||
self,
|
||||
address: int,
|
||||
length: int,
|
||||
dxl_id: int,
|
||||
value: int,
|
||||
reply: bool = True,
|
||||
error: int = 0,
|
||||
num_invalid_try: int = 0,
|
||||
) -> str:
|
||||
read_request = MockInstructionPacket.read(dxl_id, address, length)
|
||||
return_packet = MockStatusPacket.read(dxl_id, value, length, error) if reply else b""
|
||||
read_response = self._build_send_fn(return_packet, num_invalid_try)
|
||||
stub_name = f"Read_{address}_{length}_{dxl_id}_{value}_{error}"
|
||||
self.stub(
|
||||
name=stub_name,
|
||||
receive_bytes=read_request,
|
||||
send_fn=read_response,
|
||||
)
|
||||
return stub_name
|
||||
|
||||
def build_write_stub(
|
||||
self,
|
||||
address: int,
|
||||
length: int,
|
||||
dxl_id: int,
|
||||
value: int,
|
||||
reply: bool = True,
|
||||
error: int = 0,
|
||||
num_invalid_try: int = 0,
|
||||
) -> str:
|
||||
sync_read_request = MockInstructionPacket.write(dxl_id, value, address, length)
|
||||
return_packet = MockStatusPacket.build(dxl_id, params=[], length=4, error=error) if reply else b""
|
||||
stub_name = f"Write_{address}_{length}_{dxl_id}"
|
||||
self.stub(
|
||||
name=stub_name,
|
||||
receive_bytes=sync_read_request,
|
||||
send_fn=self._build_send_fn(return_packet, num_invalid_try),
|
||||
)
|
||||
return stub_name
|
||||
|
||||
def build_sync_read_stub(
|
||||
self,
|
||||
address: int,
|
||||
length: int,
|
||||
ids_values: dict[int, int],
|
||||
reply: bool = True,
|
||||
num_invalid_try: int = 0,
|
||||
) -> str:
|
||||
sync_read_request = MockInstructionPacket.sync_read(list(ids_values), address, length)
|
||||
return_packets = (
|
||||
b"".join(MockStatusPacket.read(id_, pos, length) for id_, pos in ids_values.items())
|
||||
if reply
|
||||
else b""
|
||||
)
|
||||
sync_read_response = self._build_send_fn(return_packets, num_invalid_try)
|
||||
stub_name = f"Sync_Read_{address}_{length}_" + "_".join([str(id_) for id_ in ids_values])
|
||||
self.stub(
|
||||
name=stub_name,
|
||||
receive_bytes=sync_read_request,
|
||||
send_fn=sync_read_response,
|
||||
)
|
||||
return stub_name
|
||||
|
||||
def build_sequential_sync_read_stub(
|
||||
self, address: int, length: int, ids_values: dict[int, list[int]] | None = None
|
||||
) -> str:
|
||||
sequence_length = len(next(iter(ids_values.values())))
|
||||
assert all(len(positions) == sequence_length for positions in ids_values.values())
|
||||
sync_read_request = MockInstructionPacket.sync_read(list(ids_values), address, length)
|
||||
sequential_packets = []
|
||||
for count in range(sequence_length):
|
||||
return_packets = b"".join(
|
||||
MockStatusPacket.read(id_, positions[count], length) for id_, positions in ids_values.items()
|
||||
)
|
||||
sequential_packets.append(return_packets)
|
||||
|
||||
sync_read_response = self._build_sequential_send_fn(sequential_packets)
|
||||
stub_name = f"Seq_Sync_Read_{address}_{length}_" + "_".join([str(id_) for id_ in ids_values])
|
||||
self.stub(
|
||||
name=stub_name,
|
||||
receive_bytes=sync_read_request,
|
||||
send_fn=sync_read_response,
|
||||
)
|
||||
return stub_name
|
||||
|
||||
def build_sync_write_stub(
|
||||
self, address: int, length: int, ids_values: dict[int, int], num_invalid_try: int = 0
|
||||
) -> str:
|
||||
sync_read_request = MockInstructionPacket.sync_write(ids_values, address, length)
|
||||
stub_name = f"Sync_Write_{address}_{length}_" + "_".join([str(id_) for id_ in ids_values])
|
||||
self.stub(
|
||||
name=stub_name,
|
||||
receive_bytes=sync_read_request,
|
||||
send_fn=self._build_send_fn(b"", num_invalid_try),
|
||||
)
|
||||
return stub_name
|
||||
|
||||
@staticmethod
|
||||
def _build_send_fn(packet: bytes, num_invalid_try: int = 0) -> Callable[[int], bytes]:
|
||||
def send_fn(_call_count: int) -> bytes:
|
||||
if num_invalid_try >= _call_count:
|
||||
return b""
|
||||
return packet
|
||||
|
||||
return send_fn
|
||||
|
||||
@staticmethod
|
||||
def _build_sequential_send_fn(packets: list[bytes]) -> Callable[[int], bytes]:
|
||||
def send_fn(_call_count: int) -> bytes:
|
||||
return packets[_call_count - 1]
|
||||
|
||||
return send_fn
|
||||
@@ -0,0 +1,428 @@
|
||||
import abc
|
||||
from typing import Callable
|
||||
|
||||
import scservo_sdk as scs
|
||||
import serial
|
||||
from mock_serial import MockSerial
|
||||
|
||||
from lerobot.common.motors.feetech.feetech import _split_into_byte_chunks, patch_setPacketTimeout
|
||||
|
||||
from .mock_serial_patch import WaitableStub
|
||||
|
||||
|
||||
class MockFeetechPacket(abc.ABC):
|
||||
@classmethod
|
||||
def build(cls, scs_id: int, params: list[int], length: int, *args, **kwargs) -> bytes:
|
||||
packet = cls._build(scs_id, params, length, *args, **kwargs)
|
||||
packet = cls._add_checksum(packet)
|
||||
return bytes(packet)
|
||||
|
||||
@abc.abstractclassmethod
|
||||
def _build(cls, scs_id: int, params: list[int], length: int, *args, **kwargs) -> list[int]:
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def _add_checksum(packet: list[int]) -> list[int]:
|
||||
checksum = 0
|
||||
for id_ in range(2, len(packet) - 1): # except header & checksum
|
||||
checksum += packet[id_]
|
||||
|
||||
packet[-1] = ~checksum & 0xFF
|
||||
|
||||
return packet
|
||||
|
||||
|
||||
class MockInstructionPacket(MockFeetechPacket):
|
||||
"""
|
||||
Helper class to build valid Feetech Instruction Packets.
|
||||
|
||||
Instruction Packet structure
|
||||
(from https://files.waveshare.com/upload/2/27/Communication_Protocol_User_Manual-EN%28191218-0923%29.pdf)
|
||||
|
||||
| Header | Packet ID | Length | Instruction | Params | Checksum |
|
||||
| --------- | --------- | ------ | ----------- | ----------------- | -------- |
|
||||
| 0xFF 0xFF | ID | Len | Instr | Param 1 … Param N | Sum |
|
||||
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def _build(cls, scs_id: int, params: list[int], length: int, instruction: int) -> list[int]:
|
||||
return [
|
||||
0xFF, 0xFF, # header
|
||||
scs_id, # servo id
|
||||
length, # length
|
||||
instruction, # instruction type
|
||||
*params, # data bytes
|
||||
0x00, # placeholder for checksum
|
||||
] # fmt: skip
|
||||
|
||||
@classmethod
|
||||
def ping(
|
||||
cls,
|
||||
scs_id: int,
|
||||
) -> bytes:
|
||||
"""
|
||||
Builds a "Ping" broadcast instruction.
|
||||
|
||||
No parameters required.
|
||||
"""
|
||||
return cls.build(scs_id=scs_id, params=[], length=2, instruction=scs.INST_PING)
|
||||
|
||||
@classmethod
|
||||
def read(
|
||||
cls,
|
||||
scs_id: int,
|
||||
start_address: int,
|
||||
data_length: int,
|
||||
) -> bytes:
|
||||
"""
|
||||
Builds a "Read" instruction.
|
||||
|
||||
The parameters for Read are:
|
||||
param[0] = start_address
|
||||
param[1] = data_length
|
||||
|
||||
And 'length' = 4, where:
|
||||
+1 is for instruction byte,
|
||||
+1 is for the address byte,
|
||||
+1 is for the length bytes,
|
||||
+1 is for the checksum at the end.
|
||||
"""
|
||||
params = [start_address, data_length]
|
||||
length = 4
|
||||
return cls.build(scs_id=scs_id, params=params, length=length, instruction=scs.INST_READ)
|
||||
|
||||
@classmethod
|
||||
def write(
|
||||
cls,
|
||||
scs_id: int,
|
||||
value: int,
|
||||
start_address: int,
|
||||
data_length: int,
|
||||
) -> bytes:
|
||||
"""
|
||||
Builds a "Write" instruction.
|
||||
|
||||
The parameters for Write are:
|
||||
param[0] = start_address L
|
||||
param[1] = start_address H
|
||||
param[2] = 1st Byte
|
||||
param[3] = 2nd Byte
|
||||
...
|
||||
param[1+X] = X-th Byte
|
||||
|
||||
And 'length' = data_length + 3, where:
|
||||
+1 is for instruction byte,
|
||||
+1 is for the length bytes,
|
||||
+1 is for the checksum at the end.
|
||||
"""
|
||||
data = _split_into_byte_chunks(value, data_length)
|
||||
params = [start_address, *data]
|
||||
length = data_length + 3
|
||||
return cls.build(scs_id=scs_id, params=params, length=length, instruction=scs.INST_WRITE)
|
||||
|
||||
@classmethod
|
||||
def sync_read(
|
||||
cls,
|
||||
scs_ids: list[int],
|
||||
start_address: int,
|
||||
data_length: int,
|
||||
) -> bytes:
|
||||
"""
|
||||
Builds a "Sync_Read" broadcast instruction.
|
||||
|
||||
The parameters for Sync Read are:
|
||||
param[0] = start_address
|
||||
param[1] = data_length
|
||||
param[2+] = motor IDs to read from
|
||||
|
||||
And 'length' = (number_of_params + 4), where:
|
||||
+1 is for instruction byte,
|
||||
+1 is for the address byte,
|
||||
+1 is for the length bytes,
|
||||
+1 is for the checksum at the end.
|
||||
"""
|
||||
params = [start_address, data_length, *scs_ids]
|
||||
length = len(scs_ids) + 4
|
||||
return cls.build(
|
||||
scs_id=scs.BROADCAST_ID, params=params, length=length, instruction=scs.INST_SYNC_READ
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def sync_write(
|
||||
cls,
|
||||
ids_values: dict[int, int],
|
||||
start_address: int,
|
||||
data_length: int,
|
||||
) -> bytes:
|
||||
"""
|
||||
Builds a "Sync_Write" broadcast instruction.
|
||||
|
||||
The parameters for Sync_Write are:
|
||||
param[0] = start_address
|
||||
param[1] = data_length
|
||||
param[2] = [1st motor] ID
|
||||
param[2+1] = [1st motor] 1st Byte
|
||||
param[2+2] = [1st motor] 2nd Byte
|
||||
...
|
||||
param[5+X] = [1st motor] X-th Byte
|
||||
param[6] = [2nd motor] ID
|
||||
param[6+1] = [2nd motor] 1st Byte
|
||||
param[6+2] = [2nd motor] 2nd Byte
|
||||
...
|
||||
param[6+X] = [2nd motor] X-th Byte
|
||||
|
||||
And 'length' = ((number_of_params * 1 + data_length) + 4), where:
|
||||
+1 is for instruction byte,
|
||||
+1 is for the address byte,
|
||||
+1 is for the length bytes,
|
||||
+1 is for the checksum at the end.
|
||||
"""
|
||||
data = []
|
||||
for id_, value in ids_values.items():
|
||||
split_value = _split_into_byte_chunks(value, data_length)
|
||||
data += [id_, *split_value]
|
||||
params = [start_address, data_length, *data]
|
||||
length = len(ids_values) * (1 + data_length) + 4
|
||||
return cls.build(
|
||||
scs_id=scs.BROADCAST_ID, params=params, length=length, instruction=scs.INST_SYNC_WRITE
|
||||
)
|
||||
|
||||
|
||||
class MockStatusPacket(MockFeetechPacket):
|
||||
"""
|
||||
Helper class to build valid Feetech Status Packets.
|
||||
|
||||
Status Packet structure
|
||||
(from https://files.waveshare.com/upload/2/27/Communication_Protocol_User_Manual-EN%28191218-0923%29.pdf)
|
||||
|
||||
| Header | Packet ID | Length | Error | Params | Checksum |
|
||||
| --------- | --------- | ------ | ----- | ----------------- | -------- |
|
||||
| 0xFF 0xFF | ID | Len | Err | Param 1 … Param N | Sum |
|
||||
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def _build(cls, scs_id: int, params: list[int], length: int, error: int = 0) -> list[int]:
|
||||
return [
|
||||
0xFF, 0xFF, # header
|
||||
scs_id, # servo id
|
||||
length, # length
|
||||
error, # status
|
||||
*params, # data bytes
|
||||
0x00, # placeholder for checksum
|
||||
] # fmt: skip
|
||||
|
||||
@classmethod
|
||||
def ping(cls, scs_id: int, error: int = 0) -> bytes:
|
||||
"""Builds a 'Ping' status packet.
|
||||
|
||||
Args:
|
||||
scs_id (int): ID of the servo responding.
|
||||
error (int, optional): Error to be returned. Defaults to 0 (success).
|
||||
|
||||
Returns:
|
||||
bytes: The raw 'Ping' status packet ready to be sent through serial.
|
||||
"""
|
||||
return cls.build(scs_id, params=[], length=2, error=error)
|
||||
|
||||
@classmethod
|
||||
def read(cls, scs_id: int, value: int, param_length: int, error: int = 0) -> bytes:
|
||||
"""Builds a 'Read' status packet.
|
||||
|
||||
Args:
|
||||
scs_id (int): ID of the servo responding.
|
||||
value (int): Desired value to be returned in the packet.
|
||||
param_length (int): The address length as reported in the control table.
|
||||
|
||||
Returns:
|
||||
bytes: The raw 'Sync Read' status packet ready to be sent through serial.
|
||||
"""
|
||||
params = _split_into_byte_chunks(value, param_length)
|
||||
length = param_length + 2
|
||||
return cls.build(scs_id, params=params, length=length, error=error)
|
||||
|
||||
|
||||
class MockPortHandler(scs.PortHandler):
|
||||
"""
|
||||
This class overwrite the 'setupPort' method of the Feetech PortHandler because it can specify
|
||||
baudrates that are not supported with a serial port on MacOS.
|
||||
"""
|
||||
|
||||
def setupPort(self, cflag_baud): # noqa: N802
|
||||
if self.is_open:
|
||||
self.closePort()
|
||||
|
||||
self.ser = serial.Serial(
|
||||
port=self.port_name,
|
||||
# baudrate=self.baudrate, <- This will fail on MacOS
|
||||
# parity = serial.PARITY_ODD,
|
||||
# stopbits = serial.STOPBITS_TWO,
|
||||
bytesize=serial.EIGHTBITS,
|
||||
timeout=0,
|
||||
)
|
||||
self.is_open = True
|
||||
self.ser.reset_input_buffer()
|
||||
self.tx_time_per_byte = (1000.0 / self.baudrate) * 10.0
|
||||
|
||||
return True
|
||||
|
||||
def setPacketTimeout(self, packet_length): # noqa: N802
|
||||
return patch_setPacketTimeout(self, packet_length)
|
||||
|
||||
|
||||
class MockMotors(MockSerial):
|
||||
"""
|
||||
This class will simulate physical motors by responding with valid status packets upon receiving some
|
||||
instruction packets. It is meant to test MotorsBus classes.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
@property
|
||||
def stubs(self) -> dict[str, WaitableStub]:
|
||||
return super().stubs
|
||||
|
||||
def stub(self, *, name=None, **kwargs):
|
||||
new_stub = WaitableStub(**kwargs)
|
||||
self._MockSerial__stubs[name or new_stub.receive_bytes] = new_stub
|
||||
return new_stub
|
||||
|
||||
def build_broadcast_ping_stub(self, ids: list[int] | None = None, num_invalid_try: int = 0) -> str:
|
||||
ping_request = MockInstructionPacket.ping(scs.BROADCAST_ID)
|
||||
return_packets = b"".join(MockStatusPacket.ping(id_) for id_ in ids)
|
||||
ping_response = self._build_send_fn(return_packets, num_invalid_try)
|
||||
stub_name = "Ping_" + "_".join([str(id_) for id_ in ids])
|
||||
self.stub(
|
||||
name=stub_name,
|
||||
receive_bytes=ping_request,
|
||||
send_fn=ping_response,
|
||||
)
|
||||
return stub_name
|
||||
|
||||
def build_ping_stub(self, scs_id: int, num_invalid_try: int = 0, error: int = 0) -> str:
|
||||
ping_request = MockInstructionPacket.ping(scs_id)
|
||||
return_packet = MockStatusPacket.ping(scs_id, error)
|
||||
ping_response = self._build_send_fn(return_packet, num_invalid_try)
|
||||
stub_name = f"Ping_{scs_id}_{error}"
|
||||
self.stub(
|
||||
name=stub_name,
|
||||
receive_bytes=ping_request,
|
||||
send_fn=ping_response,
|
||||
)
|
||||
return stub_name
|
||||
|
||||
def build_read_stub(
|
||||
self,
|
||||
address: int,
|
||||
length: int,
|
||||
scs_id: int,
|
||||
value: int,
|
||||
reply: bool = True,
|
||||
error: int = 0,
|
||||
num_invalid_try: int = 0,
|
||||
) -> str:
|
||||
read_request = MockInstructionPacket.read(scs_id, address, length)
|
||||
return_packet = MockStatusPacket.read(scs_id, value, length, error) if reply else b""
|
||||
read_response = self._build_send_fn(return_packet, num_invalid_try)
|
||||
stub_name = f"Read_{address}_{length}_{scs_id}_{value}_{error}"
|
||||
self.stub(
|
||||
name=stub_name,
|
||||
receive_bytes=read_request,
|
||||
send_fn=read_response,
|
||||
)
|
||||
return stub_name
|
||||
|
||||
def build_write_stub(
|
||||
self,
|
||||
address: int,
|
||||
length: int,
|
||||
scs_id: int,
|
||||
value: int,
|
||||
reply: bool = True,
|
||||
error: int = 0,
|
||||
num_invalid_try: int = 0,
|
||||
) -> str:
|
||||
sync_read_request = MockInstructionPacket.write(scs_id, value, address, length)
|
||||
return_packet = MockStatusPacket.build(scs_id, params=[], length=2, error=error) if reply else b""
|
||||
stub_name = f"Write_{address}_{length}_{scs_id}"
|
||||
self.stub(
|
||||
name=stub_name,
|
||||
receive_bytes=sync_read_request,
|
||||
send_fn=self._build_send_fn(return_packet, num_invalid_try),
|
||||
)
|
||||
return stub_name
|
||||
|
||||
def build_sync_read_stub(
|
||||
self,
|
||||
address: int,
|
||||
length: int,
|
||||
ids_values: dict[int, int],
|
||||
reply: bool = True,
|
||||
num_invalid_try: int = 0,
|
||||
) -> str:
|
||||
sync_read_request = MockInstructionPacket.sync_read(list(ids_values), address, length)
|
||||
return_packets = (
|
||||
b"".join(MockStatusPacket.read(id_, pos, length) for id_, pos in ids_values.items())
|
||||
if reply
|
||||
else b""
|
||||
)
|
||||
sync_read_response = self._build_send_fn(return_packets, num_invalid_try)
|
||||
stub_name = f"Sync_Read_{address}_{length}_" + "_".join([str(id_) for id_ in ids_values])
|
||||
self.stub(
|
||||
name=stub_name,
|
||||
receive_bytes=sync_read_request,
|
||||
send_fn=sync_read_response,
|
||||
)
|
||||
return stub_name
|
||||
|
||||
def build_sequential_sync_read_stub(
|
||||
self, address: int, length: int, ids_values: dict[int, list[int]] | None = None
|
||||
) -> str:
|
||||
sequence_length = len(next(iter(ids_values.values())))
|
||||
assert all(len(positions) == sequence_length for positions in ids_values.values())
|
||||
sync_read_request = MockInstructionPacket.sync_read(list(ids_values), address, length)
|
||||
sequential_packets = []
|
||||
for count in range(sequence_length):
|
||||
return_packets = b"".join(
|
||||
MockStatusPacket.read(id_, positions[count], length) for id_, positions in ids_values.items()
|
||||
)
|
||||
sequential_packets.append(return_packets)
|
||||
|
||||
sync_read_response = self._build_sequential_send_fn(sequential_packets)
|
||||
stub_name = f"Seq_Sync_Read_{address}_{length}_" + "_".join([str(id_) for id_ in ids_values])
|
||||
self.stub(
|
||||
name=stub_name,
|
||||
receive_bytes=sync_read_request,
|
||||
send_fn=sync_read_response,
|
||||
)
|
||||
return stub_name
|
||||
|
||||
def build_sync_write_stub(
|
||||
self, address: int, length: int, ids_values: dict[int, int], num_invalid_try: int = 0
|
||||
) -> str:
|
||||
sync_read_request = MockInstructionPacket.sync_write(ids_values, address, length)
|
||||
stub_name = f"Sync_Write_{address}_{length}_" + "_".join([str(id_) for id_ in ids_values])
|
||||
self.stub(
|
||||
name=stub_name,
|
||||
receive_bytes=sync_read_request,
|
||||
send_fn=self._build_send_fn(b"", num_invalid_try),
|
||||
)
|
||||
return stub_name
|
||||
|
||||
@staticmethod
|
||||
def _build_send_fn(packet: bytes, num_invalid_try: int = 0) -> Callable[[int], bytes]:
|
||||
def send_fn(_call_count: int) -> bytes:
|
||||
if num_invalid_try >= _call_count:
|
||||
return b""
|
||||
return packet
|
||||
|
||||
return send_fn
|
||||
|
||||
@staticmethod
|
||||
def _build_sequential_send_fn(packets: list[bytes]) -> Callable[[int], bytes]:
|
||||
def send_fn(_call_count: int) -> bytes:
|
||||
return packets[_call_count - 1]
|
||||
|
||||
return send_fn
|
||||
@@ -0,0 +1,138 @@
|
||||
# ruff: noqa: N802
|
||||
|
||||
from lerobot.common.motors.motors_bus import (
|
||||
Motor,
|
||||
MotorsBus,
|
||||
)
|
||||
|
||||
DUMMY_CTRL_TABLE_1 = {
|
||||
"Firmware_Version": (0, 1),
|
||||
"Model_Number": (1, 2),
|
||||
"Present_Position": (3, 4),
|
||||
"Goal_Position": (11, 2),
|
||||
}
|
||||
|
||||
DUMMY_CTRL_TABLE_2 = {
|
||||
"Model_Number": (0, 2),
|
||||
"Firmware_Version": (2, 1),
|
||||
"Present_Position": (3, 4),
|
||||
"Present_Velocity": (7, 4),
|
||||
"Goal_Position": (11, 4),
|
||||
"Goal_Velocity": (15, 4),
|
||||
"Lock": (19, 1),
|
||||
}
|
||||
|
||||
DUMMY_MODEL_CTRL_TABLE = {
|
||||
"model_1": DUMMY_CTRL_TABLE_1,
|
||||
"model_2": DUMMY_CTRL_TABLE_2,
|
||||
"model_3": DUMMY_CTRL_TABLE_2,
|
||||
}
|
||||
|
||||
DUMMY_BAUDRATE_TABLE = {
|
||||
0: 1_000_000,
|
||||
1: 500_000,
|
||||
2: 250_000,
|
||||
}
|
||||
|
||||
DUMMY_MODEL_BAUDRATE_TABLE = {
|
||||
"model_1": DUMMY_BAUDRATE_TABLE,
|
||||
"model_2": DUMMY_BAUDRATE_TABLE,
|
||||
"model_3": DUMMY_BAUDRATE_TABLE,
|
||||
}
|
||||
|
||||
DUMMY_ENCODING_TABLE = {
|
||||
"Present_Position": 8,
|
||||
"Goal_Position": 10,
|
||||
}
|
||||
|
||||
DUMMY_MODEL_ENCODING_TABLE = {
|
||||
"model_1": DUMMY_ENCODING_TABLE,
|
||||
"model_2": DUMMY_ENCODING_TABLE,
|
||||
"model_3": DUMMY_ENCODING_TABLE,
|
||||
}
|
||||
|
||||
DUMMY_MODEL_NUMBER_TABLE = {
|
||||
"model_1": 1234,
|
||||
"model_2": 5678,
|
||||
"model_3": 5799,
|
||||
}
|
||||
|
||||
DUMMY_MODEL_RESOLUTION_TABLE = {
|
||||
"model_1": 4096,
|
||||
"model_2": 1024,
|
||||
"model_3": 4096,
|
||||
}
|
||||
|
||||
|
||||
class MockPortHandler:
|
||||
def __init__(self, port_name):
|
||||
self.is_open: bool = False
|
||||
self.baudrate: int
|
||||
self.packet_start_time: float
|
||||
self.packet_timeout: float
|
||||
self.tx_time_per_byte: float
|
||||
self.is_using: bool = False
|
||||
self.port_name: str = port_name
|
||||
self.ser = None
|
||||
|
||||
def openPort(self):
|
||||
self.is_open = True
|
||||
return self.is_open
|
||||
|
||||
def closePort(self):
|
||||
self.is_open = False
|
||||
|
||||
def clearPort(self): ...
|
||||
def setPortName(self, port_name):
|
||||
self.port_name = port_name
|
||||
|
||||
def getPortName(self):
|
||||
return self.port_name
|
||||
|
||||
def setBaudRate(self, baudrate):
|
||||
self.baudrate: baudrate
|
||||
|
||||
def getBaudRate(self):
|
||||
return self.baudrate
|
||||
|
||||
def getBytesAvailable(self): ...
|
||||
def readPort(self, length): ...
|
||||
def writePort(self, packet): ...
|
||||
def setPacketTimeout(self, packet_length): ...
|
||||
def setPacketTimeoutMillis(self, msec): ...
|
||||
def isPacketTimeout(self): ...
|
||||
def getCurrentTime(self): ...
|
||||
def getTimeSinceStart(self): ...
|
||||
def setupPort(self, cflag_baud): ...
|
||||
def getCFlagBaud(self, baudrate): ...
|
||||
|
||||
|
||||
class MockMotorsBus(MotorsBus):
|
||||
available_baudrates = [500_000, 1_000_000]
|
||||
default_timeout = 1000
|
||||
model_baudrate_table = DUMMY_MODEL_BAUDRATE_TABLE
|
||||
model_ctrl_table = DUMMY_MODEL_CTRL_TABLE
|
||||
model_encoding_table = DUMMY_MODEL_ENCODING_TABLE
|
||||
model_number_table = DUMMY_MODEL_NUMBER_TABLE
|
||||
model_resolution_table = DUMMY_MODEL_RESOLUTION_TABLE
|
||||
normalized_data = ["Present_Position", "Goal_Position"]
|
||||
|
||||
def __init__(self, port: str, motors: dict[str, Motor]):
|
||||
super().__init__(port, motors)
|
||||
self.port_handler = MockPortHandler(port)
|
||||
|
||||
def _assert_protocol_is_compatible(self, instruction_name): ...
|
||||
def _handshake(self): ...
|
||||
def _find_single_motor(self, motor, initial_baudrate): ...
|
||||
def configure_motors(self): ...
|
||||
def is_calibrated(self): ...
|
||||
def read_calibration(self): ...
|
||||
def write_calibration(self, calibration_dict): ...
|
||||
def disable_torque(self, motors, num_retry): ...
|
||||
def _disable_torque(self, motor, model, num_retry): ...
|
||||
def enable_torque(self, motors, num_retry): ...
|
||||
def _get_half_turn_homings(self, positions): ...
|
||||
def _encode_sign(self, data_name, ids_values): ...
|
||||
def _decode_sign(self, data_name, ids_values): ...
|
||||
def _split_into_byte_chunks(self, value, length): ...
|
||||
def broadcast_ping(self, num_retry, raise_on_error): ...
|
||||
@@ -0,0 +1,112 @@
|
||||
import random
|
||||
from dataclasses import dataclass, field
|
||||
from functools import cached_property
|
||||
from typing import Any
|
||||
|
||||
from lerobot.common.cameras import CameraConfig, make_cameras_from_configs
|
||||
from lerobot.common.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
|
||||
from lerobot.common.robots import Robot, RobotConfig
|
||||
|
||||
|
||||
@RobotConfig.register_subclass("mock_robot")
|
||||
@dataclass
|
||||
class MockRobotConfig(RobotConfig):
|
||||
n_motors: int = 3
|
||||
cameras: dict[str, CameraConfig] = field(default_factory=dict)
|
||||
random_values: bool = True
|
||||
static_values: list[float] | None = None
|
||||
calibrated: bool = True
|
||||
|
||||
def __post_init__(self):
|
||||
if self.n_motors < 1:
|
||||
raise ValueError(self.n_motors)
|
||||
|
||||
if self.random_values and self.static_values is not None:
|
||||
raise ValueError("Choose either random values or static values")
|
||||
|
||||
if self.static_values is not None and len(self.static_values) != self.n_motors:
|
||||
raise ValueError("Specify the same number of static values as motors")
|
||||
|
||||
if len(self.cameras) > 0:
|
||||
raise NotImplementedError # TODO with the cameras refactor
|
||||
|
||||
|
||||
class MockRobot(Robot):
|
||||
"""Mock Robot to be used for testing."""
|
||||
|
||||
config_class = MockRobotConfig
|
||||
name = "mock_robot"
|
||||
|
||||
def __init__(self, config: MockRobotConfig):
|
||||
super().__init__(config)
|
||||
self.config = config
|
||||
self._is_connected = False
|
||||
self._is_calibrated = config.calibrated
|
||||
self.motors = [f"motor_{i + 1}" for i in range(config.n_motors)]
|
||||
self.cameras = make_cameras_from_configs(config.cameras)
|
||||
|
||||
@property
|
||||
def _motors_ft(self) -> dict[str, type]:
|
||||
return {f"{motor}.pos": float for motor in self.motors}
|
||||
|
||||
@property
|
||||
def _cameras_ft(self) -> dict[str, tuple]:
|
||||
return {
|
||||
cam: (self.config.cameras[cam].height, self.config.cameras[cam].width, 3) for cam in self.cameras
|
||||
}
|
||||
|
||||
@cached_property
|
||||
def observation_features(self) -> dict[str, type | tuple]:
|
||||
return {**self._motors_ft, **self._cameras_ft}
|
||||
|
||||
@cached_property
|
||||
def action_features(self) -> dict[str, type]:
|
||||
return self._motors_ft
|
||||
|
||||
@property
|
||||
def is_connected(self) -> bool:
|
||||
return self._is_connected
|
||||
|
||||
def connect(self, calibrate: bool = True) -> None:
|
||||
if self.is_connected:
|
||||
raise DeviceAlreadyConnectedError(f"{self} already connected")
|
||||
|
||||
self._is_connected = True
|
||||
if calibrate:
|
||||
self.calibrate()
|
||||
|
||||
@property
|
||||
def is_calibrated(self) -> bool:
|
||||
return self._is_calibrated
|
||||
|
||||
def calibrate(self) -> None:
|
||||
if not self.is_connected:
|
||||
raise DeviceNotConnectedError(f"{self} is not connected.")
|
||||
|
||||
self._is_calibrated = True
|
||||
|
||||
def configure(self) -> None:
|
||||
pass
|
||||
|
||||
def get_observation(self) -> dict[str, Any]:
|
||||
if not self.is_connected:
|
||||
raise DeviceNotConnectedError(f"{self} is not connected.")
|
||||
|
||||
if self.config.random_values:
|
||||
return {f"{motor}.pos": random.uniform(-100, 100) for motor in self.motors}
|
||||
else:
|
||||
return {
|
||||
f"{motor}.pos": val for motor, val in zip(self.motors, self.config.static_values, strict=True)
|
||||
}
|
||||
|
||||
def send_action(self, action: dict[str, Any]) -> dict[str, Any]:
|
||||
if not self.is_connected:
|
||||
raise DeviceNotConnectedError(f"{self} is not connected.")
|
||||
|
||||
return action
|
||||
|
||||
def disconnect(self) -> None:
|
||||
if not self.is_connected:
|
||||
raise DeviceNotConnectedError(f"{self} is not connected.")
|
||||
|
||||
self._is_connected = False
|
||||
@@ -0,0 +1,35 @@
|
||||
import threading
|
||||
import time
|
||||
|
||||
from mock_serial.mock_serial import Stub
|
||||
|
||||
|
||||
class WaitableStub(Stub):
|
||||
"""
|
||||
In some situations, a test might be checking if a stub has been called before `MockSerial` thread had time
|
||||
to read, match, and call the stub. In these situations, the test can fail randomly.
|
||||
|
||||
Use `wait_called()` or `wait_calls()` to block until the stub is called, avoiding race conditions.
|
||||
|
||||
Proposed fix:
|
||||
https://github.com/benthorner/mock_serial/pull/3
|
||||
"""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self._event = threading.Event()
|
||||
|
||||
def call(self):
|
||||
self._event.set()
|
||||
return super().call()
|
||||
|
||||
def wait_called(self, timeout: float = 1.0):
|
||||
return self._event.wait(timeout)
|
||||
|
||||
def wait_calls(self, min_calls: int = 1, timeout: float = 1.0):
|
||||
start = time.perf_counter()
|
||||
while time.perf_counter() - start < timeout:
|
||||
if self.calls >= min_calls:
|
||||
return self.calls
|
||||
time.sleep(0.005)
|
||||
raise TimeoutError(f"Stub not called {min_calls} times within {timeout} seconds.")
|
||||
@@ -0,0 +1,94 @@
|
||||
import random
|
||||
from dataclasses import dataclass
|
||||
from functools import cached_property
|
||||
from typing import Any
|
||||
|
||||
from lerobot.common.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
|
||||
from lerobot.common.teleoperators import Teleoperator, TeleoperatorConfig
|
||||
|
||||
|
||||
@TeleoperatorConfig.register_subclass("mock_teleop")
|
||||
@dataclass
|
||||
class MockTeleopConfig(TeleoperatorConfig):
|
||||
n_motors: int = 3
|
||||
random_values: bool = True
|
||||
static_values: list[float] | None = None
|
||||
calibrated: bool = True
|
||||
|
||||
def __post_init__(self):
|
||||
if self.n_motors < 1:
|
||||
raise ValueError(self.n_motors)
|
||||
|
||||
if self.random_values and self.static_values is not None:
|
||||
raise ValueError("Choose either random values or static values")
|
||||
|
||||
if self.static_values is not None and len(self.static_values) != self.n_motors:
|
||||
raise ValueError("Specify the same number of static values as motors")
|
||||
|
||||
|
||||
class MockTeleop(Teleoperator):
|
||||
"""Mock Teleoperator to be used for testing."""
|
||||
|
||||
config_class = MockTeleopConfig
|
||||
name = "mock_teleop"
|
||||
|
||||
def __init__(self, config: MockTeleopConfig):
|
||||
super().__init__(config)
|
||||
self.config = config
|
||||
self._is_connected = False
|
||||
self._is_calibrated = config.calibrated
|
||||
self.motors = [f"motor_{i + 1}" for i in range(config.n_motors)]
|
||||
|
||||
@cached_property
|
||||
def action_features(self) -> dict[str, type]:
|
||||
return {f"{motor}.pos": float for motor in self.motors}
|
||||
|
||||
@cached_property
|
||||
def feedback_features(self) -> dict[str, type]:
|
||||
return {f"{motor}.pos": float for motor in self.motors}
|
||||
|
||||
@property
|
||||
def is_connected(self) -> bool:
|
||||
return self._is_connected
|
||||
|
||||
def connect(self, calibrate: bool = True) -> None:
|
||||
if self.is_connected:
|
||||
raise DeviceAlreadyConnectedError(f"{self} already connected")
|
||||
|
||||
self._is_connected = True
|
||||
if calibrate:
|
||||
self.calibrate()
|
||||
|
||||
@property
|
||||
def is_calibrated(self) -> bool:
|
||||
return self._is_calibrated
|
||||
|
||||
def calibrate(self) -> None:
|
||||
if not self.is_connected:
|
||||
raise DeviceNotConnectedError(f"{self} is not connected.")
|
||||
|
||||
self._is_calibrated = True
|
||||
|
||||
def configure(self) -> None:
|
||||
pass
|
||||
|
||||
def get_action(self) -> dict[str, Any]:
|
||||
if not self.is_connected:
|
||||
raise DeviceNotConnectedError(f"{self} is not connected.")
|
||||
|
||||
if self.config.random_values:
|
||||
return {f"{motor}.pos": random.uniform(-100, 100) for motor in self.motors}
|
||||
else:
|
||||
return {
|
||||
f"{motor}.pos": val for motor, val in zip(self.motors, self.config.static_values, strict=True)
|
||||
}
|
||||
|
||||
def send_feedback(self, feedback: dict[str, Any]) -> None:
|
||||
if not self.is_connected:
|
||||
raise DeviceNotConnectedError(f"{self} is not connected.")
|
||||
|
||||
def disconnect(self) -> None:
|
||||
if not self.is_connected:
|
||||
raise DeviceNotConnectedError(f"{self} is not connected.")
|
||||
|
||||
self._is_connected = False
|
||||
@@ -1,107 +0,0 @@
|
||||
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# 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.
|
||||
"""Mocked classes and functions from dynamixel_sdk to allow for continuous integration
|
||||
and testing code logic that requires hardware and devices (e.g. robot arms, cameras)
|
||||
|
||||
Warning: These mocked versions are minimalist. They do not exactly mock every behaviors
|
||||
from the original classes and functions (e.g. return types might be None instead of boolean).
|
||||
"""
|
||||
|
||||
# from dynamixel_sdk import COMM_SUCCESS
|
||||
|
||||
DEFAULT_BAUDRATE = 9_600
|
||||
COMM_SUCCESS = 0 # tx or rx packet communication success
|
||||
|
||||
|
||||
def convert_to_bytes(value, bytes):
|
||||
# TODO(rcadene): remove need to mock `convert_to_bytes` by implemented the inverse transform
|
||||
# `convert_bytes_to_value`
|
||||
del bytes # unused
|
||||
return value
|
||||
|
||||
|
||||
def get_default_motor_values(motor_index):
|
||||
return {
|
||||
# Key (int) are from X_SERIES_CONTROL_TABLE
|
||||
7: motor_index, # ID
|
||||
8: DEFAULT_BAUDRATE, # Baud_rate
|
||||
10: 0, # Drive_Mode
|
||||
64: 0, # Torque_Enable
|
||||
# Set 2560 since calibration values for Aloha gripper is between start_pos=2499 and end_pos=3144
|
||||
# For other joints, 2560 will be autocorrected to be in calibration range
|
||||
132: 2560, # Present_Position
|
||||
}
|
||||
|
||||
|
||||
class PortHandler:
|
||||
def __init__(self, port):
|
||||
self.port = port
|
||||
# factory default baudrate
|
||||
self.baudrate = DEFAULT_BAUDRATE
|
||||
|
||||
def openPort(self): # noqa: N802
|
||||
return True
|
||||
|
||||
def closePort(self): # noqa: N802
|
||||
pass
|
||||
|
||||
def setPacketTimeoutMillis(self, timeout_ms): # noqa: N802
|
||||
del timeout_ms # unused
|
||||
|
||||
def getBaudRate(self): # noqa: N802
|
||||
return self.baudrate
|
||||
|
||||
def setBaudRate(self, baudrate): # noqa: N802
|
||||
self.baudrate = baudrate
|
||||
|
||||
|
||||
class PacketHandler:
|
||||
def __init__(self, protocol_version):
|
||||
del protocol_version # unused
|
||||
# Use packet_handler.data to communicate across Read and Write
|
||||
self.data = {}
|
||||
|
||||
|
||||
class GroupSyncRead:
|
||||
def __init__(self, port_handler, packet_handler, address, bytes):
|
||||
self.packet_handler = packet_handler
|
||||
|
||||
def addParam(self, motor_index): # noqa: N802
|
||||
# Initialize motor default values
|
||||
if motor_index not in self.packet_handler.data:
|
||||
self.packet_handler.data[motor_index] = get_default_motor_values(motor_index)
|
||||
|
||||
def txRxPacket(self): # noqa: N802
|
||||
return COMM_SUCCESS
|
||||
|
||||
def getData(self, index, address, bytes): # noqa: N802
|
||||
return self.packet_handler.data[index][address]
|
||||
|
||||
|
||||
class GroupSyncWrite:
|
||||
def __init__(self, port_handler, packet_handler, address, bytes):
|
||||
self.packet_handler = packet_handler
|
||||
self.address = address
|
||||
|
||||
def addParam(self, index, data): # noqa: N802
|
||||
# Initialize motor default values
|
||||
if index not in self.packet_handler.data:
|
||||
self.packet_handler.data[index] = get_default_motor_values(index)
|
||||
self.changeParam(index, data)
|
||||
|
||||
def txPacket(self): # noqa: N802
|
||||
return COMM_SUCCESS
|
||||
|
||||
def changeParam(self, index, data): # noqa: N802
|
||||
self.packet_handler.data[index][self.address] = data
|
||||
@@ -1,125 +0,0 @@
|
||||
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# 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.
|
||||
"""Mocked classes and functions from dynamixel_sdk to allow for continuous integration
|
||||
and testing code logic that requires hardware and devices (e.g. robot arms, cameras)
|
||||
|
||||
Warning: These mocked versions are minimalist. They do not exactly mock every behaviors
|
||||
from the original classes and functions (e.g. return types might be None instead of boolean).
|
||||
"""
|
||||
|
||||
# from dynamixel_sdk import COMM_SUCCESS
|
||||
|
||||
DEFAULT_BAUDRATE = 1_000_000
|
||||
COMM_SUCCESS = 0 # tx or rx packet communication success
|
||||
|
||||
|
||||
def convert_to_bytes(value, bytes):
|
||||
# TODO(rcadene): remove need to mock `convert_to_bytes` by implemented the inverse transform
|
||||
# `convert_bytes_to_value`
|
||||
del bytes # unused
|
||||
return value
|
||||
|
||||
|
||||
def get_default_motor_values(motor_index):
|
||||
return {
|
||||
# Key (int) are from SCS_SERIES_CONTROL_TABLE
|
||||
5: motor_index, # ID
|
||||
6: DEFAULT_BAUDRATE, # Baud_rate
|
||||
10: 0, # Drive_Mode
|
||||
21: 32, # P_Coefficient
|
||||
22: 32, # D_Coefficient
|
||||
23: 0, # I_Coefficient
|
||||
40: 0, # Torque_Enable
|
||||
41: 254, # Acceleration
|
||||
31: -2047, # Offset
|
||||
33: 0, # Mode
|
||||
55: 1, # Lock
|
||||
# Set 2560 since calibration values for Aloha gripper is between start_pos=2499 and end_pos=3144
|
||||
# For other joints, 2560 will be autocorrected to be in calibration range
|
||||
56: 2560, # Present_Position
|
||||
58: 0, # Present_Speed
|
||||
69: 0, # Present_Current
|
||||
85: 150, # Maximum_Acceleration
|
||||
}
|
||||
|
||||
|
||||
class PortHandler:
|
||||
def __init__(self, port):
|
||||
self.port = port
|
||||
# factory default baudrate
|
||||
self.baudrate = DEFAULT_BAUDRATE
|
||||
self.ser = SerialMock()
|
||||
|
||||
def openPort(self): # noqa: N802
|
||||
return True
|
||||
|
||||
def closePort(self): # noqa: N802
|
||||
pass
|
||||
|
||||
def setPacketTimeoutMillis(self, timeout_ms): # noqa: N802
|
||||
del timeout_ms # unused
|
||||
|
||||
def getBaudRate(self): # noqa: N802
|
||||
return self.baudrate
|
||||
|
||||
def setBaudRate(self, baudrate): # noqa: N802
|
||||
self.baudrate = baudrate
|
||||
|
||||
|
||||
class PacketHandler:
|
||||
def __init__(self, protocol_version):
|
||||
del protocol_version # unused
|
||||
# Use packet_handler.data to communicate across Read and Write
|
||||
self.data = {}
|
||||
|
||||
|
||||
class GroupSyncRead:
|
||||
def __init__(self, port_handler, packet_handler, address, bytes):
|
||||
self.packet_handler = packet_handler
|
||||
|
||||
def addParam(self, motor_index): # noqa: N802
|
||||
# Initialize motor default values
|
||||
if motor_index not in self.packet_handler.data:
|
||||
self.packet_handler.data[motor_index] = get_default_motor_values(motor_index)
|
||||
|
||||
def txRxPacket(self): # noqa: N802
|
||||
return COMM_SUCCESS
|
||||
|
||||
def getData(self, index, address, bytes): # noqa: N802
|
||||
return self.packet_handler.data[index][address]
|
||||
|
||||
|
||||
class GroupSyncWrite:
|
||||
def __init__(self, port_handler, packet_handler, address, bytes):
|
||||
self.packet_handler = packet_handler
|
||||
self.address = address
|
||||
|
||||
def addParam(self, index, data): # noqa: N802
|
||||
if index not in self.packet_handler.data:
|
||||
self.packet_handler.data[index] = get_default_motor_values(index)
|
||||
self.changeParam(index, data)
|
||||
|
||||
def txPacket(self): # noqa: N802
|
||||
return COMM_SUCCESS
|
||||
|
||||
def changeParam(self, index, data): # noqa: N802
|
||||
self.packet_handler.data[index][self.address] = data
|
||||
|
||||
|
||||
class SerialMock:
|
||||
def reset_output_buffer(self):
|
||||
pass
|
||||
|
||||
def reset_input_buffer(self):
|
||||
pass
|
||||
@@ -0,0 +1,400 @@
|
||||
import re
|
||||
import sys
|
||||
from typing import Generator
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from lerobot.common.motors import Motor, MotorCalibration, MotorNormMode
|
||||
from lerobot.common.motors.dynamixel import MODEL_NUMBER_TABLE, DynamixelMotorsBus
|
||||
from lerobot.common.motors.dynamixel.tables import X_SERIES_CONTROL_TABLE
|
||||
from lerobot.common.utils.encoding_utils import encode_twos_complement
|
||||
|
||||
try:
|
||||
import dynamixel_sdk as dxl
|
||||
|
||||
from tests.mocks.mock_dynamixel import MockMotors, MockPortHandler
|
||||
except (ImportError, ModuleNotFoundError):
|
||||
pytest.skip("dynamixel_sdk not available", allow_module_level=True)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def patch_port_handler():
|
||||
if sys.platform == "darwin":
|
||||
with patch.object(dxl, "PortHandler", MockPortHandler):
|
||||
yield
|
||||
else:
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_motors() -> Generator[MockMotors, None, None]:
|
||||
motors = MockMotors()
|
||||
motors.open()
|
||||
yield motors
|
||||
motors.close()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def dummy_motors() -> dict[str, Motor]:
|
||||
return {
|
||||
"dummy_1": Motor(1, "xl430-w250", MotorNormMode.RANGE_M100_100),
|
||||
"dummy_2": Motor(2, "xm540-w270", MotorNormMode.RANGE_M100_100),
|
||||
"dummy_3": Motor(3, "xl330-m077", MotorNormMode.RANGE_M100_100),
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def dummy_calibration(dummy_motors) -> dict[str, MotorCalibration]:
|
||||
drive_modes = [0, 1, 0]
|
||||
homings = [-709, -2006, 1624]
|
||||
mins = [43, 27, 145]
|
||||
maxes = [1335, 3608, 3999]
|
||||
calibration = {}
|
||||
for motor, m in dummy_motors.items():
|
||||
calibration[motor] = MotorCalibration(
|
||||
id=m.id,
|
||||
drive_mode=drive_modes[m.id - 1],
|
||||
homing_offset=homings[m.id - 1],
|
||||
range_min=mins[m.id - 1],
|
||||
range_max=maxes[m.id - 1],
|
||||
)
|
||||
return calibration
|
||||
|
||||
|
||||
@pytest.mark.skipif(sys.platform != "darwin", reason=f"No patching needed on {sys.platform=}")
|
||||
def test_autouse_patch():
|
||||
"""Ensures that the autouse fixture correctly patches dxl.PortHandler with MockPortHandler."""
|
||||
assert dxl.PortHandler is MockPortHandler
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"value, length, expected",
|
||||
[
|
||||
(0x12, 1, [0x12]),
|
||||
(0x1234, 2, [0x34, 0x12]),
|
||||
(0x12345678, 4, [0x78, 0x56, 0x34, 0x12]),
|
||||
],
|
||||
ids=[
|
||||
"1 byte",
|
||||
"2 bytes",
|
||||
"4 bytes",
|
||||
],
|
||||
) # fmt: skip
|
||||
def test__split_into_byte_chunks(value, length, expected):
|
||||
bus = DynamixelMotorsBus("", {})
|
||||
assert bus._split_into_byte_chunks(value, length) == expected
|
||||
|
||||
|
||||
def test_abc_implementation(dummy_motors):
|
||||
"""Instantiation should raise an error if the class doesn't implement abstract methods/properties."""
|
||||
DynamixelMotorsBus(port="/dev/dummy-port", motors=dummy_motors)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("id_", [1, 2, 3])
|
||||
def test_ping(id_, mock_motors, dummy_motors):
|
||||
expected_model_nb = MODEL_NUMBER_TABLE[dummy_motors[f"dummy_{id_}"].model]
|
||||
stub = mock_motors.build_ping_stub(id_, expected_model_nb)
|
||||
bus = DynamixelMotorsBus(port=mock_motors.port, motors=dummy_motors)
|
||||
bus.connect(handshake=False)
|
||||
|
||||
ping_model_nb = bus.ping(id_)
|
||||
|
||||
assert ping_model_nb == expected_model_nb
|
||||
assert mock_motors.stubs[stub].called
|
||||
|
||||
|
||||
def test_broadcast_ping(mock_motors, dummy_motors):
|
||||
models = {m.id: m.model for m in dummy_motors.values()}
|
||||
expected_model_nbs = {id_: MODEL_NUMBER_TABLE[model] for id_, model in models.items()}
|
||||
stub = mock_motors.build_broadcast_ping_stub(expected_model_nbs)
|
||||
bus = DynamixelMotorsBus(port=mock_motors.port, motors=dummy_motors)
|
||||
bus.connect(handshake=False)
|
||||
|
||||
ping_model_nbs = bus.broadcast_ping()
|
||||
|
||||
assert ping_model_nbs == expected_model_nbs
|
||||
assert mock_motors.stubs[stub].called
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"addr, length, id_, value",
|
||||
[
|
||||
(0, 1, 1, 2),
|
||||
(10, 2, 2, 999),
|
||||
(42, 4, 3, 1337),
|
||||
],
|
||||
)
|
||||
def test__read(addr, length, id_, value, mock_motors, dummy_motors):
|
||||
stub = mock_motors.build_read_stub(addr, length, id_, value)
|
||||
bus = DynamixelMotorsBus(port=mock_motors.port, motors=dummy_motors)
|
||||
bus.connect(handshake=False)
|
||||
|
||||
read_value, _, _ = bus._read(addr, length, id_)
|
||||
|
||||
assert mock_motors.stubs[stub].called
|
||||
assert read_value == value
|
||||
|
||||
|
||||
@pytest.mark.parametrize("raise_on_error", (True, False))
|
||||
def test__read_error(raise_on_error, mock_motors, dummy_motors):
|
||||
addr, length, id_, value, error = (10, 4, 1, 1337, dxl.ERRNUM_DATA_LIMIT)
|
||||
stub = mock_motors.build_read_stub(addr, length, id_, value, error=error)
|
||||
bus = DynamixelMotorsBus(port=mock_motors.port, motors=dummy_motors)
|
||||
bus.connect(handshake=False)
|
||||
|
||||
if raise_on_error:
|
||||
with pytest.raises(
|
||||
RuntimeError, match=re.escape("[RxPacketError] The data value exceeds the limit value!")
|
||||
):
|
||||
bus._read(addr, length, id_, raise_on_error=raise_on_error)
|
||||
else:
|
||||
_, _, read_error = bus._read(addr, length, id_, raise_on_error=raise_on_error)
|
||||
assert read_error == error
|
||||
|
||||
assert mock_motors.stubs[stub].called
|
||||
|
||||
|
||||
@pytest.mark.parametrize("raise_on_error", (True, False))
|
||||
def test__read_comm(raise_on_error, mock_motors, dummy_motors):
|
||||
addr, length, id_, value = (10, 4, 1, 1337)
|
||||
stub = mock_motors.build_read_stub(addr, length, id_, value, reply=False)
|
||||
bus = DynamixelMotorsBus(port=mock_motors.port, motors=dummy_motors)
|
||||
bus.connect(handshake=False)
|
||||
|
||||
if raise_on_error:
|
||||
with pytest.raises(ConnectionError, match=re.escape("[TxRxResult] There is no status packet!")):
|
||||
bus._read(addr, length, id_, raise_on_error=raise_on_error)
|
||||
else:
|
||||
_, read_comm, _ = bus._read(addr, length, id_, raise_on_error=raise_on_error)
|
||||
assert read_comm == dxl.COMM_RX_TIMEOUT
|
||||
|
||||
assert mock_motors.stubs[stub].called
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"addr, length, id_, value",
|
||||
[
|
||||
(0, 1, 1, 2),
|
||||
(10, 2, 2, 999),
|
||||
(42, 4, 3, 1337),
|
||||
],
|
||||
)
|
||||
def test__write(addr, length, id_, value, mock_motors, dummy_motors):
|
||||
stub = mock_motors.build_write_stub(addr, length, id_, value)
|
||||
bus = DynamixelMotorsBus(port=mock_motors.port, motors=dummy_motors)
|
||||
bus.connect(handshake=False)
|
||||
|
||||
comm, error = bus._write(addr, length, id_, value)
|
||||
|
||||
assert mock_motors.stubs[stub].called
|
||||
assert comm == dxl.COMM_SUCCESS
|
||||
assert error == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("raise_on_error", (True, False))
|
||||
def test__write_error(raise_on_error, mock_motors, dummy_motors):
|
||||
addr, length, id_, value, error = (10, 4, 1, 1337, dxl.ERRNUM_DATA_LIMIT)
|
||||
stub = mock_motors.build_write_stub(addr, length, id_, value, error=error)
|
||||
bus = DynamixelMotorsBus(port=mock_motors.port, motors=dummy_motors)
|
||||
bus.connect(handshake=False)
|
||||
|
||||
if raise_on_error:
|
||||
with pytest.raises(
|
||||
RuntimeError, match=re.escape("[RxPacketError] The data value exceeds the limit value!")
|
||||
):
|
||||
bus._write(addr, length, id_, value, raise_on_error=raise_on_error)
|
||||
else:
|
||||
_, write_error = bus._write(addr, length, id_, value, raise_on_error=raise_on_error)
|
||||
assert write_error == error
|
||||
|
||||
assert mock_motors.stubs[stub].called
|
||||
|
||||
|
||||
@pytest.mark.parametrize("raise_on_error", (True, False))
|
||||
def test__write_comm(raise_on_error, mock_motors, dummy_motors):
|
||||
addr, length, id_, value = (10, 4, 1, 1337)
|
||||
stub = mock_motors.build_write_stub(addr, length, id_, value, reply=False)
|
||||
bus = DynamixelMotorsBus(port=mock_motors.port, motors=dummy_motors)
|
||||
bus.connect(handshake=False)
|
||||
|
||||
if raise_on_error:
|
||||
with pytest.raises(ConnectionError, match=re.escape("[TxRxResult] There is no status packet!")):
|
||||
bus._write(addr, length, id_, value, raise_on_error=raise_on_error)
|
||||
else:
|
||||
write_comm, _ = bus._write(addr, length, id_, value, raise_on_error=raise_on_error)
|
||||
assert write_comm == dxl.COMM_RX_TIMEOUT
|
||||
|
||||
assert mock_motors.stubs[stub].called
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"addr, length, ids_values",
|
||||
[
|
||||
(0, 1, {1: 4}),
|
||||
(10, 2, {1: 1337, 2: 42}),
|
||||
(42, 4, {1: 1337, 2: 42, 3: 4016}),
|
||||
],
|
||||
ids=["1 motor", "2 motors", "3 motors"],
|
||||
)
|
||||
def test__sync_read(addr, length, ids_values, mock_motors, dummy_motors):
|
||||
stub = mock_motors.build_sync_read_stub(addr, length, ids_values)
|
||||
bus = DynamixelMotorsBus(port=mock_motors.port, motors=dummy_motors)
|
||||
bus.connect(handshake=False)
|
||||
|
||||
read_values, _ = bus._sync_read(addr, length, list(ids_values))
|
||||
|
||||
assert mock_motors.stubs[stub].called
|
||||
assert read_values == ids_values
|
||||
|
||||
|
||||
@pytest.mark.parametrize("raise_on_error", (True, False))
|
||||
def test__sync_read_comm(raise_on_error, mock_motors, dummy_motors):
|
||||
addr, length, ids_values = (10, 4, {1: 1337})
|
||||
stub = mock_motors.build_sync_read_stub(addr, length, ids_values, reply=False)
|
||||
bus = DynamixelMotorsBus(port=mock_motors.port, motors=dummy_motors)
|
||||
bus.connect(handshake=False)
|
||||
|
||||
if raise_on_error:
|
||||
with pytest.raises(ConnectionError, match=re.escape("[TxRxResult] There is no status packet!")):
|
||||
bus._sync_read(addr, length, list(ids_values), raise_on_error=raise_on_error)
|
||||
else:
|
||||
_, read_comm = bus._sync_read(addr, length, list(ids_values), raise_on_error=raise_on_error)
|
||||
assert read_comm == dxl.COMM_RX_TIMEOUT
|
||||
|
||||
assert mock_motors.stubs[stub].called
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"addr, length, ids_values",
|
||||
[
|
||||
(0, 1, {1: 4}),
|
||||
(10, 2, {1: 1337, 2: 42}),
|
||||
(42, 4, {1: 1337, 2: 42, 3: 4016}),
|
||||
],
|
||||
ids=["1 motor", "2 motors", "3 motors"],
|
||||
)
|
||||
def test__sync_write(addr, length, ids_values, mock_motors, dummy_motors):
|
||||
stub = mock_motors.build_sync_write_stub(addr, length, ids_values)
|
||||
bus = DynamixelMotorsBus(port=mock_motors.port, motors=dummy_motors)
|
||||
bus.connect(handshake=False)
|
||||
|
||||
comm = bus._sync_write(addr, length, ids_values)
|
||||
|
||||
assert mock_motors.stubs[stub].wait_called()
|
||||
assert comm == dxl.COMM_SUCCESS
|
||||
|
||||
|
||||
def test_is_calibrated(mock_motors, dummy_motors, dummy_calibration):
|
||||
drive_modes = {m.id: m.drive_mode for m in dummy_calibration.values()}
|
||||
encoded_homings = {m.id: encode_twos_complement(m.homing_offset, 4) for m in dummy_calibration.values()}
|
||||
mins = {m.id: m.range_min for m in dummy_calibration.values()}
|
||||
maxes = {m.id: m.range_max for m in dummy_calibration.values()}
|
||||
drive_modes_stub = mock_motors.build_sync_read_stub(*X_SERIES_CONTROL_TABLE["Drive_Mode"], drive_modes)
|
||||
offsets_stub = mock_motors.build_sync_read_stub(*X_SERIES_CONTROL_TABLE["Homing_Offset"], encoded_homings)
|
||||
mins_stub = mock_motors.build_sync_read_stub(*X_SERIES_CONTROL_TABLE["Min_Position_Limit"], mins)
|
||||
maxes_stub = mock_motors.build_sync_read_stub(*X_SERIES_CONTROL_TABLE["Max_Position_Limit"], maxes)
|
||||
bus = DynamixelMotorsBus(
|
||||
port=mock_motors.port,
|
||||
motors=dummy_motors,
|
||||
calibration=dummy_calibration,
|
||||
)
|
||||
bus.connect(handshake=False)
|
||||
|
||||
is_calibrated = bus.is_calibrated
|
||||
|
||||
assert is_calibrated
|
||||
assert mock_motors.stubs[drive_modes_stub].called
|
||||
assert mock_motors.stubs[offsets_stub].called
|
||||
assert mock_motors.stubs[mins_stub].called
|
||||
assert mock_motors.stubs[maxes_stub].called
|
||||
|
||||
|
||||
def test_reset_calibration(mock_motors, dummy_motors):
|
||||
write_homing_stubs = []
|
||||
write_mins_stubs = []
|
||||
write_maxes_stubs = []
|
||||
for motor in dummy_motors.values():
|
||||
write_homing_stubs.append(
|
||||
mock_motors.build_write_stub(*X_SERIES_CONTROL_TABLE["Homing_Offset"], motor.id, 0)
|
||||
)
|
||||
write_mins_stubs.append(
|
||||
mock_motors.build_write_stub(*X_SERIES_CONTROL_TABLE["Min_Position_Limit"], motor.id, 0)
|
||||
)
|
||||
write_maxes_stubs.append(
|
||||
mock_motors.build_write_stub(*X_SERIES_CONTROL_TABLE["Max_Position_Limit"], motor.id, 4095)
|
||||
)
|
||||
|
||||
bus = DynamixelMotorsBus(port=mock_motors.port, motors=dummy_motors)
|
||||
bus.connect(handshake=False)
|
||||
|
||||
bus.reset_calibration()
|
||||
|
||||
assert all(mock_motors.stubs[stub].called for stub in write_homing_stubs)
|
||||
assert all(mock_motors.stubs[stub].called for stub in write_mins_stubs)
|
||||
assert all(mock_motors.stubs[stub].called for stub in write_maxes_stubs)
|
||||
|
||||
|
||||
def test_set_half_turn_homings(mock_motors, dummy_motors):
|
||||
"""
|
||||
For this test, we assume that the homing offsets are already 0 such that
|
||||
Present_Position == Actual_Position
|
||||
"""
|
||||
current_positions = {
|
||||
1: 1337,
|
||||
2: 42,
|
||||
3: 3672,
|
||||
}
|
||||
expected_homings = {
|
||||
1: 710, # 2047 - 1337
|
||||
2: 2005, # 2047 - 42
|
||||
3: -1625, # 2047 - 3672
|
||||
}
|
||||
read_pos_stub = mock_motors.build_sync_read_stub(
|
||||
*X_SERIES_CONTROL_TABLE["Present_Position"], current_positions
|
||||
)
|
||||
write_homing_stubs = []
|
||||
for id_, homing in expected_homings.items():
|
||||
encoded_homing = encode_twos_complement(homing, 4)
|
||||
stub = mock_motors.build_write_stub(*X_SERIES_CONTROL_TABLE["Homing_Offset"], id_, encoded_homing)
|
||||
write_homing_stubs.append(stub)
|
||||
|
||||
bus = DynamixelMotorsBus(port=mock_motors.port, motors=dummy_motors)
|
||||
bus.connect(handshake=False)
|
||||
bus.reset_calibration = MagicMock()
|
||||
|
||||
bus.set_half_turn_homings()
|
||||
|
||||
bus.reset_calibration.assert_called_once()
|
||||
assert mock_motors.stubs[read_pos_stub].called
|
||||
assert all(mock_motors.stubs[stub].called for stub in write_homing_stubs)
|
||||
|
||||
|
||||
def test_record_ranges_of_motion(mock_motors, dummy_motors):
|
||||
positions = {
|
||||
1: [351, 42, 1337],
|
||||
2: [28, 3600, 2444],
|
||||
3: [4002, 2999, 146],
|
||||
}
|
||||
expected_mins = {
|
||||
"dummy_1": 42,
|
||||
"dummy_2": 28,
|
||||
"dummy_3": 146,
|
||||
}
|
||||
expected_maxes = {
|
||||
"dummy_1": 1337,
|
||||
"dummy_2": 3600,
|
||||
"dummy_3": 4002,
|
||||
}
|
||||
read_pos_stub = mock_motors.build_sequential_sync_read_stub(
|
||||
*X_SERIES_CONTROL_TABLE["Present_Position"], positions
|
||||
)
|
||||
with patch("lerobot.common.motors.motors_bus.enter_pressed", side_effect=[False, True]):
|
||||
bus = DynamixelMotorsBus(port=mock_motors.port, motors=dummy_motors)
|
||||
bus.connect(handshake=False)
|
||||
|
||||
mins, maxes = bus.record_ranges_of_motion(display_values=False)
|
||||
|
||||
assert mock_motors.stubs[read_pos_stub].calls == 3
|
||||
assert mins == expected_mins
|
||||
assert maxes == expected_maxes
|
||||
@@ -0,0 +1,443 @@
|
||||
import re
|
||||
import sys
|
||||
from typing import Generator
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from lerobot.common.motors import Motor, MotorCalibration, MotorNormMode
|
||||
from lerobot.common.motors.feetech import MODEL_NUMBER, MODEL_NUMBER_TABLE, FeetechMotorsBus
|
||||
from lerobot.common.motors.feetech.tables import STS_SMS_SERIES_CONTROL_TABLE
|
||||
from lerobot.common.utils.encoding_utils import encode_sign_magnitude
|
||||
|
||||
try:
|
||||
import scservo_sdk as scs
|
||||
|
||||
from tests.mocks.mock_feetech import MockMotors, MockPortHandler
|
||||
except (ImportError, ModuleNotFoundError):
|
||||
pytest.skip("scservo_sdk not available", allow_module_level=True)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def patch_port_handler():
|
||||
if sys.platform == "darwin":
|
||||
with patch.object(scs, "PortHandler", MockPortHandler):
|
||||
yield
|
||||
else:
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_motors() -> Generator[MockMotors, None, None]:
|
||||
motors = MockMotors()
|
||||
motors.open()
|
||||
yield motors
|
||||
motors.close()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def dummy_motors() -> dict[str, Motor]:
|
||||
return {
|
||||
"dummy_1": Motor(1, "sts3215", MotorNormMode.RANGE_M100_100),
|
||||
"dummy_2": Motor(2, "sts3215", MotorNormMode.RANGE_M100_100),
|
||||
"dummy_3": Motor(3, "sts3215", MotorNormMode.RANGE_M100_100),
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def dummy_calibration(dummy_motors) -> dict[str, MotorCalibration]:
|
||||
homings = [-709, -2006, 1624]
|
||||
mins = [43, 27, 145]
|
||||
maxes = [1335, 3608, 3999]
|
||||
calibration = {}
|
||||
for motor, m in dummy_motors.items():
|
||||
calibration[motor] = MotorCalibration(
|
||||
id=m.id,
|
||||
drive_mode=0,
|
||||
homing_offset=homings[m.id - 1],
|
||||
range_min=mins[m.id - 1],
|
||||
range_max=maxes[m.id - 1],
|
||||
)
|
||||
return calibration
|
||||
|
||||
|
||||
@pytest.mark.skipif(sys.platform != "darwin", reason=f"No patching needed on {sys.platform=}")
|
||||
def test_autouse_patch():
|
||||
"""Ensures that the autouse fixture correctly patches scs.PortHandler with MockPortHandler."""
|
||||
assert scs.PortHandler is MockPortHandler
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"protocol, value, length, expected",
|
||||
[
|
||||
(0, 0x12, 1, [0x12]),
|
||||
(1, 0x12, 1, [0x12]),
|
||||
(0, 0x1234, 2, [0x34, 0x12]),
|
||||
(1, 0x1234, 2, [0x12, 0x34]),
|
||||
(0, 0x12345678, 4, [0x78, 0x56, 0x34, 0x12]),
|
||||
(1, 0x12345678, 4, [0x56, 0x78, 0x12, 0x34]),
|
||||
],
|
||||
ids=[
|
||||
"P0: 1 byte",
|
||||
"P1: 1 byte",
|
||||
"P0: 2 bytes",
|
||||
"P1: 2 bytes",
|
||||
"P0: 4 bytes",
|
||||
"P1: 4 bytes",
|
||||
],
|
||||
) # fmt: skip
|
||||
def test__split_into_byte_chunks(protocol, value, length, expected):
|
||||
bus = FeetechMotorsBus("", {}, protocol_version=protocol)
|
||||
assert bus._split_into_byte_chunks(value, length) == expected
|
||||
|
||||
|
||||
def test_abc_implementation(dummy_motors):
|
||||
"""Instantiation should raise an error if the class doesn't implement abstract methods/properties."""
|
||||
FeetechMotorsBus(port="/dev/dummy-port", motors=dummy_motors)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("id_", [1, 2, 3])
|
||||
def test_ping(id_, mock_motors, dummy_motors):
|
||||
expected_model_nb = MODEL_NUMBER_TABLE[dummy_motors[f"dummy_{id_}"].model]
|
||||
addr, length = MODEL_NUMBER
|
||||
ping_stub = mock_motors.build_ping_stub(id_)
|
||||
mobel_nb_stub = mock_motors.build_read_stub(addr, length, id_, expected_model_nb)
|
||||
bus = FeetechMotorsBus(
|
||||
port=mock_motors.port,
|
||||
motors=dummy_motors,
|
||||
)
|
||||
bus.connect(handshake=False)
|
||||
|
||||
ping_model_nb = bus.ping(id_)
|
||||
|
||||
assert ping_model_nb == expected_model_nb
|
||||
assert mock_motors.stubs[ping_stub].called
|
||||
assert mock_motors.stubs[mobel_nb_stub].called
|
||||
|
||||
|
||||
def test_broadcast_ping(mock_motors, dummy_motors):
|
||||
models = {m.id: m.model for m in dummy_motors.values()}
|
||||
addr, length = MODEL_NUMBER
|
||||
ping_stub = mock_motors.build_broadcast_ping_stub(list(models))
|
||||
mobel_nb_stubs = []
|
||||
expected_model_nbs = {}
|
||||
for id_, model in models.items():
|
||||
model_nb = MODEL_NUMBER_TABLE[model]
|
||||
stub = mock_motors.build_read_stub(addr, length, id_, model_nb)
|
||||
expected_model_nbs[id_] = model_nb
|
||||
mobel_nb_stubs.append(stub)
|
||||
bus = FeetechMotorsBus(
|
||||
port=mock_motors.port,
|
||||
motors=dummy_motors,
|
||||
)
|
||||
bus.connect(handshake=False)
|
||||
|
||||
ping_model_nbs = bus.broadcast_ping()
|
||||
|
||||
assert ping_model_nbs == expected_model_nbs
|
||||
assert mock_motors.stubs[ping_stub].called
|
||||
assert all(mock_motors.stubs[stub].called for stub in mobel_nb_stubs)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"addr, length, id_, value",
|
||||
[
|
||||
(0, 1, 1, 2),
|
||||
(10, 2, 2, 999),
|
||||
(42, 4, 3, 1337),
|
||||
],
|
||||
)
|
||||
def test__read(addr, length, id_, value, mock_motors, dummy_motors):
|
||||
stub = mock_motors.build_read_stub(addr, length, id_, value)
|
||||
bus = FeetechMotorsBus(
|
||||
port=mock_motors.port,
|
||||
motors=dummy_motors,
|
||||
)
|
||||
bus.connect(handshake=False)
|
||||
|
||||
read_value, _, _ = bus._read(addr, length, id_)
|
||||
|
||||
assert mock_motors.stubs[stub].called
|
||||
assert read_value == value
|
||||
|
||||
|
||||
@pytest.mark.parametrize("raise_on_error", (True, False))
|
||||
def test__read_error(raise_on_error, mock_motors, dummy_motors):
|
||||
addr, length, id_, value, error = (10, 4, 1, 1337, scs.ERRBIT_VOLTAGE)
|
||||
stub = mock_motors.build_read_stub(addr, length, id_, value, error=error)
|
||||
bus = FeetechMotorsBus(
|
||||
port=mock_motors.port,
|
||||
motors=dummy_motors,
|
||||
)
|
||||
bus.connect(handshake=False)
|
||||
|
||||
if raise_on_error:
|
||||
with pytest.raises(RuntimeError, match=re.escape("[RxPacketError] Input voltage error!")):
|
||||
bus._read(addr, length, id_, raise_on_error=raise_on_error)
|
||||
else:
|
||||
_, _, read_error = bus._read(addr, length, id_, raise_on_error=raise_on_error)
|
||||
assert read_error == error
|
||||
|
||||
assert mock_motors.stubs[stub].called
|
||||
|
||||
|
||||
@pytest.mark.parametrize("raise_on_error", (True, False))
|
||||
def test__read_comm(raise_on_error, mock_motors, dummy_motors):
|
||||
addr, length, id_, value = (10, 4, 1, 1337)
|
||||
stub = mock_motors.build_read_stub(addr, length, id_, value, reply=False)
|
||||
bus = FeetechMotorsBus(
|
||||
port=mock_motors.port,
|
||||
motors=dummy_motors,
|
||||
)
|
||||
bus.connect(handshake=False)
|
||||
|
||||
if raise_on_error:
|
||||
with pytest.raises(ConnectionError, match=re.escape("[TxRxResult] There is no status packet!")):
|
||||
bus._read(addr, length, id_, raise_on_error=raise_on_error)
|
||||
else:
|
||||
_, read_comm, _ = bus._read(addr, length, id_, raise_on_error=raise_on_error)
|
||||
assert read_comm == scs.COMM_RX_TIMEOUT
|
||||
|
||||
assert mock_motors.stubs[stub].called
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"addr, length, id_, value",
|
||||
[
|
||||
(0, 1, 1, 2),
|
||||
(10, 2, 2, 999),
|
||||
(42, 4, 3, 1337),
|
||||
],
|
||||
)
|
||||
def test__write(addr, length, id_, value, mock_motors, dummy_motors):
|
||||
stub = mock_motors.build_write_stub(addr, length, id_, value)
|
||||
bus = FeetechMotorsBus(
|
||||
port=mock_motors.port,
|
||||
motors=dummy_motors,
|
||||
)
|
||||
bus.connect(handshake=False)
|
||||
|
||||
comm, error = bus._write(addr, length, id_, value)
|
||||
|
||||
assert mock_motors.stubs[stub].called
|
||||
assert comm == scs.COMM_SUCCESS
|
||||
assert error == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("raise_on_error", (True, False))
|
||||
def test__write_error(raise_on_error, mock_motors, dummy_motors):
|
||||
addr, length, id_, value, error = (10, 4, 1, 1337, scs.ERRBIT_VOLTAGE)
|
||||
stub = mock_motors.build_write_stub(addr, length, id_, value, error=error)
|
||||
bus = FeetechMotorsBus(port=mock_motors.port, motors=dummy_motors)
|
||||
bus.connect(handshake=False)
|
||||
|
||||
if raise_on_error:
|
||||
with pytest.raises(RuntimeError, match=re.escape("[RxPacketError] Input voltage error!")):
|
||||
bus._write(addr, length, id_, value, raise_on_error=raise_on_error)
|
||||
else:
|
||||
_, write_error = bus._write(addr, length, id_, value, raise_on_error=raise_on_error)
|
||||
assert write_error == error
|
||||
|
||||
assert mock_motors.stubs[stub].called
|
||||
|
||||
|
||||
@pytest.mark.parametrize("raise_on_error", (True, False))
|
||||
def test__write_comm(raise_on_error, mock_motors, dummy_motors):
|
||||
addr, length, id_, value = (10, 4, 1, 1337)
|
||||
stub = mock_motors.build_write_stub(addr, length, id_, value, reply=False)
|
||||
bus = FeetechMotorsBus(port=mock_motors.port, motors=dummy_motors)
|
||||
bus.connect(handshake=False)
|
||||
|
||||
if raise_on_error:
|
||||
with pytest.raises(ConnectionError, match=re.escape("[TxRxResult] There is no status packet!")):
|
||||
bus._write(addr, length, id_, value, raise_on_error=raise_on_error)
|
||||
else:
|
||||
write_comm, _ = bus._write(addr, length, id_, value, raise_on_error=raise_on_error)
|
||||
assert write_comm == scs.COMM_RX_TIMEOUT
|
||||
|
||||
assert mock_motors.stubs[stub].called
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"addr, length, ids_values",
|
||||
[
|
||||
(0, 1, {1: 4}),
|
||||
(10, 2, {1: 1337, 2: 42}),
|
||||
(42, 4, {1: 1337, 2: 42, 3: 4016}),
|
||||
],
|
||||
ids=["1 motor", "2 motors", "3 motors"],
|
||||
)
|
||||
def test__sync_read(addr, length, ids_values, mock_motors, dummy_motors):
|
||||
stub = mock_motors.build_sync_read_stub(addr, length, ids_values)
|
||||
bus = FeetechMotorsBus(port=mock_motors.port, motors=dummy_motors)
|
||||
bus.connect(handshake=False)
|
||||
|
||||
read_values, _ = bus._sync_read(addr, length, list(ids_values))
|
||||
|
||||
assert mock_motors.stubs[stub].called
|
||||
assert read_values == ids_values
|
||||
|
||||
|
||||
@pytest.mark.parametrize("raise_on_error", (True, False))
|
||||
def test__sync_read_comm(raise_on_error, mock_motors, dummy_motors):
|
||||
addr, length, ids_values = (10, 4, {1: 1337})
|
||||
stub = mock_motors.build_sync_read_stub(addr, length, ids_values, reply=False)
|
||||
bus = FeetechMotorsBus(port=mock_motors.port, motors=dummy_motors)
|
||||
bus.connect(handshake=False)
|
||||
|
||||
if raise_on_error:
|
||||
with pytest.raises(ConnectionError, match=re.escape("[TxRxResult] There is no status packet!")):
|
||||
bus._sync_read(addr, length, list(ids_values), raise_on_error=raise_on_error)
|
||||
else:
|
||||
_, read_comm = bus._sync_read(addr, length, list(ids_values), raise_on_error=raise_on_error)
|
||||
assert read_comm == scs.COMM_RX_TIMEOUT
|
||||
|
||||
assert mock_motors.stubs[stub].called
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"addr, length, ids_values",
|
||||
[
|
||||
(0, 1, {1: 4}),
|
||||
(10, 2, {1: 1337, 2: 42}),
|
||||
(42, 4, {1: 1337, 2: 42, 3: 4016}),
|
||||
],
|
||||
ids=["1 motor", "2 motors", "3 motors"],
|
||||
)
|
||||
def test__sync_write(addr, length, ids_values, mock_motors, dummy_motors):
|
||||
stub = mock_motors.build_sync_write_stub(addr, length, ids_values)
|
||||
bus = FeetechMotorsBus(port=mock_motors.port, motors=dummy_motors)
|
||||
bus.connect(handshake=False)
|
||||
|
||||
comm = bus._sync_write(addr, length, ids_values)
|
||||
|
||||
assert mock_motors.stubs[stub].wait_called()
|
||||
assert comm == scs.COMM_SUCCESS
|
||||
|
||||
|
||||
def test_is_calibrated(mock_motors, dummy_motors, dummy_calibration):
|
||||
mins_stubs, maxes_stubs, homings_stubs = [], [], []
|
||||
for cal in dummy_calibration.values():
|
||||
mins_stubs.append(
|
||||
mock_motors.build_read_stub(
|
||||
*STS_SMS_SERIES_CONTROL_TABLE["Min_Position_Limit"], cal.id, cal.range_min
|
||||
)
|
||||
)
|
||||
maxes_stubs.append(
|
||||
mock_motors.build_read_stub(
|
||||
*STS_SMS_SERIES_CONTROL_TABLE["Max_Position_Limit"], cal.id, cal.range_max
|
||||
)
|
||||
)
|
||||
homings_stubs.append(
|
||||
mock_motors.build_read_stub(
|
||||
*STS_SMS_SERIES_CONTROL_TABLE["Homing_Offset"],
|
||||
cal.id,
|
||||
encode_sign_magnitude(cal.homing_offset, 11),
|
||||
)
|
||||
)
|
||||
|
||||
bus = FeetechMotorsBus(
|
||||
port=mock_motors.port,
|
||||
motors=dummy_motors,
|
||||
calibration=dummy_calibration,
|
||||
)
|
||||
bus.connect(handshake=False)
|
||||
|
||||
is_calibrated = bus.is_calibrated
|
||||
|
||||
assert is_calibrated
|
||||
assert all(mock_motors.stubs[stub].called for stub in mins_stubs)
|
||||
assert all(mock_motors.stubs[stub].called for stub in maxes_stubs)
|
||||
assert all(mock_motors.stubs[stub].called for stub in homings_stubs)
|
||||
|
||||
|
||||
def test_reset_calibration(mock_motors, dummy_motors):
|
||||
write_homing_stubs = []
|
||||
write_mins_stubs = []
|
||||
write_maxes_stubs = []
|
||||
for motor in dummy_motors.values():
|
||||
write_homing_stubs.append(
|
||||
mock_motors.build_write_stub(*STS_SMS_SERIES_CONTROL_TABLE["Homing_Offset"], motor.id, 0)
|
||||
)
|
||||
write_mins_stubs.append(
|
||||
mock_motors.build_write_stub(*STS_SMS_SERIES_CONTROL_TABLE["Min_Position_Limit"], motor.id, 0)
|
||||
)
|
||||
write_maxes_stubs.append(
|
||||
mock_motors.build_write_stub(*STS_SMS_SERIES_CONTROL_TABLE["Max_Position_Limit"], motor.id, 4095)
|
||||
)
|
||||
|
||||
bus = FeetechMotorsBus(port=mock_motors.port, motors=dummy_motors)
|
||||
bus.connect(handshake=False)
|
||||
|
||||
bus.reset_calibration()
|
||||
|
||||
assert all(mock_motors.stubs[stub].called for stub in write_homing_stubs)
|
||||
assert all(mock_motors.stubs[stub].called for stub in write_mins_stubs)
|
||||
assert all(mock_motors.stubs[stub].called for stub in write_maxes_stubs)
|
||||
|
||||
|
||||
def test_set_half_turn_homings(mock_motors, dummy_motors):
|
||||
"""
|
||||
For this test, we assume that the homing offsets are already 0 such that
|
||||
Present_Position == Actual_Position
|
||||
"""
|
||||
current_positions = {
|
||||
1: 1337,
|
||||
2: 42,
|
||||
3: 3672,
|
||||
}
|
||||
expected_homings = {
|
||||
1: -710, # 1337 - 2047
|
||||
2: -2005, # 42 - 2047
|
||||
3: 1625, # 3672 - 2047
|
||||
}
|
||||
read_pos_stub = mock_motors.build_sync_read_stub(
|
||||
*STS_SMS_SERIES_CONTROL_TABLE["Present_Position"], current_positions
|
||||
)
|
||||
write_homing_stubs = []
|
||||
for id_, homing in expected_homings.items():
|
||||
encoded_homing = encode_sign_magnitude(homing, 11)
|
||||
stub = mock_motors.build_write_stub(
|
||||
*STS_SMS_SERIES_CONTROL_TABLE["Homing_Offset"], id_, encoded_homing
|
||||
)
|
||||
write_homing_stubs.append(stub)
|
||||
|
||||
bus = FeetechMotorsBus(port=mock_motors.port, motors=dummy_motors)
|
||||
bus.connect(handshake=False)
|
||||
bus.reset_calibration = MagicMock()
|
||||
|
||||
bus.set_half_turn_homings()
|
||||
|
||||
bus.reset_calibration.assert_called_once()
|
||||
assert mock_motors.stubs[read_pos_stub].called
|
||||
assert all(mock_motors.stubs[stub].called for stub in write_homing_stubs)
|
||||
|
||||
|
||||
def test_record_ranges_of_motion(mock_motors, dummy_motors):
|
||||
positions = {
|
||||
1: [351, 42, 1337],
|
||||
2: [28, 3600, 2444],
|
||||
3: [4002, 2999, 146],
|
||||
}
|
||||
expected_mins = {
|
||||
"dummy_1": 42,
|
||||
"dummy_2": 28,
|
||||
"dummy_3": 146,
|
||||
}
|
||||
expected_maxes = {
|
||||
"dummy_1": 1337,
|
||||
"dummy_2": 3600,
|
||||
"dummy_3": 4002,
|
||||
}
|
||||
stub = mock_motors.build_sequential_sync_read_stub(
|
||||
*STS_SMS_SERIES_CONTROL_TABLE["Present_Position"], positions
|
||||
)
|
||||
with patch("lerobot.common.motors.motors_bus.enter_pressed", side_effect=[False, True]):
|
||||
bus = FeetechMotorsBus(port=mock_motors.port, motors=dummy_motors)
|
||||
bus.connect(handshake=False)
|
||||
|
||||
mins, maxes = bus.record_ranges_of_motion(display_values=False)
|
||||
|
||||
assert mock_motors.stubs[stub].calls == 3
|
||||
assert mins == expected_mins
|
||||
assert maxes == expected_maxes
|
||||
@@ -1,157 +0,0 @@
|
||||
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# 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.
|
||||
"""
|
||||
Tests for physical motors and their mocked versions.
|
||||
If the physical motors 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_motors.py::test_find_port
|
||||
pytest -sx tests/test_motors.py::test_motors_bus
|
||||
```
|
||||
|
||||
Example of running test on real dynamixel motors connected to the computer:
|
||||
```bash
|
||||
pytest -sx 'tests/test_motors.py::test_motors_bus[dynamixel-False]'
|
||||
```
|
||||
|
||||
Example of running test on a mocked version of dynamixel motors:
|
||||
```bash
|
||||
pytest -sx 'tests/test_motors.py::test_motors_bus[dynamixel-True]'
|
||||
```
|
||||
"""
|
||||
|
||||
# TODO(rcadene): measure fps in nightly?
|
||||
# TODO(rcadene): test logs
|
||||
# TODO(rcadene): test calibration
|
||||
# TODO(rcadene): add compatibility with other motors bus
|
||||
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from lerobot.common.robot_devices.utils import RobotDeviceAlreadyConnectedError, RobotDeviceNotConnectedError
|
||||
from lerobot.scripts.find_motors_bus_port import find_port
|
||||
from tests.utils import TEST_MOTOR_TYPES, make_motors_bus, require_motor
|
||||
|
||||
|
||||
@pytest.mark.parametrize("motor_type, mock", TEST_MOTOR_TYPES)
|
||||
@require_motor
|
||||
def test_find_port(request, motor_type, mock):
|
||||
if mock:
|
||||
request.getfixturevalue("patch_builtins_input")
|
||||
with pytest.raises(OSError):
|
||||
find_port()
|
||||
else:
|
||||
find_port()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("motor_type, mock", TEST_MOTOR_TYPES)
|
||||
@require_motor
|
||||
def test_configure_motors_all_ids_1(request, motor_type, mock):
|
||||
if mock:
|
||||
request.getfixturevalue("patch_builtins_input")
|
||||
|
||||
if motor_type == "dynamixel":
|
||||
# see X_SERIES_BAUDRATE_TABLE
|
||||
smaller_baudrate = 9_600
|
||||
smaller_baudrate_value = 0
|
||||
elif motor_type == "feetech":
|
||||
# see SCS_SERIES_BAUDRATE_TABLE
|
||||
smaller_baudrate = 19_200
|
||||
smaller_baudrate_value = 7
|
||||
else:
|
||||
raise ValueError(motor_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(motor_type, mock=mock)
|
||||
motors_bus.connect()
|
||||
motors_bus.write("Baud_Rate", [smaller_baudrate_value] * len(motors_bus.motors))
|
||||
|
||||
motors_bus.set_bus_baudrate(smaller_baudrate)
|
||||
motors_bus.write("ID", [1] * len(motors_bus.motors))
|
||||
del motors_bus
|
||||
|
||||
# Test configure
|
||||
motors_bus = make_motors_bus(motor_type, mock=mock)
|
||||
motors_bus.connect()
|
||||
assert motors_bus.are_motors_configured()
|
||||
del motors_bus
|
||||
|
||||
|
||||
@pytest.mark.parametrize("motor_type, mock", TEST_MOTOR_TYPES)
|
||||
@require_motor
|
||||
def test_motors_bus(request, motor_type, mock):
|
||||
if mock:
|
||||
request.getfixturevalue("patch_builtins_input")
|
||||
|
||||
motors_bus = make_motors_bus(motor_type, mock=mock)
|
||||
|
||||
# Test reading and writing 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 = make_motors_bus(motor_type, mock=mock)
|
||||
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_bus.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()
|
||||
@@ -0,0 +1,342 @@
|
||||
import re
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from lerobot.common.motors.motors_bus import (
|
||||
Motor,
|
||||
MotorNormMode,
|
||||
assert_same_address,
|
||||
get_address,
|
||||
get_ctrl_table,
|
||||
)
|
||||
from tests.mocks.mock_motors_bus import (
|
||||
DUMMY_CTRL_TABLE_1,
|
||||
DUMMY_CTRL_TABLE_2,
|
||||
DUMMY_MODEL_CTRL_TABLE,
|
||||
MockMotorsBus,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def dummy_motors() -> dict[str, Motor]:
|
||||
return {
|
||||
"dummy_1": Motor(1, "model_2", MotorNormMode.RANGE_M100_100),
|
||||
"dummy_2": Motor(2, "model_3", MotorNormMode.RANGE_M100_100),
|
||||
"dummy_3": Motor(3, "model_2", MotorNormMode.RANGE_0_100),
|
||||
}
|
||||
|
||||
|
||||
def test_get_ctrl_table():
|
||||
model = "model_1"
|
||||
ctrl_table = get_ctrl_table(DUMMY_MODEL_CTRL_TABLE, model)
|
||||
assert ctrl_table == DUMMY_CTRL_TABLE_1
|
||||
|
||||
|
||||
def test_get_ctrl_table_error():
|
||||
model = "model_99"
|
||||
with pytest.raises(KeyError, match=f"Control table for {model=} not found."):
|
||||
get_ctrl_table(DUMMY_MODEL_CTRL_TABLE, model)
|
||||
|
||||
|
||||
def test_get_address():
|
||||
addr, n_bytes = get_address(DUMMY_MODEL_CTRL_TABLE, "model_1", "Firmware_Version")
|
||||
assert addr == 0
|
||||
assert n_bytes == 1
|
||||
|
||||
|
||||
def test_get_address_error():
|
||||
model = "model_1"
|
||||
data_name = "Lock"
|
||||
with pytest.raises(KeyError, match=f"Address for '{data_name}' not found in {model} control table."):
|
||||
get_address(DUMMY_MODEL_CTRL_TABLE, "model_1", data_name)
|
||||
|
||||
|
||||
def test_assert_same_address():
|
||||
models = ["model_1", "model_2"]
|
||||
assert_same_address(DUMMY_MODEL_CTRL_TABLE, models, "Present_Position")
|
||||
|
||||
|
||||
def test_assert_same_length_different_addresses():
|
||||
models = ["model_1", "model_2"]
|
||||
with pytest.raises(
|
||||
NotImplementedError,
|
||||
match=re.escape("At least two motor models use a different address"),
|
||||
):
|
||||
assert_same_address(DUMMY_MODEL_CTRL_TABLE, models, "Model_Number")
|
||||
|
||||
|
||||
def test_assert_same_address_different_length():
|
||||
models = ["model_1", "model_2"]
|
||||
with pytest.raises(
|
||||
NotImplementedError,
|
||||
match=re.escape("At least two motor models use a different bytes representation"),
|
||||
):
|
||||
assert_same_address(DUMMY_MODEL_CTRL_TABLE, models, "Goal_Position")
|
||||
|
||||
|
||||
def test__serialize_data_invalid_length():
|
||||
bus = MockMotorsBus("", {})
|
||||
with pytest.raises(NotImplementedError):
|
||||
bus._serialize_data(100, 3)
|
||||
|
||||
|
||||
def test__serialize_data_negative_numbers():
|
||||
bus = MockMotorsBus("", {})
|
||||
with pytest.raises(ValueError):
|
||||
bus._serialize_data(-1, 1)
|
||||
|
||||
|
||||
def test__serialize_data_large_number():
|
||||
bus = MockMotorsBus("", {})
|
||||
with pytest.raises(ValueError):
|
||||
bus._serialize_data(2**32, 4) # 4-byte max is 0xFFFFFFFF
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"data_name, id_, value",
|
||||
[
|
||||
("Firmware_Version", 1, 14),
|
||||
("Model_Number", 1, 5678),
|
||||
("Present_Position", 2, 1337),
|
||||
("Present_Velocity", 3, 42),
|
||||
],
|
||||
)
|
||||
def test_read(data_name, id_, value, dummy_motors):
|
||||
bus = MockMotorsBus("/dev/dummy-port", dummy_motors)
|
||||
bus.connect(handshake=False)
|
||||
addr, length = DUMMY_CTRL_TABLE_2[data_name]
|
||||
|
||||
with (
|
||||
patch.object(MockMotorsBus, "_read", return_value=(value, 0, 0)) as mock__read,
|
||||
patch.object(MockMotorsBus, "_decode_sign", return_value={id_: value}) as mock__decode_sign,
|
||||
patch.object(MockMotorsBus, "_normalize", return_value={id_: value}) as mock__normalize,
|
||||
):
|
||||
returned_value = bus.read(data_name, f"dummy_{id_}")
|
||||
|
||||
assert returned_value == value
|
||||
mock__read.assert_called_once_with(
|
||||
addr,
|
||||
length,
|
||||
id_,
|
||||
num_retry=0,
|
||||
raise_on_error=True,
|
||||
err_msg=f"Failed to read '{data_name}' on {id_=} after 1 tries.",
|
||||
)
|
||||
mock__decode_sign.assert_called_once_with(data_name, {id_: value})
|
||||
if data_name in bus.normalized_data:
|
||||
mock__normalize.assert_called_once_with({id_: value})
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"data_name, id_, value",
|
||||
[
|
||||
("Goal_Position", 1, 1337),
|
||||
("Goal_Velocity", 2, 3682),
|
||||
("Lock", 3, 1),
|
||||
],
|
||||
)
|
||||
def test_write(data_name, id_, value, dummy_motors):
|
||||
bus = MockMotorsBus("/dev/dummy-port", dummy_motors)
|
||||
bus.connect(handshake=False)
|
||||
addr, length = DUMMY_CTRL_TABLE_2[data_name]
|
||||
|
||||
with (
|
||||
patch.object(MockMotorsBus, "_write", return_value=(0, 0)) as mock__write,
|
||||
patch.object(MockMotorsBus, "_encode_sign", return_value={id_: value}) as mock__encode_sign,
|
||||
patch.object(MockMotorsBus, "_unnormalize", return_value={id_: value}) as mock__unnormalize,
|
||||
):
|
||||
bus.write(data_name, f"dummy_{id_}", value)
|
||||
|
||||
mock__write.assert_called_once_with(
|
||||
addr,
|
||||
length,
|
||||
id_,
|
||||
value,
|
||||
num_retry=0,
|
||||
raise_on_error=True,
|
||||
err_msg=f"Failed to write '{data_name}' on {id_=} with '{value}' after 1 tries.",
|
||||
)
|
||||
mock__encode_sign.assert_called_once_with(data_name, {id_: value})
|
||||
if data_name in bus.normalized_data:
|
||||
mock__unnormalize.assert_called_once_with({id_: value})
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"data_name, id_, value",
|
||||
[
|
||||
("Firmware_Version", 1, 14),
|
||||
("Model_Number", 1, 5678),
|
||||
("Present_Position", 2, 1337),
|
||||
("Present_Velocity", 3, 42),
|
||||
],
|
||||
)
|
||||
def test_sync_read_by_str(data_name, id_, value, dummy_motors):
|
||||
bus = MockMotorsBus("/dev/dummy-port", dummy_motors)
|
||||
bus.connect(handshake=False)
|
||||
addr, length = DUMMY_CTRL_TABLE_2[data_name]
|
||||
ids = [id_]
|
||||
expected_value = {f"dummy_{id_}": value}
|
||||
|
||||
with (
|
||||
patch.object(MockMotorsBus, "_sync_read", return_value=({id_: value}, 0)) as mock__sync_read,
|
||||
patch.object(MockMotorsBus, "_decode_sign", return_value={id_: value}) as mock__decode_sign,
|
||||
patch.object(MockMotorsBus, "_normalize", return_value={id_: value}) as mock__normalize,
|
||||
):
|
||||
returned_dict = bus.sync_read(data_name, f"dummy_{id_}")
|
||||
|
||||
assert returned_dict == expected_value
|
||||
mock__sync_read.assert_called_once_with(
|
||||
addr,
|
||||
length,
|
||||
ids,
|
||||
num_retry=0,
|
||||
raise_on_error=True,
|
||||
err_msg=f"Failed to sync read '{data_name}' on {ids=} after 1 tries.",
|
||||
)
|
||||
mock__decode_sign.assert_called_once_with(data_name, {id_: value})
|
||||
if data_name in bus.normalized_data:
|
||||
mock__normalize.assert_called_once_with({id_: value})
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"data_name, ids_values",
|
||||
[
|
||||
("Model_Number", {1: 5678}),
|
||||
("Present_Position", {1: 1337, 2: 42}),
|
||||
("Present_Velocity", {1: 1337, 2: 42, 3: 4016}),
|
||||
],
|
||||
ids=["1 motor", "2 motors", "3 motors"],
|
||||
)
|
||||
def test_sync_read_by_list(data_name, ids_values, dummy_motors):
|
||||
bus = MockMotorsBus("/dev/dummy-port", dummy_motors)
|
||||
bus.connect(handshake=False)
|
||||
addr, length = DUMMY_CTRL_TABLE_2[data_name]
|
||||
ids = list(ids_values)
|
||||
expected_values = {f"dummy_{id_}": val for id_, val in ids_values.items()}
|
||||
|
||||
with (
|
||||
patch.object(MockMotorsBus, "_sync_read", return_value=(ids_values, 0)) as mock__sync_read,
|
||||
patch.object(MockMotorsBus, "_decode_sign", return_value=ids_values) as mock__decode_sign,
|
||||
patch.object(MockMotorsBus, "_normalize", return_value=ids_values) as mock__normalize,
|
||||
):
|
||||
returned_dict = bus.sync_read(data_name, [f"dummy_{id_}" for id_ in ids])
|
||||
|
||||
assert returned_dict == expected_values
|
||||
mock__sync_read.assert_called_once_with(
|
||||
addr,
|
||||
length,
|
||||
ids,
|
||||
num_retry=0,
|
||||
raise_on_error=True,
|
||||
err_msg=f"Failed to sync read '{data_name}' on {ids=} after 1 tries.",
|
||||
)
|
||||
mock__decode_sign.assert_called_once_with(data_name, ids_values)
|
||||
if data_name in bus.normalized_data:
|
||||
mock__normalize.assert_called_once_with(ids_values)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"data_name, ids_values",
|
||||
[
|
||||
("Model_Number", {1: 5678, 2: 5799, 3: 5678}),
|
||||
("Present_Position", {1: 1337, 2: 42, 3: 4016}),
|
||||
("Goal_Position", {1: 4008, 2: 199, 3: 3446}),
|
||||
],
|
||||
ids=["Model_Number", "Present_Position", "Goal_Position"],
|
||||
)
|
||||
def test_sync_read_by_none(data_name, ids_values, dummy_motors):
|
||||
bus = MockMotorsBus("/dev/dummy-port", dummy_motors)
|
||||
bus.connect(handshake=False)
|
||||
addr, length = DUMMY_CTRL_TABLE_2[data_name]
|
||||
ids = list(ids_values)
|
||||
expected_values = {f"dummy_{id_}": val for id_, val in ids_values.items()}
|
||||
|
||||
with (
|
||||
patch.object(MockMotorsBus, "_sync_read", return_value=(ids_values, 0)) as mock__sync_read,
|
||||
patch.object(MockMotorsBus, "_decode_sign", return_value=ids_values) as mock__decode_sign,
|
||||
patch.object(MockMotorsBus, "_normalize", return_value=ids_values) as mock__normalize,
|
||||
):
|
||||
returned_dict = bus.sync_read(data_name)
|
||||
|
||||
assert returned_dict == expected_values
|
||||
mock__sync_read.assert_called_once_with(
|
||||
addr,
|
||||
length,
|
||||
ids,
|
||||
num_retry=0,
|
||||
raise_on_error=True,
|
||||
err_msg=f"Failed to sync read '{data_name}' on {ids=} after 1 tries.",
|
||||
)
|
||||
mock__decode_sign.assert_called_once_with(data_name, ids_values)
|
||||
if data_name in bus.normalized_data:
|
||||
mock__normalize.assert_called_once_with(ids_values)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"data_name, value",
|
||||
[
|
||||
("Goal_Position", 500),
|
||||
("Goal_Velocity", 4010),
|
||||
("Lock", 0),
|
||||
],
|
||||
)
|
||||
def test_sync_write_by_single_value(data_name, value, dummy_motors):
|
||||
bus = MockMotorsBus("/dev/dummy-port", dummy_motors)
|
||||
bus.connect(handshake=False)
|
||||
addr, length = DUMMY_CTRL_TABLE_2[data_name]
|
||||
ids_values = {m.id: value for m in dummy_motors.values()}
|
||||
|
||||
with (
|
||||
patch.object(MockMotorsBus, "_sync_write", return_value=(ids_values, 0)) as mock__sync_write,
|
||||
patch.object(MockMotorsBus, "_encode_sign", return_value=ids_values) as mock__encode_sign,
|
||||
patch.object(MockMotorsBus, "_unnormalize", return_value=ids_values) as mock__unnormalize,
|
||||
):
|
||||
bus.sync_write(data_name, value)
|
||||
|
||||
mock__sync_write.assert_called_once_with(
|
||||
addr,
|
||||
length,
|
||||
ids_values,
|
||||
num_retry=0,
|
||||
raise_on_error=True,
|
||||
err_msg=f"Failed to sync write '{data_name}' with {ids_values=} after 1 tries.",
|
||||
)
|
||||
mock__encode_sign.assert_called_once_with(data_name, ids_values)
|
||||
if data_name in bus.normalized_data:
|
||||
mock__unnormalize.assert_called_once_with(ids_values)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"data_name, ids_values",
|
||||
[
|
||||
("Goal_Position", {1: 1337, 2: 42, 3: 4016}),
|
||||
("Goal_Velocity", {1: 50, 2: 83, 3: 2777}),
|
||||
("Lock", {1: 0, 2: 0, 3: 1}),
|
||||
],
|
||||
ids=["Goal_Position", "Goal_Velocity", "Lock"],
|
||||
)
|
||||
def test_sync_write_by_value_dict(data_name, ids_values, dummy_motors):
|
||||
bus = MockMotorsBus("/dev/dummy-port", dummy_motors)
|
||||
bus.connect(handshake=False)
|
||||
addr, length = DUMMY_CTRL_TABLE_2[data_name]
|
||||
values = {f"dummy_{id_}": val for id_, val in ids_values.items()}
|
||||
|
||||
with (
|
||||
patch.object(MockMotorsBus, "_sync_write", return_value=(ids_values, 0)) as mock__sync_write,
|
||||
patch.object(MockMotorsBus, "_encode_sign", return_value=ids_values) as mock__encode_sign,
|
||||
patch.object(MockMotorsBus, "_unnormalize", return_value=ids_values) as mock__unnormalize,
|
||||
):
|
||||
bus.sync_write(data_name, values)
|
||||
|
||||
mock__sync_write.assert_called_once_with(
|
||||
addr,
|
||||
length,
|
||||
ids_values,
|
||||
num_retry=0,
|
||||
raise_on_error=True,
|
||||
err_msg=f"Failed to sync write '{data_name}' with {ids_values=} after 1 tries.",
|
||||
)
|
||||
mock__encode_sign.assert_called_once_with(data_name, ids_values)
|
||||
if data_name in bus.normalized_data:
|
||||
mock__unnormalize.assert_called_once_with(ids_values)
|
||||
@@ -37,7 +37,6 @@ def test_diffuser_scheduler(optimizer):
|
||||
"base_lrs": [0.001],
|
||||
"last_epoch": 1,
|
||||
"lr_lambdas": [None],
|
||||
"verbose": False,
|
||||
}
|
||||
assert scheduler.state_dict() == expected_state_dict
|
||||
|
||||
@@ -56,7 +55,6 @@ def test_vqbet_scheduler(optimizer):
|
||||
"base_lrs": [0.001],
|
||||
"last_epoch": 1,
|
||||
"lr_lambdas": [None],
|
||||
"verbose": False,
|
||||
}
|
||||
assert scheduler.state_dict() == expected_state_dict
|
||||
|
||||
@@ -77,7 +75,6 @@ def test_cosine_decay_with_warmup_scheduler(optimizer):
|
||||
"base_lrs": [0.001],
|
||||
"last_epoch": 1,
|
||||
"lr_lambdas": [None],
|
||||
"verbose": False,
|
||||
}
|
||||
assert scheduler.state_dict() == expected_state_dict
|
||||
|
||||
|
||||
@@ -0,0 +1,139 @@
|
||||
# !/usr/bin/env python
|
||||
|
||||
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# 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 torch
|
||||
|
||||
from lerobot.common.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
|
||||
from lerobot.common.policies.sac.reward_model.modeling_classifier import ClassifierOutput
|
||||
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
|
||||
from tests.utils import require_package
|
||||
|
||||
|
||||
def test_classifier_output():
|
||||
output = ClassifierOutput(
|
||||
logits=torch.tensor([1, 2, 3]),
|
||||
probabilities=torch.tensor([0.1, 0.2, 0.3]),
|
||||
hidden_states=None,
|
||||
)
|
||||
|
||||
assert (
|
||||
f"{output}"
|
||||
== "ClassifierOutput(logits=tensor([1, 2, 3]), probabilities=tensor([0.1000, 0.2000, 0.3000]), hidden_states=None)"
|
||||
)
|
||||
|
||||
|
||||
@require_package("transformers")
|
||||
def test_binary_classifier_with_default_params():
|
||||
from lerobot.common.policies.sac.reward_model.modeling_classifier import Classifier
|
||||
|
||||
config = RewardClassifierConfig()
|
||||
config.input_features = {
|
||||
"observation.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224)),
|
||||
}
|
||||
config.output_features = {
|
||||
"next.reward": PolicyFeature(type=FeatureType.REWARD, shape=(1,)),
|
||||
}
|
||||
config.normalization_mapping = {
|
||||
"VISUAL": NormalizationMode.IDENTITY,
|
||||
"REWARD": NormalizationMode.IDENTITY,
|
||||
}
|
||||
config.num_cameras = 1
|
||||
classifier = Classifier(config)
|
||||
|
||||
batch_size = 10
|
||||
|
||||
input = {
|
||||
"observation.image": torch.rand((batch_size, 3, 128, 128)),
|
||||
"next.reward": torch.randint(low=0, high=2, size=(batch_size,)).float(),
|
||||
}
|
||||
|
||||
images, labels = classifier.extract_images_and_labels(input)
|
||||
assert len(images) == 1
|
||||
assert images[0].shape == torch.Size([batch_size, 3, 128, 128])
|
||||
assert labels.shape == torch.Size([batch_size])
|
||||
|
||||
output = classifier.predict(images)
|
||||
|
||||
assert output is not None
|
||||
assert output.logits.size() == torch.Size([batch_size])
|
||||
assert not torch.isnan(output.logits).any(), "Tensor contains NaN values"
|
||||
assert output.probabilities.shape == torch.Size([batch_size])
|
||||
assert not torch.isnan(output.probabilities).any(), "Tensor contains NaN values"
|
||||
assert output.hidden_states.shape == torch.Size([batch_size, 256])
|
||||
assert not torch.isnan(output.hidden_states).any(), "Tensor contains NaN values"
|
||||
|
||||
|
||||
@require_package("transformers")
|
||||
def test_multiclass_classifier():
|
||||
from lerobot.common.policies.sac.reward_model.modeling_classifier import Classifier
|
||||
|
||||
num_classes = 5
|
||||
config = RewardClassifierConfig()
|
||||
config.input_features = {
|
||||
"observation.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224)),
|
||||
}
|
||||
config.output_features = {
|
||||
"next.reward": PolicyFeature(type=FeatureType.REWARD, shape=(num_classes,)),
|
||||
}
|
||||
config.num_cameras = 1
|
||||
config.num_classes = num_classes
|
||||
classifier = Classifier(config)
|
||||
|
||||
batch_size = 10
|
||||
|
||||
input = {
|
||||
"observation.image": torch.rand((batch_size, 3, 128, 128)),
|
||||
"next.reward": torch.rand((batch_size, num_classes)),
|
||||
}
|
||||
|
||||
images, labels = classifier.extract_images_and_labels(input)
|
||||
assert len(images) == 1
|
||||
assert images[0].shape == torch.Size([batch_size, 3, 128, 128])
|
||||
assert labels.shape == torch.Size([batch_size, num_classes])
|
||||
|
||||
output = classifier.predict(images)
|
||||
|
||||
assert output is not None
|
||||
assert output.logits.shape == torch.Size([batch_size, num_classes])
|
||||
assert not torch.isnan(output.logits).any(), "Tensor contains NaN values"
|
||||
assert output.probabilities.shape == torch.Size([batch_size, num_classes])
|
||||
assert not torch.isnan(output.probabilities).any(), "Tensor contains NaN values"
|
||||
assert output.hidden_states.shape == torch.Size([batch_size, 256])
|
||||
assert not torch.isnan(output.hidden_states).any(), "Tensor contains NaN values"
|
||||
|
||||
|
||||
@require_package("transformers")
|
||||
def test_default_device():
|
||||
from lerobot.common.policies.sac.reward_model.modeling_classifier import Classifier
|
||||
|
||||
config = RewardClassifierConfig()
|
||||
assert config.device == "cpu"
|
||||
|
||||
classifier = Classifier(config)
|
||||
for p in classifier.parameters():
|
||||
assert p.device == torch.device("cpu")
|
||||
|
||||
|
||||
@require_package("transformers")
|
||||
def test_explicit_device_setup():
|
||||
from lerobot.common.policies.sac.reward_model.modeling_classifier import Classifier
|
||||
|
||||
config = RewardClassifierConfig(device="cpu")
|
||||
assert config.device == "cpu"
|
||||
|
||||
classifier = Classifier(config)
|
||||
for p in classifier.parameters():
|
||||
assert p.device == torch.device("cpu")
|
||||
@@ -0,0 +1,217 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# 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 lerobot.common.policies.sac.configuration_sac import (
|
||||
ActorLearnerConfig,
|
||||
ActorNetworkConfig,
|
||||
ConcurrencyConfig,
|
||||
CriticNetworkConfig,
|
||||
PolicyConfig,
|
||||
SACConfig,
|
||||
)
|
||||
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
|
||||
|
||||
|
||||
def test_sac_config_default_initialization():
|
||||
config = SACConfig()
|
||||
|
||||
assert config.normalization_mapping == {
|
||||
"VISUAL": NormalizationMode.MEAN_STD,
|
||||
"STATE": NormalizationMode.MIN_MAX,
|
||||
"ENV": NormalizationMode.MIN_MAX,
|
||||
"ACTION": NormalizationMode.MIN_MAX,
|
||||
}
|
||||
assert config.dataset_stats == {
|
||||
"observation.image": {
|
||||
"mean": [0.485, 0.456, 0.406],
|
||||
"std": [0.229, 0.224, 0.225],
|
||||
},
|
||||
"observation.state": {
|
||||
"min": [0.0, 0.0],
|
||||
"max": [1.0, 1.0],
|
||||
},
|
||||
"action": {
|
||||
"min": [0.0, 0.0, 0.0],
|
||||
"max": [1.0, 1.0, 1.0],
|
||||
},
|
||||
}
|
||||
|
||||
# Basic parameters
|
||||
assert config.device == "cpu"
|
||||
assert config.storage_device == "cpu"
|
||||
assert config.discount == 0.99
|
||||
assert config.temperature_init == 1.0
|
||||
assert config.num_critics == 2
|
||||
|
||||
# Architecture specifics
|
||||
assert config.vision_encoder_name is None
|
||||
assert config.freeze_vision_encoder is True
|
||||
assert config.image_encoder_hidden_dim == 32
|
||||
assert config.shared_encoder is True
|
||||
assert config.num_discrete_actions is None
|
||||
assert config.image_embedding_pooling_dim == 8
|
||||
|
||||
# Training parameters
|
||||
assert config.online_steps == 1000000
|
||||
assert config.online_env_seed == 10000
|
||||
assert config.online_buffer_capacity == 100000
|
||||
assert config.offline_buffer_capacity == 100000
|
||||
assert config.async_prefetch is False
|
||||
assert config.online_step_before_learning == 100
|
||||
assert config.policy_update_freq == 1
|
||||
|
||||
# SAC algorithm parameters
|
||||
assert config.num_subsample_critics is None
|
||||
assert config.critic_lr == 3e-4
|
||||
assert config.actor_lr == 3e-4
|
||||
assert config.temperature_lr == 3e-4
|
||||
assert config.critic_target_update_weight == 0.005
|
||||
assert config.utd_ratio == 1
|
||||
assert config.state_encoder_hidden_dim == 256
|
||||
assert config.latent_dim == 256
|
||||
assert config.target_entropy is None
|
||||
assert config.use_backup_entropy is True
|
||||
assert config.grad_clip_norm == 40.0
|
||||
|
||||
# Dataset stats defaults
|
||||
expected_dataset_stats = {
|
||||
"observation.image": {
|
||||
"mean": [0.485, 0.456, 0.406],
|
||||
"std": [0.229, 0.224, 0.225],
|
||||
},
|
||||
"observation.state": {
|
||||
"min": [0.0, 0.0],
|
||||
"max": [1.0, 1.0],
|
||||
},
|
||||
"action": {
|
||||
"min": [0.0, 0.0, 0.0],
|
||||
"max": [1.0, 1.0, 1.0],
|
||||
},
|
||||
}
|
||||
assert config.dataset_stats == expected_dataset_stats
|
||||
|
||||
# Critic network configuration
|
||||
assert config.critic_network_kwargs.hidden_dims == [256, 256]
|
||||
assert config.critic_network_kwargs.activate_final is True
|
||||
assert config.critic_network_kwargs.final_activation is None
|
||||
|
||||
# Actor network configuration
|
||||
assert config.actor_network_kwargs.hidden_dims == [256, 256]
|
||||
assert config.actor_network_kwargs.activate_final is True
|
||||
|
||||
# Policy configuration
|
||||
assert config.policy_kwargs.use_tanh_squash is True
|
||||
assert config.policy_kwargs.std_min == 1e-5
|
||||
assert config.policy_kwargs.std_max == 10.0
|
||||
assert config.policy_kwargs.init_final == 0.05
|
||||
|
||||
# Discrete critic network configuration
|
||||
assert config.discrete_critic_network_kwargs.hidden_dims == [256, 256]
|
||||
assert config.discrete_critic_network_kwargs.activate_final is True
|
||||
assert config.discrete_critic_network_kwargs.final_activation is None
|
||||
|
||||
# Actor learner configuration
|
||||
assert config.actor_learner_config.learner_host == "127.0.0.1"
|
||||
assert config.actor_learner_config.learner_port == 50051
|
||||
assert config.actor_learner_config.policy_parameters_push_frequency == 4
|
||||
|
||||
# Concurrency configuration
|
||||
assert config.concurrency.actor == "threads"
|
||||
assert config.concurrency.learner == "threads"
|
||||
|
||||
assert isinstance(config.actor_network_kwargs, ActorNetworkConfig)
|
||||
assert isinstance(config.critic_network_kwargs, CriticNetworkConfig)
|
||||
assert isinstance(config.policy_kwargs, PolicyConfig)
|
||||
assert isinstance(config.actor_learner_config, ActorLearnerConfig)
|
||||
assert isinstance(config.concurrency, ConcurrencyConfig)
|
||||
|
||||
|
||||
def test_critic_network_kwargs():
|
||||
config = CriticNetworkConfig()
|
||||
assert config.hidden_dims == [256, 256]
|
||||
assert config.activate_final is True
|
||||
assert config.final_activation is None
|
||||
|
||||
|
||||
def test_actor_network_kwargs():
|
||||
config = ActorNetworkConfig()
|
||||
assert config.hidden_dims == [256, 256]
|
||||
assert config.activate_final is True
|
||||
|
||||
|
||||
def test_policy_kwargs():
|
||||
config = PolicyConfig()
|
||||
assert config.use_tanh_squash is True
|
||||
assert config.std_min == 1e-5
|
||||
assert config.std_max == 10.0
|
||||
assert config.init_final == 0.05
|
||||
|
||||
|
||||
def test_actor_learner_config():
|
||||
config = ActorLearnerConfig()
|
||||
assert config.learner_host == "127.0.0.1"
|
||||
assert config.learner_port == 50051
|
||||
assert config.policy_parameters_push_frequency == 4
|
||||
|
||||
|
||||
def test_concurrency_config():
|
||||
config = ConcurrencyConfig()
|
||||
assert config.actor == "threads"
|
||||
assert config.learner == "threads"
|
||||
|
||||
|
||||
def test_sac_config_custom_initialization():
|
||||
config = SACConfig(
|
||||
device="cpu",
|
||||
discount=0.95,
|
||||
temperature_init=0.5,
|
||||
num_critics=3,
|
||||
)
|
||||
|
||||
assert config.device == "cpu"
|
||||
assert config.discount == 0.95
|
||||
assert config.temperature_init == 0.5
|
||||
assert config.num_critics == 3
|
||||
|
||||
|
||||
def test_validate_features():
|
||||
config = SACConfig(
|
||||
input_features={"observation.state": PolicyFeature(type=FeatureType.STATE, shape=(10,))},
|
||||
output_features={"action": PolicyFeature(type=FeatureType.ACTION, shape=(3,))},
|
||||
)
|
||||
config.validate_features()
|
||||
|
||||
|
||||
def test_validate_features_missing_observation():
|
||||
config = SACConfig(
|
||||
input_features={"wrong_key": PolicyFeature(type=FeatureType.STATE, shape=(10,))},
|
||||
output_features={"action": PolicyFeature(type=FeatureType.ACTION, shape=(3,))},
|
||||
)
|
||||
with pytest.raises(
|
||||
ValueError, match="You must provide either 'observation.state' or an image observation"
|
||||
):
|
||||
config.validate_features()
|
||||
|
||||
|
||||
def test_validate_features_missing_action():
|
||||
config = SACConfig(
|
||||
input_features={"observation.state": PolicyFeature(type=FeatureType.STATE, shape=(10,))},
|
||||
output_features={"wrong_key": PolicyFeature(type=FeatureType.ACTION, shape=(3,))},
|
||||
)
|
||||
with pytest.raises(ValueError, match="You must provide 'action' in the output features"):
|
||||
config.validate_features()
|
||||
@@ -0,0 +1,541 @@
|
||||
# !/usr/bin/env python
|
||||
|
||||
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# 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 math
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch import Tensor, nn
|
||||
|
||||
from lerobot.common.policies.sac.configuration_sac import SACConfig
|
||||
from lerobot.common.policies.sac.modeling_sac import MLP, SACPolicy
|
||||
from lerobot.common.utils.random_utils import seeded_context, set_seed
|
||||
from lerobot.configs.types import FeatureType, PolicyFeature
|
||||
|
||||
try:
|
||||
import transformers # noqa: F401
|
||||
|
||||
TRANSFORMERS_AVAILABLE = True
|
||||
except ImportError:
|
||||
TRANSFORMERS_AVAILABLE = False
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def set_random_seed():
|
||||
seed = 42
|
||||
set_seed(seed)
|
||||
|
||||
|
||||
def test_mlp_with_default_args():
|
||||
mlp = MLP(input_dim=10, hidden_dims=[256, 256])
|
||||
|
||||
x = torch.randn(10)
|
||||
y = mlp(x)
|
||||
assert y.shape == (256,)
|
||||
|
||||
|
||||
def test_mlp_with_batch_dim():
|
||||
mlp = MLP(input_dim=10, hidden_dims=[256, 256])
|
||||
x = torch.randn(2, 10)
|
||||
y = mlp(x)
|
||||
assert y.shape == (2, 256)
|
||||
|
||||
|
||||
def test_forward_with_empty_hidden_dims():
|
||||
mlp = MLP(input_dim=10, hidden_dims=[])
|
||||
x = torch.randn(1, 10)
|
||||
assert mlp(x).shape == (1, 10)
|
||||
|
||||
|
||||
def test_mlp_with_dropout():
|
||||
mlp = MLP(input_dim=10, hidden_dims=[256, 256, 11], dropout_rate=0.1)
|
||||
x = torch.randn(1, 10)
|
||||
y = mlp(x)
|
||||
assert y.shape == (1, 11)
|
||||
|
||||
drop_out_layers_count = sum(isinstance(layer, nn.Dropout) for layer in mlp.net)
|
||||
assert drop_out_layers_count == 2
|
||||
|
||||
|
||||
def test_mlp_with_custom_final_activation():
|
||||
mlp = MLP(input_dim=10, hidden_dims=[256, 256], final_activation=torch.nn.Tanh())
|
||||
x = torch.randn(1, 10)
|
||||
y = mlp(x)
|
||||
assert y.shape == (1, 256)
|
||||
assert (y >= -1).all() and (y <= 1).all()
|
||||
|
||||
|
||||
def test_sac_policy_with_default_args():
|
||||
with pytest.raises(ValueError, match="should be an instance of class `PreTrainedConfig`"):
|
||||
SACPolicy()
|
||||
|
||||
|
||||
def create_dummy_state(batch_size: int, state_dim: int = 10) -> Tensor:
|
||||
return {
|
||||
"observation.state": torch.randn(batch_size, state_dim),
|
||||
}
|
||||
|
||||
|
||||
def create_dummy_with_visual_input(batch_size: int, state_dim: int = 10) -> Tensor:
|
||||
return {
|
||||
"observation.image": torch.randn(batch_size, 3, 84, 84),
|
||||
"observation.state": torch.randn(batch_size, state_dim),
|
||||
}
|
||||
|
||||
|
||||
def create_dummy_action(batch_size: int, action_dim: int = 10) -> Tensor:
|
||||
return torch.randn(batch_size, action_dim)
|
||||
|
||||
|
||||
def create_default_train_batch(
|
||||
batch_size: int = 8, state_dim: int = 10, action_dim: int = 10
|
||||
) -> dict[str, Tensor]:
|
||||
return {
|
||||
"action": create_dummy_action(batch_size, action_dim),
|
||||
"reward": torch.randn(batch_size),
|
||||
"state": create_dummy_state(batch_size, state_dim),
|
||||
"next_state": create_dummy_state(batch_size, state_dim),
|
||||
"done": torch.randn(batch_size),
|
||||
}
|
||||
|
||||
|
||||
def create_train_batch_with_visual_input(
|
||||
batch_size: int = 8, state_dim: int = 10, action_dim: int = 10
|
||||
) -> dict[str, Tensor]:
|
||||
return {
|
||||
"action": create_dummy_action(batch_size, action_dim),
|
||||
"reward": torch.randn(batch_size),
|
||||
"state": create_dummy_with_visual_input(batch_size, state_dim),
|
||||
"next_state": create_dummy_with_visual_input(batch_size, state_dim),
|
||||
"done": torch.randn(batch_size),
|
||||
}
|
||||
|
||||
|
||||
def create_observation_batch(batch_size: int = 8, state_dim: int = 10) -> dict[str, Tensor]:
|
||||
return {
|
||||
"observation.state": torch.randn(batch_size, state_dim),
|
||||
}
|
||||
|
||||
|
||||
def create_observation_batch_with_visual_input(batch_size: int = 8, state_dim: int = 10) -> dict[str, Tensor]:
|
||||
return {
|
||||
"observation.state": torch.randn(batch_size, state_dim),
|
||||
"observation.image": torch.randn(batch_size, 3, 84, 84),
|
||||
}
|
||||
|
||||
|
||||
def make_optimizers(policy: SACPolicy, has_discrete_action: bool = False) -> dict[str, torch.optim.Optimizer]:
|
||||
"""Create optimizers for the SAC policy."""
|
||||
optimizer_actor = torch.optim.Adam(
|
||||
# Handle the case of shared encoder where the encoder weights are not optimized with the actor gradient
|
||||
params=[
|
||||
p
|
||||
for n, p in policy.actor.named_parameters()
|
||||
if not policy.config.shared_encoder or not n.startswith("encoder")
|
||||
],
|
||||
lr=policy.config.actor_lr,
|
||||
)
|
||||
optimizer_critic = torch.optim.Adam(
|
||||
params=policy.critic_ensemble.parameters(),
|
||||
lr=policy.config.critic_lr,
|
||||
)
|
||||
optimizer_temperature = torch.optim.Adam(
|
||||
params=[policy.log_alpha],
|
||||
lr=policy.config.critic_lr,
|
||||
)
|
||||
|
||||
optimizers = {
|
||||
"actor": optimizer_actor,
|
||||
"critic": optimizer_critic,
|
||||
"temperature": optimizer_temperature,
|
||||
}
|
||||
|
||||
if has_discrete_action:
|
||||
optimizers["discrete_critic"] = torch.optim.Adam(
|
||||
params=policy.discrete_critic.parameters(),
|
||||
lr=policy.config.critic_lr,
|
||||
)
|
||||
|
||||
return optimizers
|
||||
|
||||
|
||||
def create_default_config(
|
||||
state_dim: int, continuous_action_dim: int, has_discrete_action: bool = False
|
||||
) -> SACConfig:
|
||||
action_dim = continuous_action_dim
|
||||
if has_discrete_action:
|
||||
action_dim += 1
|
||||
|
||||
config = SACConfig(
|
||||
input_features={"observation.state": PolicyFeature(type=FeatureType.STATE, shape=(state_dim,))},
|
||||
output_features={"action": PolicyFeature(type=FeatureType.ACTION, shape=(continuous_action_dim,))},
|
||||
dataset_stats={
|
||||
"observation.state": {
|
||||
"min": [0.0] * state_dim,
|
||||
"max": [1.0] * state_dim,
|
||||
},
|
||||
"action": {
|
||||
"min": [0.0] * continuous_action_dim,
|
||||
"max": [1.0] * continuous_action_dim,
|
||||
},
|
||||
},
|
||||
)
|
||||
config.validate_features()
|
||||
return config
|
||||
|
||||
|
||||
def create_config_with_visual_input(
|
||||
state_dim: int, continuous_action_dim: int, has_discrete_action: bool = False
|
||||
) -> SACConfig:
|
||||
config = create_default_config(
|
||||
state_dim=state_dim,
|
||||
continuous_action_dim=continuous_action_dim,
|
||||
has_discrete_action=has_discrete_action,
|
||||
)
|
||||
config.input_features["observation.image"] = PolicyFeature(type=FeatureType.VISUAL, shape=(3, 84, 84))
|
||||
config.dataset_stats["observation.image"] = {
|
||||
"mean": torch.randn(3, 1, 1),
|
||||
"std": torch.randn(3, 1, 1),
|
||||
}
|
||||
|
||||
# Let make tests a little bit faster
|
||||
config.state_encoder_hidden_dim = 32
|
||||
config.latent_dim = 32
|
||||
|
||||
config.validate_features()
|
||||
return config
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch_size,state_dim,action_dim", [(2, 6, 6), (1, 10, 10)])
|
||||
def test_sac_policy_with_default_config(batch_size: int, state_dim: int, action_dim: int):
|
||||
batch = create_default_train_batch(batch_size=batch_size, action_dim=action_dim, state_dim=state_dim)
|
||||
config = create_default_config(state_dim=state_dim, continuous_action_dim=action_dim)
|
||||
|
||||
policy = SACPolicy(config=config)
|
||||
policy.train()
|
||||
|
||||
optimizers = make_optimizers(policy)
|
||||
|
||||
cirtic_loss = policy.forward(batch, model="critic")["loss_critic"]
|
||||
assert cirtic_loss.item() is not None
|
||||
assert cirtic_loss.shape == ()
|
||||
cirtic_loss.backward()
|
||||
optimizers["critic"].step()
|
||||
|
||||
actor_loss = policy.forward(batch, model="actor")["loss_actor"]
|
||||
assert actor_loss.item() is not None
|
||||
assert actor_loss.shape == ()
|
||||
|
||||
actor_loss.backward()
|
||||
optimizers["actor"].step()
|
||||
|
||||
temperature_loss = policy.forward(batch, model="temperature")["loss_temperature"]
|
||||
assert temperature_loss.item() is not None
|
||||
assert temperature_loss.shape == ()
|
||||
|
||||
temperature_loss.backward()
|
||||
optimizers["temperature"].step()
|
||||
|
||||
policy.eval()
|
||||
with torch.no_grad():
|
||||
observation_batch = create_observation_batch(batch_size=batch_size, state_dim=state_dim)
|
||||
selected_action = policy.select_action(observation_batch)
|
||||
assert selected_action.shape == (batch_size, action_dim)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch_size,state_dim,action_dim", [(2, 6, 6), (1, 10, 10)])
|
||||
def test_sac_policy_with_visual_input(batch_size: int, state_dim: int, action_dim: int):
|
||||
config = create_config_with_visual_input(state_dim=state_dim, continuous_action_dim=action_dim)
|
||||
policy = SACPolicy(config=config)
|
||||
|
||||
batch = create_train_batch_with_visual_input(
|
||||
batch_size=batch_size, state_dim=state_dim, action_dim=action_dim
|
||||
)
|
||||
|
||||
policy.train()
|
||||
|
||||
optimizers = make_optimizers(policy)
|
||||
|
||||
cirtic_loss = policy.forward(batch, model="critic")["loss_critic"]
|
||||
assert cirtic_loss.item() is not None
|
||||
assert cirtic_loss.shape == ()
|
||||
cirtic_loss.backward()
|
||||
optimizers["critic"].step()
|
||||
|
||||
actor_loss = policy.forward(batch, model="actor")["loss_actor"]
|
||||
assert actor_loss.item() is not None
|
||||
assert actor_loss.shape == ()
|
||||
|
||||
actor_loss.backward()
|
||||
optimizers["actor"].step()
|
||||
|
||||
temperature_loss = policy.forward(batch, model="temperature")["loss_temperature"]
|
||||
assert temperature_loss.item() is not None
|
||||
assert temperature_loss.shape == ()
|
||||
|
||||
temperature_loss.backward()
|
||||
optimizers["temperature"].step()
|
||||
|
||||
policy.eval()
|
||||
with torch.no_grad():
|
||||
observation_batch = create_observation_batch_with_visual_input(
|
||||
batch_size=batch_size, state_dim=state_dim
|
||||
)
|
||||
selected_action = policy.select_action(observation_batch)
|
||||
assert selected_action.shape == (batch_size, action_dim)
|
||||
|
||||
|
||||
# Let's check best candidates for pretrained encoders
|
||||
@pytest.mark.parametrize(
|
||||
"batch_size,state_dim,action_dim,vision_encoder_name",
|
||||
[(1, 6, 6, "helper2424/resnet10"), (1, 6, 6, "facebook/convnext-base-224")],
|
||||
)
|
||||
@pytest.mark.skipif(not TRANSFORMERS_AVAILABLE, reason="Transformers are not installed")
|
||||
def test_sac_policy_with_pretrained_encoder(
|
||||
batch_size: int, state_dim: int, action_dim: int, vision_encoder_name: str
|
||||
):
|
||||
config = create_config_with_visual_input(state_dim=state_dim, continuous_action_dim=action_dim)
|
||||
config.vision_encoder_name = vision_encoder_name
|
||||
policy = SACPolicy(config=config)
|
||||
policy.train()
|
||||
|
||||
batch = create_train_batch_with_visual_input(
|
||||
batch_size=batch_size, state_dim=state_dim, action_dim=action_dim
|
||||
)
|
||||
|
||||
optimizers = make_optimizers(policy)
|
||||
|
||||
cirtic_loss = policy.forward(batch, model="critic")["loss_critic"]
|
||||
assert cirtic_loss.item() is not None
|
||||
assert cirtic_loss.shape == ()
|
||||
cirtic_loss.backward()
|
||||
optimizers["critic"].step()
|
||||
|
||||
actor_loss = policy.forward(batch, model="actor")["loss_actor"]
|
||||
assert actor_loss.item() is not None
|
||||
assert actor_loss.shape == ()
|
||||
|
||||
|
||||
def test_sac_policy_with_shared_encoder():
|
||||
batch_size = 2
|
||||
action_dim = 10
|
||||
state_dim = 10
|
||||
config = create_config_with_visual_input(state_dim=state_dim, continuous_action_dim=action_dim)
|
||||
config.shared_encoder = True
|
||||
|
||||
policy = SACPolicy(config=config)
|
||||
policy.train()
|
||||
|
||||
batch = create_train_batch_with_visual_input(
|
||||
batch_size=batch_size, state_dim=state_dim, action_dim=action_dim
|
||||
)
|
||||
|
||||
policy.train()
|
||||
|
||||
optimizers = make_optimizers(policy)
|
||||
|
||||
cirtic_loss = policy.forward(batch, model="critic")["loss_critic"]
|
||||
assert cirtic_loss.item() is not None
|
||||
assert cirtic_loss.shape == ()
|
||||
cirtic_loss.backward()
|
||||
optimizers["critic"].step()
|
||||
|
||||
actor_loss = policy.forward(batch, model="actor")["loss_actor"]
|
||||
assert actor_loss.item() is not None
|
||||
assert actor_loss.shape == ()
|
||||
|
||||
actor_loss.backward()
|
||||
optimizers["actor"].step()
|
||||
|
||||
|
||||
def test_sac_policy_with_discrete_critic():
|
||||
batch_size = 2
|
||||
continuous_action_dim = 9
|
||||
full_action_dim = continuous_action_dim + 1 # the last action is discrete
|
||||
state_dim = 10
|
||||
config = create_config_with_visual_input(
|
||||
state_dim=state_dim, continuous_action_dim=continuous_action_dim, has_discrete_action=True
|
||||
)
|
||||
|
||||
num_discrete_actions = 5
|
||||
config.num_discrete_actions = num_discrete_actions
|
||||
|
||||
policy = SACPolicy(config=config)
|
||||
policy.train()
|
||||
|
||||
batch = create_train_batch_with_visual_input(
|
||||
batch_size=batch_size, state_dim=state_dim, action_dim=full_action_dim
|
||||
)
|
||||
|
||||
policy.train()
|
||||
|
||||
optimizers = make_optimizers(policy, has_discrete_action=True)
|
||||
|
||||
cirtic_loss = policy.forward(batch, model="critic")["loss_critic"]
|
||||
assert cirtic_loss.item() is not None
|
||||
assert cirtic_loss.shape == ()
|
||||
cirtic_loss.backward()
|
||||
optimizers["critic"].step()
|
||||
|
||||
discrete_critic_loss = policy.forward(batch, model="discrete_critic")["loss_discrete_critic"]
|
||||
assert discrete_critic_loss.item() is not None
|
||||
assert discrete_critic_loss.shape == ()
|
||||
discrete_critic_loss.backward()
|
||||
optimizers["discrete_critic"].step()
|
||||
|
||||
actor_loss = policy.forward(batch, model="actor")["loss_actor"]
|
||||
assert actor_loss.item() is not None
|
||||
assert actor_loss.shape == ()
|
||||
|
||||
actor_loss.backward()
|
||||
optimizers["actor"].step()
|
||||
|
||||
policy.eval()
|
||||
with torch.no_grad():
|
||||
observation_batch = create_observation_batch_with_visual_input(
|
||||
batch_size=batch_size, state_dim=state_dim
|
||||
)
|
||||
selected_action = policy.select_action(observation_batch)
|
||||
assert selected_action.shape == (batch_size, full_action_dim)
|
||||
|
||||
discrete_actions = selected_action[:, -1].long()
|
||||
discrete_action_values = set(discrete_actions.tolist())
|
||||
|
||||
assert all(action in range(num_discrete_actions) for action in discrete_action_values), (
|
||||
f"Discrete action {discrete_action_values} is not in range({num_discrete_actions})"
|
||||
)
|
||||
|
||||
|
||||
def test_sac_policy_with_default_entropy():
|
||||
config = create_default_config(continuous_action_dim=10, state_dim=10)
|
||||
policy = SACPolicy(config=config)
|
||||
assert policy.target_entropy == -5.0
|
||||
|
||||
|
||||
def test_sac_policy_default_target_entropy_with_discrete_action():
|
||||
config = create_config_with_visual_input(state_dim=10, continuous_action_dim=6, has_discrete_action=True)
|
||||
policy = SACPolicy(config=config)
|
||||
assert policy.target_entropy == -3.0
|
||||
|
||||
|
||||
def test_sac_policy_with_predefined_entropy():
|
||||
config = create_default_config(state_dim=10, continuous_action_dim=6)
|
||||
config.target_entropy = -3.5
|
||||
|
||||
policy = SACPolicy(config=config)
|
||||
assert policy.target_entropy == pytest.approx(-3.5)
|
||||
|
||||
|
||||
def test_sac_policy_update_temperature():
|
||||
config = create_default_config(continuous_action_dim=10, state_dim=10)
|
||||
policy = SACPolicy(config=config)
|
||||
|
||||
assert policy.temperature == pytest.approx(1.0)
|
||||
policy.log_alpha.data = torch.tensor([math.log(0.1)])
|
||||
policy.update_temperature()
|
||||
assert policy.temperature == pytest.approx(0.1)
|
||||
|
||||
|
||||
def test_sac_policy_update_target_network():
|
||||
config = create_default_config(state_dim=10, continuous_action_dim=6)
|
||||
config.critic_target_update_weight = 1.0
|
||||
|
||||
policy = SACPolicy(config=config)
|
||||
policy.train()
|
||||
|
||||
for p in policy.critic_ensemble.parameters():
|
||||
p.data = torch.ones_like(p.data)
|
||||
|
||||
policy.update_target_networks()
|
||||
for p in policy.critic_target.parameters():
|
||||
assert torch.allclose(p.data, torch.ones_like(p.data)), (
|
||||
f"Target network {p.data} is not equal to {torch.ones_like(p.data)}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_critics", [1, 3])
|
||||
def test_sac_policy_with_critics_number_of_heads(num_critics: int):
|
||||
batch_size = 2
|
||||
action_dim = 10
|
||||
state_dim = 10
|
||||
config = create_config_with_visual_input(state_dim=state_dim, continuous_action_dim=action_dim)
|
||||
config.num_critics = num_critics
|
||||
|
||||
policy = SACPolicy(config=config)
|
||||
policy.train()
|
||||
|
||||
assert len(policy.critic_ensemble.critics) == num_critics
|
||||
|
||||
batch = create_train_batch_with_visual_input(
|
||||
batch_size=batch_size, state_dim=state_dim, action_dim=action_dim
|
||||
)
|
||||
|
||||
policy.train()
|
||||
|
||||
optimizers = make_optimizers(policy)
|
||||
|
||||
cirtic_loss = policy.forward(batch, model="critic")["loss_critic"]
|
||||
assert cirtic_loss.item() is not None
|
||||
assert cirtic_loss.shape == ()
|
||||
cirtic_loss.backward()
|
||||
optimizers["critic"].step()
|
||||
|
||||
|
||||
def test_sac_policy_save_and_load(tmp_path):
|
||||
root = tmp_path / "test_sac_save_and_load"
|
||||
|
||||
state_dim = 10
|
||||
action_dim = 10
|
||||
batch_size = 2
|
||||
|
||||
config = create_default_config(state_dim=state_dim, continuous_action_dim=action_dim)
|
||||
policy = SACPolicy(config=config)
|
||||
policy.eval()
|
||||
policy.save_pretrained(root)
|
||||
loaded_policy = SACPolicy.from_pretrained(root, config=config)
|
||||
loaded_policy.eval()
|
||||
|
||||
batch = create_default_train_batch(batch_size=1, state_dim=10, action_dim=10)
|
||||
|
||||
with torch.no_grad():
|
||||
with seeded_context(12):
|
||||
# Collect policy values before saving
|
||||
cirtic_loss = policy.forward(batch, model="critic")["loss_critic"]
|
||||
actor_loss = policy.forward(batch, model="actor")["loss_actor"]
|
||||
temperature_loss = policy.forward(batch, model="temperature")["loss_temperature"]
|
||||
|
||||
observation_batch = create_observation_batch(batch_size=batch_size, state_dim=state_dim)
|
||||
actions = policy.select_action(observation_batch)
|
||||
|
||||
with seeded_context(12):
|
||||
# Collect policy values after loading
|
||||
loaded_cirtic_loss = loaded_policy.forward(batch, model="critic")["loss_critic"]
|
||||
loaded_actor_loss = loaded_policy.forward(batch, model="actor")["loss_actor"]
|
||||
loaded_temperature_loss = loaded_policy.forward(batch, model="temperature")["loss_temperature"]
|
||||
|
||||
loaded_observation_batch = create_observation_batch(batch_size=batch_size, state_dim=state_dim)
|
||||
loaded_actions = loaded_policy.select_action(loaded_observation_batch)
|
||||
|
||||
assert policy.state_dict().keys() == loaded_policy.state_dict().keys()
|
||||
for k in policy.state_dict():
|
||||
assert torch.allclose(policy.state_dict()[k], loaded_policy.state_dict()[k], atol=1e-6)
|
||||
|
||||
# Compare values before and after saving and loading
|
||||
# They should be the same
|
||||
assert torch.allclose(cirtic_loss, loaded_cirtic_loss)
|
||||
assert torch.allclose(actor_loss, loaded_actor_loss)
|
||||
assert torch.allclose(temperature_loss, loaded_temperature_loss)
|
||||
assert torch.allclose(actions, loaded_actions)
|
||||
@@ -0,0 +1,208 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# 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.
|
||||
|
||||
from concurrent import futures
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch.multiprocessing import Event, Queue
|
||||
|
||||
from lerobot.common.utils.transition import Transition
|
||||
from tests.utils import require_package
|
||||
|
||||
|
||||
def create_learner_service_stub():
|
||||
import grpc
|
||||
|
||||
from lerobot.common.transport import services_pb2, services_pb2_grpc
|
||||
|
||||
class MockLearnerService(services_pb2_grpc.LearnerServiceServicer):
|
||||
def __init__(self):
|
||||
self.ready_call_count = 0
|
||||
self.should_fail = False
|
||||
|
||||
def Ready(self, request, context): # noqa: N802
|
||||
self.ready_call_count += 1
|
||||
if self.should_fail:
|
||||
context.set_code(grpc.StatusCode.UNAVAILABLE)
|
||||
context.set_details("Service unavailable")
|
||||
raise grpc.RpcError("Service unavailable")
|
||||
return services_pb2.Empty()
|
||||
|
||||
"""Fixture to start a LearnerService gRPC server and provide a connected stub."""
|
||||
|
||||
servicer = MockLearnerService()
|
||||
|
||||
# Create a gRPC server and add our servicer to it.
|
||||
server = grpc.server(futures.ThreadPoolExecutor(max_workers=4))
|
||||
services_pb2_grpc.add_LearnerServiceServicer_to_server(servicer, server)
|
||||
port = server.add_insecure_port("[::]:0") # bind to a free port chosen by OS
|
||||
server.start() # start the server (non-blocking call):contentReference[oaicite:1]{index=1}
|
||||
|
||||
# Create a client channel and stub connected to the server's port.
|
||||
channel = grpc.insecure_channel(f"localhost:{port}")
|
||||
return services_pb2_grpc.LearnerServiceStub(channel), servicer, channel, server
|
||||
|
||||
|
||||
def close_service_stub(channel, server):
|
||||
channel.close()
|
||||
server.stop(None)
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
def test_establish_learner_connection_success():
|
||||
from lerobot.scripts.rl.actor import establish_learner_connection
|
||||
|
||||
"""Test successful connection establishment."""
|
||||
stub, _servicer, channel, server = create_learner_service_stub()
|
||||
|
||||
shutdown_event = Event()
|
||||
|
||||
# Test successful connection
|
||||
result = establish_learner_connection(stub, shutdown_event, attempts=5)
|
||||
|
||||
assert result is True
|
||||
|
||||
close_service_stub(channel, server)
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
def test_establish_learner_connection_failure():
|
||||
from lerobot.scripts.rl.actor import establish_learner_connection
|
||||
|
||||
"""Test connection failure."""
|
||||
stub, servicer, channel, server = create_learner_service_stub()
|
||||
servicer.should_fail = True
|
||||
|
||||
shutdown_event = Event()
|
||||
|
||||
# Test failed connection
|
||||
with patch("time.sleep"): # Speed up the test
|
||||
result = establish_learner_connection(stub, shutdown_event, attempts=2)
|
||||
|
||||
assert result is False
|
||||
|
||||
close_service_stub(channel, server)
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
def test_push_transitions_to_transport_queue():
|
||||
from lerobot.common.transport.utils import bytes_to_transitions
|
||||
from lerobot.scripts.rl.actor import push_transitions_to_transport_queue
|
||||
from tests.transport.test_transport_utils import assert_transitions_equal
|
||||
|
||||
"""Test pushing transitions to transport queue."""
|
||||
# Create mock transitions
|
||||
transitions = []
|
||||
for i in range(3):
|
||||
transition = Transition(
|
||||
state={"observation": torch.randn(3, 64, 64), "state": torch.randn(10)},
|
||||
action=torch.randn(5),
|
||||
reward=torch.tensor(1.0 + i),
|
||||
done=torch.tensor(False),
|
||||
truncated=torch.tensor(False),
|
||||
next_state={"observation": torch.randn(3, 64, 64), "state": torch.randn(10)},
|
||||
complementary_info={"step": torch.tensor(i)},
|
||||
)
|
||||
transitions.append(transition)
|
||||
|
||||
transitions_queue = Queue()
|
||||
|
||||
# Test pushing transitions
|
||||
push_transitions_to_transport_queue(transitions, transitions_queue)
|
||||
|
||||
# Verify the data can be retrieved
|
||||
serialized_data = transitions_queue.get()
|
||||
assert isinstance(serialized_data, bytes)
|
||||
deserialized_transitions = bytes_to_transitions(serialized_data)
|
||||
assert len(deserialized_transitions) == len(transitions)
|
||||
for i, deserialized_transition in enumerate(deserialized_transitions):
|
||||
assert_transitions_equal(deserialized_transition, transitions[i])
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@pytest.mark.timeout(3) # force cross-platform watchdog
|
||||
def test_transitions_stream():
|
||||
from lerobot.scripts.rl.actor import transitions_stream
|
||||
|
||||
"""Test transitions stream functionality."""
|
||||
shutdown_event = Event()
|
||||
transitions_queue = Queue()
|
||||
|
||||
# Add test data to queue
|
||||
test_data = [b"transition_data_1", b"transition_data_2", b"transition_data_3"]
|
||||
for data in test_data:
|
||||
transitions_queue.put(data)
|
||||
|
||||
# Collect streamed data
|
||||
streamed_data = []
|
||||
stream_generator = transitions_stream(shutdown_event, transitions_queue, 0.1)
|
||||
|
||||
# Process a few items
|
||||
for i, message in enumerate(stream_generator):
|
||||
streamed_data.append(message)
|
||||
if i >= len(test_data) - 1:
|
||||
shutdown_event.set()
|
||||
break
|
||||
|
||||
# Verify we got messages
|
||||
assert len(streamed_data) == len(test_data)
|
||||
assert streamed_data[0].data == b"transition_data_1"
|
||||
assert streamed_data[1].data == b"transition_data_2"
|
||||
assert streamed_data[2].data == b"transition_data_3"
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@pytest.mark.timeout(3) # force cross-platform watchdog
|
||||
def test_interactions_stream():
|
||||
from lerobot.common.transport.utils import bytes_to_python_object, python_object_to_bytes
|
||||
from lerobot.scripts.rl.actor import interactions_stream
|
||||
|
||||
"""Test interactions stream functionality."""
|
||||
shutdown_event = Event()
|
||||
interactions_queue = Queue()
|
||||
|
||||
# Create test interaction data (similar structure to what would be sent)
|
||||
test_interactions = [
|
||||
{"episode_reward": 10.5, "step": 1, "policy_fps": 30.2},
|
||||
{"episode_reward": 15.2, "step": 2, "policy_fps": 28.7},
|
||||
{"episode_reward": 8.7, "step": 3, "policy_fps": 29.1},
|
||||
]
|
||||
|
||||
# Serialize the interaction data as it would be in practice
|
||||
test_data = [
|
||||
interactions_queue.put(python_object_to_bytes(interaction)) for interaction in test_interactions
|
||||
]
|
||||
|
||||
# Collect streamed data
|
||||
streamed_data = []
|
||||
stream_generator = interactions_stream(shutdown_event, interactions_queue, 0.1)
|
||||
|
||||
# Process the items
|
||||
for i, message in enumerate(stream_generator):
|
||||
streamed_data.append(message)
|
||||
if i >= len(test_data) - 1:
|
||||
shutdown_event.set()
|
||||
break
|
||||
|
||||
# Verify we got messages
|
||||
assert len(streamed_data) == len(test_data)
|
||||
|
||||
# Verify the messages can be deserialized back to original data
|
||||
for i, message in enumerate(streamed_data):
|
||||
deserialized_interaction = bytes_to_python_object(message.data)
|
||||
assert deserialized_interaction == test_interactions[i]
|
||||
@@ -0,0 +1,297 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# 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 socket
|
||||
import threading
|
||||
import time
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch.multiprocessing import Event, Queue
|
||||
|
||||
from lerobot.common.policies.sac.configuration_sac import SACConfig
|
||||
from lerobot.common.utils.transition import Transition
|
||||
from lerobot.configs.train import TrainRLServerPipelineConfig
|
||||
from tests.utils import require_package
|
||||
|
||||
|
||||
def create_test_transitions(count: int = 3) -> list[Transition]:
|
||||
"""Create test transitions for integration testing."""
|
||||
transitions = []
|
||||
for i in range(count):
|
||||
transition = Transition(
|
||||
state={"observation": torch.randn(3, 64, 64), "state": torch.randn(10)},
|
||||
action=torch.randn(5),
|
||||
reward=torch.tensor(1.0 + i),
|
||||
done=torch.tensor(i == count - 1), # Last transition is done
|
||||
truncated=torch.tensor(False),
|
||||
next_state={"observation": torch.randn(3, 64, 64), "state": torch.randn(10)},
|
||||
complementary_info={"step": torch.tensor(i), "episode_id": i // 2},
|
||||
)
|
||||
transitions.append(transition)
|
||||
return transitions
|
||||
|
||||
|
||||
def create_test_interactions(count: int = 3) -> list[dict]:
|
||||
"""Create test interactions for integration testing."""
|
||||
interactions = []
|
||||
for i in range(count):
|
||||
interaction = {
|
||||
"episode_reward": 10.0 + i * 5,
|
||||
"step": i * 100,
|
||||
"policy_fps": 30.0 + i,
|
||||
"intervention_rate": 0.1 * i,
|
||||
"episode_length": 200 + i * 50,
|
||||
}
|
||||
interactions.append(interaction)
|
||||
return interactions
|
||||
|
||||
|
||||
def find_free_port():
|
||||
"""Finds a free port on the local machine."""
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind(("", 0)) # Bind to port 0 to let the OS choose a free port
|
||||
s.listen(1)
|
||||
port = s.getsockname()[1]
|
||||
return port
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def cfg():
|
||||
cfg = TrainRLServerPipelineConfig()
|
||||
|
||||
port = find_free_port()
|
||||
|
||||
policy_cfg = SACConfig()
|
||||
policy_cfg.actor_learner_config.learner_host = "127.0.0.1"
|
||||
policy_cfg.actor_learner_config.learner_port = port
|
||||
policy_cfg.concurrency.actor = "threads"
|
||||
policy_cfg.concurrency.learner = "threads"
|
||||
policy_cfg.actor_learner_config.queue_get_timeout = 0.1
|
||||
|
||||
cfg.policy = policy_cfg
|
||||
|
||||
return cfg
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@pytest.mark.timeout(10) # force cross-platform watchdog
|
||||
def test_end_to_end_transitions_flow(cfg):
|
||||
from lerobot.common.transport.utils import bytes_to_transitions
|
||||
from lerobot.scripts.rl.actor import (
|
||||
establish_learner_connection,
|
||||
learner_service_client,
|
||||
push_transitions_to_transport_queue,
|
||||
send_transitions,
|
||||
)
|
||||
from lerobot.scripts.rl.learner import start_learner
|
||||
from tests.transport.test_transport_utils import assert_transitions_equal
|
||||
|
||||
"""Test complete transitions flow from actor to learner."""
|
||||
transitions_actor_queue = Queue()
|
||||
transitions_learner_queue = Queue()
|
||||
|
||||
interactions_queue = Queue()
|
||||
parameters_queue = Queue()
|
||||
shutdown_event = Event()
|
||||
|
||||
learner_thread = threading.Thread(
|
||||
target=start_learner,
|
||||
args=(parameters_queue, transitions_learner_queue, interactions_queue, shutdown_event, cfg),
|
||||
)
|
||||
learner_thread.start()
|
||||
|
||||
policy_cfg = cfg.policy
|
||||
learner_client, channel = learner_service_client(
|
||||
host=policy_cfg.actor_learner_config.learner_host, port=policy_cfg.actor_learner_config.learner_port
|
||||
)
|
||||
|
||||
assert establish_learner_connection(learner_client, shutdown_event, attempts=5)
|
||||
|
||||
send_transitions_thread = threading.Thread(
|
||||
target=send_transitions, args=(cfg, transitions_actor_queue, shutdown_event, learner_client, channel)
|
||||
)
|
||||
send_transitions_thread.start()
|
||||
|
||||
input_transitions = create_test_transitions(count=5)
|
||||
|
||||
push_transitions_to_transport_queue(input_transitions, transitions_actor_queue)
|
||||
|
||||
# Wait for learner to start
|
||||
time.sleep(0.1)
|
||||
|
||||
shutdown_event.set()
|
||||
|
||||
# Wait for learner to receive transitions
|
||||
learner_thread.join()
|
||||
send_transitions_thread.join()
|
||||
channel.close()
|
||||
|
||||
received_transitions = []
|
||||
while not transitions_learner_queue.empty():
|
||||
received_transitions.extend(bytes_to_transitions(transitions_learner_queue.get()))
|
||||
|
||||
assert len(received_transitions) == len(input_transitions)
|
||||
for i, transition in enumerate(received_transitions):
|
||||
assert_transitions_equal(transition, input_transitions[i])
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@pytest.mark.timeout(10)
|
||||
def test_end_to_end_interactions_flow(cfg):
|
||||
from lerobot.common.transport.utils import bytes_to_python_object, python_object_to_bytes
|
||||
from lerobot.scripts.rl.actor import (
|
||||
establish_learner_connection,
|
||||
learner_service_client,
|
||||
send_interactions,
|
||||
)
|
||||
from lerobot.scripts.rl.learner import start_learner
|
||||
|
||||
"""Test complete interactions flow from actor to learner."""
|
||||
# Queues for actor-learner communication
|
||||
interactions_actor_queue = Queue()
|
||||
interactions_learner_queue = Queue()
|
||||
|
||||
# Other queues required by the learner
|
||||
parameters_queue = Queue()
|
||||
transitions_learner_queue = Queue()
|
||||
|
||||
shutdown_event = Event()
|
||||
|
||||
# Start the learner in a separate thread
|
||||
learner_thread = threading.Thread(
|
||||
target=start_learner,
|
||||
args=(parameters_queue, transitions_learner_queue, interactions_learner_queue, shutdown_event, cfg),
|
||||
)
|
||||
learner_thread.start()
|
||||
|
||||
# Establish connection from actor to learner
|
||||
policy_cfg = cfg.policy
|
||||
learner_client, channel = learner_service_client(
|
||||
host=policy_cfg.actor_learner_config.learner_host, port=policy_cfg.actor_learner_config.learner_port
|
||||
)
|
||||
|
||||
assert establish_learner_connection(learner_client, shutdown_event, attempts=5)
|
||||
|
||||
# Start the actor's interaction sending process in a separate thread
|
||||
send_interactions_thread = threading.Thread(
|
||||
target=send_interactions,
|
||||
args=(cfg, interactions_actor_queue, shutdown_event, learner_client, channel),
|
||||
)
|
||||
send_interactions_thread.start()
|
||||
|
||||
# Create and push test interactions to the actor's queue
|
||||
input_interactions = create_test_interactions(count=5)
|
||||
for interaction in input_interactions:
|
||||
interactions_actor_queue.put(python_object_to_bytes(interaction))
|
||||
|
||||
# Wait for the communication to happen
|
||||
time.sleep(0.1)
|
||||
|
||||
# Signal shutdown and wait for threads to complete
|
||||
shutdown_event.set()
|
||||
learner_thread.join()
|
||||
send_interactions_thread.join()
|
||||
channel.close()
|
||||
|
||||
# Verify that the learner received the interactions
|
||||
received_interactions = []
|
||||
while not interactions_learner_queue.empty():
|
||||
received_interactions.append(bytes_to_python_object(interactions_learner_queue.get()))
|
||||
|
||||
assert len(received_interactions) == len(input_interactions)
|
||||
|
||||
# Sort by a unique key to handle potential reordering in queues
|
||||
received_interactions.sort(key=lambda x: x["step"])
|
||||
input_interactions.sort(key=lambda x: x["step"])
|
||||
|
||||
for received, expected in zip(received_interactions, input_interactions, strict=False):
|
||||
assert received == expected
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@pytest.mark.parametrize("data_size", ["small", "large"])
|
||||
@pytest.mark.timeout(10)
|
||||
def test_end_to_end_parameters_flow(cfg, data_size):
|
||||
from lerobot.common.transport.utils import bytes_to_state_dict, state_to_bytes
|
||||
from lerobot.scripts.rl.actor import establish_learner_connection, learner_service_client, receive_policy
|
||||
from lerobot.scripts.rl.learner import start_learner
|
||||
|
||||
"""Test complete parameter flow from learner to actor, with small and large data."""
|
||||
# Actor's local queue to receive params
|
||||
parameters_actor_queue = Queue()
|
||||
# Learner's queue to send params from
|
||||
parameters_learner_queue = Queue()
|
||||
|
||||
# Other queues required by the learner
|
||||
transitions_learner_queue = Queue()
|
||||
interactions_learner_queue = Queue()
|
||||
|
||||
shutdown_event = Event()
|
||||
|
||||
# Start the learner in a separate thread
|
||||
learner_thread = threading.Thread(
|
||||
target=start_learner,
|
||||
args=(
|
||||
parameters_learner_queue,
|
||||
transitions_learner_queue,
|
||||
interactions_learner_queue,
|
||||
shutdown_event,
|
||||
cfg,
|
||||
),
|
||||
)
|
||||
learner_thread.start()
|
||||
|
||||
# Establish connection from actor to learner
|
||||
policy_cfg = cfg.policy
|
||||
learner_client, channel = learner_service_client(
|
||||
host=policy_cfg.actor_learner_config.learner_host, port=policy_cfg.actor_learner_config.learner_port
|
||||
)
|
||||
|
||||
assert establish_learner_connection(learner_client, shutdown_event, attempts=5)
|
||||
|
||||
# Start the actor's parameter receiving process in a separate thread
|
||||
receive_params_thread = threading.Thread(
|
||||
target=receive_policy,
|
||||
args=(cfg, parameters_actor_queue, shutdown_event, learner_client, channel),
|
||||
)
|
||||
receive_params_thread.start()
|
||||
|
||||
# Create test parameters based on parametrization
|
||||
if data_size == "small":
|
||||
input_params = {"layer.weight": torch.randn(128, 64)}
|
||||
else: # "large"
|
||||
# CHUNK_SIZE is 2MB, so this tensor (4MB) will force chunking
|
||||
input_params = {"large_layer.weight": torch.randn(1024, 1024)}
|
||||
|
||||
# Simulate learner having new parameters to send
|
||||
parameters_learner_queue.put(state_to_bytes(input_params))
|
||||
|
||||
# Wait for the actor to receive the parameters
|
||||
time.sleep(0.1)
|
||||
|
||||
# Signal shutdown and wait for threads to complete
|
||||
shutdown_event.set()
|
||||
learner_thread.join()
|
||||
receive_params_thread.join()
|
||||
channel.close()
|
||||
|
||||
# Verify that the actor received the parameters correctly
|
||||
received_params = bytes_to_state_dict(parameters_actor_queue.get())
|
||||
|
||||
assert received_params.keys() == input_params.keys()
|
||||
for key in input_params:
|
||||
assert torch.allclose(received_params[key], input_params[key])
|
||||
@@ -0,0 +1,374 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# 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 threading
|
||||
import time
|
||||
from concurrent import futures
|
||||
from multiprocessing import Event, Queue
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.utils import require_package # our gRPC servicer class
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def learner_service_stub():
|
||||
shutdown_event = Event()
|
||||
parameters_queue = Queue()
|
||||
transitions_queue = Queue()
|
||||
interactions_queue = Queue()
|
||||
seconds_between_pushes = 1
|
||||
client, channel, server = create_learner_service_stub(
|
||||
shutdown_event, parameters_queue, transitions_queue, interactions_queue, seconds_between_pushes
|
||||
)
|
||||
|
||||
yield client # provide the stub to the test function
|
||||
|
||||
close_learner_service_stub(channel, server)
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
def create_learner_service_stub(
|
||||
shutdown_event: Event,
|
||||
parameters_queue: Queue,
|
||||
transitions_queue: Queue,
|
||||
interactions_queue: Queue,
|
||||
seconds_between_pushes: int,
|
||||
queue_get_timeout: float = 0.1,
|
||||
):
|
||||
import grpc
|
||||
|
||||
from lerobot.common.transport import services_pb2_grpc # generated from .proto
|
||||
from lerobot.scripts.rl.learner_service import LearnerService
|
||||
|
||||
"""Fixture to start a LearnerService gRPC server and provide a connected stub."""
|
||||
|
||||
servicer = LearnerService(
|
||||
shutdown_event=shutdown_event,
|
||||
parameters_queue=parameters_queue,
|
||||
seconds_between_pushes=seconds_between_pushes,
|
||||
transition_queue=transitions_queue,
|
||||
interaction_message_queue=interactions_queue,
|
||||
queue_get_timeout=queue_get_timeout,
|
||||
)
|
||||
|
||||
# Create a gRPC server and add our servicer to it.
|
||||
server = grpc.server(futures.ThreadPoolExecutor(max_workers=4))
|
||||
services_pb2_grpc.add_LearnerServiceServicer_to_server(servicer, server)
|
||||
port = server.add_insecure_port("[::]:0") # bind to a free port chosen by OS
|
||||
server.start() # start the server (non-blocking call):contentReference[oaicite:1]{index=1}
|
||||
|
||||
# Create a client channel and stub connected to the server's port.
|
||||
channel = grpc.insecure_channel(f"localhost:{port}")
|
||||
return services_pb2_grpc.LearnerServiceStub(channel), channel, server
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
def close_learner_service_stub(channel, server):
|
||||
channel.close()
|
||||
server.stop(None)
|
||||
|
||||
|
||||
@pytest.mark.timeout(3) # force cross-platform watchdog
|
||||
def test_ready_method(learner_service_stub):
|
||||
from lerobot.common.transport import services_pb2
|
||||
|
||||
"""Test the ready method of the UserService."""
|
||||
request = services_pb2.Empty()
|
||||
response = learner_service_stub.Ready(request)
|
||||
assert response == services_pb2.Empty()
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@pytest.mark.timeout(3) # force cross-platform watchdog
|
||||
def test_send_interactions():
|
||||
from lerobot.common.transport import services_pb2
|
||||
|
||||
shutdown_event = Event()
|
||||
|
||||
parameters_queue = Queue()
|
||||
transitions_queue = Queue()
|
||||
interactions_queue = Queue()
|
||||
seconds_between_pushes = 1
|
||||
client, channel, server = create_learner_service_stub(
|
||||
shutdown_event, parameters_queue, transitions_queue, interactions_queue, seconds_between_pushes
|
||||
)
|
||||
|
||||
list_of_interaction_messages = [
|
||||
services_pb2.InteractionMessage(transfer_state=services_pb2.TransferState.TRANSFER_BEGIN, data=b"1"),
|
||||
services_pb2.InteractionMessage(transfer_state=services_pb2.TransferState.TRANSFER_MIDDLE, data=b"2"),
|
||||
services_pb2.InteractionMessage(transfer_state=services_pb2.TransferState.TRANSFER_END, data=b"3"),
|
||||
services_pb2.InteractionMessage(transfer_state=services_pb2.TransferState.TRANSFER_END, data=b"4"),
|
||||
services_pb2.InteractionMessage(transfer_state=services_pb2.TransferState.TRANSFER_END, data=b"5"),
|
||||
services_pb2.InteractionMessage(transfer_state=services_pb2.TransferState.TRANSFER_BEGIN, data=b"6"),
|
||||
services_pb2.InteractionMessage(transfer_state=services_pb2.TransferState.TRANSFER_MIDDLE, data=b"7"),
|
||||
services_pb2.InteractionMessage(transfer_state=services_pb2.TransferState.TRANSFER_END, data=b"8"),
|
||||
]
|
||||
|
||||
def mock_intercations_stream():
|
||||
yield from list_of_interaction_messages
|
||||
|
||||
return services_pb2.Empty()
|
||||
|
||||
response = client.SendInteractions(mock_intercations_stream())
|
||||
assert response == services_pb2.Empty()
|
||||
|
||||
close_learner_service_stub(channel, server)
|
||||
|
||||
# Extract the data from the interactions queue
|
||||
interactions = []
|
||||
while not interactions_queue.empty():
|
||||
interactions.append(interactions_queue.get())
|
||||
|
||||
assert interactions == [b"123", b"4", b"5", b"678"]
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@pytest.mark.timeout(3) # force cross-platform watchdog
|
||||
def test_send_transitions():
|
||||
from lerobot.common.transport import services_pb2
|
||||
|
||||
"""Test the SendTransitions method with various transition data."""
|
||||
shutdown_event = Event()
|
||||
parameters_queue = Queue()
|
||||
transitions_queue = Queue()
|
||||
interactions_queue = Queue()
|
||||
seconds_between_pushes = 1
|
||||
|
||||
client, channel, server = create_learner_service_stub(
|
||||
shutdown_event, parameters_queue, transitions_queue, interactions_queue, seconds_between_pushes
|
||||
)
|
||||
|
||||
# Create test transition messages
|
||||
list_of_transition_messages = [
|
||||
services_pb2.Transition(
|
||||
transfer_state=services_pb2.TransferState.TRANSFER_BEGIN, data=b"transition_1"
|
||||
),
|
||||
services_pb2.Transition(
|
||||
transfer_state=services_pb2.TransferState.TRANSFER_MIDDLE, data=b"transition_2"
|
||||
),
|
||||
services_pb2.Transition(transfer_state=services_pb2.TransferState.TRANSFER_END, data=b"transition_3"),
|
||||
services_pb2.Transition(transfer_state=services_pb2.TransferState.TRANSFER_BEGIN, data=b"batch_1"),
|
||||
services_pb2.Transition(transfer_state=services_pb2.TransferState.TRANSFER_END, data=b"batch_2"),
|
||||
]
|
||||
|
||||
def mock_transitions_stream():
|
||||
yield from list_of_transition_messages
|
||||
|
||||
response = client.SendTransitions(mock_transitions_stream())
|
||||
assert response == services_pb2.Empty()
|
||||
|
||||
close_learner_service_stub(channel, server)
|
||||
|
||||
# Extract the data from the transitions queue
|
||||
transitions = []
|
||||
while not transitions_queue.empty():
|
||||
transitions.append(transitions_queue.get())
|
||||
|
||||
# Should have assembled the chunked data
|
||||
assert transitions == [b"transition_1transition_2transition_3", b"batch_1batch_2"]
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@pytest.mark.timeout(3) # force cross-platform watchdog
|
||||
def test_send_transitions_empty_stream():
|
||||
from lerobot.common.transport import services_pb2
|
||||
|
||||
"""Test SendTransitions with empty stream."""
|
||||
shutdown_event = Event()
|
||||
parameters_queue = Queue()
|
||||
transitions_queue = Queue()
|
||||
interactions_queue = Queue()
|
||||
seconds_between_pushes = 1
|
||||
|
||||
client, channel, server = create_learner_service_stub(
|
||||
shutdown_event, parameters_queue, transitions_queue, interactions_queue, seconds_between_pushes
|
||||
)
|
||||
|
||||
def empty_stream():
|
||||
return iter([])
|
||||
|
||||
response = client.SendTransitions(empty_stream())
|
||||
assert response == services_pb2.Empty()
|
||||
|
||||
close_learner_service_stub(channel, server)
|
||||
|
||||
# Queue should remain empty
|
||||
assert transitions_queue.empty()
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@pytest.mark.timeout(10) # force cross-platform watchdog
|
||||
def test_stream_parameters():
|
||||
import time
|
||||
|
||||
from lerobot.common.transport import services_pb2
|
||||
|
||||
"""Test the StreamParameters method."""
|
||||
shutdown_event = Event()
|
||||
parameters_queue = Queue()
|
||||
transitions_queue = Queue()
|
||||
interactions_queue = Queue()
|
||||
seconds_between_pushes = 0.2 # Short delay for testing
|
||||
|
||||
client, channel, server = create_learner_service_stub(
|
||||
shutdown_event, parameters_queue, transitions_queue, interactions_queue, seconds_between_pushes
|
||||
)
|
||||
|
||||
# Add test parameters to the queue
|
||||
test_params = [b"param_batch_1", b"param_batch_2"]
|
||||
for param in test_params:
|
||||
parameters_queue.put(param)
|
||||
|
||||
# Start streaming parameters
|
||||
request = services_pb2.Empty()
|
||||
stream = client.StreamParameters(request)
|
||||
|
||||
# Collect streamed parameters and timestamps
|
||||
received_params = []
|
||||
timestamps = []
|
||||
|
||||
for response in stream:
|
||||
received_params.append(response.data)
|
||||
timestamps.append(time.time())
|
||||
|
||||
# We should receive one last item
|
||||
break
|
||||
|
||||
parameters_queue.put(b"param_batch_3")
|
||||
|
||||
for response in stream:
|
||||
received_params.append(response.data)
|
||||
timestamps.append(time.time())
|
||||
|
||||
# We should receive only one item
|
||||
break
|
||||
|
||||
shutdown_event.set()
|
||||
close_learner_service_stub(channel, server)
|
||||
|
||||
assert received_params == [b"param_batch_2", b"param_batch_3"]
|
||||
|
||||
# Check the time difference between the two sends
|
||||
time_diff = timestamps[1] - timestamps[0]
|
||||
# Check if the time difference is close to the expected push frequency
|
||||
assert time_diff == pytest.approx(seconds_between_pushes, abs=0.1)
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@pytest.mark.timeout(3) # force cross-platform watchdog
|
||||
def test_stream_parameters_with_shutdown():
|
||||
from lerobot.common.transport import services_pb2
|
||||
|
||||
"""Test StreamParameters handles shutdown gracefully."""
|
||||
shutdown_event = Event()
|
||||
parameters_queue = Queue()
|
||||
transitions_queue = Queue()
|
||||
interactions_queue = Queue()
|
||||
seconds_between_pushes = 0.1
|
||||
queue_get_timeout = 0.001
|
||||
|
||||
client, channel, server = create_learner_service_stub(
|
||||
shutdown_event,
|
||||
parameters_queue,
|
||||
transitions_queue,
|
||||
interactions_queue,
|
||||
seconds_between_pushes,
|
||||
queue_get_timeout=queue_get_timeout,
|
||||
)
|
||||
|
||||
test_params = [b"param_batch_1", b"stop", b"param_batch_3", b"param_batch_4"]
|
||||
|
||||
# create a thread that will put the parameters in the queue
|
||||
def producer():
|
||||
for param in test_params:
|
||||
parameters_queue.put(param)
|
||||
time.sleep(0.1)
|
||||
|
||||
producer_thread = threading.Thread(target=producer)
|
||||
producer_thread.start()
|
||||
|
||||
# Start streaming
|
||||
request = services_pb2.Empty()
|
||||
stream = client.StreamParameters(request)
|
||||
|
||||
# Collect streamed parameters
|
||||
received_params = []
|
||||
|
||||
for response in stream:
|
||||
received_params.append(response.data)
|
||||
|
||||
if response.data == b"stop":
|
||||
shutdown_event.set()
|
||||
|
||||
producer_thread.join()
|
||||
close_learner_service_stub(channel, server)
|
||||
|
||||
assert received_params == [b"param_batch_1", b"stop"]
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@pytest.mark.timeout(3) # force cross-platform watchdog
|
||||
def test_stream_parameters_waits_and_retries_on_empty_queue():
|
||||
import threading
|
||||
import time
|
||||
|
||||
from lerobot.common.transport import services_pb2
|
||||
|
||||
"""Test that StreamParameters waits and retries when the queue is empty."""
|
||||
shutdown_event = Event()
|
||||
parameters_queue = Queue()
|
||||
transitions_queue = Queue()
|
||||
interactions_queue = Queue()
|
||||
seconds_between_pushes = 0.05
|
||||
queue_get_timeout = 0.01
|
||||
|
||||
client, channel, server = create_learner_service_stub(
|
||||
shutdown_event,
|
||||
parameters_queue,
|
||||
transitions_queue,
|
||||
interactions_queue,
|
||||
seconds_between_pushes,
|
||||
queue_get_timeout=queue_get_timeout,
|
||||
)
|
||||
|
||||
request = services_pb2.Empty()
|
||||
stream = client.StreamParameters(request)
|
||||
|
||||
received_params = []
|
||||
|
||||
def producer():
|
||||
# Let the consumer start and find an empty queue.
|
||||
# It will wait `seconds_between_pushes` (0.05s), then `get` will timeout after `queue_get_timeout` (0.01s).
|
||||
# Total time for the first empty loop is > 0.06s. We wait a bit longer to be safe.
|
||||
time.sleep(0.06)
|
||||
parameters_queue.put(b"param_after_wait")
|
||||
time.sleep(0.05)
|
||||
parameters_queue.put(b"param_after_wait_2")
|
||||
|
||||
producer_thread = threading.Thread(target=producer)
|
||||
producer_thread.start()
|
||||
|
||||
# The consumer will block here until the producer sends an item.
|
||||
for response in stream:
|
||||
received_params.append(response.data)
|
||||
if response.data == b"param_after_wait_2":
|
||||
break # We only need one item for this test.
|
||||
|
||||
shutdown_event.set()
|
||||
producer_thread.join()
|
||||
close_learner_service_stub(channel, server)
|
||||
|
||||
assert received_params == [b"param_after_wait", b"param_after_wait_2"]
|
||||
@@ -1,144 +0,0 @@
|
||||
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# 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.
|
||||
"""
|
||||
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_robots.py::test_robot
|
||||
```
|
||||
|
||||
Example of running test on real robots connected to the computer:
|
||||
```bash
|
||||
pytest -sx 'tests/test_robots.py::test_robot[koch-False]'
|
||||
pytest -sx 'tests/test_robots.py::test_robot[koch_bimanual-False]'
|
||||
pytest -sx 'tests/test_robots.py::test_robot[aloha-False]'
|
||||
```
|
||||
|
||||
Example of running test on a mocked version of robots:
|
||||
```bash
|
||||
pytest -sx 'tests/test_robots.py::test_robot[koch-True]'
|
||||
pytest -sx 'tests/test_robots.py::test_robot[koch_bimanual-True]'
|
||||
pytest -sx 'tests/test_robots.py::test_robot[aloha-True]'
|
||||
```
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from lerobot.common.robot_devices.robots.utils import make_robot
|
||||
from lerobot.common.robot_devices.utils import RobotDeviceAlreadyConnectedError, RobotDeviceNotConnectedError
|
||||
from tests.utils import TEST_ROBOT_TYPES, mock_calibration_dir, require_robot
|
||||
|
||||
|
||||
@pytest.mark.parametrize("robot_type, mock", TEST_ROBOT_TYPES)
|
||||
@require_robot
|
||||
def test_robot(tmp_path, request, robot_type, mock):
|
||||
# TODO(rcadene): measure fps in nightly?
|
||||
# TODO(rcadene): test logs
|
||||
# TODO(rcadene): add compatibility with other robots
|
||||
robot_kwargs = {"robot_type": robot_type, "mock": mock}
|
||||
|
||||
if robot_type == "aloha" and mock:
|
||||
# To simplify unit test, we do not rerun manual calibration for Aloha mock=True.
|
||||
# Instead, we use the files from '.cache/calibration/aloha_default'
|
||||
pass
|
||||
else:
|
||||
if mock:
|
||||
request.getfixturevalue("patch_builtins_input")
|
||||
|
||||
# Create an empty calibration directory to trigger manual calibration
|
||||
calibration_dir = tmp_path / robot_type
|
||||
mock_calibration_dir(calibration_dir)
|
||||
robot_kwargs["calibration_dir"] = calibration_dir
|
||||
|
||||
# Test using robot before connecting raises an error
|
||||
robot = make_robot(**robot_kwargs)
|
||||
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 (triggers manual calibration)
|
||||
robot = make_robot(**robot_kwargs)
|
||||
robot.connect()
|
||||
assert robot.is_connected
|
||||
|
||||
# Test connecting twice raises an error
|
||||
with pytest.raises(RobotDeviceAlreadyConnectedError):
|
||||
robot.connect()
|
||||
|
||||
# TODO(rcadene, aliberts): Test disconnecting with `__del__` instead of `disconnect`
|
||||
# del robot
|
||||
robot.disconnect()
|
||||
|
||||
# Test teleop can run
|
||||
robot = make_robot(**robot_kwargs)
|
||||
robot.connect()
|
||||
robot.teleop_step()
|
||||
|
||||
# Test data recorded during teleop are well formatted
|
||||
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
|
||||
torch.testing.assert_close(captured_observation[name], observation[name], rtol=1e-4, atol=1)
|
||||
assert captured_observation[name].shape == observation[name].shape
|
||||
|
||||
# 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
|
||||
@@ -0,0 +1,95 @@
|
||||
from contextlib import contextmanager
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from lerobot.common.robots.so100_follower import (
|
||||
SO100Follower,
|
||||
SO100FollowerConfig,
|
||||
)
|
||||
|
||||
|
||||
def _make_bus_mock() -> MagicMock:
|
||||
"""Return a bus mock with just the attributes used by the robot."""
|
||||
bus = MagicMock(name="FeetechBusMock")
|
||||
bus.is_connected = False
|
||||
|
||||
def _connect():
|
||||
bus.is_connected = True
|
||||
|
||||
def _disconnect(_disable=True):
|
||||
bus.is_connected = False
|
||||
|
||||
bus.connect.side_effect = _connect
|
||||
bus.disconnect.side_effect = _disconnect
|
||||
|
||||
@contextmanager
|
||||
def _dummy_cm():
|
||||
yield
|
||||
|
||||
bus.torque_disabled.side_effect = _dummy_cm
|
||||
|
||||
return bus
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def follower():
|
||||
bus_mock = _make_bus_mock()
|
||||
|
||||
def _bus_side_effect(*_args, **kwargs):
|
||||
bus_mock.motors = kwargs["motors"]
|
||||
motors_order: list[str] = list(bus_mock.motors)
|
||||
|
||||
bus_mock.sync_read.return_value = {motor: idx for idx, motor in enumerate(motors_order, 1)}
|
||||
bus_mock.sync_write.return_value = None
|
||||
bus_mock.write.return_value = None
|
||||
bus_mock.disable_torque.return_value = None
|
||||
bus_mock.enable_torque.return_value = None
|
||||
bus_mock.is_calibrated = True
|
||||
return bus_mock
|
||||
|
||||
with (
|
||||
patch(
|
||||
"lerobot.common.robots.so100_follower.so100_follower.FeetechMotorsBus",
|
||||
side_effect=_bus_side_effect,
|
||||
),
|
||||
patch.object(SO100Follower, "configure", lambda self: None),
|
||||
):
|
||||
cfg = SO100FollowerConfig(port="/dev/null")
|
||||
robot = SO100Follower(cfg)
|
||||
yield robot
|
||||
if robot.is_connected:
|
||||
robot.disconnect()
|
||||
|
||||
|
||||
def test_connect_disconnect(follower):
|
||||
assert not follower.is_connected
|
||||
|
||||
follower.connect()
|
||||
assert follower.is_connected
|
||||
|
||||
follower.disconnect()
|
||||
assert not follower.is_connected
|
||||
|
||||
|
||||
def test_get_observation(follower):
|
||||
follower.connect()
|
||||
obs = follower.get_observation()
|
||||
|
||||
expected_keys = {f"{m}.pos" for m in follower.bus.motors}
|
||||
assert set(obs.keys()) == expected_keys
|
||||
|
||||
for idx, motor in enumerate(follower.bus.motors, 1):
|
||||
assert obs[f"{motor}.pos"] == idx
|
||||
|
||||
|
||||
def test_send_action(follower):
|
||||
follower.connect()
|
||||
|
||||
action = {f"{m}.pos": i * 10 for i, m in enumerate(follower.bus.motors, 1)}
|
||||
returned = follower.send_action(action)
|
||||
|
||||
assert returned == action
|
||||
|
||||
goal_pos = {m: (i + 1) * 10 for i, m in enumerate(follower.bus.motors)}
|
||||
follower.bus.sync_write.assert_called_once_with("Goal_Position", goal_pos)
|
||||
@@ -45,12 +45,7 @@ def test_available_policies():
|
||||
This test verifies that the class attribute `name` for all policies is
|
||||
consistent with those listed in `lerobot/__init__.py`.
|
||||
"""
|
||||
policy_classes = [
|
||||
ACTPolicy,
|
||||
DiffusionPolicy,
|
||||
TDMPCPolicy,
|
||||
VQBeTPolicy,
|
||||
]
|
||||
policy_classes = [ACTPolicy, DiffusionPolicy, TDMPCPolicy, VQBeTPolicy]
|
||||
policies = [pol_cls.name for pol_cls in policy_classes]
|
||||
assert set(policies) == set(lerobot.available_policies), policies
|
||||
|
||||
|
||||
@@ -0,0 +1,571 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# 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 io
|
||||
from multiprocessing import Event, Queue
|
||||
from pickle import UnpicklingError
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from lerobot.common.utils.transition import Transition
|
||||
from tests.utils import require_cuda, require_package
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
def test_bytes_buffer_size_empty_buffer():
|
||||
from lerobot.common.transport.utils import bytes_buffer_size
|
||||
|
||||
"""Test with an empty buffer."""
|
||||
buffer = io.BytesIO()
|
||||
assert bytes_buffer_size(buffer) == 0
|
||||
# Ensure position is reset to beginning
|
||||
assert buffer.tell() == 0
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
def test_bytes_buffer_size_small_buffer():
|
||||
from lerobot.common.transport.utils import bytes_buffer_size
|
||||
|
||||
"""Test with a small buffer."""
|
||||
buffer = io.BytesIO(b"Hello, World!")
|
||||
assert bytes_buffer_size(buffer) == 13
|
||||
assert buffer.tell() == 0
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
def test_bytes_buffer_size_large_buffer():
|
||||
from lerobot.common.transport.utils import CHUNK_SIZE, bytes_buffer_size
|
||||
|
||||
"""Test with a large buffer."""
|
||||
data = b"x" * (CHUNK_SIZE * 2 + 1000)
|
||||
buffer = io.BytesIO(data)
|
||||
assert bytes_buffer_size(buffer) == len(data)
|
||||
assert buffer.tell() == 0
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
def test_send_bytes_in_chunks_empty_data():
|
||||
from lerobot.common.transport.utils import send_bytes_in_chunks, services_pb2
|
||||
|
||||
"""Test sending empty data."""
|
||||
message_class = services_pb2.InteractionMessage
|
||||
chunks = list(send_bytes_in_chunks(b"", message_class))
|
||||
assert len(chunks) == 0
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
def test_single_chunk_small_data():
|
||||
from lerobot.common.transport.utils import send_bytes_in_chunks, services_pb2
|
||||
|
||||
"""Test data that fits in a single chunk."""
|
||||
data = b"Some data"
|
||||
message_class = services_pb2.InteractionMessage
|
||||
chunks = list(send_bytes_in_chunks(data, message_class))
|
||||
|
||||
assert len(chunks) == 1
|
||||
assert chunks[0].data == b"Some data"
|
||||
assert chunks[0].transfer_state == services_pb2.TransferState.TRANSFER_END
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
def test_not_silent_mode():
|
||||
from lerobot.common.transport.utils import send_bytes_in_chunks, services_pb2
|
||||
|
||||
"""Test not silent mode."""
|
||||
data = b"Some data"
|
||||
message_class = services_pb2.InteractionMessage
|
||||
chunks = list(send_bytes_in_chunks(data, message_class, silent=False))
|
||||
assert len(chunks) == 1
|
||||
assert chunks[0].data == b"Some data"
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
def test_send_bytes_in_chunks_large_data():
|
||||
from lerobot.common.transport.utils import CHUNK_SIZE, send_bytes_in_chunks, services_pb2
|
||||
|
||||
"""Test sending large data."""
|
||||
data = b"x" * (CHUNK_SIZE * 2 + 1000)
|
||||
message_class = services_pb2.InteractionMessage
|
||||
chunks = list(send_bytes_in_chunks(data, message_class))
|
||||
assert len(chunks) == 3
|
||||
assert chunks[0].data == b"x" * CHUNK_SIZE
|
||||
assert chunks[0].transfer_state == services_pb2.TransferState.TRANSFER_BEGIN
|
||||
assert chunks[1].data == b"x" * CHUNK_SIZE
|
||||
assert chunks[1].transfer_state == services_pb2.TransferState.TRANSFER_MIDDLE
|
||||
assert chunks[2].data == b"x" * 1000
|
||||
assert chunks[2].transfer_state == services_pb2.TransferState.TRANSFER_END
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
def test_send_bytes_in_chunks_large_data_with_exact_chunk_size():
|
||||
from lerobot.common.transport.utils import CHUNK_SIZE, send_bytes_in_chunks, services_pb2
|
||||
|
||||
"""Test sending large data with exact chunk size."""
|
||||
data = b"x" * CHUNK_SIZE
|
||||
message_class = services_pb2.InteractionMessage
|
||||
chunks = list(send_bytes_in_chunks(data, message_class))
|
||||
assert len(chunks) == 1
|
||||
assert chunks[0].data == data
|
||||
assert chunks[0].transfer_state == services_pb2.TransferState.TRANSFER_END
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
def test_receive_bytes_in_chunks_empty_data():
|
||||
from lerobot.common.transport.utils import receive_bytes_in_chunks
|
||||
|
||||
"""Test receiving empty data."""
|
||||
queue = Queue()
|
||||
shutdown_event = Event()
|
||||
|
||||
# Empty iterator
|
||||
receive_bytes_in_chunks(iter([]), queue, shutdown_event)
|
||||
|
||||
assert queue.empty()
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
def test_receive_bytes_in_chunks_single_chunk():
|
||||
from lerobot.common.transport.utils import receive_bytes_in_chunks, services_pb2
|
||||
|
||||
"""Test receiving a single chunk message."""
|
||||
queue = Queue()
|
||||
shutdown_event = Event()
|
||||
|
||||
data = b"Single chunk data"
|
||||
chunks = [
|
||||
services_pb2.InteractionMessage(data=data, transfer_state=services_pb2.TransferState.TRANSFER_END)
|
||||
]
|
||||
|
||||
receive_bytes_in_chunks(iter(chunks), queue, shutdown_event)
|
||||
|
||||
assert queue.get(timeout=0.01) == data
|
||||
assert queue.empty()
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
def test_receive_bytes_in_chunks_single_not_end_chunk():
|
||||
from lerobot.common.transport.utils import receive_bytes_in_chunks, services_pb2
|
||||
|
||||
"""Test receiving a single chunk message."""
|
||||
queue = Queue()
|
||||
shutdown_event = Event()
|
||||
|
||||
data = b"Single chunk data"
|
||||
chunks = [
|
||||
services_pb2.InteractionMessage(data=data, transfer_state=services_pb2.TransferState.TRANSFER_MIDDLE)
|
||||
]
|
||||
|
||||
receive_bytes_in_chunks(iter(chunks), queue, shutdown_event)
|
||||
|
||||
assert queue.empty()
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
def test_receive_bytes_in_chunks_multiple_chunks():
|
||||
from lerobot.common.transport.utils import receive_bytes_in_chunks, services_pb2
|
||||
|
||||
"""Test receiving a multi-chunk message."""
|
||||
queue = Queue()
|
||||
shutdown_event = Event()
|
||||
|
||||
chunks = [
|
||||
services_pb2.InteractionMessage(
|
||||
data=b"First ", transfer_state=services_pb2.TransferState.TRANSFER_BEGIN
|
||||
),
|
||||
services_pb2.InteractionMessage(
|
||||
data=b"Middle ", transfer_state=services_pb2.TransferState.TRANSFER_MIDDLE
|
||||
),
|
||||
services_pb2.InteractionMessage(data=b"Last", transfer_state=services_pb2.TransferState.TRANSFER_END),
|
||||
]
|
||||
|
||||
receive_bytes_in_chunks(iter(chunks), queue, shutdown_event)
|
||||
|
||||
assert queue.get(timeout=0.01) == b"First Middle Last"
|
||||
assert queue.empty()
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
def test_receive_bytes_in_chunks_multiple_messages():
|
||||
from lerobot.common.transport.utils import receive_bytes_in_chunks, services_pb2
|
||||
|
||||
"""Test receiving multiple complete messages in sequence."""
|
||||
queue = Queue()
|
||||
shutdown_event = Event()
|
||||
|
||||
chunks = [
|
||||
# First message - single chunk
|
||||
services_pb2.InteractionMessage(
|
||||
data=b"Message1", transfer_state=services_pb2.TransferState.TRANSFER_END
|
||||
),
|
||||
# Second message - multi chunk
|
||||
services_pb2.InteractionMessage(
|
||||
data=b"Start2 ", transfer_state=services_pb2.TransferState.TRANSFER_BEGIN
|
||||
),
|
||||
services_pb2.InteractionMessage(
|
||||
data=b"Middle2 ", transfer_state=services_pb2.TransferState.TRANSFER_MIDDLE
|
||||
),
|
||||
services_pb2.InteractionMessage(data=b"End2", transfer_state=services_pb2.TransferState.TRANSFER_END),
|
||||
# Third message - single chunk
|
||||
services_pb2.InteractionMessage(
|
||||
data=b"Message3", transfer_state=services_pb2.TransferState.TRANSFER_END
|
||||
),
|
||||
]
|
||||
|
||||
receive_bytes_in_chunks(iter(chunks), queue, shutdown_event)
|
||||
|
||||
# Should have three messages in queue
|
||||
assert queue.get(timeout=0.01) == b"Message1"
|
||||
assert queue.get(timeout=0.01) == b"Start2 Middle2 End2"
|
||||
assert queue.get(timeout=0.01) == b"Message3"
|
||||
assert queue.empty()
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
def test_receive_bytes_in_chunks_shutdown_during_receive():
|
||||
from lerobot.common.transport.utils import receive_bytes_in_chunks, services_pb2
|
||||
|
||||
"""Test that shutdown event stops receiving mid-stream."""
|
||||
queue = Queue()
|
||||
shutdown_event = Event()
|
||||
shutdown_event.set()
|
||||
|
||||
chunks = [
|
||||
services_pb2.InteractionMessage(
|
||||
data=b"First ", transfer_state=services_pb2.TransferState.TRANSFER_BEGIN
|
||||
),
|
||||
services_pb2.InteractionMessage(
|
||||
data=b"Middle ", transfer_state=services_pb2.TransferState.TRANSFER_MIDDLE
|
||||
),
|
||||
services_pb2.InteractionMessage(data=b"Last", transfer_state=services_pb2.TransferState.TRANSFER_END),
|
||||
]
|
||||
|
||||
receive_bytes_in_chunks(iter(chunks), queue, shutdown_event)
|
||||
|
||||
assert queue.empty()
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
def test_receive_bytes_in_chunks_only_begin_chunk():
|
||||
from lerobot.common.transport.utils import receive_bytes_in_chunks, services_pb2
|
||||
|
||||
"""Test receiving only a BEGIN chunk without END."""
|
||||
queue = Queue()
|
||||
shutdown_event = Event()
|
||||
|
||||
chunks = [
|
||||
services_pb2.InteractionMessage(
|
||||
data=b"Start", transfer_state=services_pb2.TransferState.TRANSFER_BEGIN
|
||||
),
|
||||
# No END chunk
|
||||
]
|
||||
|
||||
receive_bytes_in_chunks(iter(chunks), queue, shutdown_event)
|
||||
|
||||
assert queue.empty()
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
def test_receive_bytes_in_chunks_missing_begin():
|
||||
from lerobot.common.transport.utils import receive_bytes_in_chunks, services_pb2
|
||||
|
||||
"""Test receiving chunks starting with MIDDLE instead of BEGIN."""
|
||||
queue = Queue()
|
||||
shutdown_event = Event()
|
||||
|
||||
chunks = [
|
||||
# Missing BEGIN
|
||||
services_pb2.InteractionMessage(
|
||||
data=b"Middle", transfer_state=services_pb2.TransferState.TRANSFER_MIDDLE
|
||||
),
|
||||
services_pb2.InteractionMessage(data=b"End", transfer_state=services_pb2.TransferState.TRANSFER_END),
|
||||
]
|
||||
|
||||
receive_bytes_in_chunks(iter(chunks), queue, shutdown_event)
|
||||
|
||||
# The implementation continues from where it is, so we should get partial data
|
||||
assert queue.get(timeout=0.01) == b"MiddleEnd"
|
||||
assert queue.empty()
|
||||
|
||||
|
||||
# Tests for state_to_bytes and bytes_to_state_dict
|
||||
@require_package("grpc")
|
||||
def test_state_to_bytes_empty_dict():
|
||||
from lerobot.common.transport.utils import bytes_to_state_dict, state_to_bytes
|
||||
|
||||
"""Test converting empty state dict to bytes."""
|
||||
state_dict = {}
|
||||
data = state_to_bytes(state_dict)
|
||||
reconstructed = bytes_to_state_dict(data)
|
||||
assert reconstructed == state_dict
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
def test_bytes_to_state_dict_empty_data():
|
||||
from lerobot.common.transport.utils import bytes_to_state_dict
|
||||
|
||||
"""Test converting empty data to state dict."""
|
||||
with pytest.raises(EOFError):
|
||||
bytes_to_state_dict(b"")
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
def test_state_to_bytes_simple_dict():
|
||||
from lerobot.common.transport.utils import bytes_to_state_dict, state_to_bytes
|
||||
|
||||
"""Test converting simple state dict to bytes."""
|
||||
state_dict = {
|
||||
"layer1.weight": torch.randn(10, 5),
|
||||
"layer1.bias": torch.randn(10),
|
||||
"layer2.weight": torch.randn(1, 10),
|
||||
"layer2.bias": torch.randn(1),
|
||||
}
|
||||
|
||||
data = state_to_bytes(state_dict)
|
||||
assert isinstance(data, bytes)
|
||||
assert len(data) > 0
|
||||
|
||||
reconstructed = bytes_to_state_dict(data)
|
||||
|
||||
assert len(reconstructed) == len(state_dict)
|
||||
for key in state_dict:
|
||||
assert key in reconstructed
|
||||
assert torch.allclose(state_dict[key], reconstructed[key])
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
def test_state_to_bytes_various_dtypes():
|
||||
from lerobot.common.transport.utils import bytes_to_state_dict, state_to_bytes
|
||||
|
||||
"""Test converting state dict with various tensor dtypes."""
|
||||
state_dict = {
|
||||
"float32": torch.randn(5, 5),
|
||||
"float64": torch.randn(3, 3).double(),
|
||||
"int32": torch.randint(0, 100, (4, 4), dtype=torch.int32),
|
||||
"int64": torch.randint(0, 100, (2, 2), dtype=torch.int64),
|
||||
"bool": torch.tensor([True, False, True]),
|
||||
"uint8": torch.randint(0, 255, (3, 3), dtype=torch.uint8),
|
||||
}
|
||||
|
||||
data = state_to_bytes(state_dict)
|
||||
reconstructed = bytes_to_state_dict(data)
|
||||
|
||||
for key in state_dict:
|
||||
assert reconstructed[key].dtype == state_dict[key].dtype
|
||||
if state_dict[key].dtype == torch.bool:
|
||||
assert torch.equal(state_dict[key], reconstructed[key])
|
||||
else:
|
||||
assert torch.allclose(state_dict[key], reconstructed[key])
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
def test_bytes_to_state_dict_invalid_data():
|
||||
from lerobot.common.transport.utils import bytes_to_state_dict
|
||||
|
||||
"""Test bytes_to_state_dict with invalid data."""
|
||||
with pytest.raises(UnpicklingError):
|
||||
bytes_to_state_dict(b"This is not a valid torch save file")
|
||||
|
||||
|
||||
@require_cuda
|
||||
@require_package("grpc")
|
||||
def test_state_to_bytes_various_dtypes_cuda():
|
||||
from lerobot.common.transport.utils import bytes_to_state_dict, state_to_bytes
|
||||
|
||||
"""Test converting state dict with various tensor dtypes."""
|
||||
state_dict = {
|
||||
"float32": torch.randn(5, 5).cuda(),
|
||||
"float64": torch.randn(3, 3).double().cuda(),
|
||||
"int32": torch.randint(0, 100, (4, 4), dtype=torch.int32).cuda(),
|
||||
"int64": torch.randint(0, 100, (2, 2), dtype=torch.int64).cuda(),
|
||||
"bool": torch.tensor([True, False, True]),
|
||||
"uint8": torch.randint(0, 255, (3, 3), dtype=torch.uint8),
|
||||
}
|
||||
|
||||
data = state_to_bytes(state_dict)
|
||||
reconstructed = bytes_to_state_dict(data)
|
||||
|
||||
for key in state_dict:
|
||||
assert reconstructed[key].dtype == state_dict[key].dtype
|
||||
if state_dict[key].dtype == torch.bool:
|
||||
assert torch.equal(state_dict[key], reconstructed[key])
|
||||
else:
|
||||
assert torch.allclose(state_dict[key], reconstructed[key])
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
def test_python_object_to_bytes_none():
|
||||
from lerobot.common.transport.utils import bytes_to_python_object, python_object_to_bytes
|
||||
|
||||
"""Test converting None to bytes."""
|
||||
obj = None
|
||||
data = python_object_to_bytes(obj)
|
||||
reconstructed = bytes_to_python_object(data)
|
||||
assert reconstructed is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"obj",
|
||||
[
|
||||
42,
|
||||
-123,
|
||||
3.14159,
|
||||
-2.71828,
|
||||
"Hello, World!",
|
||||
"Unicode: 你好世界 🌍",
|
||||
True,
|
||||
False,
|
||||
b"byte string",
|
||||
[],
|
||||
[1, 2, 3],
|
||||
[1, "two", 3.0, True, None],
|
||||
{},
|
||||
{"key": "value", "number": 123, "nested": {"a": 1}},
|
||||
(),
|
||||
(1, 2, 3),
|
||||
],
|
||||
)
|
||||
@require_package("grpc")
|
||||
def test_python_object_to_bytes_simple_types(obj):
|
||||
from lerobot.common.transport.utils import bytes_to_python_object, python_object_to_bytes
|
||||
|
||||
"""Test converting simple Python types."""
|
||||
data = python_object_to_bytes(obj)
|
||||
reconstructed = bytes_to_python_object(data)
|
||||
assert reconstructed == obj
|
||||
assert type(reconstructed) is type(obj)
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
def test_python_object_to_bytes_with_tensors():
|
||||
from lerobot.common.transport.utils import bytes_to_python_object, python_object_to_bytes
|
||||
|
||||
"""Test converting objects containing PyTorch tensors."""
|
||||
obj = {
|
||||
"tensor": torch.randn(5, 5),
|
||||
"list_with_tensor": [1, 2, torch.randn(3, 3), "string"],
|
||||
"nested": {
|
||||
"tensor1": torch.randn(2, 2),
|
||||
"tensor2": torch.tensor([1, 2, 3]),
|
||||
},
|
||||
}
|
||||
|
||||
data = python_object_to_bytes(obj)
|
||||
reconstructed = bytes_to_python_object(data)
|
||||
|
||||
assert torch.allclose(obj["tensor"], reconstructed["tensor"])
|
||||
assert reconstructed["list_with_tensor"][0] == 1
|
||||
assert reconstructed["list_with_tensor"][3] == "string"
|
||||
assert torch.allclose(obj["list_with_tensor"][2], reconstructed["list_with_tensor"][2])
|
||||
assert torch.allclose(obj["nested"]["tensor1"], reconstructed["nested"]["tensor1"])
|
||||
assert torch.equal(obj["nested"]["tensor2"], reconstructed["nested"]["tensor2"])
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
def test_transitions_to_bytes_empty_list():
|
||||
from lerobot.common.transport.utils import bytes_to_transitions, transitions_to_bytes
|
||||
|
||||
"""Test converting empty transitions list."""
|
||||
transitions = []
|
||||
data = transitions_to_bytes(transitions)
|
||||
reconstructed = bytes_to_transitions(data)
|
||||
assert reconstructed == transitions
|
||||
assert isinstance(reconstructed, list)
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
def test_transitions_to_bytes_single_transition():
|
||||
from lerobot.common.transport.utils import bytes_to_transitions, transitions_to_bytes
|
||||
|
||||
"""Test converting a single transition."""
|
||||
transition = Transition(
|
||||
state={"image": torch.randn(3, 64, 64), "state": torch.randn(10)},
|
||||
action=torch.randn(5),
|
||||
reward=torch.tensor(1.5),
|
||||
done=torch.tensor(False),
|
||||
next_state={"image": torch.randn(3, 64, 64), "state": torch.randn(10)},
|
||||
)
|
||||
|
||||
transitions = [transition]
|
||||
data = transitions_to_bytes(transitions)
|
||||
reconstructed = bytes_to_transitions(data)
|
||||
|
||||
assert len(reconstructed) == 1
|
||||
|
||||
assert_transitions_equal(transitions[0], reconstructed[0])
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
def assert_transitions_equal(t1: Transition, t2: Transition):
|
||||
"""Helper to assert two transitions are equal."""
|
||||
assert_observation_equal(t1["state"], t2["state"])
|
||||
assert torch.allclose(t1["action"], t2["action"])
|
||||
assert torch.allclose(t1["reward"], t2["reward"])
|
||||
assert torch.equal(t1["done"], t2["done"])
|
||||
assert_observation_equal(t1["next_state"], t2["next_state"])
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
def assert_observation_equal(o1: dict, o2: dict):
|
||||
"""Helper to assert two observations are equal."""
|
||||
assert set(o1.keys()) == set(o2.keys())
|
||||
for key in o1:
|
||||
assert torch.allclose(o1[key], o2[key])
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
def test_transitions_to_bytes_multiple_transitions():
|
||||
from lerobot.common.transport.utils import bytes_to_transitions, transitions_to_bytes
|
||||
|
||||
"""Test converting multiple transitions."""
|
||||
transitions = []
|
||||
for i in range(5):
|
||||
transition = Transition(
|
||||
state={"data": torch.randn(10)},
|
||||
action=torch.randn(3),
|
||||
reward=torch.tensor(float(i)),
|
||||
done=torch.tensor(i == 4),
|
||||
next_state={"data": torch.randn(10)},
|
||||
)
|
||||
transitions.append(transition)
|
||||
|
||||
data = transitions_to_bytes(transitions)
|
||||
reconstructed = bytes_to_transitions(data)
|
||||
|
||||
assert len(reconstructed) == len(transitions)
|
||||
for original, reconstructed_item in zip(transitions, reconstructed, strict=False):
|
||||
assert_transitions_equal(original, reconstructed_item)
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
def test_receive_bytes_in_chunks_unknown_state():
|
||||
from lerobot.common.transport.utils import receive_bytes_in_chunks
|
||||
|
||||
"""Test receive_bytes_in_chunks with an unknown transfer state."""
|
||||
|
||||
# Mock the gRPC message object, which has `transfer_state` and `data` attributes.
|
||||
class MockMessage:
|
||||
def __init__(self, transfer_state, data):
|
||||
self.transfer_state = transfer_state
|
||||
self.data = data
|
||||
|
||||
# 10 is not a valid TransferState enum value
|
||||
bad_iterator = [MockMessage(transfer_state=10, data=b"bad_data")]
|
||||
output_queue = Queue()
|
||||
shutdown_event = Event()
|
||||
|
||||
with pytest.raises(ValueError, match="Received unknown transfer state"):
|
||||
receive_bytes_in_chunks(bad_iterator, output_queue, shutdown_event)
|
||||
-147
@@ -13,20 +13,14 @@
|
||||
# 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 json
|
||||
import os
|
||||
import platform
|
||||
from functools import wraps
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from lerobot import available_cameras, available_motors, available_robots
|
||||
from lerobot.common.robot_devices.cameras.utils import Camera
|
||||
from lerobot.common.robot_devices.cameras.utils import make_camera as make_camera_device
|
||||
from lerobot.common.robot_devices.motors.utils import MotorsBus
|
||||
from lerobot.common.robot_devices.motors.utils import make_motors_bus as make_motors_bus_device
|
||||
from lerobot.common.utils.import_utils import is_package_available
|
||||
|
||||
DEVICE = os.environ.get("LEROBOT_TEST_DEVICE", "cuda") if torch.cuda.is_available() else "cpu"
|
||||
@@ -188,144 +182,3 @@ def require_package(package_name):
|
||||
return wrapper
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
def require_robot(func):
|
||||
"""
|
||||
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_robot_available fixture
|
||||
request = kwargs.get("request")
|
||||
robot_type = kwargs.get("robot_type")
|
||||
mock = kwargs.get("mock")
|
||||
|
||||
if robot_type is None:
|
||||
raise ValueError("The 'robot_type' must be an argument of the test function.")
|
||||
if request is None:
|
||||
raise ValueError("The 'request' fixture must be an argument of the test function.")
|
||||
if mock is None:
|
||||
raise ValueError("The 'mock' variable must be an argument of the test function.")
|
||||
|
||||
# Run test with a real robot. Skip test if robot connection fails.
|
||||
if not mock and not request.getfixturevalue("is_robot_available"):
|
||||
pytest.skip(f"A {robot_type} robot is not available.")
|
||||
|
||||
return func(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def require_camera(func):
|
||||
@wraps(func)
|
||||
def wrapper(*args, **kwargs):
|
||||
# Access the pytest request context to get the is_camera_available fixture
|
||||
request = kwargs.get("request")
|
||||
camera_type = kwargs.get("camera_type")
|
||||
mock = kwargs.get("mock")
|
||||
|
||||
if request is None:
|
||||
raise ValueError("The 'request' fixture must be an argument of the test function.")
|
||||
if camera_type is None:
|
||||
raise ValueError("The 'camera_type' must be an argument of the test function.")
|
||||
if mock is None:
|
||||
raise ValueError("The 'mock' variable must be an argument of the test function.")
|
||||
|
||||
if not mock and not request.getfixturevalue("is_camera_available"):
|
||||
pytest.skip(f"A {camera_type} camera is not available.")
|
||||
|
||||
return func(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def require_motor(func):
|
||||
@wraps(func)
|
||||
def wrapper(*args, **kwargs):
|
||||
# Access the pytest request context to get the is_motor_available fixture
|
||||
request = kwargs.get("request")
|
||||
motor_type = kwargs.get("motor_type")
|
||||
mock = kwargs.get("mock")
|
||||
|
||||
if request is None:
|
||||
raise ValueError("The 'request' fixture must be an argument of the test function.")
|
||||
if motor_type is None:
|
||||
raise ValueError("The 'motor_type' must be an argument of the test function.")
|
||||
if mock is None:
|
||||
raise ValueError("The 'mock' variable must be an argument of the test function.")
|
||||
|
||||
if not mock and not request.getfixturevalue("is_motor_available"):
|
||||
pytest.skip(f"A {motor_type} motor is not available.")
|
||||
|
||||
return func(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def mock_calibration_dir(calibration_dir):
|
||||
# TODO(rcadene): remove this hack
|
||||
# calibration file produced with Moss v1, but works with Koch, Koch bimanual and SO-100
|
||||
example_calib = {
|
||||
"homing_offset": [-1416, -845, 2130, 2872, 1950, -2211],
|
||||
"drive_mode": [0, 0, 1, 1, 1, 0],
|
||||
"start_pos": [1442, 843, 2166, 2849, 1988, 1835],
|
||||
"end_pos": [2440, 1869, -1106, -1848, -926, 3235],
|
||||
"calib_mode": ["DEGREE", "DEGREE", "DEGREE", "DEGREE", "DEGREE", "LINEAR"],
|
||||
"motor_names": ["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll", "gripper"],
|
||||
}
|
||||
Path(str(calibration_dir)).mkdir(parents=True, exist_ok=True)
|
||||
with open(calibration_dir / "main_follower.json", "w") as f:
|
||||
json.dump(example_calib, f)
|
||||
with open(calibration_dir / "main_leader.json", "w") as f:
|
||||
json.dump(example_calib, f)
|
||||
with open(calibration_dir / "left_follower.json", "w") as f:
|
||||
json.dump(example_calib, f)
|
||||
with open(calibration_dir / "left_leader.json", "w") as f:
|
||||
json.dump(example_calib, f)
|
||||
with open(calibration_dir / "right_follower.json", "w") as f:
|
||||
json.dump(example_calib, f)
|
||||
with open(calibration_dir / "right_leader.json", "w") as f:
|
||||
json.dump(example_calib, f)
|
||||
|
||||
|
||||
# TODO(rcadene, aliberts): remove this dark pattern that overrides
|
||||
def make_camera(camera_type: str, **kwargs) -> Camera:
|
||||
if camera_type == "opencv":
|
||||
camera_index = kwargs.pop("camera_index", OPENCV_CAMERA_INDEX)
|
||||
return make_camera_device(camera_type, camera_index=camera_index, **kwargs)
|
||||
|
||||
elif camera_type == "intelrealsense":
|
||||
serial_number = kwargs.pop("serial_number", INTELREALSENSE_SERIAL_NUMBER)
|
||||
return make_camera_device(camera_type, serial_number=serial_number, **kwargs)
|
||||
else:
|
||||
raise ValueError(f"The camera type '{camera_type}' is not valid.")
|
||||
|
||||
|
||||
# TODO(rcadene, aliberts): remove this dark pattern that overrides
|
||||
def make_motors_bus(motor_type: str, **kwargs) -> MotorsBus:
|
||||
if motor_type == "dynamixel":
|
||||
port = kwargs.pop("port", DYNAMIXEL_PORT)
|
||||
motors = kwargs.pop("motors", DYNAMIXEL_MOTORS)
|
||||
return make_motors_bus_device(motor_type, port=port, motors=motors, **kwargs)
|
||||
|
||||
elif motor_type == "feetech":
|
||||
port = kwargs.pop("port", FEETECH_PORT)
|
||||
motors = kwargs.pop("motors", FEETECH_MOTORS)
|
||||
return make_motors_bus_device(motor_type, port=port, motors=motors, **kwargs)
|
||||
|
||||
else:
|
||||
raise ValueError(f"The motor type '{motor_type}' is not valid.")
|
||||
|
||||
@@ -0,0 +1,155 @@
|
||||
import pytest
|
||||
|
||||
from lerobot.common.utils.encoding_utils import (
|
||||
decode_sign_magnitude,
|
||||
decode_twos_complement,
|
||||
encode_sign_magnitude,
|
||||
encode_twos_complement,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"value, sign_bit_index, expected",
|
||||
[
|
||||
(5, 4, 5),
|
||||
(0, 4, 0),
|
||||
(7, 3, 7),
|
||||
(-1, 4, 17),
|
||||
(-8, 4, 24),
|
||||
(-3, 3, 11),
|
||||
],
|
||||
)
|
||||
def test_encode_sign_magnitude(value, sign_bit_index, expected):
|
||||
assert encode_sign_magnitude(value, sign_bit_index) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"encoded, sign_bit_index, expected",
|
||||
[
|
||||
(5, 4, 5),
|
||||
(0, 4, 0),
|
||||
(7, 3, 7),
|
||||
(17, 4, -1),
|
||||
(24, 4, -8),
|
||||
(11, 3, -3),
|
||||
],
|
||||
)
|
||||
def test_decode_sign_magnitude(encoded, sign_bit_index, expected):
|
||||
assert decode_sign_magnitude(encoded, sign_bit_index) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"encoded, sign_bit_index",
|
||||
[
|
||||
(16, 4),
|
||||
(-9, 3),
|
||||
],
|
||||
)
|
||||
def test_encode_raises_on_overflow(encoded, sign_bit_index):
|
||||
with pytest.raises(ValueError):
|
||||
encode_sign_magnitude(encoded, sign_bit_index)
|
||||
|
||||
|
||||
def test_encode_decode_sign_magnitude():
|
||||
for sign_bit_index in range(2, 6):
|
||||
max_val = (1 << sign_bit_index) - 1
|
||||
for value in range(-max_val, max_val + 1):
|
||||
encoded = encode_sign_magnitude(value, sign_bit_index)
|
||||
decoded = decode_sign_magnitude(encoded, sign_bit_index)
|
||||
assert decoded == value, f"Failed at value={value}, index={sign_bit_index}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"value, n_bytes, expected",
|
||||
[
|
||||
(0, 1, 0),
|
||||
(5, 1, 5),
|
||||
(-1, 1, 255),
|
||||
(-128, 1, 128),
|
||||
(-2, 1, 254),
|
||||
(127, 1, 127),
|
||||
(0, 2, 0),
|
||||
(5, 2, 5),
|
||||
(-1, 2, 65_535),
|
||||
(-32_768, 2, 32_768),
|
||||
(-2, 2, 65_534),
|
||||
(32_767, 2, 32_767),
|
||||
(0, 4, 0),
|
||||
(5, 4, 5),
|
||||
(-1, 4, 4_294_967_295),
|
||||
(-2_147_483_648, 4, 2_147_483_648),
|
||||
(-2, 4, 4_294_967_294),
|
||||
(2_147_483_647, 4, 2_147_483_647),
|
||||
],
|
||||
)
|
||||
def test_encode_twos_complement(value, n_bytes, expected):
|
||||
assert encode_twos_complement(value, n_bytes) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"value, n_bytes, expected",
|
||||
[
|
||||
(0, 1, 0),
|
||||
(5, 1, 5),
|
||||
(255, 1, -1),
|
||||
(128, 1, -128),
|
||||
(254, 1, -2),
|
||||
(127, 1, 127),
|
||||
(0, 2, 0),
|
||||
(5, 2, 5),
|
||||
(65_535, 2, -1),
|
||||
(32_768, 2, -32_768),
|
||||
(65_534, 2, -2),
|
||||
(32_767, 2, 32_767),
|
||||
(0, 4, 0),
|
||||
(5, 4, 5),
|
||||
(4_294_967_295, 4, -1),
|
||||
(2_147_483_648, 4, -2_147_483_648),
|
||||
(4_294_967_294, 4, -2),
|
||||
(2_147_483_647, 4, 2_147_483_647),
|
||||
],
|
||||
)
|
||||
def test_decode_twos_complement(value, n_bytes, expected):
|
||||
assert decode_twos_complement(value, n_bytes) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"value, n_bytes",
|
||||
[
|
||||
(-129, 1),
|
||||
(128, 1),
|
||||
(-32_769, 2),
|
||||
(32_768, 2),
|
||||
(-2_147_483_649, 4),
|
||||
(2_147_483_648, 4),
|
||||
],
|
||||
)
|
||||
def test_encode_twos_complement_out_of_range(value, n_bytes):
|
||||
with pytest.raises(ValueError):
|
||||
encode_twos_complement(value, n_bytes)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"value, n_bytes",
|
||||
[
|
||||
(-128, 1),
|
||||
(-1, 1),
|
||||
(0, 1),
|
||||
(1, 1),
|
||||
(127, 1),
|
||||
(-32_768, 2),
|
||||
(-1, 2),
|
||||
(0, 2),
|
||||
(1, 2),
|
||||
(32_767, 2),
|
||||
(-2_147_483_648, 4),
|
||||
(-1, 4),
|
||||
(0, 4),
|
||||
(1, 4),
|
||||
(2_147_483_647, 4),
|
||||
],
|
||||
)
|
||||
def test_encode_decode_twos_complement(value, n_bytes):
|
||||
encoded = encode_twos_complement(value, n_bytes)
|
||||
decoded = decode_twos_complement(encoded, n_bytes)
|
||||
assert decoded == value, f"Failed at value={value}, n_bytes={n_bytes}"
|
||||
@@ -0,0 +1,112 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# 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 multiprocessing
|
||||
import os
|
||||
import signal
|
||||
import threading
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from lerobot.common.utils.process import ProcessSignalHandler
|
||||
|
||||
|
||||
# Fixture to reset shutdown_event_counter and original signal handlers before and after each test
|
||||
@pytest.fixture(autouse=True)
|
||||
def reset_globals_and_handlers():
|
||||
# Store original signal handlers
|
||||
original_handlers = {
|
||||
sig: signal.getsignal(sig)
|
||||
for sig in [signal.SIGINT, signal.SIGTERM, signal.SIGHUP, signal.SIGQUIT]
|
||||
if hasattr(signal, sig.name)
|
||||
}
|
||||
|
||||
yield
|
||||
|
||||
# Restore original signal handlers
|
||||
for sig, handler in original_handlers.items():
|
||||
signal.signal(sig, handler)
|
||||
|
||||
|
||||
def test_setup_process_handlers_event_with_threads():
|
||||
"""Test that setup_process_handlers returns the correct event type."""
|
||||
handler = ProcessSignalHandler(use_threads=True)
|
||||
shutdown_event = handler.shutdown_event
|
||||
assert isinstance(shutdown_event, threading.Event), "Should be a threading.Event"
|
||||
assert not shutdown_event.is_set(), "Event should initially be unset"
|
||||
|
||||
|
||||
def test_setup_process_handlers_event_with_processes():
|
||||
"""Test that setup_process_handlers returns the correct event type."""
|
||||
handler = ProcessSignalHandler(use_threads=False)
|
||||
shutdown_event = handler.shutdown_event
|
||||
assert isinstance(shutdown_event, type(multiprocessing.Event())), "Should be a multiprocessing.Event"
|
||||
assert not shutdown_event.is_set(), "Event should initially be unset"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_threads", [True, False])
|
||||
@pytest.mark.parametrize(
|
||||
"sig",
|
||||
[
|
||||
signal.SIGINT,
|
||||
signal.SIGTERM,
|
||||
# SIGHUP and SIGQUIT are not reliably available on all platforms (e.g. Windows)
|
||||
pytest.param(
|
||||
signal.SIGHUP,
|
||||
marks=pytest.mark.skipif(not hasattr(signal, "SIGHUP"), reason="SIGHUP not available"),
|
||||
),
|
||||
pytest.param(
|
||||
signal.SIGQUIT,
|
||||
marks=pytest.mark.skipif(not hasattr(signal, "SIGQUIT"), reason="SIGQUIT not available"),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_signal_handler_sets_event(use_threads, sig):
|
||||
"""Test that the signal handler sets the event on receiving a signal."""
|
||||
handler = ProcessSignalHandler(use_threads=use_threads)
|
||||
shutdown_event = handler.shutdown_event
|
||||
|
||||
assert handler.counter == 0
|
||||
|
||||
os.kill(os.getpid(), sig)
|
||||
|
||||
# In some environments, the signal might take a moment to be handled.
|
||||
shutdown_event.wait(timeout=1.0)
|
||||
|
||||
assert shutdown_event.is_set(), f"Event should be set after receiving signal {sig}"
|
||||
|
||||
# Ensure the internal counter was incremented
|
||||
assert handler.counter == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_threads", [True, False])
|
||||
@patch("sys.exit")
|
||||
def test_force_shutdown_on_second_signal(mock_sys_exit, use_threads):
|
||||
"""Test that a second signal triggers a force shutdown."""
|
||||
handler = ProcessSignalHandler(use_threads=use_threads)
|
||||
|
||||
os.kill(os.getpid(), signal.SIGINT)
|
||||
# Give a moment for the first signal to be processed
|
||||
import time
|
||||
|
||||
time.sleep(0.1)
|
||||
os.kill(os.getpid(), signal.SIGINT)
|
||||
|
||||
time.sleep(0.1)
|
||||
|
||||
assert handler.counter == 2
|
||||
mock_sys_exit.assert_called_once_with(1)
|
||||
@@ -0,0 +1,150 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# 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 threading
|
||||
import time
|
||||
from queue import Queue
|
||||
|
||||
from lerobot.common.utils.queue import get_last_item_from_queue
|
||||
|
||||
|
||||
def test_get_last_item_single_item():
|
||||
"""Test getting the last item when queue has only one item."""
|
||||
queue = Queue()
|
||||
queue.put("single_item")
|
||||
|
||||
result = get_last_item_from_queue(queue)
|
||||
|
||||
assert result == "single_item"
|
||||
assert queue.empty()
|
||||
|
||||
|
||||
def test_get_last_item_multiple_items():
|
||||
"""Test getting the last item when queue has multiple items."""
|
||||
queue = Queue()
|
||||
items = ["first", "second", "third", "fourth", "last"]
|
||||
|
||||
for item in items:
|
||||
queue.put(item)
|
||||
|
||||
result = get_last_item_from_queue(queue)
|
||||
|
||||
assert result == "last"
|
||||
assert queue.empty()
|
||||
|
||||
|
||||
def test_get_last_item_different_types():
|
||||
"""Test with different data types in the queue."""
|
||||
queue = Queue()
|
||||
items = [1, 2.5, "string", {"key": "value"}, [1, 2, 3], ("tuple", "data")]
|
||||
|
||||
for item in items:
|
||||
queue.put(item)
|
||||
|
||||
result = get_last_item_from_queue(queue)
|
||||
|
||||
assert result == ("tuple", "data")
|
||||
assert queue.empty()
|
||||
|
||||
|
||||
def test_get_last_item_maxsize_queue():
|
||||
"""Test with a queue that has a maximum size."""
|
||||
queue = Queue(maxsize=5)
|
||||
|
||||
# Fill the queue
|
||||
for i in range(5):
|
||||
queue.put(i)
|
||||
|
||||
# Give the queue time to fill
|
||||
time.sleep(0.1)
|
||||
|
||||
result = get_last_item_from_queue(queue)
|
||||
|
||||
assert result == 4
|
||||
assert queue.empty()
|
||||
|
||||
|
||||
def test_get_last_item_with_none_values():
|
||||
"""Test with None values in the queue."""
|
||||
queue = Queue()
|
||||
items = [1, None, 2, None, 3]
|
||||
|
||||
for item in items:
|
||||
queue.put(item)
|
||||
|
||||
# Give the queue time to fill
|
||||
time.sleep(0.1)
|
||||
|
||||
result = get_last_item_from_queue(queue)
|
||||
|
||||
assert result == 3
|
||||
assert queue.empty()
|
||||
|
||||
|
||||
def test_get_last_item_blocking_timeout():
|
||||
"""Test get_last_item_from_queue returns None on timeout."""
|
||||
queue = Queue()
|
||||
result = get_last_item_from_queue(queue, block=True, timeout=0.1)
|
||||
assert result is None
|
||||
|
||||
|
||||
def test_get_last_item_non_blocking_empty():
|
||||
"""Test get_last_item_from_queue with block=False on an empty queue returns None."""
|
||||
queue = Queue()
|
||||
result = get_last_item_from_queue(queue, block=False)
|
||||
assert result is None
|
||||
|
||||
|
||||
def test_get_last_item_non_blocking_success():
|
||||
"""Test get_last_item_from_queue with block=False on a non-empty queue."""
|
||||
queue = Queue()
|
||||
items = ["first", "second", "last"]
|
||||
for item in items:
|
||||
queue.put(item)
|
||||
|
||||
# Give the queue time to fill
|
||||
time.sleep(0.1)
|
||||
|
||||
result = get_last_item_from_queue(queue, block=False)
|
||||
assert result == "last"
|
||||
assert queue.empty()
|
||||
|
||||
|
||||
def test_get_last_item_blocking_waits_for_item():
|
||||
"""Test that get_last_item_from_queue waits for an item if block=True."""
|
||||
queue = Queue()
|
||||
result = []
|
||||
|
||||
def producer():
|
||||
queue.put("item1")
|
||||
queue.put("item2")
|
||||
|
||||
def consumer():
|
||||
# This will block until the producer puts the first item
|
||||
item = get_last_item_from_queue(queue, block=True, timeout=0.2)
|
||||
result.append(item)
|
||||
|
||||
producer_thread = threading.Thread(target=producer)
|
||||
consumer_thread = threading.Thread(target=consumer)
|
||||
|
||||
producer_thread.start()
|
||||
consumer_thread.start()
|
||||
|
||||
producer_thread.join()
|
||||
consumer_thread.join()
|
||||
|
||||
assert result == ["item2"]
|
||||
assert queue.empty()
|
||||
@@ -0,0 +1,682 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# 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 sys
|
||||
from typing import Callable
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
|
||||
from lerobot.common.utils.buffer import BatchTransition, ReplayBuffer, random_crop_vectorized
|
||||
from tests.fixtures.constants import DUMMY_REPO_ID
|
||||
|
||||
|
||||
def state_dims() -> list[str]:
|
||||
return ["observation.image", "observation.state"]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def replay_buffer() -> ReplayBuffer:
|
||||
return create_empty_replay_buffer()
|
||||
|
||||
|
||||
def clone_state(state: dict) -> dict:
|
||||
return {k: v.clone() for k, v in state.items()}
|
||||
|
||||
|
||||
def create_empty_replay_buffer(
|
||||
optimize_memory: bool = False,
|
||||
use_drq: bool = False,
|
||||
image_augmentation_function: Callable | None = None,
|
||||
) -> ReplayBuffer:
|
||||
buffer_capacity = 10
|
||||
device = "cpu"
|
||||
return ReplayBuffer(
|
||||
buffer_capacity,
|
||||
device,
|
||||
state_dims(),
|
||||
optimize_memory=optimize_memory,
|
||||
use_drq=use_drq,
|
||||
image_augmentation_function=image_augmentation_function,
|
||||
)
|
||||
|
||||
|
||||
def create_random_image() -> torch.Tensor:
|
||||
return torch.rand(3, 84, 84)
|
||||
|
||||
|
||||
def create_dummy_transition() -> dict:
|
||||
return {
|
||||
"observation.image": create_random_image(),
|
||||
"action": torch.randn(4),
|
||||
"reward": torch.tensor(1.0),
|
||||
"observation.state": torch.randn(
|
||||
10,
|
||||
),
|
||||
"done": torch.tensor(False),
|
||||
"truncated": torch.tensor(False),
|
||||
"complementary_info": {},
|
||||
}
|
||||
|
||||
|
||||
def create_dataset_from_replay_buffer(tmp_path) -> tuple[LeRobotDataset, ReplayBuffer]:
|
||||
dummy_state_1 = create_dummy_state()
|
||||
dummy_action_1 = create_dummy_action()
|
||||
|
||||
dummy_state_2 = create_dummy_state()
|
||||
dummy_action_2 = create_dummy_action()
|
||||
|
||||
dummy_state_3 = create_dummy_state()
|
||||
dummy_action_3 = create_dummy_action()
|
||||
|
||||
dummy_state_4 = create_dummy_state()
|
||||
dummy_action_4 = create_dummy_action()
|
||||
|
||||
replay_buffer = create_empty_replay_buffer()
|
||||
replay_buffer.add(dummy_state_1, dummy_action_1, 1.0, dummy_state_1, False, False)
|
||||
replay_buffer.add(dummy_state_2, dummy_action_2, 1.0, dummy_state_2, False, False)
|
||||
replay_buffer.add(dummy_state_3, dummy_action_3, 1.0, dummy_state_3, True, True)
|
||||
replay_buffer.add(dummy_state_4, dummy_action_4, 1.0, dummy_state_4, True, True)
|
||||
|
||||
root = tmp_path / "test"
|
||||
return (replay_buffer.to_lerobot_dataset(DUMMY_REPO_ID, root=root), replay_buffer)
|
||||
|
||||
|
||||
def create_dummy_state() -> dict:
|
||||
return {
|
||||
"observation.image": create_random_image(),
|
||||
"observation.state": torch.randn(
|
||||
10,
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def get_tensor_memory_consumption(tensor):
|
||||
return tensor.nelement() * tensor.element_size()
|
||||
|
||||
|
||||
def get_tensors_memory_consumption(obj, visited_addresses):
|
||||
total_size = 0
|
||||
|
||||
address = id(obj)
|
||||
if address in visited_addresses:
|
||||
return 0
|
||||
|
||||
visited_addresses.add(address)
|
||||
|
||||
if isinstance(obj, torch.Tensor):
|
||||
return get_tensor_memory_consumption(obj)
|
||||
elif isinstance(obj, (list, tuple)):
|
||||
for item in obj:
|
||||
total_size += get_tensors_memory_consumption(item, visited_addresses)
|
||||
elif isinstance(obj, dict):
|
||||
for value in obj.values():
|
||||
total_size += get_tensors_memory_consumption(value, visited_addresses)
|
||||
elif hasattr(obj, "__dict__"):
|
||||
# It's an object, we need to get the size of the attributes
|
||||
for _, attr in vars(obj).items():
|
||||
total_size += get_tensors_memory_consumption(attr, visited_addresses)
|
||||
|
||||
return total_size
|
||||
|
||||
|
||||
def get_object_memory(obj):
|
||||
# Track visited addresses to avoid infinite loops
|
||||
# and cases when two properties point to the same object
|
||||
visited_addresses = set()
|
||||
|
||||
# Get the size of the object in bytes
|
||||
total_size = sys.getsizeof(obj)
|
||||
|
||||
# Get the size of the tensor attributes
|
||||
total_size += get_tensors_memory_consumption(obj, visited_addresses)
|
||||
|
||||
return total_size
|
||||
|
||||
|
||||
def create_dummy_action() -> torch.Tensor:
|
||||
return torch.randn(4)
|
||||
|
||||
|
||||
def dict_properties() -> list:
|
||||
return ["state", "next_state"]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def dummy_state() -> dict:
|
||||
return create_dummy_state()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def next_dummy_state() -> dict:
|
||||
return create_dummy_state()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def dummy_action() -> torch.Tensor:
|
||||
return torch.randn(4)
|
||||
|
||||
|
||||
def test_empty_buffer_sample_raises_error(replay_buffer):
|
||||
assert len(replay_buffer) == 0, "Replay buffer should be empty."
|
||||
assert replay_buffer.capacity == 10, "Replay buffer capacity should be 10."
|
||||
with pytest.raises(RuntimeError, match="Cannot sample from an empty buffer"):
|
||||
replay_buffer.sample(1)
|
||||
|
||||
|
||||
def test_zero_capacity_buffer_raises_error():
|
||||
with pytest.raises(ValueError, match="Capacity must be greater than 0."):
|
||||
ReplayBuffer(0, "cpu", ["observation", "next_observation"])
|
||||
|
||||
|
||||
def test_add_transition(replay_buffer, dummy_state, dummy_action):
|
||||
replay_buffer.add(dummy_state, dummy_action, 1.0, dummy_state, False, False)
|
||||
assert len(replay_buffer) == 1, "Replay buffer should have one transition after adding."
|
||||
assert torch.equal(replay_buffer.actions[0], dummy_action), (
|
||||
"Action should be equal to the first transition."
|
||||
)
|
||||
assert replay_buffer.rewards[0] == 1.0, "Reward should be equal to the first transition."
|
||||
assert not replay_buffer.dones[0], "Done should be False for the first transition."
|
||||
assert not replay_buffer.truncateds[0], "Truncated should be False for the first transition."
|
||||
|
||||
for dim in state_dims():
|
||||
assert torch.equal(replay_buffer.states[dim][0], dummy_state[dim]), (
|
||||
"Observation should be equal to the first transition."
|
||||
)
|
||||
assert torch.equal(replay_buffer.next_states[dim][0], dummy_state[dim]), (
|
||||
"Next observation should be equal to the first transition."
|
||||
)
|
||||
|
||||
|
||||
def test_add_over_capacity():
|
||||
replay_buffer = ReplayBuffer(2, "cpu", ["observation", "next_observation"])
|
||||
dummy_state_1 = create_dummy_state()
|
||||
dummy_action_1 = create_dummy_action()
|
||||
|
||||
dummy_state_2 = create_dummy_state()
|
||||
dummy_action_2 = create_dummy_action()
|
||||
|
||||
dummy_state_3 = create_dummy_state()
|
||||
dummy_action_3 = create_dummy_action()
|
||||
|
||||
replay_buffer.add(dummy_state_1, dummy_action_1, 1.0, dummy_state_1, False, False)
|
||||
replay_buffer.add(dummy_state_2, dummy_action_2, 1.0, dummy_state_2, False, False)
|
||||
replay_buffer.add(dummy_state_3, dummy_action_3, 1.0, dummy_state_3, True, True)
|
||||
|
||||
assert len(replay_buffer) == 2, "Replay buffer should have 2 transitions after adding 3."
|
||||
|
||||
for dim in state_dims():
|
||||
assert torch.equal(replay_buffer.states[dim][0], dummy_state_3[dim]), (
|
||||
"Observation should be equal to the first transition."
|
||||
)
|
||||
assert torch.equal(replay_buffer.next_states[dim][0], dummy_state_3[dim]), (
|
||||
"Next observation should be equal to the first transition."
|
||||
)
|
||||
|
||||
assert torch.equal(replay_buffer.actions[0], dummy_action_3), (
|
||||
"Action should be equal to the last transition."
|
||||
)
|
||||
assert replay_buffer.rewards[0] == 1.0, "Reward should be equal to the last transition."
|
||||
assert replay_buffer.dones[0], "Done should be True for the first transition."
|
||||
assert replay_buffer.truncateds[0], "Truncated should be True for the first transition."
|
||||
|
||||
|
||||
def test_sample_from_empty_buffer(replay_buffer):
|
||||
with pytest.raises(RuntimeError, match="Cannot sample from an empty buffer"):
|
||||
replay_buffer.sample(1)
|
||||
|
||||
|
||||
def test_sample_with_1_transition(replay_buffer, dummy_state, next_dummy_state, dummy_action):
|
||||
replay_buffer.add(dummy_state, dummy_action, 1.0, next_dummy_state, False, False)
|
||||
got_batch_transition = replay_buffer.sample(1)
|
||||
|
||||
expected_batch_transition = BatchTransition(
|
||||
state=clone_state(dummy_state),
|
||||
action=dummy_action.clone(),
|
||||
reward=1.0,
|
||||
next_state=clone_state(next_dummy_state),
|
||||
done=False,
|
||||
truncated=False,
|
||||
)
|
||||
|
||||
for buffer_property in dict_properties():
|
||||
for k, v in expected_batch_transition[buffer_property].items():
|
||||
got_state = got_batch_transition[buffer_property][k]
|
||||
|
||||
assert got_state.shape[0] == 1, f"{k} should have 1 transition."
|
||||
assert got_state.device.type == "cpu", f"{k} should be on cpu."
|
||||
|
||||
assert torch.equal(got_state[0], v), f"{k} should be equal to the expected batch transition."
|
||||
|
||||
for key, _value in expected_batch_transition.items():
|
||||
if key in dict_properties():
|
||||
continue
|
||||
|
||||
got_value = got_batch_transition[key]
|
||||
|
||||
v_tensor = expected_batch_transition[key]
|
||||
if not isinstance(v_tensor, torch.Tensor):
|
||||
v_tensor = torch.tensor(v_tensor)
|
||||
|
||||
assert got_value.shape[0] == 1, f"{key} should have 1 transition."
|
||||
assert got_value.device.type == "cpu", f"{key} should be on cpu."
|
||||
assert torch.equal(got_value[0], v_tensor), f"{key} should be equal to the expected batch transition."
|
||||
|
||||
|
||||
def test_sample_with_batch_bigger_than_buffer_size(
|
||||
replay_buffer, dummy_state, next_dummy_state, dummy_action
|
||||
):
|
||||
replay_buffer.add(dummy_state, dummy_action, 1.0, next_dummy_state, False, False)
|
||||
got_batch_transition = replay_buffer.sample(10)
|
||||
|
||||
expected_batch_transition = BatchTransition(
|
||||
state=dummy_state,
|
||||
action=dummy_action,
|
||||
reward=1.0,
|
||||
next_state=next_dummy_state,
|
||||
done=False,
|
||||
truncated=False,
|
||||
)
|
||||
|
||||
for buffer_property in dict_properties():
|
||||
for k in expected_batch_transition[buffer_property]:
|
||||
got_state = got_batch_transition[buffer_property][k]
|
||||
|
||||
assert got_state.shape[0] == 1, f"{k} should have 1 transition."
|
||||
|
||||
for key in expected_batch_transition:
|
||||
if key in dict_properties():
|
||||
continue
|
||||
|
||||
got_value = got_batch_transition[key]
|
||||
assert got_value.shape[0] == 1, f"{key} should have 1 transition."
|
||||
|
||||
|
||||
def test_sample_batch(replay_buffer):
|
||||
dummy_state_1 = create_dummy_state()
|
||||
dummy_action_1 = create_dummy_action()
|
||||
|
||||
dummy_state_2 = create_dummy_state()
|
||||
dummy_action_2 = create_dummy_action()
|
||||
|
||||
dummy_state_3 = create_dummy_state()
|
||||
dummy_action_3 = create_dummy_action()
|
||||
|
||||
dummy_state_4 = create_dummy_state()
|
||||
dummy_action_4 = create_dummy_action()
|
||||
|
||||
replay_buffer.add(dummy_state_1, dummy_action_1, 1.0, dummy_state_1, False, False)
|
||||
replay_buffer.add(dummy_state_2, dummy_action_2, 2.0, dummy_state_2, False, False)
|
||||
replay_buffer.add(dummy_state_3, dummy_action_3, 3.0, dummy_state_3, True, True)
|
||||
replay_buffer.add(dummy_state_4, dummy_action_4, 4.0, dummy_state_4, True, True)
|
||||
|
||||
dummy_states = [dummy_state_1, dummy_state_2, dummy_state_3, dummy_state_4]
|
||||
dummy_actions = [dummy_action_1, dummy_action_2, dummy_action_3, dummy_action_4]
|
||||
|
||||
got_batch_transition = replay_buffer.sample(3)
|
||||
|
||||
for buffer_property in dict_properties():
|
||||
for k in got_batch_transition[buffer_property]:
|
||||
got_state = got_batch_transition[buffer_property][k]
|
||||
|
||||
assert got_state.shape[0] == 3, f"{k} should have 3 transition."
|
||||
|
||||
for got_state_item in got_state:
|
||||
assert any(torch.equal(got_state_item, dummy_state[k]) for dummy_state in dummy_states), (
|
||||
f"{k} should be equal to one of the dummy states."
|
||||
)
|
||||
|
||||
for got_action_item in got_batch_transition["action"]:
|
||||
assert any(torch.equal(got_action_item, dummy_action) for dummy_action in dummy_actions), (
|
||||
"Actions should be equal to the dummy actions."
|
||||
)
|
||||
|
||||
for k in got_batch_transition:
|
||||
if k in dict_properties() or k == "complementary_info":
|
||||
continue
|
||||
|
||||
got_value = got_batch_transition[k]
|
||||
assert got_value.shape[0] == 3, f"{k} should have 3 transition."
|
||||
|
||||
|
||||
def test_to_lerobot_dataset_with_empty_buffer(replay_buffer):
|
||||
with pytest.raises(ValueError, match="The replay buffer is empty. Cannot convert to a dataset."):
|
||||
replay_buffer.to_lerobot_dataset("dummy_repo")
|
||||
|
||||
|
||||
def test_to_lerobot_dataset(tmp_path):
|
||||
ds, buffer = create_dataset_from_replay_buffer(tmp_path)
|
||||
|
||||
assert len(ds) == len(buffer), "Dataset should have the same size as the Replay Buffer"
|
||||
assert ds.fps == 1, "FPS should be 1"
|
||||
assert ds.repo_id == "dummy/repo", "The dataset should have `dummy/repo` repo id"
|
||||
|
||||
for dim in state_dims():
|
||||
assert dim in ds.features
|
||||
assert ds.features[dim]["shape"] == buffer.states[dim][0].shape
|
||||
|
||||
assert ds.num_episodes == 2
|
||||
assert ds.num_frames == 4
|
||||
|
||||
for j, value in enumerate(ds):
|
||||
print(torch.equal(value["observation.image"], buffer.next_states["observation.image"][j]))
|
||||
|
||||
for i in range(len(ds)):
|
||||
for feature, value in ds[i].items():
|
||||
if feature == "action":
|
||||
assert torch.equal(value, buffer.actions[i])
|
||||
elif feature == "next.reward":
|
||||
assert torch.equal(value, buffer.rewards[i])
|
||||
elif feature == "next.done":
|
||||
assert torch.equal(value, buffer.dones[i])
|
||||
elif feature == "observation.image":
|
||||
# Tenssor -> numpy is not precise, so we have some diff there
|
||||
# TODO: Check and fix it
|
||||
torch.testing.assert_close(value, buffer.states["observation.image"][i], rtol=0.3, atol=0.003)
|
||||
elif feature == "observation.state":
|
||||
assert torch.equal(value, buffer.states["observation.state"][i])
|
||||
|
||||
|
||||
def test_from_lerobot_dataset(tmp_path):
|
||||
dummy_state_1 = create_dummy_state()
|
||||
dummy_action_1 = create_dummy_action()
|
||||
|
||||
dummy_state_2 = create_dummy_state()
|
||||
dummy_action_2 = create_dummy_action()
|
||||
|
||||
dummy_state_3 = create_dummy_state()
|
||||
dummy_action_3 = create_dummy_action()
|
||||
|
||||
dummy_state_4 = create_dummy_state()
|
||||
dummy_action_4 = create_dummy_action()
|
||||
|
||||
replay_buffer = create_empty_replay_buffer()
|
||||
replay_buffer.add(dummy_state_1, dummy_action_1, 1.0, dummy_state_1, False, False)
|
||||
replay_buffer.add(dummy_state_2, dummy_action_2, 1.0, dummy_state_2, False, False)
|
||||
replay_buffer.add(dummy_state_3, dummy_action_3, 1.0, dummy_state_3, True, True)
|
||||
replay_buffer.add(dummy_state_4, dummy_action_4, 1.0, dummy_state_4, True, True)
|
||||
|
||||
root = tmp_path / "test"
|
||||
ds = replay_buffer.to_lerobot_dataset(DUMMY_REPO_ID, root=root)
|
||||
|
||||
reconverted_buffer = ReplayBuffer.from_lerobot_dataset(
|
||||
ds, state_keys=list(state_dims()), device="cpu", capacity=replay_buffer.capacity, use_drq=False
|
||||
)
|
||||
|
||||
# Check only the part of the buffer that's actually filled with data
|
||||
assert torch.equal(
|
||||
reconverted_buffer.actions[: len(replay_buffer)],
|
||||
replay_buffer.actions[: len(replay_buffer)],
|
||||
), "Actions from converted buffer should be equal to the original replay buffer."
|
||||
assert torch.equal(
|
||||
reconverted_buffer.rewards[: len(replay_buffer)], replay_buffer.rewards[: len(replay_buffer)]
|
||||
), "Rewards from converted buffer should be equal to the original replay buffer."
|
||||
assert torch.equal(
|
||||
reconverted_buffer.dones[: len(replay_buffer)], replay_buffer.dones[: len(replay_buffer)]
|
||||
), "Dones from converted buffer should be equal to the original replay buffer."
|
||||
|
||||
# Lerobot DS haven't supported truncateds yet
|
||||
expected_truncateds = torch.zeros(len(replay_buffer)).bool()
|
||||
assert torch.equal(reconverted_buffer.truncateds[: len(replay_buffer)], expected_truncateds), (
|
||||
"Truncateds from converted buffer should be equal False"
|
||||
)
|
||||
|
||||
assert torch.equal(
|
||||
replay_buffer.states["observation.state"][: len(replay_buffer)],
|
||||
reconverted_buffer.states["observation.state"][: len(replay_buffer)],
|
||||
), "State should be the same after converting to dataset and return back"
|
||||
|
||||
for i in range(4):
|
||||
torch.testing.assert_close(
|
||||
replay_buffer.states["observation.image"][i],
|
||||
reconverted_buffer.states["observation.image"][i],
|
||||
rtol=0.4,
|
||||
atol=0.004,
|
||||
)
|
||||
|
||||
# The 2, 3 frames have done flag, so their values will be equal to the current state
|
||||
for i in range(2):
|
||||
# In the current implementation we take the next state from the `states` and ignore `next_states`
|
||||
next_index = (i + 1) % 4
|
||||
|
||||
torch.testing.assert_close(
|
||||
replay_buffer.states["observation.image"][next_index],
|
||||
reconverted_buffer.next_states["observation.image"][i],
|
||||
rtol=0.4,
|
||||
atol=0.004,
|
||||
)
|
||||
|
||||
for i in range(2, 4):
|
||||
assert torch.equal(
|
||||
replay_buffer.states["observation.state"][i],
|
||||
reconverted_buffer.next_states["observation.state"][i],
|
||||
)
|
||||
|
||||
|
||||
def test_buffer_sample_alignment():
|
||||
# Initialize buffer
|
||||
buffer = ReplayBuffer(capacity=100, device="cpu", state_keys=["state_value"], storage_device="cpu")
|
||||
|
||||
# Fill buffer with patterned data
|
||||
for i in range(100):
|
||||
signature = float(i) / 100.0
|
||||
state = {"state_value": torch.tensor([[signature]]).float()}
|
||||
action = torch.tensor([[2.0 * signature]]).float()
|
||||
reward = 3.0 * signature
|
||||
|
||||
is_end = (i + 1) % 10 == 0
|
||||
if is_end:
|
||||
next_state = {"state_value": torch.tensor([[signature]]).float()}
|
||||
done = True
|
||||
else:
|
||||
next_signature = float(i + 1) / 100.0
|
||||
next_state = {"state_value": torch.tensor([[next_signature]]).float()}
|
||||
done = False
|
||||
|
||||
buffer.add(state, action, reward, next_state, done, False)
|
||||
|
||||
# Sample and verify
|
||||
batch = buffer.sample(50)
|
||||
|
||||
for i in range(50):
|
||||
state_sig = batch["state"]["state_value"][i].item()
|
||||
action_val = batch["action"][i].item()
|
||||
reward_val = batch["reward"][i].item()
|
||||
next_state_sig = batch["next_state"]["state_value"][i].item()
|
||||
is_done = batch["done"][i].item() > 0.5
|
||||
|
||||
# Verify relationships
|
||||
assert abs(action_val - 2.0 * state_sig) < 1e-4, (
|
||||
f"Action {action_val} should be 2x state signature {state_sig}"
|
||||
)
|
||||
|
||||
assert abs(reward_val - 3.0 * state_sig) < 1e-4, (
|
||||
f"Reward {reward_val} should be 3x state signature {state_sig}"
|
||||
)
|
||||
|
||||
if is_done:
|
||||
assert abs(next_state_sig - state_sig) < 1e-4, (
|
||||
f"For done states, next_state {next_state_sig} should equal state {state_sig}"
|
||||
)
|
||||
else:
|
||||
# Either it's the next sequential state (+0.01) or same state (for episode boundaries)
|
||||
valid_next = (
|
||||
abs(next_state_sig - state_sig - 0.01) < 1e-4 or abs(next_state_sig - state_sig) < 1e-4
|
||||
)
|
||||
assert valid_next, (
|
||||
f"Next state {next_state_sig} should be either state+0.01 or same as state {state_sig}"
|
||||
)
|
||||
|
||||
|
||||
def test_memory_optimization():
|
||||
dummy_state_1 = create_dummy_state()
|
||||
dummy_action_1 = create_dummy_action()
|
||||
|
||||
dummy_state_2 = create_dummy_state()
|
||||
dummy_action_2 = create_dummy_action()
|
||||
|
||||
dummy_state_3 = create_dummy_state()
|
||||
dummy_action_3 = create_dummy_action()
|
||||
|
||||
dummy_state_4 = create_dummy_state()
|
||||
dummy_action_4 = create_dummy_action()
|
||||
|
||||
replay_buffer = create_empty_replay_buffer()
|
||||
replay_buffer.add(dummy_state_1, dummy_action_1, 1.0, dummy_state_2, False, False)
|
||||
replay_buffer.add(dummy_state_2, dummy_action_2, 1.0, dummy_state_3, False, False)
|
||||
replay_buffer.add(dummy_state_3, dummy_action_3, 1.0, dummy_state_4, False, False)
|
||||
replay_buffer.add(dummy_state_4, dummy_action_4, 1.0, dummy_state_4, True, True)
|
||||
|
||||
optimized_replay_buffer = create_empty_replay_buffer(True)
|
||||
optimized_replay_buffer.add(dummy_state_1, dummy_action_1, 1.0, dummy_state_2, False, False)
|
||||
optimized_replay_buffer.add(dummy_state_2, dummy_action_2, 1.0, dummy_state_3, False, False)
|
||||
optimized_replay_buffer.add(dummy_state_3, dummy_action_3, 1.0, dummy_state_4, False, False)
|
||||
optimized_replay_buffer.add(dummy_state_4, dummy_action_4, 1.0, None, True, True)
|
||||
|
||||
assert get_object_memory(optimized_replay_buffer) < get_object_memory(replay_buffer), (
|
||||
"Optimized replay buffer should be smaller than the original replay buffer"
|
||||
)
|
||||
|
||||
|
||||
def test_check_image_augmentations_with_drq_and_dummy_image_augmentation_function(dummy_state, dummy_action):
|
||||
def dummy_image_augmentation_function(x):
|
||||
return torch.ones_like(x) * 10
|
||||
|
||||
replay_buffer = create_empty_replay_buffer(
|
||||
use_drq=True, image_augmentation_function=dummy_image_augmentation_function
|
||||
)
|
||||
|
||||
replay_buffer.add(dummy_state, dummy_action, 1.0, dummy_state, False, False)
|
||||
|
||||
sampled_transitions = replay_buffer.sample(1)
|
||||
assert torch.all(sampled_transitions["state"]["observation.image"] == 10), (
|
||||
"Image augmentations should be applied"
|
||||
)
|
||||
assert torch.all(sampled_transitions["next_state"]["observation.image"] == 10), (
|
||||
"Image augmentations should be applied"
|
||||
)
|
||||
|
||||
|
||||
def test_check_image_augmentations_with_drq_and_default_image_augmentation_function(
|
||||
dummy_state, dummy_action
|
||||
):
|
||||
replay_buffer = create_empty_replay_buffer(use_drq=True)
|
||||
|
||||
replay_buffer.add(dummy_state, dummy_action, 1.0, dummy_state, False, False)
|
||||
|
||||
# Let's check that it doesn't fail and shapes are correct
|
||||
sampled_transitions = replay_buffer.sample(1)
|
||||
assert sampled_transitions["state"]["observation.image"].shape == (1, 3, 84, 84)
|
||||
assert sampled_transitions["next_state"]["observation.image"].shape == (1, 3, 84, 84)
|
||||
|
||||
|
||||
def test_random_crop_vectorized_basic():
|
||||
# Create a batch of 2 images with known patterns
|
||||
batch_size, channels, height, width = 2, 3, 10, 8
|
||||
images = torch.zeros((batch_size, channels, height, width))
|
||||
|
||||
# Fill with unique values for testing
|
||||
for b in range(batch_size):
|
||||
images[b] = b + 1
|
||||
|
||||
crop_size = (6, 4) # Smaller than original
|
||||
cropped = random_crop_vectorized(images, crop_size)
|
||||
|
||||
# Check output shape
|
||||
assert cropped.shape == (batch_size, channels, *crop_size)
|
||||
|
||||
# Check that values are preserved (should be either 1s or 2s for respective batches)
|
||||
assert torch.all(cropped[0] == 1)
|
||||
assert torch.all(cropped[1] == 2)
|
||||
|
||||
|
||||
def test_random_crop_vectorized_invalid_size():
|
||||
images = torch.zeros((2, 3, 10, 8))
|
||||
|
||||
# Test crop size larger than image
|
||||
with pytest.raises(ValueError, match="Requested crop size .* is bigger than the image size"):
|
||||
random_crop_vectorized(images, (12, 8))
|
||||
|
||||
with pytest.raises(ValueError, match="Requested crop size .* is bigger than the image size"):
|
||||
random_crop_vectorized(images, (10, 10))
|
||||
|
||||
|
||||
def _populate_buffer_for_async_test(capacity: int = 10) -> ReplayBuffer:
|
||||
"""Create a small buffer with deterministic 3×128×128 images and 11-D state."""
|
||||
buffer = ReplayBuffer(
|
||||
capacity=capacity,
|
||||
device="cpu",
|
||||
state_keys=["observation.image", "observation.state"],
|
||||
storage_device="cpu",
|
||||
)
|
||||
|
||||
for i in range(capacity):
|
||||
img = torch.ones(3, 128, 128) * i
|
||||
state_vec = torch.arange(11).float() + i
|
||||
state = {
|
||||
"observation.image": img,
|
||||
"observation.state": state_vec,
|
||||
}
|
||||
buffer.add(
|
||||
state=state,
|
||||
action=torch.tensor([0.0]),
|
||||
reward=0.0,
|
||||
next_state=state,
|
||||
done=False,
|
||||
truncated=False,
|
||||
)
|
||||
return buffer
|
||||
|
||||
|
||||
def test_async_iterator_shapes_basic():
|
||||
buffer = _populate_buffer_for_async_test()
|
||||
batch_size = 2
|
||||
iterator = buffer.get_iterator(batch_size=batch_size, async_prefetch=True, queue_size=1)
|
||||
batch = next(iterator)
|
||||
|
||||
images = batch["state"]["observation.image"]
|
||||
states = batch["state"]["observation.state"]
|
||||
|
||||
assert images.shape == (batch_size, 3, 128, 128)
|
||||
assert states.shape == (batch_size, 11)
|
||||
|
||||
next_images = batch["next_state"]["observation.image"]
|
||||
next_states = batch["next_state"]["observation.state"]
|
||||
|
||||
assert next_images.shape == (batch_size, 3, 128, 128)
|
||||
assert next_states.shape == (batch_size, 11)
|
||||
|
||||
|
||||
def test_async_iterator_multiple_iterations():
|
||||
buffer = _populate_buffer_for_async_test()
|
||||
batch_size = 2
|
||||
iterator = buffer.get_iterator(batch_size=batch_size, async_prefetch=True, queue_size=2)
|
||||
|
||||
for _ in range(5):
|
||||
batch = next(iterator)
|
||||
images = batch["state"]["observation.image"]
|
||||
states = batch["state"]["observation.state"]
|
||||
assert images.shape == (batch_size, 3, 128, 128)
|
||||
assert states.shape == (batch_size, 11)
|
||||
|
||||
next_images = batch["next_state"]["observation.image"]
|
||||
next_states = batch["next_state"]["observation.state"]
|
||||
assert next_images.shape == (batch_size, 3, 128, 128)
|
||||
assert next_states.shape == (batch_size, 11)
|
||||
|
||||
# Ensure iterator can be disposed without blocking
|
||||
del iterator
|
||||
Reference in New Issue
Block a user