Merge remote-tracking branch 'origin/main' into user/khalil-meftah/2026-02-16-rl-stack-refactor

# Conflicts:
#	src/lerobot/policies/__init__.py
#	src/lerobot/rl/actor.py

tests/artifacts/policies/aloha_sim_insertion_human_act_1000_steps/actions.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c2b8f8532c7a0b776de5e536b8b54e30b1a0c2e3d5cc25a2d86fe43e40ae5e8c
+oid sha256:8a31653c11eccdd4d80fd3f6a351cd54c49b8a48db1f7e9faf38fddd7900a09f
 size 515400

tests/artifacts/policies/aloha_sim_insertion_human_act_1000_steps/grad_stats.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:224b5fa4828aa88171b68c036e8919c1eae563e2113f03b6461eadf5bf8525a6
+oid sha256:75bf051698b37dcd7517ec8025a896ab5a0551a6dde5f89d0a3d5d50966e83e6
 size 31672

tests/artifacts/policies/aloha_sim_insertion_human_act_1000_steps/output_dict.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:016d2fa8fe5f58017dfd46f4632fdc19dfd751e32a2c7cde2077c6f95546d6bd
+oid sha256:88e10930a10041d50f2cf369e6813ac14618d13dad1c21bdde1ac7798611c6ba
 size 68

tests/artifacts/policies/aloha_sim_insertion_human_act_1000_steps/param_stats.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:eca0d87a699620e4fec7e68539b0be91e4cc933f6bf12032da52c182ab6f38cf
+oid sha256:89833a5ccdb7d85c83f717ff8ec68b8e822005cb8803899acaae88c578e2e3ae
 size 31672

tests/artifacts/policies/save_policy_to_safetensors.py

@@ -52,6 +52,9 @@ def get_policy_stats(ds_repo_id: str, policy_name: str, policy_kwargs: dict):
     )
 
     batch = next(iter(dataloader))
+    for key in batch:
+        if isinstance(batch[key], torch.Tensor) and batch[key].dtype == torch.uint8:
+            batch[key] = batch[key].to(dtype=torch.float32) / 255.0
     batch = preprocessor(batch)
     loss, output_dict = policy.forward(batch)
 
@@ -82,6 +85,9 @@ def get_policy_stats(ds_repo_id: str, policy_name: str, policy_kwargs: dict):
     # indicating padding (those ending with "_is_pad")
     dataset.reader.delta_indices = None
     batch = next(iter(dataloader))
+    for key in batch:
+        if isinstance(batch[key], torch.Tensor) and batch[key].dtype == torch.uint8:
+            batch[key] = batch[key].to(dtype=torch.float32) / 255.0
     obs = {}
     for k in batch:
         # TODO: regenerate the safetensors

tests/datasets/test_datasets.py

@@ -454,6 +454,35 @@ def test_tmp_video_deletion(tmp_path, empty_lerobot_dataset_factory):
     )
 
 
+def test_cleanup_interrupted_episode_removes_image_temp_dirs(tmp_path, empty_lerobot_dataset_factory):
+    """Verify interrupted episode cleanup removes temporary image directories for both image and video features."""
+    features = {
+        "image": {"dtype": "image", "shape": DUMMY_CHW, "names": ["channels", "height", "width"]},
+        "video": {"dtype": "video", "shape": DUMMY_HWC, "names": ["height", "width", "channels"]},
+    }
+    ds = empty_lerobot_dataset_factory(
+        root=tmp_path / "interrupted", features=features, streaming_encoding=False
+    )
+    # Add one frame without saving episode simulating an interruption
+    ds.add_frame(
+        {
+            "image": np.random.rand(*DUMMY_CHW),
+            "video": np.random.rand(*DUMMY_HWC),
+            "task": "Dummy task",
+        }
+    )
+    img_dir = ds.writer._get_image_file_dir(0, "image")
+    vid_img_dir = ds.writer._get_image_file_dir(0, "video")
+    # Precondition: both temp dirs exist after add_frame.
+    assert img_dir.exists()
+    assert vid_img_dir.exists()
+
+    ds.writer.cleanup_interrupted_episode(episode_index=0)
+
+    assert not img_dir.exists(), "image temp dir leaked after cleanup_interrupted_episode"
+    assert not vid_img_dir.exists(), "video temp dir leaked after cleanup_interrupted_episode"
+
+
 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"

tests/datasets/test_lerobot_dataset.py

@@ -416,6 +416,18 @@ def test_create_initial_counts_zero(tmp_path):
     assert dataset.num_frames == 0
 
 
