mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-12 04:21:45 +00:00
Merge branch 'main' into envs/support-more-args
Signed-off-by: Steven Palma <imstevenpmwork@ieee.org>
This commit is contained in:
Steven Palma
@@ -144,12 +144,18 @@ def test_async_inference_e2e(monkeypatch):
|
||||
client = RobotClient(client_config)
|
||||
assert client.start(), "Client failed initial handshake with the server"
|
||||
|
||||
# Track action chunks received without modifying RobotClient
|
||||
action_chunks_received = {"count": 0}
|
||||
# Track action chunks received and verify device type
|
||||
action_chunks_received = {"count": 0, "actions_on_cpu": True}
|
||||
original_aggregate = client._aggregate_action_queues
|
||||
|
||||
def counting_aggregate(*args, **kwargs):
|
||||
action_chunks_received["count"] += 1
|
||||
# Check that all received actions are on CPU
|
||||
if args:
|
||||
for timed_action in args[0]: # args[0] is the list of TimedAction
|
||||
action_tensor = timed_action.get_action()
|
||||
if action_tensor.device.type != "cpu":
|
||||
action_chunks_received["actions_on_cpu"] = False
|
||||
return original_aggregate(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(client, "_aggregate_action_queues", counting_aggregate)
|
||||
|
||||
@@ -62,7 +62,7 @@ class MockPolicy:
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def policy_server():
|
||||
"""Fresh `PolicyServer` instance with a stubbed-out policy model."""
|
||||
# Import only when the test actually runs (after decorator check)
|
||||
|
||||
+127
-57
@@ -20,7 +20,9 @@
|
||||
# ```
|
||||
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
@@ -28,6 +30,50 @@ from lerobot.cameras.configs import Cv2Rotation
|
||||
from lerobot.cameras.opencv import OpenCVCamera, OpenCVCameraConfig
|
||||
from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
|
||||
|
||||
RealVideoCapture = cv2.VideoCapture
|
||||
|
||||
|
||||
class MockLoopingVideoCapture:
|
||||
"""
|
||||
Wraps the real OpenCV VideoCapture.
|
||||
Motivation: cv2.VideoCapture(file.png) is only valid for one read.
|
||||
Strategy: Read the file once & return the cached frame for subsequent reads.
|
||||
Consequence: No recurrent I/O operations, but we keep the test artifacts simple.
|
||||
"""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
args_clean = [str(a) if isinstance(a, Path) else a for a in args]
|
||||
self._real_vc = RealVideoCapture(*args_clean, **kwargs)
|
||||
self._cached_frame = None
|
||||
|
||||
def read(self):
|
||||
ret, frame = self._real_vc.read()
|
||||
|
||||
if ret:
|
||||
self._cached_frame = frame
|
||||
return ret, frame
|
||||
|
||||
if not ret and self._cached_frame is not None:
|
||||
return True, self._cached_frame.copy()
|
||||
|
||||
return ret, frame
|
||||
|
||||
def __getattr__(self, name):
|
||||
return getattr(self._real_vc, name)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def patch_opencv_videocapture():
|
||||
"""
|
||||
Automatically patches cv2.VideoCapture for all tests.
|
||||
"""
|
||||
module_path = OpenCVCamera.__module__
|
||||
target = f"{module_path}.cv2.VideoCapture"
|
||||
|
||||
with patch(target, new=MockLoopingVideoCapture):
|
||||
yield
|
||||
|
||||
|
||||
# NOTE(Steven): more tests + assertions?
|
||||
TEST_ARTIFACTS_DIR = Path(__file__).parent.parent / "artifacts" / "cameras"
|
||||
DEFAULT_PNG_FILE_PATH = TEST_ARTIFACTS_DIR / "image_160x120.png"
|
||||
@@ -43,25 +89,22 @@ def test_abc_implementation():
|
||||
|
||||
|
||||
def test_connect():
|
||||
config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH)
|
||||
camera = OpenCVCamera(config)
|
||||
config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH, warmup_s=0)
|
||||
|
||||
camera.connect(warmup=False)
|
||||
|
||||
assert camera.is_connected
|
||||
with OpenCVCamera(config) as camera:
|
||||
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)
|
||||
config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH, warmup_s=0)
|
||||
|
||||
with pytest.raises(DeviceAlreadyConnectedError):
|
||||
camera.connect(warmup=False)
|
||||
with OpenCVCamera(config) as camera, pytest.raises(DeviceAlreadyConnectedError):
|
||||
camera.connect()
|
||||
|
||||
|
||||
def test_connect_invalid_camera_path():
|
||||
config = OpenCVCameraConfig(index_or_path="nonexistent/camera.png")
|
||||
|
||||
camera = OpenCVCamera(config)
|
||||
|
||||
with pytest.raises(ConnectionError):
|
||||
@@ -74,27 +117,25 @@ def test_invalid_width_connect():
|
||||
width=99999, # Invalid width to trigger error
|
||||
height=480,
|
||||
)
|
||||
camera = OpenCVCamera(config)
|
||||
|
||||
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)
|
||||
config = OpenCVCameraConfig(index_or_path=index_or_path, warmup_s=0)
|
||||
|
||||
img = camera.read()
|
||||
|
||||
assert isinstance(img, np.ndarray)
|
||||
with OpenCVCamera(config) as camera:
|
||||
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)
|
||||
|
||||
camera = OpenCVCamera(config)
|
||||
with pytest.raises(DeviceNotConnectedError):
|
||||
_ = camera.read()
|
||||
|
||||
@@ -119,32 +160,22 @@ def test_disconnect_before_connect():
|
||||
|
||||
@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)
|
||||
config = OpenCVCameraConfig(index_or_path=index_or_path, warmup_s=0)
|
||||
|
||||
try:
|
||||
with OpenCVCamera(config) as camera:
|
||||
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)
|
||||
config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH, warmup_s=0)
|
||||
|
||||
try:
|
||||
with pytest.raises(TimeoutError):
|
||||
camera.async_read(timeout_ms=0)
|
||||
finally:
|
||||
if camera.is_connected:
|
||||
camera.disconnect()
|
||||
with OpenCVCamera(config) as camera, pytest.raises(TimeoutError):
|
||||
camera.async_read(timeout_ms=0) # consumes any available frame by then
|
||||
camera.async_read(timeout_ms=0) # request immediately another one
|
||||
|
||||
|
||||
def test_async_read_before_connect():
|
||||
@@ -155,6 +186,50 @@ def test_async_read_before_connect():
|
||||
_ = camera.async_read()
|
||||
|
||||
|
||||
def test_read_latest():
|
||||
config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH, warmup_s=0)
|
||||
|
||||
with OpenCVCamera(config) as camera:
|
||||
# ensure at least one fresh frame is captured
|
||||
frame = camera.read()
|
||||
latest = camera.read_latest()
|
||||
|
||||
assert isinstance(latest, np.ndarray)
|
||||
assert latest.shape == frame.shape
|
||||
|
||||
|
||||
def test_read_latest_before_connect():
|
||||
config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH)
|
||||
|
||||
camera = OpenCVCamera(config)
|
||||
with pytest.raises(DeviceNotConnectedError):
|
||||
_ = camera.read_latest()
|
||||
|
||||
|
||||
def test_read_latest_high_frequency():
|
||||
config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH, warmup_s=0)
|
||||
|
||||
with OpenCVCamera(config) as camera:
|
||||
# prime to ensure frames are available
|
||||
ref = camera.read()
|
||||
|
||||
for _ in range(20):
|
||||
latest = camera.read_latest()
|
||||
assert isinstance(latest, np.ndarray)
|
||||
assert latest.shape == ref.shape
|
||||
|
||||
|
||||
def test_read_latest_too_old():
|
||||
config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH, warmup_s=0)
|
||||
|
||||
with OpenCVCamera(config) as camera:
|
||||
# prime to ensure frames are available
|
||||
_ = camera.read()
|
||||
|
||||
with pytest.raises(TimeoutError):
|
||||
_ = camera.read_latest(max_age_ms=0) # immediately too old
|
||||
|
||||
|
||||
def test_fourcc_configuration():
|
||||
"""Test FourCC configuration validation and application."""
|
||||
|
||||
@@ -181,18 +256,15 @@ def test_fourcc_configuration():
|
||||
|
||||
def test_fourcc_with_camera():
|
||||
"""Test FourCC functionality with actual camera connection."""
|
||||
config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH, fourcc="MJPG")
|
||||
camera = OpenCVCamera(config)
|
||||
config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH, fourcc="MJPG", warmup_s=0)
|
||||
|
||||
# Connect should work with MJPG specified
|
||||
camera.connect(warmup=False)
|
||||
assert camera.is_connected
|
||||
with OpenCVCamera(config) as camera:
|
||||
assert camera.is_connected
|
||||
|
||||
# Read should work normally
|
||||
img = camera.read()
|
||||
assert isinstance(img, np.ndarray)
|
||||
|
||||
camera.disconnect()
|
||||
# Read should work normally
|
||||
img = camera.read()
|
||||
assert isinstance(img, np.ndarray)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("index_or_path", TEST_IMAGE_PATHS, ids=TEST_IMAGE_SIZES)
|
||||
@@ -211,18 +283,16 @@ def test_rotation(rotation, index_or_path):
|
||||
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)
|
||||
config = OpenCVCameraConfig(index_or_path=index_or_path, rotation=rotation, warmup_s=0)
|
||||
with OpenCVCamera(config) as camera:
|
||||
img = camera.read()
|
||||
assert isinstance(img, np.ndarray)
|
||||
|
||||
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)
|
||||
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)
|
||||
|
||||
@@ -20,6 +20,8 @@ from unittest.mock import MagicMock, patch
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("reachy2_sdk")
|
||||
|
||||
from lerobot.cameras.reachy2_camera import Reachy2Camera, Reachy2CameraConfig
|
||||
from lerobot.utils.errors import DeviceNotConnectedError
|
||||
|
||||
@@ -127,24 +129,12 @@ def test_async_read(camera):
|
||||
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()
|
||||
|
||||
|
||||
def test_async_read_timeout(camera):
|
||||
camera.connect()
|
||||
try:
|
||||
with pytest.raises(TimeoutError):
|
||||
camera.async_read(timeout_ms=0)
|
||||
finally:
|
||||
if camera.is_connected:
|
||||
camera.disconnect()
|
||||
|
||||
|
||||
def test_read_before_connect(camera):
|
||||
with pytest.raises(DeviceNotConnectedError):
|
||||
_ = camera.read()
|
||||
@@ -160,6 +150,44 @@ def test_async_read_before_connect(camera):
|
||||
_ = camera.async_read()
|
||||
|
||||
|
||||
def test_read_latest(camera):
|
||||
camera.connect()
|
||||
|
||||
frame = camera.read()
|
||||
latest = camera.read_latest()
|
||||
|
||||
assert isinstance(latest, np.ndarray)
|
||||
assert latest.shape == frame.shape
|
||||
|
||||
|
||||
def test_read_latest_before_connect(camera):
|
||||
# camera fixture yields an unconnected camera instance
|
||||
with pytest.raises(DeviceNotConnectedError):
|
||||
_ = camera.read_latest()
|
||||
|
||||
|
||||
def test_read_latest_high_frequency(camera):
|
||||
camera.connect()
|
||||
|
||||
# prime to ensure frames are available
|
||||
ref = camera.read()
|
||||
|
||||
for _ in range(20):
|
||||
latest = camera.read_latest()
|
||||
assert isinstance(latest, np.ndarray)
|
||||
assert latest.shape == ref.shape
|
||||
|
||||
|
||||
def test_read_latest_too_old(camera):
|
||||
camera.connect()
|
||||
|
||||
# prime to ensure frames are available
|
||||
_ = camera.read()
|
||||
|
||||
with pytest.raises(TimeoutError):
|
||||
_ = camera.read_latest(max_age_ms=0) # immediately too old
|
||||
|
||||
|
||||
def test_wrong_camera_name():
|
||||
with pytest.raises(ValueError):
|
||||
_ = Reachy2CameraConfig(name="wrong-name", image_type="left")
|
||||
|
||||
@@ -62,19 +62,15 @@ def test_abc_implementation():
|
||||
|
||||
|
||||
def test_connect():
|
||||
config = RealSenseCameraConfig(serial_number_or_name="042")
|
||||
camera = RealSenseCamera(config)
|
||||
config = RealSenseCameraConfig(serial_number_or_name="042", warmup_s=0)
|
||||
|
||||
camera.connect(warmup=False)
|
||||
assert camera.is_connected
|
||||
with RealSenseCamera(config) as camera:
|
||||
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):
|
||||
config = RealSenseCameraConfig(serial_number_or_name="042", warmup_s=0)
|
||||
with RealSenseCamera(config) as camera, pytest.raises(DeviceAlreadyConnectedError):
|
||||
camera.connect(warmup=False)
|
||||
|
||||
|
||||
@@ -96,12 +92,10 @@ def test_invalid_width_connect():
|
||||
|
||||
|
||||
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)
|
||||
config = RealSenseCameraConfig(serial_number_or_name="042", width=640, height=480, fps=30, warmup_s=0)
|
||||
with RealSenseCamera(config) as camera:
|
||||
img = camera.read()
|
||||
assert isinstance(img, np.ndarray)
|
||||
|
||||
|
||||
# TODO(Steven): Fix this test for the latest version of pyrealsense2.
|
||||
@@ -142,32 +136,21 @@ def test_disconnect_before_connect():
|
||||
|
||||
|
||||
def test_async_read():
|
||||
config = RealSenseCameraConfig(serial_number_or_name="042", width=640, height=480, fps=30)
|
||||
camera = RealSenseCamera(config)
|
||||
camera.connect(warmup=False)
|
||||
config = RealSenseCameraConfig(serial_number_or_name="042", width=640, height=480, fps=30, warmup_s=0)
|
||||
|
||||
try:
|
||||
with RealSenseCamera(config) as camera:
|
||||
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()
|
||||
config = RealSenseCameraConfig(serial_number_or_name="042", width=640, height=480, fps=30, warmup_s=0)
|
||||
with RealSenseCamera(config) as camera, pytest.raises(TimeoutError):
|
||||
camera.async_read(timeout_ms=0) # consumes any available frame by then
|
||||
camera.async_read(timeout_ms=0) # request immediately another one
|
||||
|
||||
|
||||
def test_async_read_before_connect():
|
||||
@@ -178,6 +161,47 @@ def test_async_read_before_connect():
|
||||
_ = camera.async_read()
|
||||
|
||||
|
||||
def test_read_latest():
|
||||
config = RealSenseCameraConfig(serial_number_or_name="042", width=640, height=480, fps=30, warmup_s=0)
|
||||
with RealSenseCamera(config) as camera:
|
||||
img = camera.read()
|
||||
latest = camera.read_latest()
|
||||
|
||||
assert isinstance(latest, np.ndarray)
|
||||
assert latest.shape == img.shape
|
||||
|
||||
|
||||
def test_read_latest_high_frequency():
|
||||
config = RealSenseCameraConfig(serial_number_or_name="042", width=640, height=480, fps=30, warmup_s=0)
|
||||
with RealSenseCamera(config) as camera:
|
||||
# prime with one read to ensure frames are available
|
||||
ref = camera.read()
|
||||
|
||||
for _ in range(20):
|
||||
latest = camera.read_latest()
|
||||
assert isinstance(latest, np.ndarray)
|
||||
assert latest.shape == ref.shape
|
||||
|
||||
|
||||
def test_read_latest_before_connect():
|
||||
config = RealSenseCameraConfig(serial_number_or_name="042")
|
||||
camera = RealSenseCamera(config)
|
||||
|
||||
with pytest.raises(DeviceNotConnectedError):
|
||||
_ = camera.read_latest()
|
||||
|
||||
|
||||
def test_read_latest_too_old():
|
||||
config = RealSenseCameraConfig(serial_number_or_name="042")
|
||||
|
||||
with RealSenseCamera(config) as camera:
|
||||
# prime to ensure frames are available
|
||||
_ = camera.read()
|
||||
|
||||
with pytest.raises(TimeoutError):
|
||||
_ = camera.read_latest(max_age_ms=0) # immediately too old
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"rotation",
|
||||
[
|
||||
@@ -189,18 +213,16 @@ def test_async_read_before_connect():
|
||||
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)
|
||||
config = RealSenseCameraConfig(serial_number_or_name="042", rotation=rotation, warmup_s=0)
|
||||
with RealSenseCamera(config) as camera:
|
||||
img = camera.read()
|
||||
assert isinstance(img, np.ndarray)
|
||||
|
||||
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)
|
||||
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)
|
||||
|
||||
@@ -28,7 +28,6 @@ pytest_plugins = [
|
||||
"tests.fixtures.files",
|
||||
"tests.fixtures.hub",
|
||||
"tests.fixtures.optimizers",
|
||||
"tests.plugins.reachy2_sdk",
|
||||
]
|
||||
|
||||
|
||||
|
||||
@@ -16,6 +16,7 @@
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import datasets
|
||||
import torch
|
||||
|
||||
from lerobot.datasets.aggregate import aggregate_datasets
|
||||
@@ -380,3 +381,236 @@ def test_video_timestamps_regression(tmp_path, lerobot_dataset_factory):
|
||||
for key in aggr_ds.meta.video_keys:
|
||||
assert key in item, f"Video key {key} missing from item {i}"
|
||||
assert item[key].shape[0] == 3, f"Expected 3 channels for video key {key}"
|
||||
|
||||
|
||||
def assert_image_schema_preserved(aggr_ds):
|
||||
"""Test that HuggingFace Image feature schema is preserved in aggregated parquet files.
|
||||
|
||||
This verifies the fix for a bug where image columns were written with a generic
|
||||
struct schema {'bytes': Value('binary'), 'path': Value('string')} instead of
|
||||
the proper Image() feature type, causing HuggingFace Hub viewer to display
|
||||
raw dict objects instead of image thumbnails.
|
||||
"""
|
||||
image_keys = aggr_ds.meta.image_keys
|
||||
if not image_keys:
|
||||
return
|
||||
|
||||
# Check that parquet files have proper Image schema
|
||||
data_dir = aggr_ds.root / "data"
|
||||
parquet_files = list(data_dir.rglob("*.parquet"))
|
||||
assert len(parquet_files) > 0, "No parquet files found in aggregated dataset"
|
||||
|
||||
for parquet_file in parquet_files:
|
||||
# Load with HuggingFace datasets to check schema
|
||||
ds = datasets.Dataset.from_parquet(str(parquet_file))
|
||||
|
||||
for image_key in image_keys:
|
||||
feature = ds.features.get(image_key)
|
||||
assert feature is not None, f"Image key '{image_key}' not found in parquet schema"
|
||||
assert isinstance(feature, datasets.Image), (
|
||||
f"Image key '{image_key}' should have Image() feature type, "
|
||||
f"but got {type(feature).__name__}: {feature}. "
|
||||
"This indicates image schema was not preserved during aggregation."
|
||||
)
|
||||
|
||||
|
||||
def assert_image_frames_integrity(aggr_ds, ds_0, ds_1):
|
||||
"""Test that image frames are correctly preserved after aggregation."""
|
||||
image_keys = aggr_ds.meta.image_keys
|
||||
if not image_keys:
|
||||
return
|
||||
|
||||
def images_equal(img1, img2):
|
||||
return torch.allclose(img1, img2)
|
||||
|
||||
# Test the section corresponding to the first dataset (ds_0)
|
||||
for i in range(len(ds_0)):
|
||||
assert aggr_ds[i]["index"] == i, (
|
||||
f"Frame index at position {i} should be {i}, but got {aggr_ds[i]['index']}"
|
||||
)
|
||||
for key in image_keys:
|
||||
assert images_equal(aggr_ds[i][key], ds_0[i][key]), (
|
||||
f"Image frames at position {i} should be equal between aggregated and ds_0"
|
||||
)
|
||||
|
||||
# Test the section corresponding to the second dataset (ds_1)
|
||||
for i in range(len(ds_0), len(ds_0) + len(ds_1)):
|
||||
assert aggr_ds[i]["index"] == i, (
|
||||
f"Frame index at position {i} should be {i}, but got {aggr_ds[i]['index']}"
|
||||
)
|
||||
for key in image_keys:
|
||||
assert images_equal(aggr_ds[i][key], ds_1[i - len(ds_0)][key]), (
|
||||
f"Image frames at position {i} should be equal between aggregated and ds_1"
|
||||
)
|
||||
|
||||
|
||||
def test_aggregate_image_datasets(tmp_path, lerobot_dataset_factory):
|
||||
"""Test aggregation of image-based datasets preserves HuggingFace Image schema.
|
||||
|
||||
This test specifically verifies that:
|
||||
1. Image-based datasets can be aggregated correctly
|
||||
2. The HuggingFace Image() feature type is preserved in parquet files
|
||||
3. Image data integrity is maintained across aggregation
|
||||
4. Images can be properly decoded after aggregation
|
||||
|
||||
This catches the bug where to_parquet_with_hf_images() was not passing
|
||||
the features schema, causing image columns to be written as generic
|
||||
struct types instead of Image() types.
|
||||
"""
|
||||
ds_0_num_frames = 50
|
||||
ds_1_num_frames = 75
|
||||
ds_0_num_episodes = 2
|
||||
ds_1_num_episodes = 3
|
||||
|
||||
# Create two image-based datasets (use_videos=False)
|
||||
ds_0 = lerobot_dataset_factory(
|
||||
root=tmp_path / "image_0",
|
||||
repo_id=f"{DUMMY_REPO_ID}_image_0",
|
||||
total_episodes=ds_0_num_episodes,
|
||||
total_frames=ds_0_num_frames,
|
||||
use_videos=False, # Image-based dataset
|
||||
)
|
||||
ds_1 = lerobot_dataset_factory(
|
||||
root=tmp_path / "image_1",
|
||||
repo_id=f"{DUMMY_REPO_ID}_image_1",
|
||||
total_episodes=ds_1_num_episodes,
|
||||
total_frames=ds_1_num_frames,
|
||||
use_videos=False, # Image-based dataset
|
||||
)
|
||||
|
||||
# Verify source datasets have image keys
|
||||
assert len(ds_0.meta.image_keys) > 0, "ds_0 should have image keys"
|
||||
assert len(ds_1.meta.image_keys) > 0, "ds_1 should have image keys"
|
||||
|
||||
# Aggregate the datasets
|
||||
aggregate_datasets(
|
||||
repo_ids=[ds_0.repo_id, ds_1.repo_id],
|
||||
roots=[ds_0.root, ds_1.root],
|
||||
aggr_repo_id=f"{DUMMY_REPO_ID}_image_aggr",
|
||||
aggr_root=tmp_path / "image_aggr",
|
||||
)
|
||||
|
||||
# Load the aggregated dataset
|
||||
with (
|
||||
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
|
||||
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
|
||||
):
|
||||
mock_get_safe_version.return_value = "v3.0"
|
||||
mock_snapshot_download.return_value = str(tmp_path / "image_aggr")
|
||||
aggr_ds = LeRobotDataset(f"{DUMMY_REPO_ID}_image_aggr", root=tmp_path / "image_aggr")
|
||||
|
||||
# Verify aggregated dataset has image keys
|
||||
assert len(aggr_ds.meta.image_keys) > 0, "Aggregated dataset should have image keys"
|
||||
assert aggr_ds.meta.image_keys == ds_0.meta.image_keys, "Image keys should match source datasets"
|
||||
|
||||
# Run standard aggregation assertions
|
||||
expected_total_episodes = ds_0_num_episodes + ds_1_num_episodes
|
||||
expected_total_frames = ds_0_num_frames + ds_1_num_frames
|
||||
|
||||
assert_episode_and_frame_counts(aggr_ds, expected_total_episodes, expected_total_frames)
|
||||
assert_dataset_content_integrity(aggr_ds, ds_0, ds_1)
|
||||
assert_metadata_consistency(aggr_ds, ds_0, ds_1)
|
||||
assert_episode_indices_updated_correctly(aggr_ds, ds_0, ds_1)
|
||||
|
||||
# Image-specific assertions
|
||||
assert_image_schema_preserved(aggr_ds)
|
||||
assert_image_frames_integrity(aggr_ds, ds_0, ds_1)
|
||||
|
||||
# Verify images can be accessed and have correct shape
|
||||
sample_item = aggr_ds[0]
|
||||
for image_key in aggr_ds.meta.image_keys:
|
||||
img = sample_item[image_key]
|
||||
assert isinstance(img, torch.Tensor), f"Image {image_key} should be a tensor"
|
||||
assert img.dim() == 3, f"Image {image_key} should have 3 dimensions (C, H, W)"
|
||||
assert img.shape[0] == 3, f"Image {image_key} should have 3 channels"
|
||||
|
||||
assert_dataset_iteration_works(aggr_ds)
|
||||
|
||||
|
||||
def test_aggregate_already_merged_dataset(tmp_path, lerobot_dataset_factory):
|
||||
"""Regression test for aggregating a dataset that is itself a result of a previous merge.
|
||||
|
||||
This test reproduces the bug where merging datasets with multiple parquet files
|
||||
(e.g., from a previous merge with file rotation) would cause FileNotFoundError
|
||||
because metadata file indices were incorrectly preserved instead of being mapped
|
||||
to their actual destination files.
|
||||
|
||||
The fix adds src_to_dst tracking in aggregate_data() to correctly map source
|
||||
file indices to destination file indices.
|
||||
"""
|
||||
# Step 1: Create datasets A and B
|
||||
ds_a = lerobot_dataset_factory(
|
||||
root=tmp_path / "ds_a",
|
||||
repo_id=f"{DUMMY_REPO_ID}_a",
|
||||
total_episodes=4,
|
||||
total_frames=200,
|
||||
)
|
||||
ds_b = lerobot_dataset_factory(
|
||||
root=tmp_path / "ds_b",
|
||||
repo_id=f"{DUMMY_REPO_ID}_b",
|
||||
total_episodes=4,
|
||||
total_frames=200,
|
||||
)
|
||||
|
||||
# Step 2: Merge A+B into AB with small file size to force multiple files
|
||||
aggregate_datasets(
|
||||
repo_ids=[ds_a.repo_id, ds_b.repo_id],
|
||||
roots=[ds_a.root, ds_b.root],
|
||||
aggr_repo_id=f"{DUMMY_REPO_ID}_ab",
|
||||
aggr_root=tmp_path / "ds_ab",
|
||||
data_files_size_in_mb=0.01, # Force file rotation
|
||||
)
|
||||
|
||||
with (
|
||||
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
|
||||
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
|
||||
):
|
||||
mock_get_safe_version.return_value = "v3.0"
|
||||
mock_snapshot_download.return_value = str(tmp_path / "ds_ab")
|
||||
ds_ab = LeRobotDataset(f"{DUMMY_REPO_ID}_ab", root=tmp_path / "ds_ab")
|
||||
|
||||
# Verify AB has multiple data files (file rotation occurred)
|
||||
ab_data_files = list((tmp_path / "ds_ab" / "data").rglob("*.parquet"))
|
||||
assert len(ab_data_files) > 1, "First merge should create multiple parquet files"
|
||||
|
||||
# Step 3: Create dataset C
|
||||
ds_c = lerobot_dataset_factory(
|
||||
root=tmp_path / "ds_c",
|
||||
repo_id=f"{DUMMY_REPO_ID}_c",
|
||||
total_episodes=2,
|
||||
total_frames=100,
|
||||
)
|
||||
|
||||
# Step 4: Merge AB+C into final - THIS IS WHERE THE BUG OCCURRED
|
||||
aggregate_datasets(
|
||||
repo_ids=[ds_ab.repo_id, ds_c.repo_id],
|
||||
roots=[ds_ab.root, ds_c.root],
|
||||
aggr_repo_id=f"{DUMMY_REPO_ID}_abc",
|
||||
aggr_root=tmp_path / "ds_abc",
|
||||
)
|
||||
|
||||
with (
|
||||
patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version,
|
||||
patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download,
|
||||
):
|
||||
mock_get_safe_version.return_value = "v3.0"
|
||||
mock_snapshot_download.return_value = str(tmp_path / "ds_abc")
|
||||
ds_abc = LeRobotDataset(f"{DUMMY_REPO_ID}_abc", root=tmp_path / "ds_abc")
|
||||
|
||||
# Step 5: Verify all data files referenced in metadata actually exist
|
||||
for ep_idx in range(ds_abc.num_episodes):
|
||||
data_file_path = ds_abc.root / ds_abc.meta.get_data_file_path(ep_idx)
|
||||
assert data_file_path.exists(), (
|
||||
f"Episode {ep_idx} references non-existent file: {data_file_path}\n"
|
||||
"This indicates the src_to_dst mapping fix is not working correctly."
|
||||
)
|
||||
|
||||
# Step 6: Verify we can iterate through the entire dataset without FileNotFoundError
|
||||
expected_episodes = ds_a.num_episodes + ds_b.num_episodes + ds_c.num_episodes
|
||||
expected_frames = ds_a.num_frames + ds_b.num_frames + ds_c.num_frames
|
||||
|
||||
assert ds_abc.num_episodes == expected_episodes
|
||||
assert ds_abc.num_frames == expected_frames
|
||||
|
||||
# This would raise FileNotFoundError before the fix
|
||||
assert_dataset_iteration_works(ds_abc)
|
||||
|
||||
@@ -26,10 +26,11 @@ from lerobot.datasets.dataset_tools import (
|
||||
delete_episodes,
|
||||
merge_datasets,
|
||||
modify_features,
|
||||
modify_tasks,
|
||||
remove_feature,
|
||||
split_dataset,
|
||||
)
|
||||
from lerobot.scripts.lerobot_edit_dataset import convert_dataset_to_videos
|
||||
from lerobot.scripts.lerobot_edit_dataset import convert_image_to_video_dataset
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
@@ -1050,7 +1051,175 @@ def test_modify_features_preserves_file_structure(sample_dataset, tmp_path):
|
||||
assert "reward" in modified_dataset.meta.features
|
||||
|
||||
|
||||
def test_convert_dataset_to_videos(tmp_path):
|
||||
def test_modify_tasks_single_task_for_all(sample_dataset):
|
||||
"""Test setting a single task for all episodes."""
|
||||
new_task = "Pick up the cube and place it"
|
||||
|
||||
modified_dataset = modify_tasks(sample_dataset, new_task=new_task)
|
||||
|
||||
# Verify all episodes have the new task
|
||||
assert len(modified_dataset.meta.tasks) == 1
|
||||
assert new_task in modified_dataset.meta.tasks.index
|
||||
|
||||
# Verify task_index is 0 for all frames (only one task)
|
||||
for i in range(len(modified_dataset)):
|
||||
item = modified_dataset[i]
|
||||
assert item["task_index"].item() == 0
|
||||
assert item["task"] == new_task
|
||||
|
||||
|
||||
def test_modify_tasks_episode_specific(sample_dataset):
|
||||
"""Test setting different tasks for specific episodes."""
|
||||
episode_tasks = {
|
||||
0: "Task A",
|
||||
1: "Task B",
|
||||
2: "Task A",
|
||||
3: "Task C",
|
||||
4: "Task B",
|
||||
}
|
||||
|
||||
modified_dataset = modify_tasks(sample_dataset, episode_tasks=episode_tasks)
|
||||
|
||||
# Verify correct number of unique tasks
|
||||
unique_tasks = set(episode_tasks.values())
|
||||
assert len(modified_dataset.meta.tasks) == len(unique_tasks)
|
||||
|
||||
# Verify each episode has the correct task
|
||||
for ep_idx, expected_task in episode_tasks.items():
|
||||
ep_data = modified_dataset.meta.episodes[ep_idx]
|
||||
assert ep_data["tasks"][0] == expected_task
|
||||
|
||||
|
||||
def test_modify_tasks_default_with_overrides(sample_dataset):
|
||||
"""Test setting a default task with specific overrides."""
|
||||
default_task = "Default task"
|
||||
override_task = "Special task"
|
||||
episode_tasks = {2: override_task, 4: override_task}
|
||||
|
||||
modified_dataset = modify_tasks(
|
||||
sample_dataset,
|
||||
new_task=default_task,
|
||||
episode_tasks=episode_tasks,
|
||||
)
|
||||
|
||||
# Verify correct number of unique tasks
|
||||
assert len(modified_dataset.meta.tasks) == 2
|
||||
assert default_task in modified_dataset.meta.tasks.index
|
||||
assert override_task in modified_dataset.meta.tasks.index
|
||||
|
||||
# Verify episodes have correct tasks
|
||||
for ep_idx in range(5):
|
||||
ep_data = modified_dataset.meta.episodes[ep_idx]
|
||||
if ep_idx in episode_tasks:
|
||||
assert ep_data["tasks"][0] == override_task
|
||||
else:
|
||||
assert ep_data["tasks"][0] == default_task
|
||||
|
||||
|
||||
def test_modify_tasks_no_task_specified(sample_dataset):
|
||||
"""Test error when no task is specified."""
|
||||
with pytest.raises(ValueError, match="Must specify at least one of new_task or episode_tasks"):
|
||||
modify_tasks(sample_dataset)
|
||||
|
||||
|
||||
def test_modify_tasks_invalid_episode_indices(sample_dataset):
|
||||
"""Test error with invalid episode indices."""
|
||||
with pytest.raises(ValueError, match="Invalid episode indices"):
|
||||
modify_tasks(sample_dataset, episode_tasks={10: "Task", 20: "Task"})
|
||||
|
||||
|
||||
def test_modify_tasks_updates_info_json(sample_dataset):
|
||||
"""Test that total_tasks is updated in info.json."""
|
||||
episode_tasks = {0: "Task A", 1: "Task B", 2: "Task C", 3: "Task A", 4: "Task B"}
|
||||
|
||||
modified_dataset = modify_tasks(sample_dataset, episode_tasks=episode_tasks)
|
||||
|
||||
# Verify total_tasks is updated
|
||||
assert modified_dataset.meta.total_tasks == 3
|
||||
|
||||
|
||||
def test_modify_tasks_preserves_other_metadata(sample_dataset):
|
||||
"""Test that modifying tasks preserves other metadata."""
|
||||
original_frames = sample_dataset.meta.total_frames
|
||||
original_episodes = sample_dataset.meta.total_episodes
|
||||
original_fps = sample_dataset.meta.fps
|
||||
|
||||
modified_dataset = modify_tasks(sample_dataset, new_task="New task")
|
||||
|
||||
# Verify other metadata is preserved
|
||||
assert modified_dataset.meta.total_frames == original_frames
|
||||
assert modified_dataset.meta.total_episodes == original_episodes
|
||||
assert modified_dataset.meta.fps == original_fps
|
||||
|
||||
|
||||
def test_modify_tasks_task_index_correct(sample_dataset):
|
||||
"""Test that task_index values are correct in data files."""
|
||||
# Create tasks that will have predictable indices (sorted alphabetically)
|
||||
episode_tasks = {
|
||||
0: "Alpha task", # Will be index 0
|
||||
1: "Beta task", # Will be index 1
|
||||
2: "Alpha task", # Will be index 0
|
||||
3: "Gamma task", # Will be index 2
|
||||
4: "Beta task", # Will be index 1
|
||||
}
|
||||
|
||||
modified_dataset = modify_tasks(sample_dataset, episode_tasks=episode_tasks)
|
||||
|
||||
# Verify task indices are correct
|
||||
task_to_expected_idx = {
|
||||
"Alpha task": 0,
|
||||
"Beta task": 1,
|
||||
"Gamma task": 2,
|
||||
}
|
||||
|
||||
for i in range(len(modified_dataset)):
|
||||
item = modified_dataset[i]
|
||||
ep_idx = item["episode_index"].item()
|
||||
expected_task = episode_tasks[ep_idx]
|
||||
expected_idx = task_to_expected_idx[expected_task]
|
||||
assert item["task_index"].item() == expected_idx
|
||||
assert item["task"] == expected_task
|
||||
|
||||
|
||||
def test_modify_tasks_in_place(sample_dataset):
|
||||
"""Test that modify_tasks modifies the dataset in-place."""
|
||||
original_root = sample_dataset.root
|
||||
|
||||
modified_dataset = modify_tasks(sample_dataset, new_task="New task")
|
||||
|
||||
# Verify same instance is returned and root is unchanged
|
||||
assert modified_dataset is sample_dataset
|
||||
assert modified_dataset.root == original_root
|
||||
|
||||
|
||||
def test_modify_tasks_keeps_original_when_not_overridden(sample_dataset):
|
||||
"""Test that original tasks are kept when using episode_tasks without new_task."""
|
||||
from lerobot.datasets.utils import load_episodes
|
||||
|
||||
# Ensure episodes metadata is loaded
|
||||
if sample_dataset.meta.episodes is None:
|
||||
sample_dataset.meta.episodes = load_episodes(sample_dataset.meta.root)
|
||||
|
||||
# Get original tasks for episodes not being overridden
|
||||
original_task_ep0 = sample_dataset.meta.episodes[0]["tasks"][0]
|
||||
original_task_ep1 = sample_dataset.meta.episodes[1]["tasks"][0]
|
||||
|
||||
# Only override episodes 2, 3, 4
|
||||
episode_tasks = {2: "New Task A", 3: "New Task B", 4: "New Task A"}
|
||||
|
||||
modified_dataset = modify_tasks(sample_dataset, episode_tasks=episode_tasks)
|
||||
|
||||
# Verify original tasks are kept for episodes 0 and 1
|
||||
assert modified_dataset.meta.episodes[0]["tasks"][0] == original_task_ep0
|
||||
assert modified_dataset.meta.episodes[1]["tasks"][0] == original_task_ep1
|
||||
|
||||
# Verify new tasks for overridden episodes
|
||||
assert modified_dataset.meta.episodes[2]["tasks"][0] == "New Task A"
|
||||
assert modified_dataset.meta.episodes[3]["tasks"][0] == "New Task B"
|
||||
assert modified_dataset.meta.episodes[4]["tasks"][0] == "New Task A"
|
||||
|
||||
|
||||
def test_convert_image_to_video_dataset(tmp_path):
|
||||
"""Test converting lerobot/pusht_image dataset to video format."""
|
||||
from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
||||
|
||||
@@ -1071,7 +1240,7 @@ def test_convert_dataset_to_videos(tmp_path):
|
||||
assert "observation.image" in source_dataset.meta.features
|
||||
|
||||
# Convert to video dataset (only first 2 episodes for speed)
|
||||
video_dataset = convert_dataset_to_videos(
|
||||
video_dataset = convert_image_to_video_dataset(
|
||||
dataset=source_dataset,
|
||||
output_dir=output_dir,
|
||||
repo_id="lerobot/pusht_video",
|
||||
@@ -1113,7 +1282,7 @@ def test_convert_dataset_to_videos(tmp_path):
|
||||
shutil.rmtree(output_dir)
|
||||
|
||||
|
||||
def test_convert_dataset_to_videos_subset_episodes(tmp_path):
|
||||
def test_convert_image_to_video_dataset_subset_episodes(tmp_path):
|
||||
"""Test converting only specific episodes from lerobot/pusht_image to video format."""
|
||||
from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
||||
|
||||
@@ -1132,7 +1301,7 @@ def test_convert_dataset_to_videos_subset_episodes(tmp_path):
|
||||
# Convert only episode 0 to video (subset of loaded episodes)
|
||||
episode_indices = [0]
|
||||
|
||||
video_dataset = convert_dataset_to_videos(
|
||||
video_dataset = convert_image_to_video_dataset(
|
||||
dataset=source_dataset,
|
||||
output_dir=output_dir,
|
||||
repo_id="lerobot/pusht_video_subset",
|
||||
|
||||
@@ -33,6 +33,7 @@ from lerobot.datasets.image_writer import image_array_to_pil_image
|
||||
from lerobot.datasets.lerobot_dataset import (
|
||||
LeRobotDataset,
|
||||
MultiLeRobotDataset,
|
||||
_encode_video_worker,
|
||||
)
|
||||
from lerobot.datasets.utils import (
|
||||
DEFAULT_CHUNK_SIZE,
|
||||
@@ -43,6 +44,7 @@ from lerobot.datasets.utils import (
|
||||
hf_transform_to_torch,
|
||||
hw_to_dataset_features,
|
||||
)
|
||||
from lerobot.datasets.video_utils import VALID_VIDEO_CODECS
|
||||
from lerobot.envs.factory import make_env_config
|
||||
from lerobot.policies.factory import make_policy_config
|
||||
from lerobot.robots import make_robot_from_config
|
||||
@@ -350,6 +352,65 @@ def test_image_array_to_pil_image_wrong_range_float_0_255():
|
||||
image_array_to_pil_image(image)
|
||||
|
||||
|
||||
def test_tmp_image_deletion(tmp_path, empty_lerobot_dataset_factory):
|
||||
"""Verify temporary image directories are removed for image features after saving episode."""
|
||||
# Image feature: images should be deleted after saving episode
|
||||
image_key = "image"
|
||||
features_image = {
|
||||
image_key: {"dtype": "image", "shape": DUMMY_CHW, "names": ["channels", "height", "width"]}
|
||||
}
|
||||
ds_img = empty_lerobot_dataset_factory(root=tmp_path / "img", features=features_image)
|
||||
ds_img.add_frame({"image": np.random.rand(*DUMMY_CHW), "task": "Dummy task"})
|
||||
ds_img.save_episode()
|
||||
img_dir = ds_img._get_image_file_dir(0, image_key)
|
||||
assert not img_dir.exists(), "Temporary image directory should be removed for image features"
|
||||
|
||||
|
||||
def test_tmp_video_deletion(tmp_path, empty_lerobot_dataset_factory):
|
||||
"""Verify temporary image directories are removed for video encoding when `batch_encoding_size == 1`."""
|
||||
# Video feature: when batch_encoding_size == 1 temporary images should be deleted
|
||||
vid_key = "video"
|
||||
features_video = {
|
||||
vid_key: {"dtype": "video", "shape": DUMMY_CHW, "names": ["channels", "height", "width"]}
|
||||
}
|
||||
|
||||
ds_vid = empty_lerobot_dataset_factory(root=tmp_path / "vid", features=features_video)
|
||||
ds_vid.batch_encoding_size = 1
|
||||
ds_vid.add_frame({vid_key: np.random.rand(*DUMMY_CHW), "task": "Dummy task"})
|
||||
ds_vid.save_episode()
|
||||
vid_img_dir = ds_vid._get_image_file_dir(0, vid_key)
|
||||
assert not vid_img_dir.exists(), (
|
||||
"Temporary image directory should be removed when batch_encoding_size == 1"
|
||||
)
|
||||
|
||||
|
||||
def test_tmp_mixed_deletion(tmp_path, empty_lerobot_dataset_factory):
|
||||
"""Verify temporary image directories are removed appropriately when both image and video features are present."""
|
||||
image_key = "image"
|
||||
vid_key = "video"
|
||||
features_mixed = {
|
||||
image_key: {"dtype": "image", "shape": DUMMY_CHW, "names": ["channels", "height", "width"]},
|
||||
vid_key: {"dtype": "video", "shape": DUMMY_HWC, "names": ["height", "width", "channels"]},
|
||||
}
|
||||
ds_mixed = empty_lerobot_dataset_factory(
|
||||
root=tmp_path / "mixed", features=features_mixed, batch_encoding_size=2, streaming_encoding=False
|
||||
)
|
||||
ds_mixed.add_frame(
|
||||
{
|
||||
"image": np.random.rand(*DUMMY_CHW),
|
||||
"video": np.random.rand(*DUMMY_HWC),
|
||||
"task": "Dummy task",
|
||||
}
|
||||
)
|
||||
ds_mixed.save_episode()
|
||||
img_dir = ds_mixed._get_image_file_dir(0, image_key)
|
||||
vid_img_dir = ds_mixed._get_image_file_dir(0, vid_key)
|
||||
assert not img_dir.exists(), "Temporary image directory should be removed for image features"
|
||||
assert vid_img_dir.exists(), (
|
||||
"Temporary image directory should not be removed for video features when batch_encoding_size == 2"
|
||||
)
|
||||
|
||||
|
||||
# TODO(aliberts):
|
||||
# - [ ] test various attributes & state from init and create
|
||||
# - [ ] test init with episodes and check num_frames
|
||||
@@ -1292,3 +1353,303 @@ def test_frames_in_current_file_calculation(tmp_path, empty_lerobot_dataset_fact
|
||||
frame = loaded_dataset[idx]
|
||||
expected_ep = idx // frames_per_episode
|
||||
assert frame["episode_index"].item() == expected_ep
|
||||
|
||||
|
||||
def test_encode_video_worker_forwards_vcodec(tmp_path):
|
||||
"""Test that _encode_video_worker correctly forwards the vcodec parameter to encode_video_frames."""
|
||||
from unittest.mock import patch
|
||||
|
||||
from lerobot.datasets.utils import DEFAULT_IMAGE_PATH
|
||||
|
||||
# Create the expected directory structure
|
||||
video_key = "observation.images.laptop"
|
||||
episode_index = 0
|
||||
frame_index = 0
|
||||
|
||||
fpath = DEFAULT_IMAGE_PATH.format(
|
||||
image_key=video_key, episode_index=episode_index, frame_index=frame_index
|
||||
)
|
||||
img_dir = tmp_path / Path(fpath).parent
|
||||
img_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Create a dummy image file
|
||||
dummy_img = Image.new("RGB", (64, 64), color="red")
|
||||
dummy_img.save(img_dir / "frame-000000.png")
|
||||
|
||||
# Track what vcodec was passed to encode_video_frames
|
||||
captured_kwargs = {}
|
||||
|
||||
def mock_encode_video_frames(imgs_dir, video_path, fps, **kwargs):
|
||||
captured_kwargs.update(kwargs)
|
||||
# Create a dummy output file so the worker doesn't fail
|
||||
Path(video_path).parent.mkdir(parents=True, exist_ok=True)
|
||||
Path(video_path).touch()
|
||||
|
||||
with patch("lerobot.datasets.lerobot_dataset.encode_video_frames", side_effect=mock_encode_video_frames):
|
||||
# Test with h264 codec
|
||||
_encode_video_worker(video_key, episode_index, tmp_path, fps=30, vcodec="h264")
|
||||
|
||||
assert "vcodec" in captured_kwargs
|
||||
assert captured_kwargs["vcodec"] == "h264"
|
||||
|
||||
|
||||
def test_encode_video_worker_default_vcodec(tmp_path):
|
||||
"""Test that _encode_video_worker uses libsvtav1 as the default codec."""
|
||||
from unittest.mock import patch
|
||||
|
||||
from lerobot.datasets.utils import DEFAULT_IMAGE_PATH
|
||||
|
||||
# Create the expected directory structure
|
||||
video_key = "observation.images.laptop"
|
||||
episode_index = 0
|
||||
frame_index = 0
|
||||
|
||||
fpath = DEFAULT_IMAGE_PATH.format(
|
||||
image_key=video_key, episode_index=episode_index, frame_index=frame_index
|
||||
)
|
||||
img_dir = tmp_path / Path(fpath).parent
|
||||
img_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Create a dummy image file
|
||||
dummy_img = Image.new("RGB", (64, 64), color="red")
|
||||
dummy_img.save(img_dir / "frame-000000.png")
|
||||
|
||||
# Track what vcodec was passed to encode_video_frames
|
||||
captured_kwargs = {}
|
||||
|
||||
def mock_encode_video_frames(imgs_dir, video_path, fps, **kwargs):
|
||||
captured_kwargs.update(kwargs)
|
||||
# Create a dummy output file so the worker doesn't fail
|
||||
Path(video_path).parent.mkdir(parents=True, exist_ok=True)
|
||||
Path(video_path).touch()
|
||||
|
||||
with patch("lerobot.datasets.lerobot_dataset.encode_video_frames", side_effect=mock_encode_video_frames):
|
||||
# Test with default codec (no vcodec specified)
|
||||
_encode_video_worker(video_key, episode_index, tmp_path, fps=30)
|
||||
|
||||
assert "vcodec" in captured_kwargs
|
||||
assert captured_kwargs["vcodec"] == "libsvtav1"
|
||||
|
||||
|
||||
def test_lerobot_dataset_vcodec_validation():
|
||||
"""Test that LeRobotDataset validates the vcodec parameter."""
|
||||
# Test that invalid vcodec raises ValueError
|
||||
with pytest.raises(ValueError, match="Invalid vcodec"):
|
||||
LeRobotDataset.__new__(LeRobotDataset) # bypass __init__ to test validation directly
|
||||
# Actually test via create since it's easier
|
||||
LeRobotDataset.create(
|
||||
repo_id="test/invalid_codec",
|
||||
fps=30,
|
||||
features={"observation.state": {"dtype": "float32", "shape": (2,), "names": ["x", "y"]}},
|
||||
vcodec="invalid_codec",
|
||||
)
|
||||
|
||||
|
||||
def test_valid_video_codecs_constant():
|
||||
"""Test that VALID_VIDEO_CODECS contains the expected codecs."""
|
||||
assert "h264" in VALID_VIDEO_CODECS
|
||||
assert "hevc" in VALID_VIDEO_CODECS
|
||||
assert "libsvtav1" in VALID_VIDEO_CODECS
|
||||
assert "auto" in VALID_VIDEO_CODECS
|
||||
assert "h264_videotoolbox" in VALID_VIDEO_CODECS
|
||||
assert "h264_nvenc" in VALID_VIDEO_CODECS
|
||||
assert len(VALID_VIDEO_CODECS) == 10
|
||||
|
||||
|
||||
def test_delta_timestamps_with_episodes_filter(tmp_path, empty_lerobot_dataset_factory):
|
||||
"""Regression test for bug where delta_timestamps incorrectly marked all frames as padded when using episodes filter.
|
||||
|
||||
The bug occurred because _get_query_indices was using the relative index (idx) in the filtered dataset
|
||||
instead of the absolute index when comparing against episode boundaries (ep_start, ep_end).
|
||||
"""
|
||||
features = {
|
||||
"observation.state": {"dtype": "float32", "shape": (2,), "names": ["x", "y"]},
|
||||
"action": {"dtype": "float32", "shape": (2,), "names": ["vx", "vy"]},
|
||||
}
|
||||
|
||||
dataset = empty_lerobot_dataset_factory(root=tmp_path / "test", features=features, use_videos=False)
|
||||
|
||||
# Create 3 episodes with 10 frames each
|
||||
frames_per_episode = 10
|
||||
for ep_idx in range(3):
|
||||
for frame_idx in range(frames_per_episode):
|
||||
dataset.add_frame(
|
||||
{
|
||||
"observation.state": torch.tensor([ep_idx, frame_idx], dtype=torch.float32),
|
||||
"action": torch.randn(2),
|
||||
"task": f"task_{ep_idx}",
|
||||
}
|
||||
)
|
||||
dataset.save_episode()
|
||||
dataset.finalize()
|
||||
|
||||
# Load only episode 1 (middle episode) with delta_timestamps
|
||||
delta_ts = {"observation.state": [0.0]} # Just the current frame
|
||||
filtered_dataset = LeRobotDataset(
|
||||
dataset.repo_id,
|
||||
root=dataset.root,
|
||||
episodes=[1],
|
||||
delta_timestamps=delta_ts,
|
||||
)
|
||||
|
||||
# Verify the filtered dataset has the correct length
|
||||
assert len(filtered_dataset) == frames_per_episode
|
||||
|
||||
# Check that no frames are marked as padded (since delta=0 should always be valid)
|
||||
for idx in range(len(filtered_dataset)):
|
||||
frame = filtered_dataset[idx]
|
||||
assert frame["observation.state_is_pad"].item() is False, f"Frame {idx} incorrectly marked as padded"
|
||||
# Verify we're getting data from episode 1
|
||||
assert frame["episode_index"].item() == 1
|
||||
|
||||
|
||||
def test_delta_timestamps_padding_at_episode_boundaries(tmp_path, empty_lerobot_dataset_factory):
|
||||
"""Test that delta_timestamps correctly marks padding at episode boundaries when using episodes filter."""
|
||||
features = {
|
||||
"observation.state": {"dtype": "float32", "shape": (2,), "names": ["x", "y"]},
|
||||
"action": {"dtype": "float32", "shape": (2,), "names": ["vx", "vy"]},
|
||||
}
|
||||
|
||||
dataset = empty_lerobot_dataset_factory(
|
||||
root=tmp_path / "test", features=features, use_videos=False, fps=10
|
||||
)
|
||||
|
||||
# Create 3 episodes with 5 frames each
|
||||
frames_per_episode = 5
|
||||
for ep_idx in range(3):
|
||||
for frame_idx in range(frames_per_episode):
|
||||
dataset.add_frame(
|
||||
{
|
||||
"observation.state": torch.tensor([ep_idx, frame_idx], dtype=torch.float32),
|
||||
"action": torch.randn(2),
|
||||
"task": f"task_{ep_idx}",
|
||||
}
|
||||
)
|
||||
dataset.save_episode()
|
||||
dataset.finalize()
|
||||
|
||||
# Load only episode 1 with delta_timestamps that go beyond episode boundaries
|
||||
# fps=10, so 0.1s = 1 frame offset
|
||||
delta_ts = {"observation.state": [-0.2, -0.1, 0.0, 0.1, 0.2]} # -2, -1, 0, +1, +2 frames
|
||||
filtered_dataset = LeRobotDataset(
|
||||
dataset.repo_id,
|
||||
root=dataset.root,
|
||||
episodes=[1],
|
||||
delta_timestamps=delta_ts,
|
||||
tolerance_s=0.04, # Slightly less than half a frame at 10fps
|
||||
)
|
||||
|
||||
assert len(filtered_dataset) == frames_per_episode
|
||||
|
||||
# Check padding at the start of the episode (first frame)
|
||||
first_frame = filtered_dataset[0]
|
||||
is_pad = first_frame["observation.state_is_pad"].tolist()
|
||||
# At frame 0 of episode 1: delta -2 and -1 should be padded, 0, +1, +2 should not
|
||||
assert is_pad == [True, True, False, False, False], f"First frame padding incorrect: {is_pad}"
|
||||
|
||||
# Check middle frame (no padding expected)
|
||||
mid_frame = filtered_dataset[2]
|
||||
is_pad = mid_frame["observation.state_is_pad"].tolist()
|
||||
assert is_pad == [False, False, False, False, False], f"Middle frame padding incorrect: {is_pad}"
|
||||
|
||||
# Check padding at the end of the episode (last frame)
|
||||
last_frame = filtered_dataset[4]
|
||||
is_pad = last_frame["observation.state_is_pad"].tolist()
|
||||
# At frame 4 of episode 1: delta -2, -1, 0 should not be padded, +1, +2 should be
|
||||
assert is_pad == [False, False, False, True, True], f"Last frame padding incorrect: {is_pad}"
|
||||
|
||||
|
||||
def test_delta_timestamps_multiple_episodes_filter(tmp_path, empty_lerobot_dataset_factory):
|
||||
"""Test delta_timestamps with multiple non-consecutive episodes selected."""
|
||||
features = {
|
||||
"observation.state": {"dtype": "float32", "shape": (2,), "names": ["x", "y"]},
|
||||
}
|
||||
|
||||
dataset = empty_lerobot_dataset_factory(
|
||||
root=tmp_path / "test", features=features, use_videos=False, fps=10
|
||||
)
|
||||
|
||||
# Create 5 episodes with 5 frames each
|
||||
frames_per_episode = 5
|
||||
for ep_idx in range(5):
|
||||
for frame_idx in range(frames_per_episode):
|
||||
dataset.add_frame(
|
||||
{
|
||||
"observation.state": torch.tensor([ep_idx, frame_idx], dtype=torch.float32),
|
||||
"task": f"task_{ep_idx}",
|
||||
}
|
||||
)
|
||||
dataset.save_episode()
|
||||
dataset.finalize()
|
||||
|
||||
# Load episodes 1 and 3 (non-consecutive)
|
||||
delta_ts = {"observation.state": [0.0]}
|
||||
filtered_dataset = LeRobotDataset(
|
||||
dataset.repo_id,
|
||||
root=dataset.root,
|
||||
episodes=[1, 3],
|
||||
delta_timestamps=delta_ts,
|
||||
)
|
||||
|
||||
assert len(filtered_dataset) == 2 * frames_per_episode
|
||||
|
||||
# All frames should have valid (non-padded) data for delta=0
|
||||
for idx in range(len(filtered_dataset)):
|
||||
frame = filtered_dataset[idx]
|
||||
assert frame["observation.state_is_pad"].item() is False
|
||||
|
||||
# Verify we're getting the correct episodes
|
||||
episode_indices = [filtered_dataset[i]["episode_index"].item() for i in range(len(filtered_dataset))]
|
||||
expected_episodes = [1] * frames_per_episode + [3] * frames_per_episode
|
||||
assert episode_indices == expected_episodes
|
||||
|
||||
|
||||
def test_delta_timestamps_query_returns_correct_values(tmp_path, empty_lerobot_dataset_factory):
|
||||
"""Test that delta_timestamps returns the correct observation values, not just correct padding."""
|
||||
features = {
|
||||
"observation.state": {"dtype": "float32", "shape": (1,), "names": ["x"]},
|
||||
}
|
||||
|
||||
dataset = empty_lerobot_dataset_factory(
|
||||
root=tmp_path / "test", features=features, use_videos=False, fps=10
|
||||
)
|
||||
|
||||
# Create 2 episodes with known values
|
||||
# Episode 0: frames with values 0, 1, 2, 3, 4
|
||||
# Episode 1: frames with values 10, 11, 12, 13, 14
|
||||
frames_per_episode = 5
|
||||
for ep_idx in range(2):
|
||||
for frame_idx in range(frames_per_episode):
|
||||
value = ep_idx * 10 + frame_idx
|
||||
dataset.add_frame(
|
||||
{
|
||||
"observation.state": torch.tensor([value], dtype=torch.float32),
|
||||
"task": f"task_{ep_idx}",
|
||||
}
|
||||
)
|
||||
dataset.save_episode()
|
||||
dataset.finalize()
|
||||
|
||||
# Load episode 1 with delta that looks at previous frame
|
||||
delta_ts = {"observation.state": [-0.1, 0.0]} # Previous frame and current frame
|
||||
filtered_dataset = LeRobotDataset(
|
||||
dataset.repo_id,
|
||||
root=dataset.root,
|
||||
episodes=[1],
|
||||
delta_timestamps=delta_ts,
|
||||
tolerance_s=0.04,
|
||||
)
|
||||
|
||||
# Check frame 2 of episode 1 (which has absolute index 7, value 12)
|
||||
frame = filtered_dataset[2]
|
||||
state_values = frame["observation.state"].tolist()
|
||||
# Should get [11, 12] - the previous and current values within episode 1
|
||||
assert state_values == [11.0, 12.0], f"Expected [11.0, 12.0], got {state_values}"
|
||||
|
||||
# Check first frame - previous frame should be clamped to episode start (padded)
|
||||
first_frame = filtered_dataset[0]
|
||||
state_values = first_frame["observation.state"].tolist()
|
||||
is_pad = first_frame["observation.state_is_pad"].tolist()
|
||||
# Previous frame is outside episode, so it's clamped to first frame and marked as padded
|
||||
assert state_values == [10.0, 10.0], f"Expected [10.0, 10.0], got {state_values}"
|
||||
assert is_pad == [True, False], f"Expected [True, False], got {is_pad}"
|
||||
|
||||
@@ -390,6 +390,30 @@ def test_sharpness_jitter_invalid_range_max_smaller():
|
||||
SharpnessJitter((2.0, 0.1))
|
||||
|
||||
|
||||
def test_make_transform_from_config_with_v2_resize(img_tensor_factory):
|
||||
img_tensor = img_tensor_factory()
|
||||
tf_cfg = ImageTransformConfig(type="Resize", kwargs={"size": (32, 32)})
|
||||
tf = make_transform_from_config(tf_cfg)
|
||||
assert isinstance(tf, v2.Resize)
|
||||
output = tf(img_tensor)
|
||||
assert output.shape[-2:] == (32, 32)
|
||||
|
||||
|
||||
def test_make_transform_from_config_with_v2_identity(img_tensor_factory):
|
||||
img_tensor = img_tensor_factory()
|
||||
tf_cfg = ImageTransformConfig(type="Identity", kwargs={})
|
||||
tf = make_transform_from_config(tf_cfg)
|
||||
assert isinstance(tf, v2.Identity)
|
||||
output = tf(img_tensor)
|
||||
assert output.shape == img_tensor.shape
|
||||
|
||||
|
||||
def test_make_transform_from_config_invalid_type():
|
||||
tf_cfg = ImageTransformConfig(type="NotARealTransform", kwargs={})
|
||||
with pytest.raises(ValueError, match="not valid"):
|
||||
make_transform_from_config(tf_cfg)
|
||||
|
||||
|
||||
def test_save_all_transforms(img_tensor_factory, tmp_path):
|
||||
img_tensor = img_tensor_factory()
|
||||
tf_cfg = ImageTransformsConfig(enable=True)
|
||||
|
||||
@@ -0,0 +1,730 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 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 streaming video encoding and hardware-accelerated encoding."""
|
||||
|
||||
import queue
|
||||
import threading
|
||||
from unittest.mock import patch
|
||||
|
||||
import av
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from lerobot.datasets.video_utils import (
|
||||
VALID_VIDEO_CODECS,
|
||||
StreamingVideoEncoder,
|
||||
_CameraEncoderThread,
|
||||
_get_codec_options,
|
||||
detect_available_hw_encoders,
|
||||
resolve_vcodec,
|
||||
)
|
||||
from lerobot.utils.constants import OBS_IMAGES
|
||||
|
||||
# ─── _get_codec_options tests ───
|
||||
|
||||
|
||||
class TestGetCodecOptions:
|
||||
def test_libsvtav1_defaults(self):
|
||||
opts = _get_codec_options("libsvtav1")
|
||||
assert opts["g"] == "2"
|
||||
assert opts["crf"] == "30"
|
||||
assert opts["preset"] == "12"
|
||||
|
||||
def test_libsvtav1_custom_preset(self):
|
||||
opts = _get_codec_options("libsvtav1", preset=8)
|
||||
assert opts["preset"] == "8"
|
||||
|
||||
def test_h264_options(self):
|
||||
opts = _get_codec_options("h264", g=10, crf=23)
|
||||
assert opts["g"] == "10"
|
||||
assert opts["crf"] == "23"
|
||||
assert "preset" not in opts
|
||||
|
||||
def test_videotoolbox_options(self):
|
||||
opts = _get_codec_options("h264_videotoolbox", g=2, crf=30)
|
||||
assert opts["g"] == "2"
|
||||
# CRF 30 maps to quality = max(1, min(100, 100 - 30*2)) = 40
|
||||
assert opts["q:v"] == "40"
|
||||
assert "crf" not in opts
|
||||
|
||||
def test_nvenc_options(self):
|
||||
opts = _get_codec_options("h264_nvenc", g=2, crf=25)
|
||||
assert opts["rc"] == "constqp"
|
||||
assert opts["qp"] == "25"
|
||||
assert "crf" not in opts
|
||||
# NVENC doesn't support g
|
||||
assert "g" not in opts
|
||||
|
||||
def test_vaapi_options(self):
|
||||
opts = _get_codec_options("h264_vaapi", crf=28)
|
||||
assert opts["qp"] == "28"
|
||||
|
||||
def test_qsv_options(self):
|
||||
opts = _get_codec_options("h264_qsv", crf=25)
|
||||
assert opts["global_quality"] == "25"
|
||||
|
||||
def test_no_g_no_crf(self):
|
||||
opts = _get_codec_options("h264", g=None, crf=None)
|
||||
assert "g" not in opts
|
||||
assert "crf" not in opts
|
||||
|
||||
|
||||
# ─── HW encoder detection tests ───
|
||||
|
||||
|
||||
class TestHWEncoderDetection:
|
||||
def test_detect_available_hw_encoders_returns_list(self):
|
||||
result = detect_available_hw_encoders()
|
||||
assert isinstance(result, list)
|
||||
|
||||
def test_detect_available_hw_encoders_only_valid(self):
|
||||
from lerobot.datasets.video_utils import HW_ENCODERS
|
||||
|
||||
result = detect_available_hw_encoders()
|
||||
for encoder in result:
|
||||
assert encoder in HW_ENCODERS
|
||||
|
||||
def test_resolve_vcodec_passthrough(self):
|
||||
assert resolve_vcodec("libsvtav1") == "libsvtav1"
|
||||
assert resolve_vcodec("h264") == "h264"
|
||||
|
||||
def test_resolve_vcodec_auto_fallback(self):
|
||||
"""When no HW encoders are available, auto should fall back to libsvtav1."""
|
||||
with patch("lerobot.datasets.video_utils.detect_available_hw_encoders", return_value=[]):
|
||||
assert resolve_vcodec("auto") == "libsvtav1"
|
||||
|
||||
def test_resolve_vcodec_auto_picks_hw(self):
|
||||
"""When a HW encoder is available, auto should pick it."""
|
||||
with patch(
|
||||
"lerobot.datasets.video_utils.detect_available_hw_encoders",
|
||||
return_value=["h264_videotoolbox"],
|
||||
):
|
||||
assert resolve_vcodec("auto") == "h264_videotoolbox"
|
||||
|
||||
def test_resolve_vcodec_auto_returns_valid(self):
|
||||
"""Test that resolve_vcodec('auto') returns a known valid codec."""
|
||||
result = resolve_vcodec("auto")
|
||||
assert result in VALID_VIDEO_CODECS
|
||||
|
||||
def test_hw_encoder_names_accepted_in_validation(self):
|
||||
"""Test that HW encoder names pass validation in VALID_VIDEO_CODECS."""
|
||||
assert "auto" in VALID_VIDEO_CODECS
|
||||
assert "h264_videotoolbox" in VALID_VIDEO_CODECS
|
||||
assert "h264_nvenc" in VALID_VIDEO_CODECS
|
||||
|
||||
def test_resolve_vcodec_invalid_raises(self):
|
||||
"""Test that resolve_vcodec raises ValueError for invalid codecs."""
|
||||
with pytest.raises(ValueError, match="Invalid vcodec"):
|
||||
resolve_vcodec("not_a_real_codec")
|
||||
|
||||
|
||||
# ─── _CameraEncoderThread tests ───
|
||||
|
||||
|
||||
class TestCameraEncoderThread:
|
||||
def test_encodes_valid_mp4(self, tmp_path):
|
||||
"""Test that the encoder thread creates a valid MP4 file with correct frame count."""
|
||||
num_frames = 30
|
||||
height, width = 64, 96
|
||||
fps = 30
|
||||
video_path = tmp_path / "test_output" / "test.mp4"
|
||||
|
||||
frame_queue: queue.Queue = queue.Queue(maxsize=60)
|
||||
result_queue: queue.Queue = queue.Queue(maxsize=1)
|
||||
stop_event = threading.Event()
|
||||
|
||||
encoder_thread = _CameraEncoderThread(
|
||||
video_path=video_path,
|
||||
fps=fps,
|
||||
vcodec="libsvtav1",
|
||||
pix_fmt="yuv420p",
|
||||
g=2,
|
||||
crf=30,
|
||||
preset=13,
|
||||
frame_queue=frame_queue,
|
||||
result_queue=result_queue,
|
||||
stop_event=stop_event,
|
||||
)
|
||||
encoder_thread.start()
|
||||
|
||||
# Feed frames (HWC uint8)
|
||||
for _ in range(num_frames):
|
||||
frame = np.random.randint(0, 255, (height, width, 3), dtype=np.uint8)
|
||||
frame_queue.put(frame)
|
||||
|
||||
# Send sentinel
|
||||
frame_queue.put(None)
|
||||
encoder_thread.join(timeout=60)
|
||||
assert not encoder_thread.is_alive()
|
||||
|
||||
# Check result
|
||||
status, data = result_queue.get(timeout=5)
|
||||
assert status == "ok"
|
||||
assert data is not None # Stats should be returned
|
||||
assert "mean" in data
|
||||
assert "std" in data
|
||||
assert "min" in data
|
||||
assert "max" in data
|
||||
assert "count" in data
|
||||
|
||||
# Verify the MP4 file is valid
|
||||
assert video_path.exists()
|
||||
with av.open(str(video_path)) as container:
|
||||
stream = container.streams.video[0]
|
||||
# The frame count should match
|
||||
total_frames = sum(1 for _ in container.decode(stream))
|
||||
assert total_frames == num_frames
|
||||
|
||||
def test_handles_chw_input(self, tmp_path):
|
||||
"""Test that CHW format input is handled correctly."""
|
||||
num_frames = 5
|
||||
fps = 30
|
||||
video_path = tmp_path / "test_chw" / "test.mp4"
|
||||
|
||||
frame_queue: queue.Queue = queue.Queue(maxsize=60)
|
||||
result_queue: queue.Queue = queue.Queue(maxsize=1)
|
||||
stop_event = threading.Event()
|
||||
|
||||
encoder_thread = _CameraEncoderThread(
|
||||
video_path=video_path,
|
||||
fps=fps,
|
||||
vcodec="libsvtav1",
|
||||
pix_fmt="yuv420p",
|
||||
g=2,
|
||||
crf=30,
|
||||
preset=13,
|
||||
frame_queue=frame_queue,
|
||||
result_queue=result_queue,
|
||||
stop_event=stop_event,
|
||||
)
|
||||
encoder_thread.start()
|
||||
|
||||
# Feed CHW frames
|
||||
for _ in range(num_frames):
|
||||
frame = np.random.randint(0, 255, (3, 64, 96), dtype=np.uint8)
|
||||
frame_queue.put(frame)
|
||||
|
||||
frame_queue.put(None)
|
||||
encoder_thread.join(timeout=60)
|
||||
|
||||
status, _ = result_queue.get(timeout=5)
|
||||
assert status == "ok"
|
||||
assert video_path.exists()
|
||||
|
||||
def test_stop_event_cancellation(self, tmp_path):
|
||||
"""Test that setting the stop event causes the thread to exit."""
|
||||
fps = 30
|
||||
video_path = tmp_path / "test_cancel" / "test.mp4"
|
||||
|
||||
frame_queue: queue.Queue = queue.Queue(maxsize=60)
|
||||
result_queue: queue.Queue = queue.Queue(maxsize=1)
|
||||
stop_event = threading.Event()
|
||||
|
||||
encoder_thread = _CameraEncoderThread(
|
||||
video_path=video_path,
|
||||
fps=fps,
|
||||
vcodec="libsvtav1",
|
||||
pix_fmt="yuv420p",
|
||||
g=2,
|
||||
crf=30,
|
||||
preset=13,
|
||||
frame_queue=frame_queue,
|
||||
result_queue=result_queue,
|
||||
stop_event=stop_event,
|
||||
)
|
||||
encoder_thread.start()
|
||||
|
||||
# Feed a few frames
|
||||
for _ in range(3):
|
||||
frame = np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8)
|
||||
frame_queue.put(frame)
|
||||
|
||||
# Signal stop instead of sending sentinel
|
||||
stop_event.set()
|
||||
encoder_thread.join(timeout=10)
|
||||
assert not encoder_thread.is_alive()
|
||||
|
||||
|
||||
# ─── StreamingVideoEncoder tests ───
|
||||
|
||||
|
||||
class TestStreamingVideoEncoder:
|
||||
def test_single_camera_episode(self, tmp_path):
|
||||
"""Test encoding a single camera episode."""
|
||||
encoder = StreamingVideoEncoder(fps=30, vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13)
|
||||
|
||||
video_keys = [f"{OBS_IMAGES}.laptop"]
|
||||
encoder.start_episode(video_keys, tmp_path)
|
||||
|
||||
num_frames = 20
|
||||
for _ in range(num_frames):
|
||||
frame = np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8)
|
||||
encoder.feed_frame(f"{OBS_IMAGES}.laptop", frame)
|
||||
|
||||
results = encoder.finish_episode()
|
||||
assert f"{OBS_IMAGES}.laptop" in results
|
||||
|
||||
mp4_path, stats = results[f"{OBS_IMAGES}.laptop"]
|
||||
assert mp4_path.exists()
|
||||
assert stats is not None
|
||||
|
||||
# Verify frame count
|
||||
with av.open(str(mp4_path)) as container:
|
||||
stream = container.streams.video[0]
|
||||
total_frames = sum(1 for _ in container.decode(stream))
|
||||
assert total_frames == num_frames
|
||||
|
||||
encoder.close()
|
||||
|
||||
def test_multi_camera_episode(self, tmp_path):
|
||||
"""Test encoding multiple cameras simultaneously."""
|
||||
encoder = StreamingVideoEncoder(fps=30, vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30)
|
||||
|
||||
video_keys = [f"{OBS_IMAGES}.laptop", f"{OBS_IMAGES}.phone"]
|
||||
encoder.start_episode(video_keys, tmp_path)
|
||||
|
||||
num_frames = 15
|
||||
for _ in range(num_frames):
|
||||
frame0 = np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8)
|
||||
frame1 = np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8)
|
||||
encoder.feed_frame(video_keys[0], frame0)
|
||||
encoder.feed_frame(video_keys[1], frame1)
|
||||
|
||||
results = encoder.finish_episode()
|
||||
|
||||
for key in video_keys:
|
||||
assert key in results
|
||||
mp4_path, stats = results[key]
|
||||
assert mp4_path.exists()
|
||||
assert stats is not None
|
||||
|
||||
encoder.close()
|
||||
|
||||
def test_sequential_episodes(self, tmp_path):
|
||||
"""Test that multiple sequential episodes work correctly."""
|
||||
encoder = StreamingVideoEncoder(fps=30, vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30)
|
||||
video_keys = [f"{OBS_IMAGES}.cam"]
|
||||
|
||||
for ep in range(3):
|
||||
encoder.start_episode(video_keys, tmp_path)
|
||||
num_frames = 10 + ep * 5
|
||||
for _ in range(num_frames):
|
||||
frame = np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8)
|
||||
encoder.feed_frame(f"{OBS_IMAGES}.cam", frame)
|
||||
results = encoder.finish_episode()
|
||||
|
||||
mp4_path, stats = results[f"{OBS_IMAGES}.cam"]
|
||||
assert mp4_path.exists()
|
||||
|
||||
with av.open(str(mp4_path)) as container:
|
||||
stream = container.streams.video[0]
|
||||
total_frames = sum(1 for _ in container.decode(stream))
|
||||
assert total_frames == num_frames
|
||||
|
||||
encoder.close()
|
||||
|
||||
def test_cancel_episode(self, tmp_path):
|
||||
"""Test that canceling an episode cleans up properly."""
|
||||
encoder = StreamingVideoEncoder(fps=30, vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30)
|
||||
video_keys = [f"{OBS_IMAGES}.cam"]
|
||||
|
||||
encoder.start_episode(video_keys, tmp_path)
|
||||
|
||||
for _ in range(5):
|
||||
frame = np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8)
|
||||
encoder.feed_frame(f"{OBS_IMAGES}.cam", frame)
|
||||
|
||||
encoder.cancel_episode()
|
||||
|
||||
# Should be able to start a new episode after cancel
|
||||
encoder.start_episode(video_keys, tmp_path)
|
||||
for _ in range(5):
|
||||
frame = np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8)
|
||||
encoder.feed_frame(f"{OBS_IMAGES}.cam", frame)
|
||||
results = encoder.finish_episode()
|
||||
|
||||
assert f"{OBS_IMAGES}.cam" in results
|
||||
encoder.close()
|
||||
|
||||
def test_feed_without_start_raises(self, tmp_path):
|
||||
"""Test that feeding frames without starting an episode raises."""
|
||||
encoder = StreamingVideoEncoder(fps=30, vcodec="libsvtav1", pix_fmt="yuv420p")
|
||||
with pytest.raises(RuntimeError, match="No active episode"):
|
||||
encoder.feed_frame("cam", np.zeros((64, 96, 3), dtype=np.uint8))
|
||||
encoder.close()
|
||||
|
||||
def test_finish_without_start_raises(self, tmp_path):
|
||||
"""Test that finishing without starting raises."""
|
||||
encoder = StreamingVideoEncoder(fps=30, vcodec="libsvtav1", pix_fmt="yuv420p")
|
||||
with pytest.raises(RuntimeError, match="No active episode"):
|
||||
encoder.finish_episode()
|
||||
encoder.close()
|
||||
|
||||
def test_close_is_idempotent(self, tmp_path):
|
||||
"""Test that close() can be called multiple times safely."""
|
||||
encoder = StreamingVideoEncoder(fps=30, vcodec="libsvtav1", pix_fmt="yuv420p")
|
||||
encoder.close()
|
||||
encoder.close() # Should not raise
|
||||
|
||||
def test_video_duration_matches_frame_count(self, tmp_path):
|
||||
"""Test that encoded video duration matches num_frames / fps."""
|
||||
encoder = StreamingVideoEncoder(fps=30, vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13)
|
||||
video_keys = [f"{OBS_IMAGES}.cam"]
|
||||
encoder.start_episode(video_keys, tmp_path)
|
||||
|
||||
num_frames = 90 # 3 seconds at 30fps
|
||||
for _ in range(num_frames):
|
||||
frame = np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8)
|
||||
encoder.feed_frame(f"{OBS_IMAGES}.cam", frame)
|
||||
|
||||
results = encoder.finish_episode()
|
||||
mp4_path, _ = results[f"{OBS_IMAGES}.cam"]
|
||||
|
||||
expected_duration = num_frames / 30.0 # 3.0 seconds
|
||||
|
||||
with av.open(str(mp4_path)) as container:
|
||||
stream = container.streams.video[0]
|
||||
total_frames = sum(1 for _ in container.decode(stream))
|
||||
if stream.duration is not None:
|
||||
actual_duration = float(stream.duration * stream.time_base)
|
||||
else:
|
||||
actual_duration = float(container.duration / av.time_base)
|
||||
|
||||
assert total_frames == num_frames
|
||||
# Allow small tolerance for duration due to codec framing
|
||||
assert abs(actual_duration - expected_duration) < 0.5, (
|
||||
f"Video duration {actual_duration:.2f}s != expected {expected_duration:.2f}s"
|
||||
)
|
||||
|
||||
encoder.close()
|
||||
|
||||
def test_multi_camera_start_episode_called_once(self, tmp_path):
|
||||
"""Test that with multiple cameras, no frames are lost due to double start_episode."""
|
||||
encoder = StreamingVideoEncoder(fps=30, vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30)
|
||||
|
||||
video_keys = [f"{OBS_IMAGES}.cam1", f"{OBS_IMAGES}.cam2"]
|
||||
encoder.start_episode(video_keys, tmp_path)
|
||||
|
||||
num_frames = 30
|
||||
for _ in range(num_frames):
|
||||
frame0 = np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8)
|
||||
frame1 = np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8)
|
||||
encoder.feed_frame(video_keys[0], frame0)
|
||||
encoder.feed_frame(video_keys[1], frame1)
|
||||
|
||||
results = encoder.finish_episode()
|
||||
|
||||
# Both cameras should have all frames
|
||||
for key in video_keys:
|
||||
mp4_path, stats = results[key]
|
||||
assert mp4_path.exists()
|
||||
with av.open(str(mp4_path)) as container:
|
||||
stream = container.streams.video[0]
|
||||
total_frames = sum(1 for _ in container.decode(stream))
|
||||
assert total_frames == num_frames, (
|
||||
f"Camera {key}: expected {num_frames} frames, got {total_frames}"
|
||||
)
|
||||
|
||||
encoder.close()
|
||||
|
||||
def test_encoder_threads_passed_to_thread(self, tmp_path):
|
||||
"""Test that encoder_threads is stored and passed through to encoder threads."""
|
||||
encoder = StreamingVideoEncoder(
|
||||
fps=30, vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, encoder_threads=2
|
||||
)
|
||||
assert encoder.encoder_threads == 2
|
||||
|
||||
video_keys = [f"{OBS_IMAGES}.cam"]
|
||||
encoder.start_episode(video_keys, tmp_path)
|
||||
|
||||
# Verify the thread received the encoder_threads value
|
||||
thread = encoder._threads[f"{OBS_IMAGES}.cam"]
|
||||
assert thread.encoder_threads == 2
|
||||
|
||||
# Feed some frames and finish to ensure it works end-to-end
|
||||
num_frames = 10
|
||||
for _ in range(num_frames):
|
||||
frame = np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8)
|
||||
encoder.feed_frame(f"{OBS_IMAGES}.cam", frame)
|
||||
|
||||
results = encoder.finish_episode()
|
||||
mp4_path, stats = results[f"{OBS_IMAGES}.cam"]
|
||||
assert mp4_path.exists()
|
||||
assert stats is not None
|
||||
|
||||
with av.open(str(mp4_path)) as container:
|
||||
stream = container.streams.video[0]
|
||||
total_frames = sum(1 for _ in container.decode(stream))
|
||||
assert total_frames == num_frames
|
||||
|
||||
encoder.close()
|
||||
|
||||
def test_encoder_threads_none_by_default(self, tmp_path):
|
||||
"""Test that encoder_threads defaults to None (codec auto-detect)."""
|
||||
encoder = StreamingVideoEncoder(fps=30, vcodec="libsvtav1", pix_fmt="yuv420p")
|
||||
assert encoder.encoder_threads is None
|
||||
encoder.close()
|
||||
|
||||
def test_graceful_frame_dropping(self, tmp_path):
|
||||
"""Test that full queue drops frames instead of crashing."""
|
||||
encoder = StreamingVideoEncoder(
|
||||
fps=30, vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13, queue_maxsize=1
|
||||
)
|
||||
video_keys = [f"{OBS_IMAGES}.cam"]
|
||||
encoder.start_episode(video_keys, tmp_path)
|
||||
|
||||
# Feed many frames quickly - with queue_maxsize=1, some will be dropped
|
||||
num_frames = 50
|
||||
for _ in range(num_frames):
|
||||
frame = np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8)
|
||||
encoder.feed_frame(f"{OBS_IMAGES}.cam", frame)
|
||||
|
||||
# Should not raise - frames are dropped gracefully
|
||||
results = encoder.finish_episode()
|
||||
assert f"{OBS_IMAGES}.cam" in results
|
||||
|
||||
mp4_path, _ = results[f"{OBS_IMAGES}.cam"]
|
||||
assert mp4_path.exists()
|
||||
|
||||
# Some frames should have been dropped (queue was tiny)
|
||||
dropped = encoder._dropped_frames.get(f"{OBS_IMAGES}.cam", 0)
|
||||
# We can't guarantee drops but can verify no crash occurred
|
||||
assert dropped >= 0
|
||||
|
||||
encoder.close()
|
||||
|
||||
|
||||
# ─── Integration tests with LeRobotDataset ───
|
||||
|
||||
|
||||
class TestStreamingEncoderIntegration:
|
||||
def test_add_frame_save_episode_streaming(self, tmp_path):
|
||||
"""Full integration test: add_frame -> save_episode with streaming encoding."""
|
||||
from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
||||
|
||||
features = {
|
||||
"observation.images.cam": {
|
||||
"dtype": "video",
|
||||
"shape": (64, 96, 3),
|
||||
"names": ["height", "width", "channels"],
|
||||
},
|
||||
"action": {"dtype": "float32", "shape": (6,), "names": ["j1", "j2", "j3", "j4", "j5", "j6"]},
|
||||
}
|
||||
|
||||
dataset = LeRobotDataset.create(
|
||||
repo_id="test/streaming",
|
||||
fps=30,
|
||||
features=features,
|
||||
root=tmp_path / "streaming_test",
|
||||
use_videos=True,
|
||||
streaming_encoding=True,
|
||||
)
|
||||
|
||||
assert dataset._streaming_encoder is not None
|
||||
|
||||
num_frames = 20
|
||||
for _ in range(num_frames):
|
||||
frame = {
|
||||
"observation.images.cam": np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8),
|
||||
"action": np.random.randn(6).astype(np.float32),
|
||||
"task": "test task",
|
||||
}
|
||||
dataset.add_frame(frame)
|
||||
|
||||
dataset.save_episode()
|
||||
|
||||
# Verify dataset metadata
|
||||
assert dataset.meta.total_episodes == 1
|
||||
assert dataset.meta.total_frames == num_frames
|
||||
|
||||
# Verify stats exist for the video key
|
||||
assert dataset.meta.stats is not None
|
||||
assert "observation.images.cam" in dataset.meta.stats
|
||||
assert "action" in dataset.meta.stats
|
||||
|
||||
dataset.finalize()
|
||||
|
||||
def test_streaming_disabled_creates_pngs(self, tmp_path):
|
||||
"""Test that disabling streaming encoding falls back to PNG path."""
|
||||
from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
||||
|
||||
features = {
|
||||
"observation.images.cam": {
|
||||
"dtype": "video",
|
||||
"shape": (64, 96, 3),
|
||||
"names": ["height", "width", "channels"],
|
||||
},
|
||||
"action": {"dtype": "float32", "shape": (6,), "names": ["j1", "j2", "j3", "j4", "j5", "j6"]},
|
||||
}
|
||||
|
||||
dataset = LeRobotDataset.create(
|
||||
repo_id="test/no_streaming",
|
||||
fps=30,
|
||||
features=features,
|
||||
root=tmp_path / "no_streaming_test",
|
||||
use_videos=True,
|
||||
streaming_encoding=False,
|
||||
)
|
||||
|
||||
assert dataset._streaming_encoder is None
|
||||
|
||||
num_frames = 5
|
||||
for _ in range(num_frames):
|
||||
frame = {
|
||||
"observation.images.cam": np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8),
|
||||
"action": np.random.randn(6).astype(np.float32),
|
||||
"task": "test task",
|
||||
}
|
||||
dataset.add_frame(frame)
|
||||
|
||||
# With streaming disabled, PNG files should be written
|
||||
images_dir = dataset.root / "images"
|
||||
assert images_dir.exists()
|
||||
|
||||
dataset.save_episode()
|
||||
dataset.finalize()
|
||||
|
||||
def test_multi_episode_streaming(self, tmp_path):
|
||||
"""Test recording multiple episodes with streaming encoding."""
|
||||
from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
||||
|
||||
features = {
|
||||
"observation.images.cam": {
|
||||
"dtype": "video",
|
||||
"shape": (64, 96, 3),
|
||||
"names": ["height", "width", "channels"],
|
||||
},
|
||||
"action": {"dtype": "float32", "shape": (2,), "names": ["j1", "j2"]},
|
||||
}
|
||||
|
||||
dataset = LeRobotDataset.create(
|
||||
repo_id="test/multi_ep",
|
||||
fps=30,
|
||||
features=features,
|
||||
root=tmp_path / "multi_ep_test",
|
||||
use_videos=True,
|
||||
streaming_encoding=True,
|
||||
)
|
||||
|
||||
for ep in range(3):
|
||||
num_frames = 10 + ep * 5
|
||||
for _ in range(num_frames):
|
||||
frame = {
|
||||
"observation.images.cam": np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8),
|
||||
"action": np.random.randn(2).astype(np.float32),
|
||||
"task": f"task_{ep}",
|
||||
}
|
||||
dataset.add_frame(frame)
|
||||
dataset.save_episode()
|
||||
|
||||
assert dataset.meta.total_episodes == 3
|
||||
assert dataset.meta.total_frames == 10 + 15 + 20
|
||||
|
||||
dataset.finalize()
|
||||
|
||||
def test_clear_episode_buffer_cancels_streaming(self, tmp_path):
|
||||
"""Test that clearing episode buffer cancels streaming encoding."""
|
||||
from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
||||
|
||||
features = {
|
||||
"observation.images.cam": {
|
||||
"dtype": "video",
|
||||
"shape": (64, 96, 3),
|
||||
"names": ["height", "width", "channels"],
|
||||
},
|
||||
"action": {"dtype": "float32", "shape": (2,), "names": ["j1", "j2"]},
|
||||
}
|
||||
|
||||
dataset = LeRobotDataset.create(
|
||||
repo_id="test/cancel",
|
||||
fps=30,
|
||||
features=features,
|
||||
root=tmp_path / "cancel_test",
|
||||
use_videos=True,
|
||||
streaming_encoding=True,
|
||||
)
|
||||
|
||||
# Add some frames
|
||||
for _ in range(5):
|
||||
frame = {
|
||||
"observation.images.cam": np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8),
|
||||
"action": np.random.randn(2).astype(np.float32),
|
||||
"task": "task",
|
||||
}
|
||||
dataset.add_frame(frame)
|
||||
|
||||
# Cancel and re-record
|
||||
dataset.clear_episode_buffer()
|
||||
|
||||
# Record a new episode
|
||||
for _ in range(10):
|
||||
frame = {
|
||||
"observation.images.cam": np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8),
|
||||
"action": np.random.randn(2).astype(np.float32),
|
||||
"task": "task",
|
||||
}
|
||||
dataset.add_frame(frame)
|
||||
dataset.save_episode()
|
||||
|
||||
assert dataset.meta.total_episodes == 1
|
||||
assert dataset.meta.total_frames == 10
|
||||
|
||||
dataset.finalize()
|
||||
|
||||
def test_multi_camera_add_frame_streaming(self, tmp_path):
|
||||
"""Test that start_episode is called once with multiple video keys."""
|
||||
from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
||||
|
||||
features = {
|
||||
"observation.images.cam1": {
|
||||
"dtype": "video",
|
||||
"shape": (64, 96, 3),
|
||||
"names": ["height", "width", "channels"],
|
||||
},
|
||||
"observation.images.cam2": {
|
||||
"dtype": "video",
|
||||
"shape": (64, 96, 3),
|
||||
"names": ["height", "width", "channels"],
|
||||
},
|
||||
"action": {"dtype": "float32", "shape": (2,), "names": ["j1", "j2"]},
|
||||
}
|
||||
|
||||
dataset = LeRobotDataset.create(
|
||||
repo_id="test/multi_cam",
|
||||
fps=30,
|
||||
features=features,
|
||||
root=tmp_path / "multi_cam_test",
|
||||
use_videos=True,
|
||||
streaming_encoding=True,
|
||||
)
|
||||
|
||||
num_frames = 15
|
||||
for _ in range(num_frames):
|
||||
frame = {
|
||||
"observation.images.cam1": np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8),
|
||||
"observation.images.cam2": np.random.randint(0, 255, (64, 96, 3), dtype=np.uint8),
|
||||
"action": np.random.randn(2).astype(np.float32),
|
||||
"task": "test task",
|
||||
}
|
||||
dataset.add_frame(frame)
|
||||
|
||||
dataset.save_episode()
|
||||
|
||||
assert dataset.meta.total_episodes == 1
|
||||
assert dataset.meta.total_frames == num_frames
|
||||
|
||||
dataset.finalize()
|
||||
@@ -0,0 +1,190 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 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 subtask functionality in LeRobotDataset.
|
||||
|
||||
These tests verify that:
|
||||
- Subtask information is correctly loaded from datasets that have subtask data
|
||||
- The __getitem__ method correctly adds subtask strings to returned items
|
||||
- Subtask handling gracefully handles missing data
|
||||
"""
|
||||
|
||||
import pandas as pd
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
||||
|
||||
|
||||
class TestSubtaskDataset:
|
||||
"""Tests for subtask handling in LeRobotDataset."""
|
||||
|
||||
@pytest.fixture
|
||||
def subtask_dataset(self):
|
||||
"""Load the test subtask dataset from the hub."""
|
||||
# Use lerobot/pusht-subtask dataset with episode 1
|
||||
return LeRobotDataset(
|
||||
repo_id="lerobot/pusht-subtask",
|
||||
episodes=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],
|
||||
)
|
||||
|
||||
def test_subtask_dataset_loads(self, subtask_dataset):
|
||||
"""Test that the subtask dataset loads successfully."""
|
||||
assert subtask_dataset is not None
|
||||
assert len(subtask_dataset) > 0
|
||||
|
||||
def test_subtask_metadata_loaded(self, subtask_dataset):
|
||||
"""Test that subtask metadata is loaded when present in dataset."""
|
||||
# The dataset should have subtasks metadata loaded
|
||||
assert subtask_dataset.meta.subtasks is not None
|
||||
assert isinstance(subtask_dataset.meta.subtasks, pd.DataFrame)
|
||||
|
||||
def test_subtask_index_in_features(self, subtask_dataset):
|
||||
"""Test that subtask_index is a feature when dataset has subtasks."""
|
||||
assert "subtask_index" in subtask_dataset.features
|
||||
|
||||
def test_getitem_returns_subtask_string(self, subtask_dataset):
|
||||
"""Test that __getitem__ correctly adds subtask string to returned item."""
|
||||
item = subtask_dataset[0]
|
||||
|
||||
# Subtask should be present in the returned item
|
||||
assert "subtask" in item
|
||||
assert isinstance(item["subtask"], str)
|
||||
assert len(item["subtask"]) > 0 # Should not be empty
|
||||
|
||||
def test_getitem_has_subtask_index(self, subtask_dataset):
|
||||
"""Test that __getitem__ includes subtask_index."""
|
||||
item = subtask_dataset[0]
|
||||
|
||||
assert "subtask_index" in item
|
||||
assert isinstance(item["subtask_index"], torch.Tensor)
|
||||
|
||||
def test_subtask_index_maps_to_valid_subtask(self, subtask_dataset):
|
||||
"""Test that subtask_index correctly maps to a subtask in metadata."""
|
||||
item = subtask_dataset[0]
|
||||
|
||||
subtask_idx = item["subtask_index"].item()
|
||||
subtask_from_metadata = subtask_dataset.meta.subtasks.iloc[subtask_idx].name
|
||||
|
||||
assert item["subtask"] == subtask_from_metadata
|
||||
|
||||
def test_all_items_have_subtask(self, subtask_dataset):
|
||||
"""Test that all items in the dataset have subtask information."""
|
||||
for i in range(min(len(subtask_dataset), 5)): # Check first 5 items
|
||||
item = subtask_dataset[i]
|
||||
assert "subtask" in item
|
||||
assert isinstance(item["subtask"], str)
|
||||
|
||||
def test_task_and_subtask_coexist(self, subtask_dataset):
|
||||
"""Test that both task and subtask are present in returned items."""
|
||||
item = subtask_dataset[0]
|
||||
|
||||
# Both task and subtask should be present
|
||||
assert "task" in item
|
||||
assert "subtask" in item
|
||||
assert isinstance(item["task"], str)
|
||||
assert isinstance(item["subtask"], str)
|
||||
|
||||
|
||||
class TestSubtaskDatasetMissing:
|
||||
"""Tests for graceful handling when subtask data is missing."""
|
||||
|
||||
@pytest.fixture
|
||||
def dataset_without_subtasks(self, tmp_path, empty_lerobot_dataset_factory):
|
||||
"""Create a dataset without subtask information."""
|
||||
features = {"state": {"dtype": "float32", "shape": (2,), "names": None}}
|
||||
dataset = empty_lerobot_dataset_factory(root=tmp_path / "no_subtask", features=features)
|
||||
|
||||
# Add some frames and save
|
||||
for _ in range(5):
|
||||
dataset.add_frame({"state": torch.randn(2), "task": "Test task"})
|
||||
dataset.save_episode()
|
||||
dataset.finalize()
|
||||
|
||||
# Reload the dataset
|
||||
return LeRobotDataset(dataset.repo_id, root=dataset.root)
|
||||
|
||||
def test_no_subtask_in_features(self, dataset_without_subtasks):
|
||||
"""Test that subtask_index is not in features when not provided."""
|
||||
assert "subtask_index" not in dataset_without_subtasks.features
|
||||
|
||||
def test_getitem_without_subtask(self, dataset_without_subtasks):
|
||||
"""Test that __getitem__ works when subtask is not present."""
|
||||
item = dataset_without_subtasks[0]
|
||||
|
||||
# Item should still be retrievable
|
||||
assert item is not None
|
||||
assert "state" in item
|
||||
assert "task" in item
|
||||
|
||||
# Subtask should NOT be present
|
||||
assert "subtask" not in item
|
||||
|
||||
def test_subtasks_metadata_is_none(self, dataset_without_subtasks):
|
||||
"""Test that subtasks metadata is None when not present."""
|
||||
assert dataset_without_subtasks.meta.subtasks is None
|
||||
|
||||
|
||||
class TestSubtaskEdgeCases:
|
||||
"""Edge case tests for subtask handling."""
|
||||
|
||||
def test_subtask_with_multiple_episodes(self):
|
||||
"""Test subtask handling with multiple episodes if available."""
|
||||
try:
|
||||
dataset = LeRobotDataset(
|
||||
repo_id="lerobot/pusht-subtask",
|
||||
episodes=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],
|
||||
)
|
||||
except Exception:
|
||||
pytest.skip("Could not load test-subtask dataset")
|
||||
|
||||
# Check first and last items have valid subtasks
|
||||
first_item = dataset[0]
|
||||
last_item = dataset[len(dataset) - 1]
|
||||
|
||||
assert "subtask" in first_item
|
||||
assert "subtask" in last_item
|
||||
assert isinstance(first_item["subtask"], str)
|
||||
assert isinstance(last_item["subtask"], str)
|
||||
|
||||
def test_subtask_index_consistency(self):
|
||||
"""Test that same subtask_index returns same subtask string."""
|
||||
try:
|
||||
dataset = LeRobotDataset(
|
||||
repo_id="lerobot/pusht-subtask",
|
||||
episodes=[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],
|
||||
)
|
||||
except Exception:
|
||||
pytest.skip("Could not load test-subtask dataset")
|
||||
|
||||
if len(dataset) < 2:
|
||||
pytest.skip("Dataset too small for this test")
|
||||
|
||||
# Collect subtask_index to subtask mappings
|
||||
subtask_map = {}
|
||||
for i in range(min(len(dataset), 10)):
|
||||
item = dataset[i]
|
||||
idx = item["subtask_index"].item()
|
||||
subtask = item["subtask"]
|
||||
|
||||
if idx in subtask_map:
|
||||
# Same index should always return same subtask
|
||||
assert subtask_map[idx] == subtask, (
|
||||
f"Inconsistent subtask for index {idx}: '{subtask_map[idx]}' vs '{subtask}'"
|
||||
)
|
||||
else:
|
||||
subtask_map[idx] = subtask
|
||||
@@ -17,12 +17,12 @@
|
||||
import random
|
||||
from dataclasses import dataclass, field
|
||||
from functools import cached_property
|
||||
from typing import Any
|
||||
|
||||
from lerobot.cameras import CameraConfig, make_cameras_from_configs
|
||||
from lerobot.motors.motors_bus import Motor, MotorNormMode
|
||||
from lerobot.processor import RobotAction, RobotObservation
|
||||
from lerobot.robots import Robot, RobotConfig
|
||||
from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
|
||||
from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
|
||||
from tests.mocks.mock_motors_bus import MockMotorsBus
|
||||
|
||||
|
||||
@@ -98,10 +98,8 @@ class MockRobot(Robot):
|
||||
def is_connected(self) -> bool:
|
||||
return self._is_connected
|
||||
|
||||
@check_if_already_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()
|
||||
@@ -110,19 +108,15 @@ class MockRobot(Robot):
|
||||
def is_calibrated(self) -> bool:
|
||||
return self._is_calibrated
|
||||
|
||||
@check_if_not_connected
|
||||
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.")
|
||||
|
||||
@check_if_not_connected
|
||||
def get_observation(self) -> RobotObservation:
|
||||
if self.config.random_values:
|
||||
return {f"{motor}.pos": random.uniform(-100, 100) for motor in self.motors}
|
||||
else:
|
||||
@@ -130,14 +124,10 @@ class MockRobot(Robot):
|
||||
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.")
|
||||
|
||||
@check_if_not_connected
|
||||
def send_action(self, action: RobotAction) -> RobotAction:
|
||||
return action
|
||||
|
||||
@check_if_not_connected
|
||||
def disconnect(self) -> None:
|
||||
if not self.is_connected:
|
||||
raise DeviceNotConnectedError(f"{self} is not connected.")
|
||||
|
||||
self._is_connected = False
|
||||
|
||||
@@ -19,8 +19,9 @@ from dataclasses import dataclass
|
||||
from functools import cached_property
|
||||
from typing import Any
|
||||
|
||||
from lerobot.processor import RobotAction
|
||||
from lerobot.teleoperators import Teleoperator, TeleoperatorConfig
|
||||
from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
|
||||
from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
|
||||
|
||||
|
||||
@TeleoperatorConfig.register_subclass("mock_teleop")
|
||||
@@ -67,10 +68,8 @@ class MockTeleop(Teleoperator):
|
||||
def is_connected(self) -> bool:
|
||||
return self._is_connected
|
||||
|
||||
@check_if_already_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()
|
||||
@@ -79,19 +78,15 @@ class MockTeleop(Teleoperator):
|
||||
def is_calibrated(self) -> bool:
|
||||
return self._is_calibrated
|
||||
|
||||
@check_if_not_connected
|
||||
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.")
|
||||
|
||||
@check_if_not_connected
|
||||
def get_action(self) -> RobotAction:
|
||||
if self.config.random_values:
|
||||
return {f"{motor}.pos": random.uniform(-100, 100) for motor in self.motors}
|
||||
else:
|
||||
@@ -99,12 +94,9 @@ class MockTeleop(Teleoperator):
|
||||
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.")
|
||||
@check_if_not_connected
|
||||
def send_feedback(self, feedback: dict[str, Any]) -> None: ...
|
||||
|
||||
@check_if_not_connected
|
||||
def disconnect(self) -> None:
|
||||
if not self.is_connected:
|
||||
raise DeviceNotConnectedError(f"{self} is not connected.")
|
||||
|
||||
self._is_connected = False
|
||||
|
||||
@@ -0,0 +1,66 @@
|
||||
"""Minimal test script for Damiao motor with ID 3."""
|
||||
|
||||
import pytest
|
||||
|
||||
from lerobot.utils.import_utils import _can_available
|
||||
|
||||
if not _can_available:
|
||||
pytest.skip("python-can not available", allow_module_level=True)
|
||||
|
||||
from lerobot.motors import Motor
|
||||
from lerobot.motors.damiao import DamiaoMotorsBus
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="Requires physical Damiao motor and CAN interface")
|
||||
def test_damiao_motor():
|
||||
motors = {
|
||||
"joint_3": Motor(
|
||||
id=0x03,
|
||||
model="damiao",
|
||||
norm_mode="degrees",
|
||||
motor_type_str="dm4310",
|
||||
recv_id=0x13,
|
||||
),
|
||||
}
|
||||
|
||||
bus = DamiaoMotorsBus(port="can0", motors=motors)
|
||||
|
||||
try:
|
||||
print("Connecting...")
|
||||
bus.connect()
|
||||
print("✓ Connected")
|
||||
|
||||
print("Enabling torque...")
|
||||
bus.enable_torque()
|
||||
print("✓ Torque enabled")
|
||||
|
||||
print("Reading all states...")
|
||||
states = bus.sync_read_all_states()
|
||||
print(f"✓ States: {states}")
|
||||
|
||||
print("Reading position...")
|
||||
positions = bus.sync_read("Present_Position")
|
||||
print(f"✓ Position: {positions}")
|
||||
|
||||
print("Testing MIT control batch...")
|
||||
current_pos = states["joint_3"]["position"]
|
||||
commands = {"joint_3": (10.0, 0.5, current_pos, 0.0, 0.0)}
|
||||
bus._mit_control_batch(commands)
|
||||
print("✓ MIT control batch sent")
|
||||
|
||||
print("Disabling torque...")
|
||||
bus.disable_torque()
|
||||
print("✓ Torque disabled")
|
||||
|
||||
print("Setting zero position...")
|
||||
bus.set_zero_position()
|
||||
print("✓ Zero position set")
|
||||
|
||||
finally:
|
||||
print("Disconnecting...")
|
||||
bus.disconnect(disable_torque=True)
|
||||
print("✓ Disconnected")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_damiao_motor()
|
||||
@@ -11,6 +11,8 @@
|
||||
# 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 packaging.version import Version
|
||||
from torch.optim.lr_scheduler import LambdaLR
|
||||
|
||||
from lerobot.optim.schedulers import (
|
||||
@@ -38,6 +40,10 @@ def test_diffuser_scheduler(optimizer):
|
||||
"last_epoch": 1,
|
||||
"lr_lambdas": [None],
|
||||
}
|
||||
|
||||
if Version(torch.__version__) >= Version("2.8"):
|
||||
expected_state_dict["_is_initial"] = False
|
||||
|
||||
assert scheduler.state_dict() == expected_state_dict
|
||||
|
||||
|
||||
@@ -56,6 +62,10 @@ def test_vqbet_scheduler(optimizer):
|
||||
"last_epoch": 1,
|
||||
"lr_lambdas": [None],
|
||||
}
|
||||
|
||||
if Version(torch.__version__) >= Version("2.8"):
|
||||
expected_state_dict["_is_initial"] = False
|
||||
|
||||
assert scheduler.state_dict() == expected_state_dict
|
||||
|
||||
|
||||
@@ -76,6 +86,10 @@ def test_cosine_decay_with_warmup_scheduler(optimizer):
|
||||
"last_epoch": 1,
|
||||
"lr_lambdas": [None],
|
||||
}
|
||||
|
||||
if Version(torch.__version__) >= Version("2.8"):
|
||||
expected_state_dict["_is_initial"] = False
|
||||
|
||||
assert scheduler.state_dict() == expected_state_dict
|
||||
|
||||
|
||||
|
||||
@@ -1,46 +0,0 @@
|
||||
#!/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
|
||||
import types
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
|
||||
def _install_reachy2_sdk_stub():
|
||||
sdk = types.ModuleType("reachy2_sdk")
|
||||
sdk.__path__ = []
|
||||
sdk.ReachySDK = MagicMock(name="ReachySDK")
|
||||
|
||||
media = types.ModuleType("reachy2_sdk.media")
|
||||
media.__path__ = []
|
||||
camera = types.ModuleType("reachy2_sdk.media.camera")
|
||||
camera.CameraView = MagicMock(name="CameraView")
|
||||
camera_manager = types.ModuleType("reachy2_sdk.media.camera_manager")
|
||||
camera_manager.CameraManager = MagicMock(name="CameraManager")
|
||||
|
||||
sdk.media = media
|
||||
media.camera = camera
|
||||
media.camera_manager = camera_manager
|
||||
|
||||
# Register in sys.modules
|
||||
sys.modules.setdefault("reachy2_sdk", sdk)
|
||||
sys.modules.setdefault("reachy2_sdk.media", media)
|
||||
sys.modules.setdefault("reachy2_sdk.media.camera", camera)
|
||||
sys.modules.setdefault("reachy2_sdk.media.camera_manager", camera_manager)
|
||||
|
||||
|
||||
def pytest_sessionstart(session):
|
||||
_install_reachy2_sdk_stub()
|
||||
@@ -0,0 +1,504 @@
|
||||
#!/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.
|
||||
|
||||
"""Test script to verify PI0Fast policy integration with LeRobot vs the original implementation"""
|
||||
# ruff: noqa: E402
|
||||
|
||||
import os
|
||||
import random
|
||||
from copy import deepcopy
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
pytest.importorskip("transformers")
|
||||
pytest.importorskip("scipy")
|
||||
pytestmark = pytest.mark.skipif(
|
||||
os.environ.get("CI") == "true" or os.environ.get("GITHUB_ACTIONS") == "true",
|
||||
reason="This test requires accepting the model license",
|
||||
)
|
||||
|
||||
from lerobot.policies.pi0_fast.configuration_pi0_fast import PI0FastConfig
|
||||
from lerobot.policies.pi0_fast.modeling_pi0_fast import PI0FastPolicy
|
||||
from lerobot.policies.pi0_fast.processor_pi0_fast import make_pi0_fast_pre_post_processors
|
||||
from lerobot.processor import PolicyAction, PolicyProcessorPipeline # noqa: E402
|
||||
from lerobot.utils.constants import (
|
||||
ACTION_TOKEN_MASK,
|
||||
ACTION_TOKENS,
|
||||
OBS_IMAGES,
|
||||
OBS_LANGUAGE_ATTENTION_MASK,
|
||||
OBS_LANGUAGE_TOKENS,
|
||||
OBS_STATE,
|
||||
) # noqa: E402
|
||||
from tests.utils import require_cuda # noqa: E402
|
||||
|
||||
# Constants
|
||||
DUMMY_ACTION_DIM = 7
|
||||
DUMMY_STATE_DIM = 20
|
||||
IMAGE_HEIGHT = 224
|
||||
IMAGE_WIDTH = 224
|
||||
NUM_VIEWS = 2 # Number of camera views
|
||||
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
MODEL_PATH_LEROBOT = "lerobot/pi0fast-base"
|
||||
|
||||
# Expected action token shape: (batch_size, max_decoding_steps)
|
||||
EXPECTED_ACTION_TOKENS_SHAPE = (1, 2)
|
||||
|
||||
# Expected first 5 action tokens (for reproducibility check)
|
||||
EXPECTED_ACTION_TOKENS_FIRST_5 = torch.tensor([255657, 255362])
|
||||
|
||||
# Expected actions after detokenization
|
||||
EXPECTED_ACTIONS_SHAPE = (1, 2, 32) # (batch_size, n_action_steps, action_dim)
|
||||
EXPECTED_ACTIONS_MEAN = 0.04419417306780815
|
||||
EXPECTED_ACTIONS_STD = 0.26231569051742554
|
||||
EXPECTED_ACTIONS_FIRST_5 = torch.tensor([0.0000, 1.4849, 0.0000, 0.0000, 0.0000])
|
||||
|
||||
|
||||
def set_seed_all(seed: int):
|
||||
"""Set random seed for all RNG sources to ensure reproducibility."""
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
|
||||
# Set deterministic behavior
|
||||
torch.backends.cudnn.deterministic = True
|
||||
torch.backends.cudnn.benchmark = False
|
||||
torch.use_deterministic_algorithms(True, warn_only=True)
|
||||
|
||||
|
||||
def instantiate_lerobot_pi0_fast(
|
||||
from_pretrained: bool = False,
|
||||
model_path: str = MODEL_PATH_LEROBOT,
|
||||
) -> tuple[
|
||||
Any, # Policy
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction],
|
||||
]:
|
||||
"""Instantiate LeRobot PI0Fast policy with preprocessor and postprocessor."""
|
||||
if from_pretrained:
|
||||
policy = PI0FastPolicy.from_pretrained(
|
||||
pretrained_name_or_path=model_path,
|
||||
strict=True,
|
||||
)
|
||||
policy.config.validate_action_token_prefix = False
|
||||
policy.config.max_action_tokens = 2
|
||||
policy.config.max_decoding_steps = 2
|
||||
policy.config.chunk_size = 2
|
||||
policy.config.n_action_steps = 2
|
||||
else:
|
||||
config = PI0FastConfig(
|
||||
n_action_steps=2,
|
||||
max_action_dim=DUMMY_ACTION_DIM,
|
||||
max_state_dim=DUMMY_STATE_DIM,
|
||||
device=DEVICE,
|
||||
validate_action_token_prefix=False,
|
||||
max_action_tokens=2,
|
||||
max_decoding_steps=2,
|
||||
chunk_size=2,
|
||||
)
|
||||
policy = PI0FastPolicy(config)
|
||||
|
||||
policy.to(DEVICE)
|
||||
policy.config.device = DEVICE
|
||||
preprocessor, postprocessor = make_pi0_fast_pre_post_processors(
|
||||
config=policy.config,
|
||||
dataset_stats=None, # Pass None for dataset_stats to disable normalization
|
||||
)
|
||||
|
||||
return policy, preprocessor, postprocessor
|
||||
|
||||
|
||||
def create_dummy_data(device=DEVICE):
|
||||
"""Create dummy data for testing both implementations."""
|
||||
batch_size = 1
|
||||
prompt = "Pick up the red block and place it in the bin"
|
||||
|
||||
# Create random RGB images in [0, 255] uint8 range (as PIL images would be)
|
||||
# Then convert to [0, 1] float32 range for LeRobot
|
||||
def fake_rgb(h, w):
|
||||
arr = np.random.randint(0, 255, (h, w, 3), dtype=np.uint8)
|
||||
t = torch.from_numpy(arr).permute(2, 0, 1) # CHW
|
||||
return t
|
||||
|
||||
batch = {
|
||||
f"{OBS_IMAGES}.base_0_rgb": torch.stack(
|
||||
[fake_rgb(IMAGE_HEIGHT, IMAGE_WIDTH) for _ in range(batch_size)]
|
||||
).to(device),
|
||||
f"{OBS_IMAGES}.left_wrist_0_rgb": torch.stack(
|
||||
[fake_rgb(IMAGE_HEIGHT, IMAGE_WIDTH) for _ in range(batch_size)]
|
||||
).to(device),
|
||||
f"{OBS_IMAGES}.right_wrist_0_rgb": torch.stack(
|
||||
[fake_rgb(IMAGE_HEIGHT, IMAGE_WIDTH) for _ in range(batch_size)]
|
||||
).to(device),
|
||||
OBS_STATE: torch.randn(batch_size, DUMMY_STATE_DIM, dtype=torch.float32, device=device),
|
||||
"task": [prompt for _ in range(batch_size)],
|
||||
}
|
||||
|
||||
return batch
|
||||
|
||||
|
||||
# Pytest fixtures
|
||||
@pytest.fixture(scope="module")
|
||||
def pi0_fast_components():
|
||||
"""Fixture to instantiate and provide all PI0Fast components for tests."""
|
||||
print(f"\nTesting with DEVICE='{DEVICE}'")
|
||||
print("\n[Setup] Instantiating LeRobot PI0Fast policy...")
|
||||
policy_obj, preprocessor_obj, postprocessor_obj = instantiate_lerobot_pi0_fast(from_pretrained=True)
|
||||
print("Model loaded successfully")
|
||||
yield policy_obj, preprocessor_obj, postprocessor_obj
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def policy(pi0_fast_components):
|
||||
"""Fixture to provide the PI0Fast policy for tests."""
|
||||
return pi0_fast_components[0]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def preprocessor(pi0_fast_components):
|
||||
"""Fixture to provide the PI0Fast preprocessor for tests."""
|
||||
return pi0_fast_components[1]
|
||||
|
||||
|
||||
@require_cuda
|
||||
def test_pi0_fast_preprocessor_alignment(policy, preprocessor):
|
||||
"""Test that LeRobot PI0Fast preprocessor produces expected outputs."""
|
||||
print("\n" + "=" * 80)
|
||||
print("Test: PI0Fast Preprocessor Outputs")
|
||||
print("=" * 80)
|
||||
|
||||
set_seed_all(42)
|
||||
|
||||
print("\nCreating dummy data...")
|
||||
batch = create_dummy_data()
|
||||
|
||||
print("\n[LeRobot] Preprocessing...")
|
||||
lerobot_observation = preprocessor(deepcopy(batch))
|
||||
|
||||
print("\nVerifying preprocessor outputs:")
|
||||
print("-" * 80)
|
||||
|
||||
# Expected keys from PI0Fast preprocessing
|
||||
expected_keys = [
|
||||
"observation.images.base_0_rgb",
|
||||
"observation.images.left_wrist_0_rgb",
|
||||
"observation.images.right_wrist_0_rgb",
|
||||
"observation.state",
|
||||
"observation.language_tokens",
|
||||
"observation.language_attention_mask",
|
||||
]
|
||||
|
||||
for key in expected_keys:
|
||||
if key in lerobot_observation:
|
||||
shape = tuple(lerobot_observation[key].shape)
|
||||
print(f"\nKey: {key}")
|
||||
print(f"Shape: {shape}")
|
||||
print(f"Dtype: {lerobot_observation[key].dtype}")
|
||||
else:
|
||||
print(f"\nKey '{key}' not found in inputs!")
|
||||
|
||||
# Check language tokens shape
|
||||
if "observation.language_tokens" in lerobot_observation:
|
||||
lang_tokens = lerobot_observation["observation.language_tokens"]
|
||||
print(f"\nLanguage tokens shape: {lang_tokens.shape}")
|
||||
# Should have batch dimension and max_length from tokenizer
|
||||
assert lang_tokens.dim() == 2, f"Expected 2D tensor, got {lang_tokens.dim()}D"
|
||||
|
||||
print("\nPreprocessor outputs verified!")
|
||||
|
||||
|
||||
@require_cuda
|
||||
def test_pi0_fast_action_generation(policy, preprocessor):
|
||||
"""Test PI0Fast LeRobot implementation generates expected actions."""
|
||||
print("\n" + "=" * 80)
|
||||
print("Test: PI0Fast Action Generation Against Expected Values")
|
||||
print("=" * 80)
|
||||
|
||||
set_seed_all(42)
|
||||
|
||||
print("\nCreating dummy data...")
|
||||
batch = create_dummy_data()
|
||||
|
||||
print("\n[LeRobot] Running inference...")
|
||||
lerobot_observation = preprocessor(deepcopy(batch))
|
||||
|
||||
# Reset seed for inference
|
||||
torch.manual_seed(42)
|
||||
with torch.no_grad():
|
||||
lerobot_actions = policy.predict_action_chunk(lerobot_observation)
|
||||
lerobot_actions = lerobot_actions.float().cpu()
|
||||
|
||||
print(f"LeRobot actions shape: {lerobot_actions.shape}")
|
||||
print(f"LeRobot actions mean: {lerobot_actions.mean().item():.6f}")
|
||||
print(f"LeRobot actions std: {lerobot_actions.std().item():.6f}")
|
||||
print(f"LeRobot actions first 5: {lerobot_actions[0, 0, :5]}")
|
||||
|
||||
print("\nExpected values (from original PI0Fast):")
|
||||
print(f"Expected actions shape: {EXPECTED_ACTIONS_SHAPE}")
|
||||
print(f"Expected actions mean: {EXPECTED_ACTIONS_MEAN:.6f}")
|
||||
print(f"Expected actions std: {EXPECTED_ACTIONS_STD:.6f}")
|
||||
print(f"Expected actions first 5: {EXPECTED_ACTIONS_FIRST_5}")
|
||||
|
||||
print("\nAction Comparison:")
|
||||
print("-" * 80)
|
||||
|
||||
# Compare shapes
|
||||
actual_shape = tuple(lerobot_actions.shape)
|
||||
print(f"Actual shape: {actual_shape}")
|
||||
|
||||
assert actual_shape == EXPECTED_ACTIONS_SHAPE, (
|
||||
f"Shape mismatch: {actual_shape} vs {EXPECTED_ACTIONS_SHAPE}"
|
||||
)
|
||||
print(f"Shape matches: {actual_shape}")
|
||||
|
||||
# Compare statistics
|
||||
actual_mean = lerobot_actions.mean().item()
|
||||
actual_std = lerobot_actions.std().item()
|
||||
|
||||
print(f"\nMean: {actual_mean:.6f} (expected: {EXPECTED_ACTIONS_MEAN:.6f})")
|
||||
print(f"Std: {actual_std:.6f} (expected: {EXPECTED_ACTIONS_STD:.6f})")
|
||||
|
||||
# Compare first 5 actions
|
||||
actual_first_5 = lerobot_actions[0, 0, :5]
|
||||
print("\nFirst 5 actions comparison:")
|
||||
print(f" Actual: {actual_first_5}")
|
||||
print(f" Expected: {EXPECTED_ACTIONS_FIRST_5}")
|
||||
|
||||
first_5_diff = torch.abs(actual_first_5 - EXPECTED_ACTIONS_FIRST_5)
|
||||
print(f" Max diff: {first_5_diff.max().item():.6e}")
|
||||
print(f" Mean diff: {first_5_diff.mean().item():.6e}")
|
||||
|
||||
# Check with different tolerances
|
||||
tolerances = [1e-5, 1e-4, 1e-3, 1e-2]
|
||||
for tol in tolerances:
|
||||
is_close = torch.allclose(actual_first_5, EXPECTED_ACTIONS_FIRST_5, atol=tol)
|
||||
status = "Success" if is_close else "Failure"
|
||||
print(f"{status}: First 5 actions close (atol={tol}): {is_close}")
|
||||
|
||||
# Assert with reasonable tolerance
|
||||
tolerance = 1e-3
|
||||
assert torch.allclose(actual_first_5, EXPECTED_ACTIONS_FIRST_5, atol=tolerance), (
|
||||
f"First 5 actions differ by more than tolerance ({tolerance})"
|
||||
)
|
||||
print(f"\nSuccess: Actions match expected values within tolerance ({tolerance})!")
|
||||
|
||||
print("\nAction generation test completed (values printed for reference)!")
|
||||
|
||||
|
||||
@require_cuda
|
||||
def test_pi0_fast_inference_reproducibility(policy, preprocessor):
|
||||
"""Test that PI0Fast inference is reproducible with the same seed."""
|
||||
print("\n" + "=" * 80)
|
||||
print("Test: PI0Fast Inference Reproducibility")
|
||||
print("=" * 80)
|
||||
|
||||
print("\nCreating dummy data...")
|
||||
batch = create_dummy_data()
|
||||
|
||||
# First inference
|
||||
print("\n[Run 1] Running inference...")
|
||||
set_seed_all(42)
|
||||
lerobot_observation = preprocessor(deepcopy(batch))
|
||||
with torch.no_grad():
|
||||
actions_1 = policy.predict_action_chunk(lerobot_observation)
|
||||
actions_1 = actions_1.float().cpu()
|
||||
|
||||
# Second inference with same seed
|
||||
print("\n[Run 2] Running inference with same seed...")
|
||||
set_seed_all(42)
|
||||
lerobot_observation = preprocessor(deepcopy(batch))
|
||||
with torch.no_grad():
|
||||
actions_2 = policy.predict_action_chunk(lerobot_observation)
|
||||
actions_2 = actions_2.float().cpu()
|
||||
|
||||
print("\nComparing two runs:")
|
||||
print("-" * 80)
|
||||
if torch.allclose(actions_1, actions_2, atol=1e-8):
|
||||
print("Inference is perfectly reproducible!")
|
||||
else:
|
||||
diff = torch.abs(actions_1 - actions_2)
|
||||
print("Small differences detected:")
|
||||
print(f" Max diff: {diff.max().item():.6e}")
|
||||
print(f" Mean diff: {diff.mean().item():.6e}")
|
||||
|
||||
assert torch.allclose(actions_1, actions_2, atol=1e-6), "Inference should be reproducible!"
|
||||
|
||||
print("\nInference is reproducible!")
|
||||
|
||||
|
||||
@require_cuda
|
||||
def test_pi0_fast_forward_pass_logits(policy, preprocessor):
|
||||
"""Test PI0Fast forward pass and compare logits against expected values."""
|
||||
print("\n" + "=" * 80)
|
||||
print("Test: PI0Fast Forward Pass Logits")
|
||||
print("=" * 80)
|
||||
|
||||
set_seed_all(42)
|
||||
|
||||
print("\nCreating dummy data with action tokens...")
|
||||
batch = create_dummy_data()
|
||||
|
||||
# Preprocess the batch
|
||||
lerobot_observation = preprocessor(deepcopy(batch))
|
||||
|
||||
# For forward pass, we need action tokens
|
||||
# Create dummy action tokens for testing
|
||||
batch_size = 1
|
||||
max_action_tokens = policy.config.max_action_tokens
|
||||
|
||||
# Create dummy action tokens (in practice, these come from the FAST tokenizer)
|
||||
dummy_action_tokens = torch.randint(
|
||||
0, 1000, (batch_size, max_action_tokens), dtype=torch.long, device=DEVICE
|
||||
)
|
||||
dummy_action_masks = torch.ones(batch_size, max_action_tokens, dtype=torch.bool, device=DEVICE)
|
||||
|
||||
# Add action tokens to the observation
|
||||
lerobot_observation[ACTION_TOKENS] = dummy_action_tokens
|
||||
lerobot_observation[ACTION_TOKEN_MASK] = dummy_action_masks
|
||||
|
||||
print("\n[LeRobot] Running forward pass...")
|
||||
policy.train()
|
||||
with torch.no_grad():
|
||||
loss, loss_dict = policy.forward(lerobot_observation)
|
||||
|
||||
print(f"Loss: {loss.item():.6f}")
|
||||
print(f"FAST Loss: {loss_dict['ce_loss']:.6f}")
|
||||
|
||||
print("\nForward pass completed successfully!")
|
||||
print(f"Loss value: {loss.item():.6f}")
|
||||
|
||||
# The loss should be a positive value
|
||||
assert loss.item() > 0, "Loss should be positive"
|
||||
assert not torch.isnan(loss), "Loss should not be NaN"
|
||||
assert not torch.isinf(loss), "Loss should not be infinite"
|
||||
|
||||
print("\nForward pass test passed!")
|
||||
|
||||
|
||||
@require_cuda
|
||||
def test_pi0_fast_action_token_sampling(policy, preprocessor):
|
||||
"""Test PI0Fast action token sampling (autoregressive decoding)."""
|
||||
print("\n" + "=" * 80)
|
||||
print("Test: PI0Fast Action Token Sampling")
|
||||
print("=" * 80)
|
||||
|
||||
set_seed_all(42)
|
||||
|
||||
print("\nCreating dummy data...")
|
||||
batch = create_dummy_data()
|
||||
|
||||
print("\n[LeRobot] Preprocessing...")
|
||||
lerobot_observation = preprocessor(deepcopy(batch))
|
||||
|
||||
# Prepare inputs for model
|
||||
images, img_masks = policy._preprocess_images(lerobot_observation)
|
||||
tokens = lerobot_observation[OBS_LANGUAGE_TOKENS]
|
||||
masks = lerobot_observation[OBS_LANGUAGE_ATTENTION_MASK]
|
||||
|
||||
print("\n[LeRobot] Sampling action tokens...")
|
||||
torch.manual_seed(42)
|
||||
with torch.no_grad():
|
||||
action_tokens = policy.model.sample_actions_fast(
|
||||
images,
|
||||
img_masks,
|
||||
tokens,
|
||||
masks,
|
||||
max_decoding_steps=2,
|
||||
temperature=0.0, # Greedy decoding for reproducibility
|
||||
)
|
||||
|
||||
print(f"Action tokens shape: {action_tokens.shape}")
|
||||
print(f"Action tokens first 10: {action_tokens[0, :10].tolist()}")
|
||||
|
||||
print("\nExpected values (from original PI0Fast):")
|
||||
print(f"Expected shape: {EXPECTED_ACTION_TOKENS_SHAPE}")
|
||||
print(f"Expected first 5: {EXPECTED_ACTION_TOKENS_FIRST_5.tolist()}")
|
||||
|
||||
# Verify shape
|
||||
actual_shape = tuple(action_tokens.shape)
|
||||
print(f"\nActual shape: {actual_shape}")
|
||||
|
||||
assert actual_shape == EXPECTED_ACTION_TOKENS_SHAPE, (
|
||||
f"Shape mismatch: {actual_shape} vs {EXPECTED_ACTION_TOKENS_SHAPE}"
|
||||
)
|
||||
|
||||
# Compare first 5 tokens
|
||||
actual_first_5 = action_tokens[0, :5].cpu()
|
||||
assert torch.equal(actual_first_5, EXPECTED_ACTION_TOKENS_FIRST_5), (
|
||||
f"First 5 tokens mismatch: {actual_first_5} vs {EXPECTED_ACTION_TOKENS_FIRST_5}"
|
||||
)
|
||||
|
||||
print("\nAction token sampling test completed!")
|
||||
|
||||
|
||||
@require_cuda
|
||||
def test_pi0_fast_detokenization(policy, preprocessor):
|
||||
"""Test PI0Fast action detokenization (FAST decoding)."""
|
||||
print("\n" + "=" * 80)
|
||||
print("Test: PI0Fast Action Detokenization")
|
||||
print("=" * 80)
|
||||
|
||||
set_seed_all(42)
|
||||
|
||||
print("\nCreating dummy data...")
|
||||
batch = create_dummy_data()
|
||||
|
||||
print("\n[LeRobot] Preprocessing...")
|
||||
lerobot_observation = preprocessor(deepcopy(batch))
|
||||
|
||||
# Prepare inputs for model
|
||||
images, img_masks = policy._preprocess_images(lerobot_observation)
|
||||
tokens = lerobot_observation[OBS_LANGUAGE_TOKENS]
|
||||
masks = lerobot_observation[OBS_LANGUAGE_ATTENTION_MASK]
|
||||
|
||||
print("\n[LeRobot] Sampling action tokens...")
|
||||
torch.manual_seed(42)
|
||||
with torch.no_grad():
|
||||
action_tokens = policy.model.sample_actions_fast(
|
||||
images,
|
||||
img_masks,
|
||||
tokens,
|
||||
masks,
|
||||
max_decoding_steps=2,
|
||||
temperature=0.0,
|
||||
)
|
||||
|
||||
print(f"Action tokens shape: {action_tokens.shape}")
|
||||
|
||||
# Detokenize
|
||||
print("\n[LeRobot] Detokenizing action tokens...")
|
||||
action_horizon = policy.config.n_action_steps
|
||||
action_dim = policy.config.output_features["action"].shape[0]
|
||||
|
||||
try:
|
||||
continuous_actions = policy.detokenize_actions(
|
||||
action_tokens, action_horizon=action_horizon, action_dim=action_dim
|
||||
)
|
||||
print(f"Continuous actions shape: {continuous_actions.shape}")
|
||||
print(f"Continuous actions mean: {continuous_actions.mean().item():.6f}")
|
||||
print(f"Continuous actions std: {continuous_actions.std().item():.6f}")
|
||||
print(f"Continuous actions first 5: {continuous_actions[0, 0, :5]}")
|
||||
print("\nDetokenization successful!")
|
||||
except Exception as e:
|
||||
print(f"\nDetokenization failed with error: {e}")
|
||||
print("This may be expected if the action tokens are not valid FAST tokens.")
|
||||
print("The test will pass as long as the sampling works correctly.")
|
||||
@@ -441,12 +441,13 @@ def test_sac_policy_with_predefined_entropy():
|
||||
|
||||
|
||||
def test_sac_policy_update_temperature():
|
||||
"""Test that temperature property is always in sync with log_alpha."""
|
||||
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()
|
||||
# Temperature property automatically reflects log_alpha changes
|
||||
assert policy.temperature == pytest.approx(0.1)
|
||||
|
||||
|
||||
|
||||
@@ -21,9 +21,10 @@ import os
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
# Skip if openpi or transformers is not available
|
||||
# Skip if required dependencies are not available
|
||||
pytest.importorskip("peft")
|
||||
pytest.importorskip("transformers==4.49.0")
|
||||
pytest.importorskip("transformers")
|
||||
pytest.importorskip("torchdiffeq")
|
||||
|
||||
# Skip this entire module in CI
|
||||
pytestmark = pytest.mark.skipif(
|
||||
|
||||
@@ -17,7 +17,7 @@
|
||||
import json
|
||||
import tempfile
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
@@ -1884,7 +1884,7 @@ class FeatureContractAddStep(ProcessorStep):
|
||||
"""Adds a PolicyFeature"""
|
||||
|
||||
key: str = "a"
|
||||
value: PolicyFeature = PolicyFeature(type=FeatureType.STATE, shape=(1,))
|
||||
value: PolicyFeature = field(default_factory=lambda: PolicyFeature(type=FeatureType.STATE, shape=(1,)))
|
||||
|
||||
def __call__(self, transition: EnvTransition) -> EnvTransition:
|
||||
return transition
|
||||
|
||||
@@ -27,7 +27,14 @@ import torch
|
||||
from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
|
||||
from lerobot.processor import DataProcessorPipeline, TokenizerProcessorStep, TransitionKey
|
||||
from lerobot.processor.converters import create_transition, identity_transition
|
||||
from lerobot.utils.constants import ACTION, OBS_IMAGE, OBS_LANGUAGE, OBS_STATE
|
||||
from lerobot.utils.constants import (
|
||||
ACTION,
|
||||
OBS_IMAGE,
|
||||
OBS_LANGUAGE,
|
||||
OBS_LANGUAGE_SUBTASK_ATTENTION_MASK,
|
||||
OBS_LANGUAGE_SUBTASK_TOKENS,
|
||||
OBS_STATE,
|
||||
)
|
||||
from tests.utils import require_package
|
||||
|
||||
|
||||
@@ -1038,3 +1045,459 @@ def test_simulated_accelerate_scenario():
|
||||
# MockTokenizer squeezes single-item batches, so shape is (max_length,) not (1, max_length)
|
||||
assert tokens.shape == (10,) # MockTokenizer behavior for single string in list
|
||||
assert attention_mask.shape == (10,)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Tests for get_subtask method
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@require_package("transformers")
|
||||
def test_get_subtask_missing_key():
|
||||
"""Test get_subtask returns None when subtask key is missing from complementary_data."""
|
||||
mock_tokenizer = MockTokenizer(vocab_size=100)
|
||||
processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)
|
||||
|
||||
transition = create_transition(
|
||||
observation={"state": torch.tensor([1.0, 2.0])},
|
||||
action=torch.tensor([0.1, 0.2]),
|
||||
complementary_data={"task": "main task"}, # No "subtask" key
|
||||
)
|
||||
|
||||
result = processor.get_subtask(transition)
|
||||
assert result is None
|
||||
|
||||
|
||||
@require_package("transformers")
|
||||
def test_get_subtask_none_value():
|
||||
"""Test get_subtask returns None when subtask value is None."""
|
||||
mock_tokenizer = MockTokenizer(vocab_size=100)
|
||||
processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)
|
||||
|
||||
transition = create_transition(
|
||||
observation={"state": torch.tensor([1.0, 2.0])},
|
||||
action=torch.tensor([0.1, 0.2]),
|
||||
complementary_data={"task": "main task", "subtask": None},
|
||||
)
|
||||
|
||||
result = processor.get_subtask(transition)
|
||||
assert result is None
|
||||
|
||||
|
||||
@require_package("transformers")
|
||||
def test_get_subtask_none_complementary_data():
|
||||
"""Test get_subtask returns None when complementary_data is None."""
|
||||
mock_tokenizer = MockTokenizer(vocab_size=100)
|
||||
processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)
|
||||
|
||||
transition = create_transition(
|
||||
observation={"state": torch.tensor([1.0, 2.0])},
|
||||
action=torch.tensor([0.1, 0.2]),
|
||||
complementary_data=None, # No complementary data
|
||||
)
|
||||
|
||||
result = processor.get_subtask(transition)
|
||||
assert result is None
|
||||
|
||||
|
||||
@require_package("transformers")
|
||||
def test_get_subtask_string():
|
||||
"""Test get_subtask returns list with single string when subtask is a string."""
|
||||
mock_tokenizer = MockTokenizer(vocab_size=100)
|
||||
processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)
|
||||
|
||||
transition = create_transition(
|
||||
observation={"state": torch.tensor([1.0, 2.0])},
|
||||
action=torch.tensor([0.1, 0.2]),
|
||||
complementary_data={"task": "main task", "subtask": "pick up the cube"},
|
||||
)
|
||||
|
||||
result = processor.get_subtask(transition)
|
||||
assert result == ["pick up the cube"]
|
||||
assert isinstance(result, list)
|
||||
assert len(result) == 1
|
||||
|
||||
|
||||
@require_package("transformers")
|
||||
def test_get_subtask_list_of_strings():
|
||||
"""Test get_subtask returns the list when subtask is already a list of strings."""
|
||||
mock_tokenizer = MockTokenizer(vocab_size=100)
|
||||
processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)
|
||||
|
||||
subtask_list = ["pick up", "move to target", "place down"]
|
||||
transition = create_transition(
|
||||
observation={"state": torch.tensor([1.0, 2.0])},
|
||||
action=torch.tensor([0.1, 0.2]),
|
||||
complementary_data={"task": "main task", "subtask": subtask_list},
|
||||
)
|
||||
|
||||
result = processor.get_subtask(transition)
|
||||
assert result == subtask_list
|
||||
assert isinstance(result, list)
|
||||
assert len(result) == 3
|
||||
|
||||
|
||||
@require_package("transformers")
|
||||
def test_get_subtask_unsupported_type_integer():
|
||||
"""Test get_subtask returns None when subtask is an unsupported type (integer)."""
|
||||
mock_tokenizer = MockTokenizer(vocab_size=100)
|
||||
processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)
|
||||
|
||||
transition = create_transition(
|
||||
observation={"state": torch.tensor([1.0, 2.0])},
|
||||
action=torch.tensor([0.1, 0.2]),
|
||||
complementary_data={"task": "main task", "subtask": 123},
|
||||
)
|
||||
|
||||
result = processor.get_subtask(transition)
|
||||
assert result is None
|
||||
|
||||
|
||||
@require_package("transformers")
|
||||
def test_get_subtask_unsupported_type_mixed_list():
|
||||
"""Test get_subtask returns None when subtask is a list with mixed types."""
|
||||
mock_tokenizer = MockTokenizer(vocab_size=100)
|
||||
processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)
|
||||
|
||||
transition = create_transition(
|
||||
observation={"state": torch.tensor([1.0, 2.0])},
|
||||
action=torch.tensor([0.1, 0.2]),
|
||||
complementary_data={"task": "main task", "subtask": ["valid string", 123, "another string"]},
|
||||
)
|
||||
|
||||
result = processor.get_subtask(transition)
|
||||
assert result is None
|
||||
|
||||
|
||||
@require_package("transformers")
|
||||
def test_get_subtask_unsupported_type_dict():
|
||||
"""Test get_subtask returns None when subtask is a dictionary."""
|
||||
mock_tokenizer = MockTokenizer(vocab_size=100)
|
||||
processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)
|
||||
|
||||
transition = create_transition(
|
||||
observation={"state": torch.tensor([1.0, 2.0])},
|
||||
action=torch.tensor([0.1, 0.2]),
|
||||
complementary_data={"task": "main task", "subtask": {"key": "value"}},
|
||||
)
|
||||
|
||||
result = processor.get_subtask(transition)
|
||||
assert result is None
|
||||
|
||||
|
||||
@require_package("transformers")
|
||||
def test_get_subtask_empty_string():
|
||||
"""Test get_subtask with empty string returns list with empty string."""
|
||||
mock_tokenizer = MockTokenizer(vocab_size=100)
|
||||
processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)
|
||||
|
||||
transition = create_transition(
|
||||
observation={"state": torch.tensor([1.0, 2.0])},
|
||||
action=torch.tensor([0.1, 0.2]),
|
||||
complementary_data={"task": "main task", "subtask": ""},
|
||||
)
|
||||
|
||||
result = processor.get_subtask(transition)
|
||||
assert result == [""]
|
||||
|
||||
|
||||
@require_package("transformers")
|
||||
def test_get_subtask_empty_list():
|
||||
"""Test get_subtask with empty list returns empty list."""
|
||||
mock_tokenizer = MockTokenizer(vocab_size=100)
|
||||
processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)
|
||||
|
||||
transition = create_transition(
|
||||
observation={"state": torch.tensor([1.0, 2.0])},
|
||||
action=torch.tensor([0.1, 0.2]),
|
||||
complementary_data={"task": "main task", "subtask": []},
|
||||
)
|
||||
|
||||
result = processor.get_subtask(transition)
|
||||
assert result == []
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Tests for subtask tokenization in observation method
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@require_package("transformers")
|
||||
def test_subtask_tokenization_when_present():
|
||||
"""Test that subtask is tokenized and added to observation when present."""
|
||||
mock_tokenizer = MockTokenizer(vocab_size=100)
|
||||
processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=8)
|
||||
|
||||
transition = create_transition(
|
||||
observation={"state": torch.tensor([1.0, 2.0])},
|
||||
action=torch.tensor([0.1, 0.2]),
|
||||
complementary_data={"task": "main task", "subtask": "pick up the red cube"},
|
||||
)
|
||||
|
||||
result = processor(transition)
|
||||
|
||||
# Check that subtask tokens were added to observation
|
||||
observation = result[TransitionKey.OBSERVATION]
|
||||
assert OBS_LANGUAGE_SUBTASK_TOKENS in observation
|
||||
assert OBS_LANGUAGE_SUBTASK_ATTENTION_MASK in observation
|
||||
|
||||
# Check token structure
|
||||
subtask_tokens = observation[OBS_LANGUAGE_SUBTASK_TOKENS]
|
||||
subtask_attention_mask = observation[OBS_LANGUAGE_SUBTASK_ATTENTION_MASK]
|
||||
assert isinstance(subtask_tokens, torch.Tensor)
|
||||
assert isinstance(subtask_attention_mask, torch.Tensor)
|
||||
assert subtask_tokens.shape == (8,)
|
||||
assert subtask_attention_mask.shape == (8,)
|
||||
assert subtask_attention_mask.dtype == torch.bool
|
||||
|
||||
|
||||
@require_package("transformers")
|
||||
def test_subtask_tokenization_not_added_when_none():
|
||||
"""Test that subtask tokens are NOT added to observation when subtask is None."""
|
||||
mock_tokenizer = MockTokenizer(vocab_size=100)
|
||||
processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=8)
|
||||
|
||||
transition = create_transition(
|
||||
observation={"state": torch.tensor([1.0, 2.0])},
|
||||
action=torch.tensor([0.1, 0.2]),
|
||||
complementary_data={"task": "main task"}, # No subtask
|
||||
)
|
||||
|
||||
result = processor(transition)
|
||||
|
||||
# Check that subtask tokens were NOT added to observation
|
||||
observation = result[TransitionKey.OBSERVATION]
|
||||
assert OBS_LANGUAGE_SUBTASK_TOKENS not in observation
|
||||
assert OBS_LANGUAGE_SUBTASK_ATTENTION_MASK not in observation
|
||||
|
||||
# But main task tokens should still be present
|
||||
assert f"{OBS_LANGUAGE}.tokens" in observation
|
||||
assert f"{OBS_LANGUAGE}.attention_mask" in observation
|
||||
|
||||
|
||||
@require_package("transformers")
|
||||
def test_subtask_tokenization_not_added_when_subtask_value_is_none():
|
||||
"""Test that subtask tokens are NOT added when subtask value is explicitly None."""
|
||||
mock_tokenizer = MockTokenizer(vocab_size=100)
|
||||
processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=8)
|
||||
|
||||
transition = create_transition(
|
||||
observation={"state": torch.tensor([1.0, 2.0])},
|
||||
action=torch.tensor([0.1, 0.2]),
|
||||
complementary_data={"task": "main task", "subtask": None},
|
||||
)
|
||||
|
||||
result = processor(transition)
|
||||
|
||||
# Check that subtask tokens were NOT added to observation
|
||||
observation = result[TransitionKey.OBSERVATION]
|
||||
assert OBS_LANGUAGE_SUBTASK_TOKENS not in observation
|
||||
assert OBS_LANGUAGE_SUBTASK_ATTENTION_MASK not in observation
|
||||
|
||||
|
||||
@require_package("transformers")
|
||||
def test_subtask_tokenization_list_of_strings():
|
||||
"""Test subtask tokenization with list of strings."""
|
||||
mock_tokenizer = MockTokenizer(vocab_size=100)
|
||||
processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=8)
|
||||
|
||||
transition = create_transition(
|
||||
observation={"state": torch.tensor([1.0, 2.0])},
|
||||
action=torch.tensor([0.1, 0.2]),
|
||||
complementary_data={"task": "main task", "subtask": ["pick up", "place down"]},
|
||||
)
|
||||
|
||||
result = processor(transition)
|
||||
|
||||
# Check that subtask tokens were added to observation
|
||||
observation = result[TransitionKey.OBSERVATION]
|
||||
assert OBS_LANGUAGE_SUBTASK_TOKENS in observation
|
||||
assert OBS_LANGUAGE_SUBTASK_ATTENTION_MASK in observation
|
||||
|
||||
# Check token structure for batch
|
||||
subtask_tokens = observation[OBS_LANGUAGE_SUBTASK_TOKENS]
|
||||
subtask_attention_mask = observation[OBS_LANGUAGE_SUBTASK_ATTENTION_MASK]
|
||||
assert subtask_tokens.shape == (2, 8) # batch_size=2, seq_len=8
|
||||
assert subtask_attention_mask.shape == (2, 8)
|
||||
|
||||
|
||||
@require_package("transformers")
|
||||
def test_subtask_tokenization_device_cpu():
|
||||
"""Test that subtask tokens are on CPU when other tensors are on CPU."""
|
||||
mock_tokenizer = MockTokenizer(vocab_size=100)
|
||||
processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)
|
||||
|
||||
# Create transition with CPU tensors
|
||||
observation = {OBS_STATE: torch.randn(10)} # CPU tensor
|
||||
action = torch.randn(5) # CPU tensor
|
||||
transition = create_transition(
|
||||
observation=observation,
|
||||
action=action,
|
||||
complementary_data={"task": "main task", "subtask": "pick up cube"},
|
||||
)
|
||||
|
||||
result = processor(transition)
|
||||
|
||||
# Check that subtask tokens are on CPU
|
||||
subtask_tokens = result[TransitionKey.OBSERVATION][OBS_LANGUAGE_SUBTASK_TOKENS]
|
||||
subtask_attention_mask = result[TransitionKey.OBSERVATION][OBS_LANGUAGE_SUBTASK_ATTENTION_MASK]
|
||||
|
||||
assert subtask_tokens.device.type == "cpu"
|
||||
assert subtask_attention_mask.device.type == "cpu"
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
|
||||
@require_package("transformers")
|
||||
def test_subtask_tokenization_device_cuda():
|
||||
"""Test that subtask tokens are moved to CUDA when other tensors are on CUDA."""
|
||||
mock_tokenizer = MockTokenizer(vocab_size=100)
|
||||
processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)
|
||||
|
||||
# Create transition with CUDA tensors
|
||||
observation = {OBS_STATE: torch.randn(10).cuda()} # CUDA tensor
|
||||
action = torch.randn(5).cuda() # CUDA tensor
|
||||
transition = create_transition(
|
||||
observation=observation,
|
||||
action=action,
|
||||
complementary_data={"task": "main task", "subtask": "pick up cube"},
|
||||
)
|
||||
|
||||
result = processor(transition)
|
||||
|
||||
# Check that subtask tokens are on CUDA
|
||||
subtask_tokens = result[TransitionKey.OBSERVATION][OBS_LANGUAGE_SUBTASK_TOKENS]
|
||||
subtask_attention_mask = result[TransitionKey.OBSERVATION][OBS_LANGUAGE_SUBTASK_ATTENTION_MASK]
|
||||
|
||||
assert subtask_tokens.device.type == "cuda"
|
||||
assert subtask_attention_mask.device.type == "cuda"
|
||||
|
||||
|
||||
@require_package("transformers")
|
||||
def test_subtask_tokenization_preserves_other_observation_data():
|
||||
"""Test that subtask tokenization preserves other observation data."""
|
||||
mock_tokenizer = MockTokenizer(vocab_size=100)
|
||||
processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)
|
||||
|
||||
original_state = torch.tensor([1.0, 2.0, 3.0])
|
||||
transition = create_transition(
|
||||
observation={"state": original_state.clone()},
|
||||
action=torch.tensor([0.1, 0.2]),
|
||||
complementary_data={"task": "main task", "subtask": "pick up cube"},
|
||||
)
|
||||
|
||||
result = processor(transition)
|
||||
observation = result[TransitionKey.OBSERVATION]
|
||||
|
||||
# Check that original observation data is preserved
|
||||
assert torch.equal(observation["state"], original_state)
|
||||
|
||||
# Check that both task and subtask tokens are present
|
||||
assert f"{OBS_LANGUAGE}.tokens" in observation
|
||||
assert f"{OBS_LANGUAGE}.attention_mask" in observation
|
||||
assert OBS_LANGUAGE_SUBTASK_TOKENS in observation
|
||||
assert OBS_LANGUAGE_SUBTASK_ATTENTION_MASK in observation
|
||||
|
||||
|
||||
@require_package("transformers")
|
||||
def test_subtask_attention_mask_dtype():
|
||||
"""Test that subtask attention mask has correct dtype (bool)."""
|
||||
mock_tokenizer = MockTokenizer(vocab_size=100)
|
||||
processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)
|
||||
|
||||
transition = create_transition(
|
||||
observation={"state": torch.tensor([1.0, 2.0])},
|
||||
action=torch.tensor([0.1, 0.2]),
|
||||
complementary_data={"task": "main task", "subtask": "pick up cube"},
|
||||
)
|
||||
|
||||
result = processor(transition)
|
||||
observation = result[TransitionKey.OBSERVATION]
|
||||
|
||||
subtask_attention_mask = observation[OBS_LANGUAGE_SUBTASK_ATTENTION_MASK]
|
||||
assert subtask_attention_mask.dtype == torch.bool
|
||||
|
||||
|
||||
@require_package("transformers")
|
||||
def test_subtask_tokenization_deterministic():
|
||||
"""Test that subtask tokenization is deterministic for the same input."""
|
||||
mock_tokenizer = MockTokenizer(vocab_size=100)
|
||||
processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)
|
||||
|
||||
transition = create_transition(
|
||||
observation={"state": torch.tensor([1.0, 2.0])},
|
||||
action=torch.tensor([0.1, 0.2]),
|
||||
complementary_data={"task": "main task", "subtask": "consistent subtask"},
|
||||
)
|
||||
|
||||
result1 = processor(transition)
|
||||
result2 = processor(transition)
|
||||
|
||||
subtask_tokens1 = result1[TransitionKey.OBSERVATION][OBS_LANGUAGE_SUBTASK_TOKENS]
|
||||
subtask_tokens2 = result2[TransitionKey.OBSERVATION][OBS_LANGUAGE_SUBTASK_TOKENS]
|
||||
subtask_mask1 = result1[TransitionKey.OBSERVATION][OBS_LANGUAGE_SUBTASK_ATTENTION_MASK]
|
||||
subtask_mask2 = result2[TransitionKey.OBSERVATION][OBS_LANGUAGE_SUBTASK_ATTENTION_MASK]
|
||||
|
||||
# Results should be identical
|
||||
assert torch.equal(subtask_tokens1, subtask_tokens2)
|
||||
assert torch.equal(subtask_mask1, subtask_mask2)
|
||||
|
||||
|
||||
@require_package("transformers")
|
||||
@patch("lerobot.processor.tokenizer_processor.AutoTokenizer")
|
||||
def test_subtask_tokenization_integration_with_pipeline(mock_auto_tokenizer):
|
||||
"""Test subtask tokenization works correctly with DataProcessorPipeline."""
|
||||
mock_tokenizer = MockTokenizer(vocab_size=100)
|
||||
mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer
|
||||
|
||||
tokenizer_processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer", max_length=6)
|
||||
robot_processor = DataProcessorPipeline(
|
||||
[tokenizer_processor], to_transition=identity_transition, to_output=identity_transition
|
||||
)
|
||||
|
||||
transition = create_transition(
|
||||
observation={"state": torch.tensor([1.0, 2.0])},
|
||||
action=torch.tensor([0.1, 0.2]),
|
||||
complementary_data={"task": "main task", "subtask": "subtask instruction"},
|
||||
)
|
||||
|
||||
result = robot_processor(transition)
|
||||
|
||||
# Check that observation exists and both tokenizations were applied
|
||||
assert TransitionKey.OBSERVATION in result
|
||||
observation = result[TransitionKey.OBSERVATION]
|
||||
|
||||
# Check task tokens
|
||||
assert f"{OBS_LANGUAGE}.tokens" in observation
|
||||
assert f"{OBS_LANGUAGE}.attention_mask" in observation
|
||||
|
||||
# Check subtask tokens
|
||||
assert OBS_LANGUAGE_SUBTASK_TOKENS in observation
|
||||
assert OBS_LANGUAGE_SUBTASK_ATTENTION_MASK in observation
|
||||
|
||||
# Check shapes
|
||||
assert observation[f"{OBS_LANGUAGE}.tokens"].shape == (6,)
|
||||
assert observation[OBS_LANGUAGE_SUBTASK_TOKENS].shape == (6,)
|
||||
|
||||
|
||||
@require_package("transformers")
|
||||
def test_subtask_not_added_for_unsupported_types():
|
||||
"""Test that subtask tokens are not added when subtask has unsupported type."""
|
||||
mock_tokenizer = MockTokenizer(vocab_size=100)
|
||||
processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=8)
|
||||
|
||||
# Test with integer subtask
|
||||
transition = create_transition(
|
||||
observation={"state": torch.tensor([1.0, 2.0])},
|
||||
action=torch.tensor([0.1, 0.2]),
|
||||
complementary_data={"task": "main task", "subtask": 123},
|
||||
)
|
||||
|
||||
result = processor(transition)
|
||||
observation = result[TransitionKey.OBSERVATION]
|
||||
|
||||
# Subtask tokens should NOT be added for unsupported types
|
||||
assert OBS_LANGUAGE_SUBTASK_TOKENS not in observation
|
||||
assert OBS_LANGUAGE_SUBTASK_ATTENTION_MASK not in observation
|
||||
|
||||
# But main task tokens should still be present
|
||||
assert f"{OBS_LANGUAGE}.tokens" in observation
|
||||
|
||||
@@ -64,7 +64,7 @@ def close_service_stub(channel, server):
|
||||
server.stop(None)
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def test_establish_learner_connection_success():
|
||||
from lerobot.rl.actor import establish_learner_connection
|
||||
|
||||
@@ -81,7 +81,7 @@ def test_establish_learner_connection_success():
|
||||
close_service_stub(channel, server)
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def test_establish_learner_connection_failure():
|
||||
from lerobot.rl.actor import establish_learner_connection
|
||||
|
||||
@@ -100,7 +100,7 @@ def test_establish_learner_connection_failure():
|
||||
close_service_stub(channel, server)
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def test_push_transitions_to_transport_queue():
|
||||
from lerobot.rl.actor import push_transitions_to_transport_queue
|
||||
from lerobot.transport.utils import bytes_to_transitions
|
||||
@@ -135,7 +135,7 @@ def test_push_transitions_to_transport_queue():
|
||||
assert_transitions_equal(deserialized_transition, transitions[i])
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
@pytest.mark.timeout(3) # force cross-platform watchdog
|
||||
def test_transitions_stream():
|
||||
from lerobot.rl.actor import transitions_stream
|
||||
@@ -167,7 +167,7 @@ def test_transitions_stream():
|
||||
assert streamed_data[2].data == b"transition_data_3"
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
@pytest.mark.timeout(3) # force cross-platform watchdog
|
||||
def test_interactions_stream():
|
||||
from lerobot.rl.actor import interactions_stream
|
||||
|
||||
@@ -88,7 +88,7 @@ def cfg():
|
||||
return cfg
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
@pytest.mark.timeout(10) # force cross-platform watchdog
|
||||
def test_end_to_end_transitions_flow(cfg):
|
||||
from lerobot.rl.actor import (
|
||||
@@ -150,7 +150,7 @@ def test_end_to_end_transitions_flow(cfg):
|
||||
assert_transitions_equal(transition, input_transitions[i])
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
@pytest.mark.timeout(10)
|
||||
def test_end_to_end_interactions_flow(cfg):
|
||||
from lerobot.rl.actor import (
|
||||
@@ -223,7 +223,7 @@ def test_end_to_end_interactions_flow(cfg):
|
||||
assert received == expected
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
@pytest.mark.parametrize("data_size", ["small", "large"])
|
||||
@pytest.mark.timeout(10)
|
||||
def test_end_to_end_parameters_flow(cfg, data_size):
|
||||
|
||||
@@ -39,7 +39,7 @@ def learner_service_stub():
|
||||
close_learner_service_stub(channel, server)
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def create_learner_service_stub(
|
||||
shutdown_event: Event,
|
||||
parameters_queue: Queue,
|
||||
@@ -75,7 +75,7 @@ def create_learner_service_stub(
|
||||
return services_pb2_grpc.LearnerServiceStub(channel), channel, server
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def close_learner_service_stub(channel, server):
|
||||
channel.close()
|
||||
server.stop(None)
|
||||
@@ -91,7 +91,7 @@ def test_ready_method(learner_service_stub):
|
||||
assert response == services_pb2.Empty()
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
@pytest.mark.timeout(3) # force cross-platform watchdog
|
||||
def test_send_interactions():
|
||||
from lerobot.transport import services_pb2
|
||||
@@ -135,7 +135,7 @@ def test_send_interactions():
|
||||
assert interactions == [b"123", b"4", b"5", b"678"]
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
@pytest.mark.timeout(3) # force cross-platform watchdog
|
||||
def test_send_transitions():
|
||||
from lerobot.transport import services_pb2
|
||||
@@ -181,7 +181,7 @@ def test_send_transitions():
|
||||
assert transitions == [b"transition_1transition_2transition_3", b"batch_1batch_2"]
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
@pytest.mark.timeout(3) # force cross-platform watchdog
|
||||
def test_send_transitions_empty_stream():
|
||||
from lerobot.transport import services_pb2
|
||||
@@ -209,7 +209,7 @@ def test_send_transitions_empty_stream():
|
||||
assert transitions_queue.empty()
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
@pytest.mark.timeout(10) # force cross-platform watchdog
|
||||
def test_stream_parameters():
|
||||
import time
|
||||
@@ -267,7 +267,7 @@ def test_stream_parameters():
|
||||
assert time_diff == pytest.approx(seconds_between_pushes, abs=0.1)
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
@pytest.mark.timeout(3) # force cross-platform watchdog
|
||||
def test_stream_parameters_with_shutdown():
|
||||
from lerobot.transport import services_pb2
|
||||
@@ -319,7 +319,7 @@ def test_stream_parameters_with_shutdown():
|
||||
assert received_params == [b"param_batch_1", b"stop"]
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
@pytest.mark.timeout(3) # force cross-platform watchdog
|
||||
def test_stream_parameters_waits_and_retries_on_empty_queue():
|
||||
import threading
|
||||
|
||||
@@ -19,6 +19,8 @@ from unittest.mock import MagicMock, patch
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("reachy2_sdk")
|
||||
|
||||
from lerobot.robots.reachy2 import (
|
||||
REACHY2_ANTENNAS_JOINTS,
|
||||
REACHY2_L_ARM_JOINTS,
|
||||
@@ -140,6 +142,7 @@ def _make_reachy2_camera_mock(*args, **kwargs):
|
||||
cam.connect = MagicMock()
|
||||
cam.disconnect = MagicMock()
|
||||
cam.async_read = MagicMock(side_effect=lambda: np.zeros((height, width, 3), dtype=np.uint8))
|
||||
cam.read_latest = MagicMock(side_effect=lambda: np.zeros((height, width, 3), dtype=np.uint8))
|
||||
return cam
|
||||
|
||||
|
||||
|
||||
@@ -19,7 +19,7 @@ from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from lerobot.robots.so100_follower import (
|
||||
from lerobot.robots.so_follower import (
|
||||
SO100Follower,
|
||||
SO100FollowerConfig,
|
||||
)
|
||||
@@ -66,7 +66,7 @@ def follower():
|
||||
|
||||
with (
|
||||
patch(
|
||||
"lerobot.robots.so100_follower.so100_follower.FeetechMotorsBus",
|
||||
"lerobot.robots.so_follower.so_follower.FeetechMotorsBus",
|
||||
side_effect=_bus_side_effect,
|
||||
),
|
||||
patch.object(SO100Follower, "configure", lambda self: None),
|
||||
|
||||
@@ -0,0 +1,74 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 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 draccus
|
||||
import pytest
|
||||
|
||||
from lerobot.scripts.lerobot_edit_dataset import (
|
||||
ConvertImageToVideoConfig,
|
||||
DeleteEpisodesConfig,
|
||||
EditDatasetConfig,
|
||||
InfoConfig,
|
||||
MergeConfig,
|
||||
ModifyTasksConfig,
|
||||
OperationConfig,
|
||||
RemoveFeatureConfig,
|
||||
SplitConfig,
|
||||
)
|
||||
|
||||
|
||||
def parse_cfg(cli_args: list[str]) -> EditDatasetConfig:
|
||||
"""Helper to parse CLI args into an EditDatasetConfig via draccus."""
|
||||
return draccus.parse(EditDatasetConfig, args=cli_args)
|
||||
|
||||
|
||||
class TestOperationTypeParsing:
|
||||
"""Test that --operation.type correctly selects the right config subclass."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"type_name, expected_cls",
|
||||
[
|
||||
("delete_episodes", DeleteEpisodesConfig),
|
||||
("split", SplitConfig),
|
||||
("merge", MergeConfig),
|
||||
("remove_feature", RemoveFeatureConfig),
|
||||
("modify_tasks", ModifyTasksConfig),
|
||||
("convert_image_to_video", ConvertImageToVideoConfig),
|
||||
("info", InfoConfig),
|
||||
],
|
||||
)
|
||||
def test_operation_type_resolves_correct_class(self, type_name, expected_cls):
|
||||
cfg = parse_cfg(["--repo_id", "test/repo", "--operation.type", type_name])
|
||||
assert isinstance(cfg.operation, expected_cls), (
|
||||
f"Expected {expected_cls.__name__}, got {type(cfg.operation).__name__}"
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"type_name, expected_cls",
|
||||
[
|
||||
("delete_episodes", DeleteEpisodesConfig),
|
||||
("split", SplitConfig),
|
||||
("merge", MergeConfig),
|
||||
("remove_feature", RemoveFeatureConfig),
|
||||
("modify_tasks", ModifyTasksConfig),
|
||||
("convert_image_to_video", ConvertImageToVideoConfig),
|
||||
("info", InfoConfig),
|
||||
],
|
||||
)
|
||||
def test_get_choice_name_roundtrips(self, type_name, expected_cls):
|
||||
cfg = parse_cfg(["--repo_id", "test/repo", "--operation.type", type_name])
|
||||
resolved_name = OperationConfig.get_choice_name(type(cfg.operation))
|
||||
assert resolved_name == type_name
|
||||
@@ -0,0 +1,235 @@
|
||||
import importlib
|
||||
import os
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from safetensors.torch import load_file
|
||||
|
||||
from .utils import require_package
|
||||
|
||||
# Skip this entire module in CI
|
||||
pytestmark = pytest.mark.skipif(
|
||||
os.environ.get("CI") == "true" or os.environ.get("GITHUB_ACTIONS") == "true",
|
||||
reason="This test requires peft and is very slow, not meant for CI",
|
||||
)
|
||||
|
||||
|
||||
def run_command(cmd, module, args):
|
||||
module = importlib.import_module(f"lerobot.scripts.{module}")
|
||||
with patch("sys.argv", [cmd] + args):
|
||||
module.main()
|
||||
|
||||
|
||||
def lerobot_train(args):
|
||||
return run_command(cmd="lerobot-train", module="lerobot_train", args=args)
|
||||
|
||||
|
||||
def lerobot_record(args):
|
||||
return run_command(cmd="lerobot-record", module="lerobot_record", args=args)
|
||||
|
||||
|
||||
def resolve_model_id_for_peft_training(policy_type):
|
||||
"""PEFT training needs pretrained models, this finds the pretrained model of a policy type for PEFT training."""
|
||||
if policy_type == "smolvla":
|
||||
return "lerobot/smolvla_base"
|
||||
|
||||
raise ValueError(f"No pretrained model known for {policy_type}. PEFT training will not work.")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("policy_type", ["smolvla"])
|
||||
@require_package("peft")
|
||||
def test_peft_training_push_to_hub_works(policy_type, tmp_path):
|
||||
"""Ensure that push to hub stores PEFT only the adapter, not the full model weights."""
|
||||
output_dir = tmp_path / f"output_{policy_type}"
|
||||
upload_folder_contents = set()
|
||||
|
||||
model_id = resolve_model_id_for_peft_training(policy_type)
|
||||
|
||||
def mock_upload_folder(*args, **kwargs):
|
||||
folder_path = kwargs["folder_path"]
|
||||
# we include more than is actually uploaded since we ignore {allow,ignore}_patterns of upload_folders()
|
||||
upload_folder_contents.update(os.listdir(folder_path))
|
||||
return MagicMock()
|
||||
|
||||
with (
|
||||
patch("huggingface_hub.HfApi.create_repo"),
|
||||
patch("huggingface_hub.HfApi.upload_folder", mock_upload_folder),
|
||||
):
|
||||
lerobot_train(
|
||||
[
|
||||
f"--policy.path={model_id}",
|
||||
"--policy.push_to_hub=true",
|
||||
"--policy.repo_id=foo/bar",
|
||||
"--policy.input_features=null",
|
||||
"--policy.output_features=null",
|
||||
"--peft.method=LORA",
|
||||
"--dataset.repo_id=lerobot/pusht",
|
||||
"--dataset.episodes=[0, 1]",
|
||||
"--steps=1",
|
||||
f"--output_dir={output_dir}",
|
||||
]
|
||||
)
|
||||
|
||||
assert "adapter_model.safetensors" in upload_folder_contents
|
||||
assert "config.json" in upload_folder_contents
|
||||
assert "adapter_config.json" in upload_folder_contents
|
||||
|
||||
|
||||
@pytest.mark.parametrize("policy_type", ["smolvla"])
|
||||
@require_package("peft")
|
||||
def test_peft_training_works(policy_type, tmp_path):
|
||||
"""Check whether the standard case of fine-tuning a (partially) pre-trained policy with PEFT works."""
|
||||
output_dir = tmp_path / f"output_{policy_type}"
|
||||
model_id = resolve_model_id_for_peft_training(policy_type)
|
||||
|
||||
lerobot_train(
|
||||
[
|
||||
f"--policy.path={model_id}",
|
||||
"--policy.push_to_hub=false",
|
||||
"--policy.input_features=null",
|
||||
"--policy.output_features=null",
|
||||
"--peft.method=LORA",
|
||||
"--dataset.repo_id=lerobot/pusht",
|
||||
"--dataset.episodes=[0, 1]",
|
||||
"--steps=1",
|
||||
f"--output_dir={output_dir}",
|
||||
]
|
||||
)
|
||||
|
||||
policy_dir = output_dir / "checkpoints" / "last" / "pretrained_model"
|
||||
|
||||
for file in ["adapter_config.json", "adapter_model.safetensors", "config.json"]:
|
||||
assert (policy_dir / file).exists()
|
||||
|
||||
# This is the default case where we train a pre-trained policy from scratch with new data.
|
||||
# We assume that we target policy-specific modules but fully fine-tune action and state projections
|
||||
# so these must be part of the trained state dict.
|
||||
state_dict = load_file(policy_dir / "adapter_model.safetensors")
|
||||
|
||||
adapted_keys = [
|
||||
"state_proj",
|
||||
"action_in_proj",
|
||||
"action_out_proj",
|
||||
"action_time_mlp_in",
|
||||
"action_time_mlp_out",
|
||||
]
|
||||
|
||||
found_keys = [
|
||||
module_key
|
||||
for module_key in adapted_keys
|
||||
for state_dict_key in state_dict
|
||||
if f".{module_key}." in state_dict_key
|
||||
]
|
||||
|
||||
assert set(found_keys) == set(adapted_keys)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("policy_type", ["smolvla"])
|
||||
@require_package("peft")
|
||||
def test_peft_training_params_are_fewer(policy_type, tmp_path):
|
||||
"""Check whether the standard case of fine-tuning a (partially) pre-trained policy with PEFT works."""
|
||||
output_dir = tmp_path / f"output_{policy_type}"
|
||||
model_id = resolve_model_id_for_peft_training(policy_type)
|
||||
|
||||
def dummy_update_policy(
|
||||
train_metrics, policy, batch, optimizer, grad_clip_norm: float, accelerator, **kwargs
|
||||
):
|
||||
params_total = sum(p.numel() for p in policy.parameters())
|
||||
params_trainable = sum(p.numel() for p in policy.parameters() if p.requires_grad)
|
||||
|
||||
assert params_total > params_trainable
|
||||
|
||||
return train_metrics, {}
|
||||
|
||||
with patch("lerobot.scripts.lerobot_train.update_policy", dummy_update_policy):
|
||||
lerobot_train(
|
||||
[
|
||||
f"--policy.path={model_id}",
|
||||
"--policy.push_to_hub=false",
|
||||
"--policy.input_features=null",
|
||||
"--policy.output_features=null",
|
||||
"--peft.method=LORA",
|
||||
"--dataset.repo_id=lerobot/pusht",
|
||||
"--dataset.episodes=[0, 1]",
|
||||
"--steps=1",
|
||||
f"--output_dir={output_dir}",
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
class DummyRobot:
|
||||
name = "dummy"
|
||||
cameras = []
|
||||
action_features = {"foo": 1.0, "bar": 2.0}
|
||||
observation_features = {"obs1": 1.0, "obs2": 2.0}
|
||||
is_connected = True
|
||||
|
||||
def connect(self, *args):
|
||||
pass
|
||||
|
||||
def disconnect(self):
|
||||
pass
|
||||
|
||||
|
||||
def dummy_make_robot_from_config(*args, **kwargs):
|
||||
return DummyRobot()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("policy_type", ["smolvla"])
|
||||
@require_package("peft")
|
||||
def test_peft_record_loads_policy(policy_type, tmp_path):
|
||||
"""Train a policy with PEFT and attempt to load it with `lerobot-record`."""
|
||||
from peft import PeftModel
|
||||
|
||||
output_dir = tmp_path / f"output_{policy_type}"
|
||||
model_id = resolve_model_id_for_peft_training(policy_type)
|
||||
|
||||
lerobot_train(
|
||||
[
|
||||
f"--policy.path={model_id}",
|
||||
"--policy.push_to_hub=false",
|
||||
"--policy.input_features=null",
|
||||
"--policy.output_features=null",
|
||||
"--peft.method=LORA",
|
||||
"--dataset.repo_id=lerobot/pusht",
|
||||
"--dataset.episodes=[0, 1]",
|
||||
"--steps=1",
|
||||
f"--output_dir={output_dir}",
|
||||
]
|
||||
)
|
||||
|
||||
policy_dir = output_dir / "checkpoints" / "last" / "pretrained_model"
|
||||
dataset_dir = tmp_path / "eval_pusht"
|
||||
single_task = "move the table"
|
||||
loaded_policy = None
|
||||
|
||||
def dummy_record_loop(*args, **kwargs):
|
||||
nonlocal loaded_policy
|
||||
|
||||
if "dataset" not in kwargs:
|
||||
return
|
||||
|
||||
dataset = kwargs["dataset"]
|
||||
dataset.add_frame({"task": single_task})
|
||||
loaded_policy = kwargs["policy"]
|
||||
|
||||
with (
|
||||
patch("lerobot.scripts.lerobot_record.make_robot_from_config", dummy_make_robot_from_config),
|
||||
# disable record loop since we're only interested in successful loading of the policy.
|
||||
patch("lerobot.scripts.lerobot_record.record_loop", dummy_record_loop),
|
||||
# disable speech output
|
||||
patch("lerobot.utils.utils.say"),
|
||||
):
|
||||
lerobot_record(
|
||||
[
|
||||
f"--policy.path={policy_dir}",
|
||||
"--robot.type=so101_follower",
|
||||
"--robot.port=/dev/null",
|
||||
"--dataset.repo_id=lerobot/eval_pusht",
|
||||
f'--dataset.single_task="{single_task}"',
|
||||
f"--dataset.root={dataset_dir}",
|
||||
"--dataset.push_to_hub=false",
|
||||
]
|
||||
)
|
||||
|
||||
assert isinstance(loaded_policy, PeftModel)
|
||||
@@ -26,7 +26,7 @@ from lerobot.utils.transition import Transition
|
||||
from tests.utils import require_cuda, require_package
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def test_bytes_buffer_size_empty_buffer():
|
||||
from lerobot.transport.utils import bytes_buffer_size
|
||||
|
||||
@@ -37,7 +37,7 @@ def test_bytes_buffer_size_empty_buffer():
|
||||
assert buffer.tell() == 0
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def test_bytes_buffer_size_small_buffer():
|
||||
from lerobot.transport.utils import bytes_buffer_size
|
||||
|
||||
@@ -47,7 +47,7 @@ def test_bytes_buffer_size_small_buffer():
|
||||
assert buffer.tell() == 0
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def test_bytes_buffer_size_large_buffer():
|
||||
from lerobot.transport.utils import CHUNK_SIZE, bytes_buffer_size
|
||||
|
||||
@@ -58,7 +58,7 @@ def test_bytes_buffer_size_large_buffer():
|
||||
assert buffer.tell() == 0
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def test_send_bytes_in_chunks_empty_data():
|
||||
from lerobot.transport.utils import send_bytes_in_chunks, services_pb2
|
||||
|
||||
@@ -68,7 +68,7 @@ def test_send_bytes_in_chunks_empty_data():
|
||||
assert len(chunks) == 0
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def test_single_chunk_small_data():
|
||||
from lerobot.transport.utils import send_bytes_in_chunks, services_pb2
|
||||
|
||||
@@ -82,7 +82,7 @@ def test_single_chunk_small_data():
|
||||
assert chunks[0].transfer_state == services_pb2.TransferState.TRANSFER_END
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def test_not_silent_mode():
|
||||
from lerobot.transport.utils import send_bytes_in_chunks, services_pb2
|
||||
|
||||
@@ -94,7 +94,7 @@ def test_not_silent_mode():
|
||||
assert chunks[0].data == b"Some data"
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def test_send_bytes_in_chunks_large_data():
|
||||
from lerobot.transport.utils import CHUNK_SIZE, send_bytes_in_chunks, services_pb2
|
||||
|
||||
@@ -111,7 +111,7 @@ def test_send_bytes_in_chunks_large_data():
|
||||
assert chunks[2].transfer_state == services_pb2.TransferState.TRANSFER_END
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def test_send_bytes_in_chunks_large_data_with_exact_chunk_size():
|
||||
from lerobot.transport.utils import CHUNK_SIZE, send_bytes_in_chunks, services_pb2
|
||||
|
||||
@@ -124,7 +124,7 @@ def test_send_bytes_in_chunks_large_data_with_exact_chunk_size():
|
||||
assert chunks[0].transfer_state == services_pb2.TransferState.TRANSFER_END
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def test_receive_bytes_in_chunks_empty_data():
|
||||
from lerobot.transport.utils import receive_bytes_in_chunks
|
||||
|
||||
@@ -138,7 +138,7 @@ def test_receive_bytes_in_chunks_empty_data():
|
||||
assert queue.empty()
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def test_receive_bytes_in_chunks_single_chunk():
|
||||
from lerobot.transport.utils import receive_bytes_in_chunks, services_pb2
|
||||
|
||||
@@ -157,7 +157,7 @@ def test_receive_bytes_in_chunks_single_chunk():
|
||||
assert queue.empty()
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def test_receive_bytes_in_chunks_single_not_end_chunk():
|
||||
from lerobot.transport.utils import receive_bytes_in_chunks, services_pb2
|
||||
|
||||
@@ -175,7 +175,7 @@ def test_receive_bytes_in_chunks_single_not_end_chunk():
|
||||
assert queue.empty()
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def test_receive_bytes_in_chunks_multiple_chunks():
|
||||
from lerobot.transport.utils import receive_bytes_in_chunks, services_pb2
|
||||
|
||||
@@ -199,7 +199,7 @@ def test_receive_bytes_in_chunks_multiple_chunks():
|
||||
assert queue.empty()
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def test_receive_bytes_in_chunks_multiple_messages():
|
||||
from lerobot.transport.utils import receive_bytes_in_chunks, services_pb2
|
||||
|
||||
@@ -235,7 +235,7 @@ def test_receive_bytes_in_chunks_multiple_messages():
|
||||
assert queue.empty()
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def test_receive_bytes_in_chunks_shutdown_during_receive():
|
||||
from lerobot.transport.utils import receive_bytes_in_chunks, services_pb2
|
||||
|
||||
@@ -259,7 +259,7 @@ def test_receive_bytes_in_chunks_shutdown_during_receive():
|
||||
assert queue.empty()
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def test_receive_bytes_in_chunks_only_begin_chunk():
|
||||
from lerobot.transport.utils import receive_bytes_in_chunks, services_pb2
|
||||
|
||||
@@ -279,7 +279,7 @@ def test_receive_bytes_in_chunks_only_begin_chunk():
|
||||
assert queue.empty()
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def test_receive_bytes_in_chunks_missing_begin():
|
||||
from lerobot.transport.utils import receive_bytes_in_chunks, services_pb2
|
||||
|
||||
@@ -303,7 +303,7 @@ def test_receive_bytes_in_chunks_missing_begin():
|
||||
|
||||
|
||||
# Tests for state_to_bytes and bytes_to_state_dict
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def test_state_to_bytes_empty_dict():
|
||||
from lerobot.transport.utils import bytes_to_state_dict, state_to_bytes
|
||||
|
||||
@@ -314,7 +314,7 @@ def test_state_to_bytes_empty_dict():
|
||||
assert reconstructed == state_dict
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def test_bytes_to_state_dict_empty_data():
|
||||
from lerobot.transport.utils import bytes_to_state_dict
|
||||
|
||||
@@ -323,7 +323,7 @@ def test_bytes_to_state_dict_empty_data():
|
||||
bytes_to_state_dict(b"")
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def test_state_to_bytes_simple_dict():
|
||||
from lerobot.transport.utils import bytes_to_state_dict, state_to_bytes
|
||||
|
||||
@@ -347,7 +347,7 @@ def test_state_to_bytes_simple_dict():
|
||||
assert torch.allclose(state_dict[key], reconstructed[key])
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def test_state_to_bytes_various_dtypes():
|
||||
from lerobot.transport.utils import bytes_to_state_dict, state_to_bytes
|
||||
|
||||
@@ -372,7 +372,7 @@ def test_state_to_bytes_various_dtypes():
|
||||
assert torch.allclose(state_dict[key], reconstructed[key])
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def test_bytes_to_state_dict_invalid_data():
|
||||
from lerobot.transport.utils import bytes_to_state_dict
|
||||
|
||||
@@ -382,7 +382,7 @@ def test_bytes_to_state_dict_invalid_data():
|
||||
|
||||
|
||||
@require_cuda
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def test_state_to_bytes_various_dtypes_cuda():
|
||||
from lerobot.transport.utils import bytes_to_state_dict, state_to_bytes
|
||||
|
||||
@@ -407,7 +407,7 @@ def test_state_to_bytes_various_dtypes_cuda():
|
||||
assert torch.allclose(state_dict[key], reconstructed[key])
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def test_python_object_to_bytes_none():
|
||||
from lerobot.transport.utils import bytes_to_python_object, python_object_to_bytes
|
||||
|
||||
@@ -439,7 +439,7 @@ def test_python_object_to_bytes_none():
|
||||
(1, 2, 3),
|
||||
],
|
||||
)
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def test_python_object_to_bytes_simple_types(obj):
|
||||
from lerobot.transport.utils import bytes_to_python_object, python_object_to_bytes
|
||||
|
||||
@@ -450,7 +450,7 @@ def test_python_object_to_bytes_simple_types(obj):
|
||||
assert type(reconstructed) is type(obj)
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def test_python_object_to_bytes_with_tensors():
|
||||
from lerobot.transport.utils import bytes_to_python_object, python_object_to_bytes
|
||||
|
||||
@@ -475,7 +475,7 @@ def test_python_object_to_bytes_with_tensors():
|
||||
assert torch.equal(obj["nested"]["tensor2"], reconstructed["nested"]["tensor2"])
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def test_transitions_to_bytes_empty_list():
|
||||
from lerobot.transport.utils import bytes_to_transitions, transitions_to_bytes
|
||||
|
||||
@@ -487,7 +487,7 @@ def test_transitions_to_bytes_empty_list():
|
||||
assert isinstance(reconstructed, list)
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def test_transitions_to_bytes_single_transition():
|
||||
from lerobot.transport.utils import bytes_to_transitions, transitions_to_bytes
|
||||
|
||||
@@ -509,7 +509,7 @@ def test_transitions_to_bytes_single_transition():
|
||||
assert_transitions_equal(transitions[0], reconstructed[0])
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def assert_transitions_equal(t1: Transition, t2: Transition):
|
||||
"""Helper to assert two transitions are equal."""
|
||||
assert_observation_equal(t1["state"], t2["state"])
|
||||
@@ -519,7 +519,7 @@ def assert_transitions_equal(t1: Transition, t2: Transition):
|
||||
assert_observation_equal(t1["next_state"], t2["next_state"])
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def assert_observation_equal(o1: dict, o2: dict):
|
||||
"""Helper to assert two observations are equal."""
|
||||
assert set(o1.keys()) == set(o2.keys())
|
||||
@@ -527,7 +527,7 @@ def assert_observation_equal(o1: dict, o2: dict):
|
||||
assert torch.allclose(o1[key], o2[key])
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def test_transitions_to_bytes_multiple_transitions():
|
||||
from lerobot.transport.utils import bytes_to_transitions, transitions_to_bytes
|
||||
|
||||
@@ -551,7 +551,7 @@ def test_transitions_to_bytes_multiple_transitions():
|
||||
assert_transitions_equal(original, reconstructed_item)
|
||||
|
||||
|
||||
@require_package("grpc")
|
||||
@require_package("grpcio", "grpc")
|
||||
def test_receive_bytes_in_chunks_unknown_state():
|
||||
from lerobot.transport.utils import receive_bytes_in_chunks
|
||||
|
||||
|
||||
+2
-2
@@ -167,7 +167,7 @@ def require_package_arg(func):
|
||||
return wrapper
|
||||
|
||||
|
||||
def require_package(package_name):
|
||||
def require_package(package_name, import_name=None):
|
||||
"""
|
||||
Decorator that skips the test if the specified package is not installed.
|
||||
"""
|
||||
@@ -175,7 +175,7 @@ def require_package(package_name):
|
||||
def decorator(func):
|
||||
@wraps(func)
|
||||
def wrapper(*args, **kwargs):
|
||||
if not is_package_available(package_name):
|
||||
if not is_package_available(pkg_name=package_name, import_name=import_name):
|
||||
pytest.skip(f"{package_name} not installed")
|
||||
return func(*args, **kwargs)
|
||||
|
||||
|
||||
@@ -82,6 +82,20 @@ def test_save_checkpoint(mock_save_training_state, tmp_path, optimizer):
|
||||
mock_save_training_state.assert_called_once()
|
||||
|
||||
|
||||
@patch("lerobot.utils.train_utils.save_training_state")
|
||||
def test_save_checkpoint_peft(mock_save_training_state, tmp_path, optimizer):
|
||||
policy = Mock()
|
||||
policy.config = Mock()
|
||||
policy.config.save_pretrained = Mock()
|
||||
cfg = Mock()
|
||||
cfg.use_peft = True
|
||||
save_checkpoint(tmp_path, 10, cfg, policy, optimizer)
|
||||
policy.save_pretrained.assert_called_once()
|
||||
cfg.save_pretrained.assert_called_once()
|
||||
policy.config.save_pretrained.assert_called_once()
|
||||
mock_save_training_state.assert_called_once()
|
||||
|
||||
|
||||
def test_save_training_state(tmp_path, optimizer, scheduler):
|
||||
save_training_state(tmp_path, 10, optimizer, scheduler)
|
||||
assert (tmp_path / TRAINING_STATE_DIR).is_dir()
|
||||
|
||||
@@ -41,7 +41,10 @@ def mock_rerun(monkeypatch):
|
||||
def __init__(self, arr):
|
||||
self.arr = arr
|
||||
|
||||
def dummy_log(key, obj, **kwargs):
|
||||
def dummy_log(key, obj=None, **kwargs):
|
||||
# Accept either positional `obj` or keyword `entity` and record remaining kwargs.
|
||||
if obj is None and "entity" in kwargs:
|
||||
obj = kwargs.pop("entity")
|
||||
calls.append((key, obj, kwargs))
|
||||
|
||||
dummy_rr = SimpleNamespace(
|
||||
|
||||
Reference in New Issue
Block a user