Merge branch 'feat/training_time_rtc' into openarms_wallx_rebased_3

tests/async_inference/test_e2e.py

@@ -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)

tests/async_inference/test_policy_server.py

@@ -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)

tests/cameras/test_reachy2_camera.py

@@ -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()

tests/conftest.py

@@ -28,7 +28,6 @@ pytest_plugins = [
     "tests.fixtures.files",
     "tests.fixtures.hub",
     "tests.fixtures.optimizers",
-    "tests.plugins.reachy2_sdk",
 ]
 
 

tests/datasets/test_aggregate.py

@@ -16,6 +16,7 @@
 
 from unittest.mock import patch
 
+import datasets
 import torch
 
 from lerobot.datasets.aggregate import aggregate_datasets
@@ -380,3 +381,147 @@ 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)

tests/datasets/test_dataset_tools.py

@@ -29,7 +29,7 @@ from lerobot.datasets.dataset_tools import (
     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 +1050,7 @@ 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_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 +1071,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 +1113,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 +1132,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",

tests/datasets/test_datasets.py

@@ -31,8 +31,10 @@ from lerobot.configs.train import TrainPipelineConfig
 from lerobot.datasets.factory import make_dataset
 from lerobot.datasets.image_writer import image_array_to_pil_image
 from lerobot.datasets.lerobot_dataset import (
+    VALID_VIDEO_CODECS,
     LeRobotDataset,
     MultiLeRobotDataset,
+    _encode_video_worker,
 )
 from lerobot.datasets.utils import (
     DEFAULT_CHUNK_SIZE,
@@ -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
+    )
+    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,300 @@ 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 len(VALID_VIDEO_CODECS) == 3
+
+
+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}"

tests/mocks/mock_robot.py

@@ -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

tests/mocks/mock_teleop.py

@@ -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

tests/plugins/reachy2_sdk.py

@@ -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()

tests/policies/pi0_fast/test_pi0_fast_original_vs_lerobot.py

@@ -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.")

tests/policies/rtc/test_training_time_rtc.py

@@ -0,0 +1,50 @@
+#!/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.
+
+"""Tests for training-time RTC helpers."""
+
+import torch
+
+from lerobot.configs.types import RTCTrainingDelayDistribution
+from lerobot.policies.rtc.configuration_rtc import RTCTrainingConfig
+from lerobot.policies.rtc.training_time import apply_rtc_training_time, sample_rtc_delay
+
+
+def test_rtc_training_config_defaults():
+    config = RTCTrainingConfig()
+    assert config.enabled is False
+    assert config.min_delay == 0
+    assert config.max_delay == 0
+    assert config.delay_distribution == RTCTrainingDelayDistribution.UNIFORM
+    assert config.exp_decay == 1.0
+
+
+def test_sample_rtc_delay_uniform_range():
+    cfg = RTCTrainingConfig(enabled=True, min_delay=1, max_delay=4)
+    delays = sample_rtc_delay(cfg, batch_size=100, device=torch.device("cpu"))
+    assert delays.min().item() >= 1
+    assert delays.max().item() <= 4
+
+
+def test_apply_rtc_training_time_prefix_mask():
+    time = torch.tensor([0.5])
+    delays = torch.tensor([2])
+    time_tokens, postfix_mask = apply_rtc_training_time(time, delays, seq_len=4)
+    assert time_tokens.shape == (1, 4)
+    assert postfix_mask.shape == (1, 4)
+    # Delay=2 means the first two steps are prefix (time forced to 0.0) and only the last two are postfix.
+    assert torch.allclose(time_tokens[0], torch.tensor([0.0, 0.0, 0.5, 0.5]))
+    assert torch.equal(postfix_mask[0], torch.tensor([False, False, True, True]))

tests/processor/test_pipeline.py

@@ -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

tests/rl/test_actor.py

@@ -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

tests/rl/test_actor_learner.py

@@ -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):

tests/rl/test_learner_service.py

@@ -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

tests/robots/test_reachy2.py

@@ -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,

tests/robots/test_so100_follower.py

@@ -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),

tests/test_cli_peft.py

@@ -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)

tests/transport/test_transport_utils.py

@@ -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
 

tests/utils.py

@@ -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)
 

tests/utils/test_train_utils.py

@@ -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()

tests/utils/test_visualization_utils.py

@@ -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(