+def test_create_propagates_video_files_size_in_mb(tmp_path):
+    """video_files_size_in_mb passed to create() is reflected in the dataset metadata."""
+    dataset = LeRobotDataset.create(
+        repo_id=DUMMY_REPO_ID,
+        fps=DEFAULT_FPS,
+        features=SIMPLE_FEATURES,
+        root=tmp_path / "ds",
+        video_files_size_in_mb=42.0,
+    )
+    assert dataset.meta.video_files_size_in_mb == 42.0
+
+
 def test_add_frame_works_in_write_mode(tmp_path):
     """add_frame() succeeds on a dataset created via create()."""
     dataset = LeRobotDataset.create(

tests/envs/test_robotwin.py

@@ -0,0 +1,282 @@
+#!/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.
+"""Unit tests for the RoboTwin 2.0 Gymnasium wrapper.
+
+These tests mock out the SAPIEN-based RoboTwin runtime (task modules +
+YAML config loader) so they run without the full RoboTwin installation
+(SAPIEN, CuRobo, mplib, asset downloads, etc.).
+"""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from unittest.mock import MagicMock, patch
+
+import gymnasium as gym
+import numpy as np
+import pytest
+
+from lerobot.envs.robotwin import (
+    ACTION_DIM,
+    ROBOTWIN_CAMERA_NAMES,
+    ROBOTWIN_TASKS,
+    RoboTwinEnv,
+    create_robotwin_envs,
+)
+
+# ---------------------------------------------------------------------------
+# Fixtures / helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_mock_task_env(
+    height: int = 240,
+    width: int = 320,
+    cameras: tuple[str, ...] = ROBOTWIN_CAMERA_NAMES,
+) -> MagicMock:
+    """Return a mock that mimics the RoboTwin task class API.
+
+    RoboTwin's real get_obs returns
+        {"observation": {cam: {"rgb": img}}, "joint_action": {"vector": np.ndarray}, ...}
+    so the mock follows the same nested shape.
+    """
+    obs_dict = {
+        "observation": {cam: {"rgb": np.zeros((height, width, 3), dtype=np.uint8)} for cam in cameras},
+        "joint_action": {"vector": np.zeros(ACTION_DIM, dtype=np.float32)},
+        "endpose": {},
+    }
+
+    mock = MagicMock()
+    mock.get_obs.return_value = obs_dict
+    mock.setup_demo.return_value = None
+    mock.take_action.return_value = None
+    mock.eval_success = False
+    mock.check_success.return_value = False
+    mock.close_env.return_value = None
+    return mock
+
+
+@contextmanager
+def _patch_runtime(mock_task_instance: MagicMock):
+    """Patch both the task-class loader and the YAML config loader so the
+    env can construct + reset without a real RoboTwin install."""
+    task_cls = MagicMock(return_value=mock_task_instance)
+    fake_setup = {
+        "head_camera_h": 240,
+        "head_camera_w": 320,
+        "left_embodiment_config": {},
+        "right_embodiment_config": {},
+        "left_robot_file": "",
+        "right_robot_file": "",
+        "dual_arm_embodied": True,
+        "render_freq": 0,
+        "task_name": "beat_block_hammer",
+        "task_config": "demo_clean",
+    }
+    with (
+        patch("lerobot.envs.robotwin._load_robotwin_task", return_value=task_cls),
+        patch("lerobot.envs.robotwin._load_robotwin_setup_kwargs", return_value=fake_setup),
+    ):
+        yield
+
+
+# ---------------------------------------------------------------------------
+# RoboTwinEnv unit tests
+# ---------------------------------------------------------------------------
+
+
+class TestRoboTwinEnv:
+    def test_observation_space_shape(self):
+        """observation_space should have the configured h×w×3 for every camera."""
+        h, w = 240, 320
+        env = RoboTwinEnv(
+            task_name="beat_block_hammer",
+            observation_height=h,
+            observation_width=w,
+            camera_names=["head_camera", "left_camera"],
+        )
+        pixels_space = env.observation_space["pixels"]
+        assert pixels_space["head_camera"].shape == (h, w, 3)
+        assert pixels_space["left_camera"].shape == (h, w, 3)
+        assert "right_camera" not in pixels_space
+
+    def test_action_space(self):
+        env = RoboTwinEnv(task_name="beat_block_hammer")
+        assert env.action_space.shape == (ACTION_DIM,)
+        assert env.action_space.dtype == np.float32
+
+    def test_reset_returns_correct_obs_keys(self):
+        mock_task = _make_mock_task_env()
+        env = RoboTwinEnv(task_name="beat_block_hammer")
+        with _patch_runtime(mock_task):
+            obs, info = env.reset()
+
+        assert "pixels" in obs
+        for cam in ROBOTWIN_CAMERA_NAMES:
+            assert cam in obs["pixels"], f"Missing camera '{cam}' in obs"
+        assert "agent_pos" in obs
+        assert obs["agent_pos"].shape == (ACTION_DIM,)
+        assert info["is_success"] is False
+
+    def test_reset_calls_setup_demo(self):
+        mock_task = _make_mock_task_env()
+        env = RoboTwinEnv(task_name="beat_block_hammer")
+        with _patch_runtime(mock_task):
+            env.reset(seed=42)
+        # setup_demo receives the full YAML-derived kwargs plus seed + is_test;
+        # we only assert the caller-provided bits.
+        assert mock_task.setup_demo.call_count == 1
+        call_kwargs = mock_task.setup_demo.call_args.kwargs
+        assert call_kwargs["seed"] == 42
+        assert call_kwargs["is_test"] is True
+
+    def test_step_returns_correct_types(self):
+        mock_task = _make_mock_task_env()
+        env = RoboTwinEnv(task_name="beat_block_hammer")
+        action = np.zeros(ACTION_DIM, dtype=np.float32)
+        with _patch_runtime(mock_task):
+            env.reset()
+            obs, reward, terminated, truncated, info = env.step(action)
+
+        assert isinstance(obs, dict)
+        assert isinstance(reward, float)
+        assert isinstance(terminated, bool)
+        assert isinstance(truncated, bool)
+        assert isinstance(info, dict)
+
+    def test_step_wrong_action_shape_raises(self):
+        mock_task = _make_mock_task_env()
+        env = RoboTwinEnv(task_name="beat_block_hammer")
+        bad_action = np.zeros(7, dtype=np.float32)  # wrong dim
+        with _patch_runtime(mock_task):
+            env.reset()
+            with pytest.raises(ValueError, match="Expected 1-D action"):
+                env.step(bad_action)
+
+    def test_success_terminates_episode(self):
+        mock_task = _make_mock_task_env()
+        mock_task.check_success.return_value = True
+        env = RoboTwinEnv(task_name="beat_block_hammer")
+        action = np.zeros(ACTION_DIM, dtype=np.float32)
+        with _patch_runtime(mock_task):
+            env.reset()
+            _, _, terminated, _, info = env.step(action)
+        assert terminated is True
+        assert info["is_success"] is True
+
+    def test_truncation_after_episode_length(self):
+        mock_task = _make_mock_task_env()
+        env = RoboTwinEnv(task_name="beat_block_hammer", episode_length=2)
+        action = np.zeros(ACTION_DIM, dtype=np.float32)
+        with _patch_runtime(mock_task):
+            env.reset()
+            env.step(action)  # step 1
+            _, _, _, truncated, _ = env.step(action)  # step 2 → truncated
+        assert truncated is True
+
+    def test_close_calls_close_env(self):
+        mock_task = _make_mock_task_env()
+        env = RoboTwinEnv(task_name="beat_block_hammer")
+        with _patch_runtime(mock_task):
+            env.reset()
+            env.close()
+        mock_task.close_env.assert_called_once()
+
+    def test_black_frame_for_missing_camera(self):
+        """If a camera key is absent from get_obs(), a black frame is returned."""
+        # Mock exposes only head_camera; we ask for both head_camera + left_camera.
+        mock_task = _make_mock_task_env(height=10, width=10, cameras=("head_camera",))
+        env = RoboTwinEnv(
+            task_name="beat_block_hammer",
+            camera_names=["head_camera", "left_camera"],
+            observation_height=10,
+            observation_width=10,
+        )
+        with _patch_runtime(mock_task):
+            obs, _ = env.reset()
+        assert obs["pixels"]["left_camera"].shape == (10, 10, 3)
+        assert obs["pixels"]["left_camera"].sum() == 0
+
+    def test_task_and_task_description_attributes(self):
+        env = RoboTwinEnv(task_name="beat_block_hammer")
+        assert env.task == "beat_block_hammer"
+        assert isinstance(env.task_description, str)
+
+    def test_deferred_init_env_is_none_before_reset(self):
+        env = RoboTwinEnv(task_name="beat_block_hammer")
+        assert env._env is None  # noqa: SLF001  (testing internal state)
+
+
+# ---------------------------------------------------------------------------
+# create_robotwin_envs tests
+# ---------------------------------------------------------------------------
+
+
+class TestCreateRoboTwinEnvs:
+    def test_returns_correct_structure(self):
+        mock_task = _make_mock_task_env()
+        with _patch_runtime(mock_task):
+            envs = create_robotwin_envs(
+                task="beat_block_hammer",
+                n_envs=1,
+                env_cls=gym.vector.SyncVectorEnv,
+            )
+        assert "beat_block_hammer" in envs
+        assert 0 in envs["beat_block_hammer"]
+        assert isinstance(envs["beat_block_hammer"][0], gym.vector.SyncVectorEnv)
+
+    def test_multi_task(self):
+        mock_task = _make_mock_task_env()
+        with _patch_runtime(mock_task):
+            envs = create_robotwin_envs(
+                task="beat_block_hammer,click_bell",
+                n_envs=1,
+                env_cls=gym.vector.SyncVectorEnv,
+            )
+        assert set(envs.keys()) == {"beat_block_hammer", "click_bell"}
+
+    def test_unknown_task_raises(self):
+        with pytest.raises(ValueError, match="Unknown RoboTwin tasks"):
+            create_robotwin_envs(
+                task="not_a_real_task",
+                n_envs=1,
+                env_cls=gym.vector.SyncVectorEnv,
+            )
+
+    def test_invalid_n_envs_raises(self):
+        with pytest.raises(ValueError, match="n_envs must be a positive int"):
+            create_robotwin_envs(
+                task="beat_block_hammer",
+                n_envs=0,
+                env_cls=gym.vector.SyncVectorEnv,
+            )
+
+
+# ---------------------------------------------------------------------------
+# ROBOTWIN_TASKS list
+# ---------------------------------------------------------------------------
+
+
+def test_task_list_not_empty():
+    assert len(ROBOTWIN_TASKS) >= 50
+
+
+def test_all_tasks_are_strings():
+    assert all(isinstance(t, str) and t for t in ROBOTWIN_TASKS)
+
+
+def test_no_duplicate_tasks():
+    assert len(ROBOTWIN_TASKS) == len(set(ROBOTWIN_TASKS))

