Merge branch 'feat/mirror' into openarms_wallx_rebased_3

tests/cameras/test_opencv.py

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

tests/cameras/test_reachy2_camera.py

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

tests/cameras/test_realsense.py

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

tests/datasets/test_aggregate.py

@@ -525,3 +525,92 @@ def test_aggregate_image_datasets(tmp_path, lerobot_dataset_factory):
         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)

tests/datasets/test_dataset_tools.py

@@ -26,6 +26,7 @@ from lerobot.datasets.dataset_tools import (
     delete_episodes,
     merge_datasets,
     modify_features,
+    modify_tasks,
     remove_feature,
     split_dataset,
 )
@@ -1050,6 +1051,174 @@ def test_modify_features_preserves_file_structure(sample_dataset, tmp_path):
         assert "reward" in modified_dataset.meta.features
 
 
+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

tests/datasets/test_image_transforms.py

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

tests/datasets/test_subtask_dataset.py

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

tests/policies/test_delta_actions.py

@@ -0,0 +1,344 @@
+"""Tests for delta action transforms — full pipeline validation.
+
+Tests the complete flow matching OpenPI:
+  raw actions → DeltaActions → Normalize(delta_stats) → model → Unnormalize → AbsoluteActions
+
+Uses real dataset: lerobot-data-collection/dagger_final_1_21
+"""
+
+import numpy as np
+import pytest
+import torch
+
+from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+from lerobot.datasets.compute_stats import get_feature_stats
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.processor import TransitionKey, batch_to_transition
+from lerobot.processor.delta_action_processor import (
+    AbsoluteActionsProcessorStep,
+    DeltaActionsProcessorStep,
+    to_absolute_actions,
+    to_delta_actions,
+)
+from lerobot.processor.normalize_processor import NormalizerProcessorStep, UnnormalizerProcessorStep
+from lerobot.utils.constants import ACTION, OBS_STATE
+
+CHUNK_SIZE = 10
+REPO_ID = "lerobot-data-collection/dagger_final_1_21"
+
+
+@pytest.fixture(scope="module")
+def dataset():
+    return LeRobotDataset(REPO_ID, episodes=[0])
+
+
+@pytest.fixture(scope="module")
+def action_dim(dataset):
+    return dataset.meta.features["action"]["shape"][0]
+
+
+def _build_action_chunks(dataset, chunk_size, max_chunks=50):
+    """Build action chunks from hf_dataset, like the training script does."""
+    hf = dataset.hf_dataset
+    total = len(hf)
+    all_ep = torch.tensor([int(hf[i]["episode_index"]) for i in range(total)])
+    chunks, states = [], []
+    for i in range(total - chunk_size + 1):
+        if all_ep[i] != all_ep[i + chunk_size - 1]:
+            continue
+        chunk_actions = torch.stack([hf[i + k]["action"] for k in range(chunk_size)]).float()
+        state = hf[i]["observation.state"].float()
+        chunks.append(chunk_actions)
+        states.append(state)
+        if len(chunks) >= max_chunks:
+            break
+    assert len(chunks) > 0, f"No valid chunks found. total={total}, ep_indices={all_ep.tolist()}"
+    return torch.stack(chunks), torch.stack(states)
+
+
+def _compute_delta_chunk_stats(action_chunks, states, mask):
+    all_deltas = []
+    for actions, state in zip(action_chunks, states):
+        delta = to_delta_actions(actions.unsqueeze(0), state.unsqueeze(0), mask).squeeze(0)
+        all_deltas.append(delta.numpy())
+    all_delta = np.concatenate(all_deltas, axis=0)
+    return get_feature_stats(all_delta, axis=0, keepdims=all_delta.ndim == 1)
+
+
+# --- Basic roundtrip tests ---
+
+def test_roundtrip_3d(action_dim):
+    actions = torch.randn(4, CHUNK_SIZE, action_dim)
+    state = torch.randn(4, action_dim)
+    mask = [True] * action_dim
+    recovered = to_absolute_actions(to_delta_actions(actions, state, mask), state, mask)
+    torch.testing.assert_close(recovered, actions)
+
+
+def test_roundtrip_2d(action_dim):
+    actions = torch.randn(4, action_dim)
+    state = torch.randn(4, action_dim)
+    mask = [True] * action_dim
+    recovered = to_absolute_actions(to_delta_actions(actions, state, mask), state, mask)
+    torch.testing.assert_close(recovered, actions)
+
+
+def test_no_mutation(action_dim):
+    actions = torch.randn(2, CHUNK_SIZE, action_dim)
+    original = actions.clone()
+    state = torch.randn(2, action_dim)
+    to_delta_actions(actions, state, [True] * action_dim)
+    torch.testing.assert_close(actions, original)
+
+
+def test_exclude_joints_supports_partial_name_matching():
+    names = [
+        "right_joint_1.pos",
+        "right_gripper.pos",
+        "left_joint_1.pos",
+        "left_gripper.pos",
+    ]
+    step = DeltaActionsProcessorStep(enabled=True, exclude_joints=["gripper"], action_names=names)
+    assert step._build_mask(len(names)) == [True, False, True, False]
+
+
+# --- Chunk-level delta stats test ---
+
+def test_chunk_stats_have_larger_std_than_frame_stats(dataset, action_dim):
+    """Chunk-level delta stats should have larger std than per-frame delta stats."""
+    action_chunks, states = _build_action_chunks(dataset, CHUNK_SIZE)
+    mask = [True] * action_dim
+
+    chunk_stats = _compute_delta_chunk_stats(action_chunks, states, mask)
+
+    # Per-frame stats
+    hf = dataset.hf_dataset
+    n = min(500, len(hf))
+    frame_actions = torch.stack([hf[i]["action"] for i in range(n)]).float()
+    frame_states = torch.stack([hf[i]["observation.state"] for i in range(n)]).float()
+    frame_deltas = to_delta_actions(frame_actions, frame_states, mask).numpy()
+    frame_stats = get_feature_stats(frame_deltas, axis=0, keepdims=frame_deltas.ndim == 1)
+
+    assert chunk_stats["std"].mean() >= frame_stats["std"].mean(), (
+        f"Chunk std ({chunk_stats['std'].mean():.4f}) should be >= "
+        f"frame std ({frame_stats['std'].mean():.4f})"
+    )
+
+
+# --- Full pipeline roundtrip: delta → normalize → unnormalize → absolute ---
+
+def test_full_pipeline_roundtrip(dataset, action_dim):
+    """Test the complete OpenPI pipeline: delta → normalize → unnormalize → absolute."""
+    action_chunks, states = _build_action_chunks(dataset, CHUNK_SIZE)
+    mask = [True] * action_dim
+
+    delta_stats = _compute_delta_chunk_stats(action_chunks, states, mask)
+    stats = {ACTION: {k: v for k, v in delta_stats.items()}}
+
+    features = {ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(action_dim,))}
+    norm_map = {FeatureType.ACTION: NormalizationMode.MEAN_STD}
+
+    delta_step = DeltaActionsProcessorStep(enabled=True)
+    normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
+    unnormalizer = UnnormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
+    absolute_step = AbsoluteActionsProcessorStep(enabled=True, delta_step=delta_step)
+
+    original_actions = action_chunks[0].unsqueeze(0)
+    state = states[0].unsqueeze(0)
+
+    batch = {ACTION: original_actions, OBS_STATE: state}
+    transition = batch_to_transition(batch)
+
+    # Forward: delta → normalize
+    t1 = delta_step(transition)
+    t2 = normalizer(t1)
+
+    normalized_action = t2[TransitionKey.ACTION]
+    assert normalized_action.abs().mean() < 10, (
+        f"Normalized actions should be in reasonable range, got mean abs {normalized_action.abs().mean():.2f}"
+    )
+
+    # Reverse: unnormalize → absolute
+    t3 = unnormalizer(t2)
+    t4 = absolute_step(t3)
+
+    recovered_actions = t4[TransitionKey.ACTION]
+    torch.testing.assert_close(recovered_actions, original_actions, atol=1e-4, rtol=1e-4)
+
+
+def test_normalized_delta_values_are_reasonable(dataset, action_dim):
+    """With correct chunk stats, normalized delta actions should be in a reasonable range."""
+    action_chunks, states = _build_action_chunks(dataset, CHUNK_SIZE)
+    mask = [True] * action_dim
+
+    delta_stats = _compute_delta_chunk_stats(action_chunks, states, mask)
+    mean = torch.tensor(delta_stats["mean"]).float()
+    std = torch.tensor(delta_stats["std"]).float()
+
+    all_normalized = []
+    for actions, state in zip(action_chunks, states):
+        delta = to_delta_actions(actions.unsqueeze(0), state.unsqueeze(0), mask).squeeze(0)
+        normalized = (delta - mean) / (std + 1e-6)
+        all_normalized.append(normalized)
+
+    all_normalized = torch.cat(all_normalized, dim=0)
+
+    pct_in_range = (all_normalized.abs() < 5).float().mean()
+    assert pct_in_range > 0.9, (
+        f"Only {pct_in_range*100:.1f}% of normalized values in [-5, 5], expected >90%"
+    )
+
+    assert all_normalized.mean().abs() < 1.0, (
+        f"Mean of normalized deltas is {all_normalized.mean():.2f}, expected near 0"
+    )
+
+
+def test_processor_step_roundtrip(dataset, action_dim):
+    """DeltaActionsProcessorStep applies delta; to_absolute_actions recovers original."""
+    hf = dataset.hf_dataset
+    batch = {
+        ACTION: torch.stack([hf[i]["action"] for i in range(4)]),
+        OBS_STATE: torch.stack([hf[i]["observation.state"] for i in range(4)]),
+    }
+    original_actions = batch[ACTION].clone()
+    transition = batch_to_transition(batch)
+
+    step = DeltaActionsProcessorStep(enabled=True)
+    delta_transition = step(transition)
+    assert not torch.allclose(delta_transition[TransitionKey.ACTION], original_actions)
+
+    state = transition[TransitionKey.OBSERVATION][OBS_STATE]
+    mask = [True] * action_dim
+    recovered = to_absolute_actions(delta_transition[TransitionKey.ACTION], state, mask)
+    torch.testing.assert_close(recovered, original_actions)
+
+
+def test_processor_step_disabled_is_noop(dataset, action_dim):
+    """enabled=False should be a no-op."""
+    hf = dataset.hf_dataset
+    batch = {
+        ACTION: torch.stack([hf[i]["action"] for i in range(2)]),
+        OBS_STATE: torch.stack([hf[i]["observation.state"] for i in range(2)]),
+    }
+    original = batch[ACTION].clone()
+    transition = batch_to_transition(batch)
+    result = DeltaActionsProcessorStep(enabled=False)(transition)
+    torch.testing.assert_close(result[TransitionKey.ACTION], original)
+
+
+# --- Training batch shape validation ---
+
+def test_delta_with_action_chunks(dataset, action_dim):
+    """Verify delta works correctly with (B, chunk_size, action_dim) shaped actions."""
+    action_chunks, states = _build_action_chunks(dataset, CHUNK_SIZE)
+
+    # Simulate a training batch: actions=(B, chunk_size, action_dim), state=(B, state_dim)
+    batch_actions = action_chunks[:4]  # (4, chunk_size, action_dim)
+    batch_states = states[:4]  # (4, state_dim)
+
+    mask = [True] * action_dim
+    delta = to_delta_actions(batch_actions, batch_states, mask)
+
+    # First action in each chunk should be close to zero (action[t] - state[t] ≈ small)
+    first_deltas = delta[:, 0, :]  # (B, action_dim)
+    assert first_deltas.abs().mean() < delta.abs().mean(), (
+        f"First action in chunk should have smaller delta than average. "
+        f"First: {first_deltas.abs().mean():.4f}, Average: {delta.abs().mean():.4f}"
+    )
+
+    # Later actions should have larger deltas
+    last_deltas = delta[:, -1, :]  # (B, action_dim)
+    assert last_deltas.abs().mean() >= first_deltas.abs().mean(), (
+        f"Last action in chunk should have >= delta than first. "
+        f"Last: {last_deltas.abs().mean():.4f}, First: {first_deltas.abs().mean():.4f}"
+    )
+
+    # Roundtrip
+    recovered = to_absolute_actions(delta, batch_states, mask)
+    torch.testing.assert_close(recovered, batch_actions)
+
+
+def test_delta_stats_match_actual_data_distribution(dataset, action_dim):
+    """Verify computed stats match the actual delta distribution."""
+    action_chunks, states = _build_action_chunks(dataset, CHUNK_SIZE)
+    mask = [True] * action_dim
+
+    # Compute stats like the training script does
+    delta_stats = _compute_delta_chunk_stats(action_chunks, states, mask)
+
+    # Also compute directly
+    all_deltas = []
+    for actions, state in zip(action_chunks, states):
+        delta = to_delta_actions(actions.unsqueeze(0), state.unsqueeze(0), mask).squeeze(0)
+        all_deltas.append(delta)
+    all_deltas_tensor = torch.cat(all_deltas, dim=0)
+
+    # Compare mean
+    actual_mean = all_deltas_tensor.mean(dim=0).numpy()
+    np.testing.assert_allclose(delta_stats["mean"], actual_mean, atol=0.01)
+
+    # Compare std
+    actual_std = all_deltas_tensor.std(dim=0).numpy()
+    np.testing.assert_allclose(delta_stats["std"], actual_std, atol=0.1)
+
+    # Verify q01 < mean < q99
+    assert (delta_stats["q01"] < delta_stats["mean"]).all(), "q01 should be < mean"
+    assert (delta_stats["mean"] < delta_stats["q99"]).all(), "mean should be < q99"
+
+
+def test_quantile_normalization_roundtrip(dataset, action_dim):
+    """Full roundtrip with QUANTILES normalization (what OpenPI uses for pi05)."""
+    action_chunks, states = _build_action_chunks(dataset, CHUNK_SIZE)
+    mask = [True] * action_dim
+
+    delta_stats = _compute_delta_chunk_stats(action_chunks, states, mask)
+    stats = {ACTION: {k: v for k, v in delta_stats.items()}}
+
+    features = {ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(action_dim,))}
+    norm_map = {FeatureType.ACTION: NormalizationMode.QUANTILES}
+
+    delta_step = DeltaActionsProcessorStep(enabled=True)
+    normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
+    unnormalizer = UnnormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
+    absolute_step = AbsoluteActionsProcessorStep(enabled=True, delta_step=delta_step)
+
+    original_actions = action_chunks[0].unsqueeze(0)
+    state = states[0].unsqueeze(0)
+
+    batch = {ACTION: original_actions, OBS_STATE: state}
+    transition = batch_to_transition(batch)
+
+    # Forward: delta → quantile normalize
+    t1 = delta_step(transition)
+    t2 = normalizer(t1)
+
+    normalized = t2[TransitionKey.ACTION]
+    # Most values should be in [-1, 1] with quantile normalization
+    pct_in_range = (normalized.abs() < 2).float().mean()
+    assert pct_in_range > 0.5, (
+        f"Only {pct_in_range*100:.1f}% in [-2, 2] after quantile norm, expected >50%"
+    )
+
+    # Reverse: unnormalize → absolute
+    t3 = unnormalizer(t2)
+    t4 = absolute_step(t3)
+
+    recovered = t4[TransitionKey.ACTION]
+    torch.testing.assert_close(recovered, original_actions, atol=1e-3, rtol=1e-3)
+
+
+def test_state_not_modified_by_delta(dataset, action_dim):
+    """State should never be modified by the delta processor."""
+    hf = dataset.hf_dataset
+    batch = {
+        ACTION: torch.stack([hf[i]["action"] for i in range(4)]),
+        OBS_STATE: torch.stack([hf[i]["observation.state"] for i in range(4)]),
+    }
+    original_state = batch[OBS_STATE].clone()
+    transition = batch_to_transition(batch)
+
+    step = DeltaActionsProcessorStep(enabled=True)
+    result = step(transition)
+
+    result_state = result[TransitionKey.OBSERVATION][OBS_STATE]
+    torch.testing.assert_close(result_state, original_state)

tests/policies/test_sac_policy.py

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

tests/processor/test_tokenizer_processor.py

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

tests/scripts/test_edit_dataset_parsing.py

@@ -0,0 +1,71 @@
+#!/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,
+    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),
+        ],
+    )
+    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),
+        ],
+    )
+    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