tests/policies/rtc/test_action_interpolator.py

@@ -17,9 +17,9 @@
 import pytest
 import torch
 
-from lerobot.policies.rtc.action_interpolator import ActionInterpolator
 from lerobot.policies.rtc.action_queue import ActionQueue
 from lerobot.policies.rtc.configuration_rtc import RTCConfig
+from lerobot.utils.action_interpolator import ActionInterpolator
 
 # ====================== Fixtures ======================
 

tests/policies/rtc/test_configuration_rtc.py

@@ -26,7 +26,7 @@ def test_rtc_config_default_initialization():
     """Test RTCConfig initializes with default values."""
     config = RTCConfig()
 
-    assert config.enabled is False
+    assert config.enabled is True
     assert config.prefix_attention_schedule == RTCAttentionSchedule.LINEAR
     assert config.max_guidance_weight == 10.0
     assert config.execution_horizon == 10

tests/policies/rtc/test_rtc_relative_actions.py

@@ -22,7 +22,7 @@ from lerobot.configs.types import (
     PolicyFeature,
     RTCAttentionSchedule,
 )
-from lerobot.processor import TransitionKey, batch_to_transition
+from lerobot.processor import TransitionKey, batch_to_transition, create_transition
 from lerobot.processor.normalize_processor import NormalizerProcessorStep, UnnormalizerProcessorStep
 from lerobot.processor.relative_action_processor import (
     AbsoluteActionsProcessorStep,
@@ -52,6 +52,9 @@ _rtc_debug_mod = _import_rtc_module("lerobot.policies.rtc.debug_tracker", "debug
 _rtc_mod = _import_rtc_module("lerobot.policies.rtc.modeling_rtc", "modeling_rtc.py")
 RTCProcessor = _rtc_mod.RTCProcessor
 
+_rtc_relative_mod = _import_rtc_module("lerobot.policies.rtc.relative", "relative.py")
+reanchor_relative_rtc_prefix = _rtc_relative_mod.reanchor_relative_rtc_prefix
+
 ACTION_DIM = 6
 CHUNK_SIZE = 50
 EXECUTION_HORIZON = 10
@@ -187,7 +190,7 @@ class TestRTCDenoiseWithRelativeLeftovers:
 
 
 class TestFullPipelineRelativeRTC:
-    """End-to-end test of the RTC + relative actions pipeline matching eval_with_real_robot.py flow."""
+    """End-to-end test of the RTC + relative actions pipeline matching lerobot-rollout flow."""
 
     def test_preprocessor_caches_state_for_postprocessor(self):
         """Preprocessor's relative step should cache state so postprocessor can convert back."""
@@ -218,7 +221,9 @@ class TestFullPipelineRelativeRTC:
 
     def test_roundtrip_with_identity_normalization(self):
         """Actions → relative → normalize → [model] → unnormalize → absolute should recover originals.
-        Using mean=0, std=1 normalization (identity)."""
+
+        Using mean=0, std=1 normalization (identity).
+        """
         relative_step, normalizer, unnormalizer, absolute_step = _make_relative_pipeline()
 
         state = torch.randn(1, ACTION_DIM)
@@ -240,7 +245,7 @@ class TestFullPipelineRelativeRTC:
         torch.testing.assert_close(recovered, actions, atol=1e-5, rtol=1e-5)
 
     def test_eval_loop_simulation(self):
-        """Simulate the eval_with_real_robot.py loop with relative actions.
+        """Simulate the lerobot-rollout loop with relative actions.
 
         Iteration 1: No leftovers → model generates relative actions → store for RTC
         Iteration 2: Use leftovers as RTC guidance → model generates new relative actions
@@ -400,13 +405,113 @@ class TestStateRebasingApproximation:
         assert error_excluded < 1e-6, f"Excluded joint should have zero error, got {error_excluded}"
 
 
+class TestRTCReanchoringWithStateNormalizer:
+    """RTC re-anchoring under non-identity OBS_STATE normalization."""
+
+    @staticmethod
+    def _build_normalizer_with_state_stats():
+        """Build a relative-action preprocessor with non-trivial OBS_STATE stats."""
+        features = {
+            ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(ACTION_DIM,)),
+            OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(ACTION_DIM,)),
+        }
+        norm_map = {
+            FeatureType.ACTION: NormalizationMode.MEAN_STD,
+            FeatureType.STATE: NormalizationMode.MEAN_STD,
+        }
+        stats = {
+            ACTION: {
+                "mean": torch.zeros(ACTION_DIM).numpy(),
+                "std": (0.5 * torch.ones(ACTION_DIM)).numpy(),
+            },
+            OBS_STATE: {
+                "mean": (5.0 * torch.ones(ACTION_DIM)).numpy(),
+                "std": (2.0 * torch.ones(ACTION_DIM)).numpy(),
+            },
+        }
+        relative_step = RelativeActionsProcessorStep(enabled=True)
+        normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
+        return relative_step, normalizer
+
+    def test_reanchor_with_raw_state_matches_normalize_of_absolute_minus_state(self):
+        """Reanchoring with the raw cached state yields ``normalize(prev_actions_absolute - raw_state)``."""
+        relative_step, normalizer = self._build_normalizer_with_state_stats()
+
+        raw_state = torch.tensor([[7.0, 8.0, 9.0, 10.0, 11.0, 12.0]])
+        relative_step(batch_to_transition({OBS_STATE: raw_state.clone()}))
+
+        prev_actions_absolute = torch.tensor([[2.0, 3.0, 4.0, 5.0, 6.0, 7.0]] * 5)
+
+        result = reanchor_relative_rtc_prefix(
+            prev_actions_absolute=prev_actions_absolute,
+            current_state=relative_step.get_cached_state(),
+            relative_step=relative_step,
+            normalizer_step=normalizer,
+            policy_device="cpu",
+        )
+
+        expected_relative = to_relative_actions(prev_actions_absolute, raw_state, [True] * ACTION_DIM)
+        expected = normalizer(create_transition(action=expected_relative))[TransitionKey.ACTION]
+        torch.testing.assert_close(result, expected, atol=1e-5, rtol=1e-5)
+
+    def test_reanchor_with_normalized_state_produces_wrong_result(self):
+        """Reanchoring with raw vs. normalized state produces meaningfully different outputs."""
+        relative_step, normalizer = self._build_normalizer_with_state_stats()
+
+        raw_state = torch.tensor([[7.0, 8.0, 9.0, 10.0, 11.0, 12.0]])
+        relative_step(batch_to_transition({OBS_STATE: raw_state.clone()}))
+
+        normalized_obs = normalizer(batch_to_transition({OBS_STATE: raw_state.clone()}))
+        normalized_state = normalized_obs[TransitionKey.OBSERVATION][OBS_STATE]
+        assert not torch.allclose(normalized_state, raw_state)
+
+        prev_actions_absolute = torch.tensor([[2.0, 3.0, 4.0, 5.0, 6.0, 7.0]] * 5)
+
+        result_raw = reanchor_relative_rtc_prefix(
+            prev_actions_absolute=prev_actions_absolute,
+            current_state=raw_state,
+            relative_step=relative_step,
+            normalizer_step=normalizer,
+            policy_device="cpu",
+        )
+        result_normalized = reanchor_relative_rtc_prefix(
+            prev_actions_absolute=prev_actions_absolute,
+            current_state=normalized_state,
+            relative_step=relative_step,
+            normalizer_step=normalizer,
+            policy_device="cpu",
+        )
+
+        max_abs_diff = (result_raw - result_normalized).abs().max()
+        assert max_abs_diff > 0.5, (
+            f"Raw and normalized state produced near-identical outputs (max diff {max_abs_diff:.4f}); "
+            "OBS_STATE stats are too close to identity to be sensitive."
+        )
+
+    def test_engine_pipeline_cached_state_is_raw_after_full_preprocess(self):
+        """``get_cached_state()`` returns raw OBS_STATE after the full preprocessor pipeline runs."""
+        relative_step, normalizer = self._build_normalizer_with_state_stats()
+
+        raw_state = torch.tensor([[7.0, 8.0, 9.0, 10.0, 11.0, 12.0]])
+
+        transition = batch_to_transition({OBS_STATE: raw_state.clone()})
+        transition = relative_step(transition)
+        preprocessed = normalizer(transition)
+
+        cached = relative_step.get_cached_state()
+        torch.testing.assert_close(cached, raw_state, atol=1e-6, rtol=1e-6)
+
+        post_normalize_state = preprocessed[TransitionKey.OBSERVATION][OBS_STATE]
+        assert not torch.allclose(cached, post_normalize_state, atol=1e-3)
+
+
 def _detect_relative_actions(preprocessor) -> bool:
-    """Mirror of the helper in eval_with_real_robot.py for testing without importing it."""
+    """Mirror of the helper in lerobot-rollout for testing without importing it."""
     return any(isinstance(step, RelativeActionsProcessorStep) and step.enabled for step in preprocessor.steps)
 
 
 class TestDetectRelativeActions:
-    """Test the _detect_relative_actions helper logic used by eval_with_real_robot.py."""
+    """Test the _detect_relative_actions helper logic used by lerobot-rollout."""
 
     def test_detects_enabled_relative_step(self):
         class FakePipeline:

tests/policies/test_policies.py

@@ -196,6 +196,8 @@ def test_policy(ds_repo_id, env_name, env_kwargs, policy_name, policy_kwargs):
 
     for key in batch:
         if isinstance(batch[key], torch.Tensor):
+            if batch[key].dtype == torch.uint8:
+                batch[key] = batch[key].to(dtype=torch.float32) / 255.0
             batch[key] = batch[key].to(DEVICE, non_blocking=True)
 
     # Test updating the policy (and test that it does not mutate the batch)

tests/teleoperators/test_reachy2_teleoperator.py

@@ -18,6 +18,11 @@ from unittest.mock import MagicMock, patch
 
 import pytest
 
+from lerobot.utils.import_utils import is_package_available
+
+if not is_package_available("reachy2_sdk"):
+    pytest.skip("reachy2_sdk not available", allow_module_level=True)
+
 from lerobot.teleoperators.reachy2_teleoperator import (
     REACHY2_ANTENNAS_JOINTS,
     REACHY2_L_ARM_JOINTS,

tests/test_cli_peft.py

@@ -24,10 +24,6 @@ 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":
@@ -155,81 +151,3 @@ def test_peft_training_params_are_fewer(policy_type, tmp_path):
                 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"])
-@skip_if_package_missing("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/test_control_robot.py

@@ -21,8 +21,9 @@ import pytest
 pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
 pytest.importorskip("deepdiff", reason="deepdiff is required (install lerobot[hardware])")
 
+from lerobot.configs.dataset import DatasetRecordConfig
 from lerobot.scripts.lerobot_calibrate import CalibrateConfig, calibrate
-from lerobot.scripts.lerobot_record import DatasetRecordConfig, RecordConfig, record
+from lerobot.scripts.lerobot_record import RecordConfig, record
 from lerobot.scripts.lerobot_replay import DatasetReplayConfig, ReplayConfig, replay
 from lerobot.scripts.lerobot_teleoperate import TeleoperateConfig, teleoperate
 from tests.fixtures.constants import DUMMY_REPO_ID

tests/test_robomme_env.py

@@ -0,0 +1,232 @@
+# 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.
+"""Unit tests for the RoboMME env wrapper and config.
+
+RoboMME requires Linux + ManiSkill (Vulkan/SAPIEN), so tests that touch the
+env wrapper mock the ``robomme`` package. Tests that only exercise the
+dataclass config run without any mocking.
+"""
+
+from __future__ import annotations
+
+import sys
+from types import ModuleType
+from unittest.mock import MagicMock
+
+import numpy as np
+
+
+def _install_robomme_stub():
+    """Register a minimal stub for the ``robomme`` package on sys.modules."""
+    stub = ModuleType("robomme")
+    wrapper_stub = ModuleType("robomme.env_record_wrapper")
+
+    class FakeBuilder:
+        def __init__(self, **kwargs):
+            pass
+
+        def make_env_for_episode(self, episode_idx: int, max_steps: int):
+            env = MagicMock()
+            obs = {
+                "front_rgb_list": [np.zeros((256, 256, 3), dtype=np.uint8)],
+                "wrist_rgb_list": [np.zeros((256, 256, 3), dtype=np.uint8)],
+                "joint_state_list": [np.zeros(7, dtype=np.float32)],
+                "gripper_state_list": [np.zeros(2, dtype=np.float32)],
+            }
+            env.reset.return_value = (obs, {"status": "ongoing", "task_goal": "pick the cube"})
+            env.step.return_value = (obs, 0.0, False, False, {"status": "ongoing", "task_goal": ""})
+            return env
+
+    wrapper_stub.BenchmarkEnvBuilder = FakeBuilder
+    stub.env_record_wrapper = wrapper_stub
+    sys.modules["robomme"] = stub
+    sys.modules["robomme.env_record_wrapper"] = wrapper_stub
+
+
+def _uninstall_robomme_stub():
+    sys.modules.pop("robomme", None)
+    sys.modules.pop("robomme.env_record_wrapper", None)
+
+
+# ---------------------------------------------------------------------------
+# Config tests (no sim required)
+# ---------------------------------------------------------------------------
+
+
+def test_robomme_env_config_defaults():
+    from lerobot.envs.configs import RoboMMEEnv
+
+    cfg = RoboMMEEnv()
+    assert cfg.task == "PickXtimes"
+    assert cfg.fps == 10
+    assert cfg.episode_length == 300
+    assert cfg.action_space == "joint_angle"
+    assert cfg.dataset_split == "test"
+    assert cfg.task_ids is None
+
+
+def test_robomme_env_config_type():
+    from lerobot.envs.configs import RoboMMEEnv
+
+    cfg = RoboMMEEnv()
+    assert cfg.type == "robomme"
+
+
+def test_robomme_features_map():
+    from lerobot.envs.configs import RoboMMEEnv
+    from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE
+
+    cfg = RoboMMEEnv()
+    assert cfg.features_map[ACTION] == ACTION
+    assert cfg.features_map["pixels/image"] == f"{OBS_IMAGES}.image"
+    assert cfg.features_map["pixels/wrist_image"] == f"{OBS_IMAGES}.wrist_image"
+    assert cfg.features_map["agent_pos"] == OBS_STATE
+
+
+def test_robomme_features_action_dim_joint_angle():
+    from lerobot.envs.configs import RoboMMEEnv
+    from lerobot.utils.constants import ACTION
+
+    cfg = RoboMMEEnv(action_space="joint_angle")
+    assert cfg.features[ACTION].shape == (8,)
+
+
+def test_robomme_features_action_dim_ee_pose():
+    """`ee_pose` uses a 7-D action; __post_init__ sets the correct shape."""
+    from lerobot.envs.configs import RoboMMEEnv
+    from lerobot.utils.constants import ACTION
+
+    cfg = RoboMMEEnv(action_space="ee_pose")
+    assert cfg.features[ACTION].shape == (7,)
+
+
+# ---------------------------------------------------------------------------
+# Obs conversion (pure Python, no sim)
+# ---------------------------------------------------------------------------
+
+
+def test_convert_obs_list_format():
+    """_convert_obs takes the last element from list-format obs fields and
+    emits a nested ``pixels`` dict (image, wrist_image) plus ``agent_pos``.
+
+    The nested layout is required so ``preprocess_observation()`` in
+    ``envs/utils.py`` maps each camera to ``observation.images.<cam>``.
+    """
+    _install_robomme_stub()
+    try:
+        from lerobot.envs.robomme import RoboMMEGymEnv
+
+        env = RoboMMEGymEnv.__new__(RoboMMEGymEnv)
+
+        front = np.full((256, 256, 3), 42, dtype=np.uint8)
+        wrist = np.full((256, 256, 3), 7, dtype=np.uint8)
+        joints = np.arange(7, dtype=np.float32)
+        gripper = np.array([0.5, 0.5], dtype=np.float32)
+
+        obs_raw = {
+            "front_rgb_list": [np.zeros_like(front), front],
+            "wrist_rgb_list": [np.zeros_like(wrist), wrist],
+            "joint_state_list": [np.zeros(7, dtype=np.float32), joints],
+            "gripper_state_list": [np.zeros(2, dtype=np.float32), gripper],
+        }
+
+        result = env._convert_obs(obs_raw)
+        np.testing.assert_array_equal(result["pixels"]["image"], front)
+        np.testing.assert_array_equal(result["pixels"]["wrist_image"], wrist)
+        assert result["agent_pos"].shape == (8,)
+        np.testing.assert_array_almost_equal(result["agent_pos"][:7], joints)
+        assert result["agent_pos"][7] == gripper[0]
+    finally:
+        _uninstall_robomme_stub()
+
+
+def test_convert_obs_array_format():
+    """_convert_obs also handles non-list (direct array) obs."""
+    _install_robomme_stub()
+    try:
+        from lerobot.envs.robomme import RoboMMEGymEnv
+
+        env = RoboMMEGymEnv.__new__(RoboMMEGymEnv)
+
+        front = np.zeros((256, 256, 3), dtype=np.uint8)
+        obs_raw = {
+            "front_rgb_list": front,
+            "wrist_rgb_list": front,
+            "joint_state_list": np.zeros(7, dtype=np.float32),
+            "gripper_state_list": np.zeros(2, dtype=np.float32),
+        }
+        result = env._convert_obs(obs_raw)
+        assert result["pixels"]["image"].shape == (256, 256, 3)
+        assert result["pixels"]["wrist_image"].shape == (256, 256, 3)
+        assert result["agent_pos"].shape == (8,)
+    finally:
+        _uninstall_robomme_stub()
+
+
+# ---------------------------------------------------------------------------
+# create_robomme_envs (mocked sim)
+# ---------------------------------------------------------------------------
+
+
+def test_create_robomme_envs_returns_correct_structure():
+    """Single task -> {task_name: {task_id: VectorEnv}} with one entry per task_id."""
+    _install_robomme_stub()
+    try:
+        from lerobot.envs.robomme import create_robomme_envs
+
+        env_cls = MagicMock(return_value=MagicMock())
+        result = create_robomme_envs(
+            task="PickXtimes",
+            n_envs=1,
+            task_ids=[0, 1],
+            env_cls=env_cls,
+        )
+
+        assert "PickXtimes" in result
+        assert 0 in result["PickXtimes"]
+        assert 1 in result["PickXtimes"]
+        assert env_cls.call_count == 2
+    finally:
+        _uninstall_robomme_stub()
+
+
+def test_create_robomme_envs_multi_task():
+    """Comma-separated task list produces one suite per task."""
+    _install_robomme_stub()
+    try:
+        from lerobot.envs.robomme import create_robomme_envs
+
+        env_cls = MagicMock(return_value=MagicMock())
+        result = create_robomme_envs(
+            task="PickXtimes,BinFill,StopCube",
+            n_envs=1,
+            env_cls=env_cls,
+        )
+
+        assert set(result.keys()) == {"PickXtimes", "BinFill", "StopCube"}
+    finally:
+        _uninstall_robomme_stub()
+
+
+def test_create_robomme_envs_raises_on_invalid_env_cls():
+    _install_robomme_stub()
+    try:
+        import pytest
+
+        from lerobot.envs.robomme import create_robomme_envs
+
+        with pytest.raises(ValueError, match="env_cls must be a callable"):
+            create_robomme_envs(task="PickXtimes", n_envs=1, env_cls=None)
+    finally:
+        _uninstall_robomme_stub()

tests/test_rollout.py

@@ -0,0 +1,345 @@
+# 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.
+
+"""Minimal tests for the rollout module's public API."""
+
+from __future__ import annotations
+
+import dataclasses
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
+
+# ---------------------------------------------------------------------------
+# Import smoke tests
+# ---------------------------------------------------------------------------
+
+
+def test_rollout_top_level_imports():
+    import lerobot.rollout
+
+    for name in lerobot.rollout.__all__:
+        assert hasattr(lerobot.rollout, name), f"Missing export: {name}"
+
+
+def test_inference_submodule_imports():
+    import lerobot.rollout.inference
+
+    for name in lerobot.rollout.inference.__all__:
+        assert hasattr(lerobot.rollout.inference, name), f"Missing export: {name}"
+
+
+def test_strategies_submodule_imports():
+    import lerobot.rollout.strategies
+
+    for name in lerobot.rollout.strategies.__all__:
+        assert hasattr(lerobot.rollout.strategies, name), f"Missing export: {name}"
+
+
+# ---------------------------------------------------------------------------
+# Config tests
+# ---------------------------------------------------------------------------
+
+
+def test_strategy_config_types():
+    from lerobot.rollout import (
+        BaseStrategyConfig,
+        DAggerStrategyConfig,
+        HighlightStrategyConfig,
+        SentryStrategyConfig,
+    )
+
+    assert BaseStrategyConfig().type == "base"
+    assert SentryStrategyConfig().type == "sentry"
+    assert HighlightStrategyConfig().type == "highlight"
+    assert DAggerStrategyConfig().type == "dagger"
+
+
+def test_dagger_config_invalid_input_device():
+    from lerobot.rollout import DAggerStrategyConfig
+
+    with pytest.raises(ValueError, match="input_device must be 'keyboard' or 'pedal'"):
+        DAggerStrategyConfig(input_device="joystick")
+
+
+def test_dagger_config_defaults():
+    from lerobot.rollout import DAggerStrategyConfig
+
+    cfg = DAggerStrategyConfig()
+    assert cfg.num_episodes is None
+    assert cfg.record_autonomous is False
+    assert cfg.input_device == "keyboard"
+
+
+def test_inference_config_types():
+    from lerobot.rollout import RTCInferenceConfig, SyncInferenceConfig
+
+    assert SyncInferenceConfig().type == "sync"
+
+    rtc = RTCInferenceConfig()
+    assert rtc.type == "rtc"
+    assert rtc.queue_threshold == 30
+    assert rtc.rtc is not None
+
+
+def test_sentry_config_defaults():
+    from lerobot.rollout import SentryStrategyConfig
+
+    cfg = SentryStrategyConfig()
+    assert cfg.upload_every_n_episodes == 5
+    assert cfg.target_video_file_size_mb is None
+
+
+# ---------------------------------------------------------------------------
+# RolloutRingBuffer
+# ---------------------------------------------------------------------------
+
+
+def test_ring_buffer_append_and_eviction():
+    from lerobot.rollout.ring_buffer import RolloutRingBuffer
+
+    buf = RolloutRingBuffer(max_seconds=0.5, max_memory_mb=100.0, fps=10.0)
+    # max_frames = 5
+    for i in range(8):
+        buf.append({"val": i})
+    assert len(buf) == 5
+
+
+def test_ring_buffer_drain():
+    from lerobot.rollout.ring_buffer import RolloutRingBuffer
+
+    buf = RolloutRingBuffer(max_seconds=1.0, max_memory_mb=100.0, fps=10.0)
+    for i in range(3):
+        buf.append({"val": i})
+    frames = buf.drain()
+    assert len(frames) == 3
+    assert len(buf) == 0
+    assert buf.estimated_bytes == 0
+
+
+def test_ring_buffer_clear():
+    from lerobot.rollout.ring_buffer import RolloutRingBuffer
+
+    buf = RolloutRingBuffer(max_seconds=1.0, max_memory_mb=100.0, fps=10.0)
+    buf.append({"val": 1})
+    buf.clear()
+    assert len(buf) == 0
+    assert buf.estimated_bytes == 0
+
+
+def test_ring_buffer_tensor_bytes():
+    from lerobot.rollout.ring_buffer import RolloutRingBuffer
+
+    buf = RolloutRingBuffer(max_seconds=1.0, max_memory_mb=100.0, fps=10.0)
+    t = torch.zeros(100, dtype=torch.float32)  # 400 bytes
+    buf.append({"tensor": t})
+    assert buf.estimated_bytes >= 400
+
+
+# ---------------------------------------------------------------------------
+# ThreadSafeRobot
+# ---------------------------------------------------------------------------
+
+
+def test_thread_safe_robot_delegates():
+    from lerobot.rollout.robot_wrapper import ThreadSafeRobot
+    from tests.mocks.mock_robot import MockRobot, MockRobotConfig
+
+    robot = MockRobot(MockRobotConfig(n_motors=3))
+    robot.connect()
+    wrapper = ThreadSafeRobot(robot)
+
+    obs = wrapper.get_observation()
+    assert "motor_1.pos" in obs
+    assert "motor_2.pos" in obs
+    assert "motor_3.pos" in obs
+
+    action = {"motor_1.pos": 0.0, "motor_2.pos": 1.0, "motor_3.pos": 2.0}
+    result = wrapper.send_action(action)
+    assert result == action
+
+    robot.disconnect()
+
+
+def test_thread_safe_robot_properties():
+    from lerobot.rollout.robot_wrapper import ThreadSafeRobot
+    from tests.mocks.mock_robot import MockRobot, MockRobotConfig
+
+    robot = MockRobot(MockRobotConfig(n_motors=3))
+    robot.connect()
+    wrapper = ThreadSafeRobot(robot)
+
+    assert wrapper.name == "mock_robot"
+    assert "motor_1.pos" in wrapper.observation_features
+    assert "motor_1.pos" in wrapper.action_features
+    assert wrapper.is_connected is True
+    assert wrapper.inner is robot
+
+    robot.disconnect()
+
+
+# ---------------------------------------------------------------------------
+# Strategy factory
+# ---------------------------------------------------------------------------
+
+
+def test_create_strategy_dispatches():
+    from lerobot.rollout import (
+        BaseStrategy,
+        BaseStrategyConfig,
+        DAggerStrategy,
+        DAggerStrategyConfig,
+        SentryStrategy,
+        SentryStrategyConfig,
+        create_strategy,
+    )
+
+    assert isinstance(create_strategy(BaseStrategyConfig()), BaseStrategy)
+    assert isinstance(create_strategy(SentryStrategyConfig()), SentryStrategy)
+    assert isinstance(create_strategy(DAggerStrategyConfig()), DAggerStrategy)
+
+
+def test_create_strategy_unknown_raises():
+    from lerobot.rollout import create_strategy
+
+    cfg = MagicMock()
+    cfg.type = "bogus"
+    with pytest.raises(ValueError, match="Unknown strategy type"):
+        create_strategy(cfg)
+
+
+# ---------------------------------------------------------------------------
+# Inference factory
+# ---------------------------------------------------------------------------
+
+
+def test_create_inference_engine_sync():
+    from lerobot.rollout import SyncInferenceConfig, SyncInferenceEngine, create_inference_engine
+
+    engine = create_inference_engine(
+        SyncInferenceConfig(),
+        policy=MagicMock(),
+        preprocessor=MagicMock(),
+        postprocessor=MagicMock(),
+        robot_wrapper=MagicMock(robot_type="mock"),
+        hw_features={},
+        dataset_features={},
+        ordered_action_keys=["k"],
+        task="test",
+        fps=30.0,
+        device="cpu",
+    )
+    assert isinstance(engine, SyncInferenceEngine)
+
+
+# ---------------------------------------------------------------------------
+# Pure functions
+# ---------------------------------------------------------------------------
+
+
+def test_estimate_max_episode_seconds_no_video():
+    from lerobot.rollout.strategies import estimate_max_episode_seconds
+
+    assert estimate_max_episode_seconds({}, fps=30.0) == 300.0
+
+
+def test_estimate_max_episode_seconds_with_video():
+    from lerobot.rollout.strategies import estimate_max_episode_seconds
+
+    features = {"cam": {"dtype": "video", "shape": (480, 640, 3)}}
+    result = estimate_max_episode_seconds(features, fps=30.0)
+    assert result > 0
+    # With a real camera, duration should differ from the fallback
+    assert result != 300.0
+
+
+def test_safe_push_to_hub():
+    from lerobot.rollout.strategies import safe_push_to_hub
+
+    ds = MagicMock()
+    ds.num_episodes = 0
+    assert safe_push_to_hub(ds) is False
+    ds.push_to_hub.assert_not_called()
+
+    ds.num_episodes = 5
+    assert safe_push_to_hub(ds, tags=["test"]) is True
+    ds.push_to_hub.assert_called_once_with(tags=["test"], private=False)
+
+
+# ---------------------------------------------------------------------------
+# DAgger state machine
+# ---------------------------------------------------------------------------
+
+
+def test_dagger_full_transition_cycle():
+    from lerobot.rollout.strategies import DAggerEvents, DAggerPhase
+
+    events = DAggerEvents()
+    assert events.phase == DAggerPhase.AUTONOMOUS
+
+    # AUTONOMOUS -> PAUSED
+    events.request_transition("pause_resume")
+    old, new = events.consume_transition()
+    assert (old, new) == (DAggerPhase.AUTONOMOUS, DAggerPhase.PAUSED)
+
+    # PAUSED -> CORRECTING
+    events.request_transition("correction")
+    old, new = events.consume_transition()
+    assert (old, new) == (DAggerPhase.PAUSED, DAggerPhase.CORRECTING)
+
+    # CORRECTING -> PAUSED
+    events.request_transition("correction")
+    old, new = events.consume_transition()
+    assert (old, new) == (DAggerPhase.CORRECTING, DAggerPhase.PAUSED)
+
+    # PAUSED -> AUTONOMOUS
+    events.request_transition("pause_resume")
+    old, new = events.consume_transition()
+    assert (old, new) == (DAggerPhase.PAUSED, DAggerPhase.AUTONOMOUS)
+
+
+def test_dagger_invalid_transition_ignored():
+    from lerobot.rollout.strategies import DAggerEvents, DAggerPhase
+
+    events = DAggerEvents()
+    events.request_transition("correction")  # Not valid from AUTONOMOUS
+    assert events.consume_transition() is None
+    assert events.phase == DAggerPhase.AUTONOMOUS
+
+
+def test_dagger_events_reset():
+    from lerobot.rollout.strategies import DAggerEvents, DAggerPhase
+
+    events = DAggerEvents()
+    events.request_transition("pause_resume")
+    events.consume_transition()  # -> PAUSED
+    events.upload_requested.set()
+    events.reset()
+    assert events.phase == DAggerPhase.AUTONOMOUS
+    assert not events.upload_requested.is_set()
+
+
+# ---------------------------------------------------------------------------
+# Context dataclass
+# ---------------------------------------------------------------------------
+
+
+def test_rollout_context_fields():
+    from lerobot.rollout import RolloutContext
+
+    field_names = {f.name for f in dataclasses.fields(RolloutContext)}
+    assert field_names == {"runtime", "hardware", "policy", "processors", "data"}

tests/utils/test_process.py

@@ -24,7 +24,7 @@ import pytest
 
 pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
 
-from lerobot.rl.process import ProcessSignalHandler  # noqa: E402
+from lerobot.utils.process import ProcessSignalHandler  # noqa: E402
 
 
 # Fixture to reset shutdown_event_counter and original signal handlers before and after each test