Merge branch 'main' into feat/bump_rerun

Signed-off-by: Steven Palma <imstevenpmwork@ieee.org>

tests/annotations/_helpers.py

@@ -0,0 +1,58 @@
+#!/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.
+"""Helpers shared across annotation-pipeline tests."""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from lerobot.annotations.steerable_pipeline.vlm_client import StubVlmClient
+
+
+def make_canned_responder(
+    responses_by_marker: dict[str, Any],
+    default: Any = None,
+) -> StubVlmClient:
+    """Return a stub that picks a response by inspecting the user prompt.
+
+    For each call the responder examines the last user-message text and
+    returns the response keyed by the first marker substring it contains.
+    Falls back to ``default`` if no marker matches.
+    """
+
+    def responder(messages: list[dict[str, Any]]) -> Any:
+        last_user_text = ""
+        for message in messages:
+            if message.get("role") != "user":
+                continue
+            content = message.get("content")
+            if isinstance(content, str):
+                last_user_text = content
+            elif isinstance(content, list):
+                for block in content:
+                    if isinstance(block, dict) and block.get("type") == "text":
+                        last_user_text = block.get("text", "")
+        for marker, response in responses_by_marker.items():
+            if marker in last_user_text:
+                return response
+        return default
+
+    return StubVlmClient(responder=responder)
+
+
+def encode_vqa_answer(payload: dict[str, Any]) -> str:
+    return json.dumps(payload, sort_keys=True)

tests/annotations/conftest.py

@@ -0,0 +1,58 @@
+#!/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.
+"""Shared fixtures for annotation-pipeline tests.
+
+The on-disk dataset builder lives with the other dataset factories in
+``tests/fixtures/dataset_factories.py`` (:func:`build_annotation_dataset`);
+these fixtures only wire it into pytest.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+
+# ``build_annotation_dataset`` pulls in ``lerobot.datasets`` (HF ``datasets``
+# + ``pandas``, only in the ``dataset`` extra), so it's imported lazily inside
+# each fixture — this conftest stays importable without that extra. The test
+# modules ``pytest.importorskip("datasets")`` so they skip rather than error.
+
+
+@pytest.fixture
+def fixture_dataset_root(tmp_path: Path) -> Path:
+    """A tiny dataset with two episodes, 12 frames each at 10 fps."""
+    from tests.fixtures.dataset_factories import build_annotation_dataset
+
+    return build_annotation_dataset(
+        tmp_path / "ds",
+        episode_specs=[
+            (0, 12, "Could you tidy the kitchen please?"),
+            (1, 12, "Please clean up the kitchen"),
+        ],
+        fps=10,
+    )
+
+
+@pytest.fixture
+def single_episode_root(tmp_path: Path) -> Path:
+    from tests.fixtures.dataset_factories import build_annotation_dataset
+
+    return build_annotation_dataset(
+        tmp_path / "ds_one",
+        episode_specs=[(0, 30, "Pour water from the bottle into the cup.")],
+        fps=10,
+    )

tests/annotations/run_e2e_smoke.py

@@ -0,0 +1,116 @@
+#!/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.
+"""Opt-in E2E smoke run for ``make annotation-e2e``.
+
+Builds the shared annotation fixture (:func:`build_annotation_dataset`),
+runs the full annotation pipeline against it with a stub VLM, and prints a
+short report. This is intentionally not a pytest test — it exercises the
+CLI plumbing — but it reuses the same on-disk dataset builder as the pytest
+fixtures so there is no duplicated fixture code.
+"""
+
+from __future__ import annotations
+
+import sys
+import tempfile
+from pathlib import Path
+
+from lerobot.annotations.steerable_pipeline.config import AnnotationPipelineConfig
+from lerobot.annotations.steerable_pipeline.executor import Executor
+from lerobot.annotations.steerable_pipeline.modules import (
+    GeneralVqaModule,
+    InterjectionsAndSpeechModule,
+    PlanSubtasksMemoryModule,
+)
+from lerobot.annotations.steerable_pipeline.validator import StagingValidator
+from lerobot.annotations.steerable_pipeline.vlm_client import StubVlmClient
+from lerobot.annotations.steerable_pipeline.writer import LanguageColumnsWriter
+from tests.fixtures.dataset_factories import build_annotation_dataset
+
+
+def _stub_responder(messages):
+    text = ""
+    for m in messages:
+        if m.get("role") == "user":
+            content = m.get("content")
+            if isinstance(content, list):
+                for block in content:
+                    if isinstance(block, dict) and block.get("type") == "text":
+                        text = block.get("text", "")
+            elif isinstance(content, str):
+                text = content
+    if "atomic subtasks" in text:
+        return {
+            "subtasks": [
+                {"text": "grasp the bottle", "start": 0.0, "end": 1.0},
+                {"text": "pour into the cup", "start": 1.0, "end": 2.0},
+                {"text": "place the bottle down", "start": 2.0, "end": 3.0},
+            ]
+        }
+    if "compressed semantic memory" in text:
+        return {"memory": "poured once"}
+    if "acknowledgement the robot" in text:
+        return {"text": "Sure."}
+    if "compact interjection" in text:
+        return {"interjection": "use less water", "speech": "Using less water."}
+    if "frame-grounded visual question" in text:
+        return {"question": "How many cups?", "answer": {"label": "cup", "count": 1}}
+    return None
+
+
+def main() -> int:
+    with tempfile.TemporaryDirectory() as tmp:
+        root = build_annotation_dataset(
+            Path(tmp) / "ds",
+            episode_specs=[(0, 30, "Pour water into the cup.")],
+            fps=10,
+        )
+        vlm = StubVlmClient(responder=_stub_responder)
+        cfg = AnnotationPipelineConfig()
+        executor = Executor(
+            config=cfg,
+            plan=PlanSubtasksMemoryModule(vlm=vlm, config=cfg.plan),
+            interjections=InterjectionsAndSpeechModule(vlm=vlm, config=cfg.interjections, seed=cfg.seed),
+            vqa=GeneralVqaModule(vlm=vlm, config=cfg.vqa, seed=cfg.seed),
+            writer=LanguageColumnsWriter(),
+            validator=StagingValidator(),
+        )
+        summary = executor.run(root)
+        print(f"phases={[(p.name, p.episodes_processed) for p in summary.phases]}")
+        print(f"validation: {summary.validation_report.summary()}")
+        print(f"shards rewritten: {len(summary.written_paths)}")
+
+        # Assert the interjection code path actually fired — otherwise a stale
+        # canned-VLM marker would silently produce zero interjections and this
+        # smoke run would still "pass" by only printing.
+        import pyarrow.parquet as pq  # noqa: PLC0415
+
+        events = [
+            r
+            for shard in summary.written_paths
+            for ev in pq.read_table(shard).column("language_events").to_pylist()
+            for r in ev
+        ]
+        n_interjections = sum(1 for r in events if r.get("style") == "interjection")
+        n_speech = sum(1 for r in events if r.get("style") is None and r.get("role") == "assistant")
+        print(f"interjections={n_interjections} speech_atoms={n_speech}")
+        assert n_interjections > 0, "no interjection rows produced — check the interjection prompt marker"
+        assert n_speech > 0, "no speech tool-call atoms produced — check the speech prompt marker"
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

tests/annotations/test_frames.py

@@ -0,0 +1,247 @@
+#!/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.
+"""Unit tests for :class:`VideoFrameProvider` method bindings.
+
+These were prompted by a real regression: ``video_for_episode`` was once
+indented one level too deep so it ended up nested *inside* a module-level
+helper (after that function's ``return`` statement) — silently dead code
+that meant production runs with ``use_video_url=False`` would
+``AttributeError`` on ``self.frame_provider.video_for_episode(...)``. The
+existing module tests didn't catch it because they exercise stub providers.
+
+The tests below assert on the class itself (not on an instance), so a
+future reindent regression flips them to red without needing a real
+LeRobot dataset on disk.
+"""
+
+from __future__ import annotations
+
+import shutil
+import subprocess
+from pathlib import Path
+
+import pytest
+import torch
+
+pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
+
+from lerobot.annotations.steerable_pipeline.frames import VideoFrameProvider  # noqa: E402
+
+
+class _FakeMeta:
+    """Minimal metadata stub exposing ``video_keys`` / ``camera_keys``."""
+
+    def __init__(self, video_keys: list[str], image_keys: list[str], video_path: Path | None = None) -> None:
+        self.video_keys = video_keys
+        self.camera_keys = [*video_keys, *image_keys]
+        self.depth_keys = []
+        self._video_path = video_path
+        self.episodes = {0: {f"videos/{key}/from_timestamp": 0.0 for key in video_keys}}
+
+    def get_video_file_path(self, episode_index: int, camera_key: str) -> Path:
+        return self._video_path
+
+
+def test_default_camera_key_skips_image_only_cameras(tmp_path: Path, monkeypatch) -> None:
+    """The default camera must be a *video* key — image-stored cameras have no
+    ``videos/<key>/from_timestamp`` and would KeyError in the clip/decode path.
+
+    Regression: a dataset whose first ``camera_keys`` entry was an image-stored
+    camera (e.g. ``observation.images.wrist``) crashed at clip extraction.
+    """
+    fake = _FakeMeta(
+        video_keys=["observation.images.robot0_agentview_right"],
+        image_keys=["observation.images.wrist"],
+    )
+    import lerobot.datasets.dataset_metadata as meta_mod
+
+    monkeypatch.setattr(meta_mod, "LeRobotDatasetMetadata", lambda *a, **k: fake, raising=True)
+    provider = VideoFrameProvider(root=tmp_path)
+    assert provider.camera_key == "observation.images.robot0_agentview_right"
+    assert "observation.images.wrist" not in provider.camera_keys
+
+
+def test_video_for_episode_is_a_method_of_videoframeprovider():
+    """``video_for_episode`` must be a bound method, not nested dead code."""
+    assert callable(getattr(VideoFrameProvider, "video_for_episode", None))
+
+
+def test_episode_clip_path_is_a_method_of_videoframeprovider():
+    """``episode_clip_path`` is now a method (was a free function reaching
+    into ``provider._meta`` from outside the class)."""
+    assert callable(getattr(VideoFrameProvider, "episode_clip_path", None))
+
+
+def test_videoframeprovider_has_a_lock_for_concurrent_use():
+    """A ``ThreadPoolExecutor`` runs the plan / interjections / vqa phases
+    concurrently; the cache + warn-flag accesses must be guarded.
+    """
+    import threading
+
+    # Fresh-instance check via a minimal fake to avoid touching the hub.
+    # The lock is declared with ``init=False`` and has a default factory,
+    # so a constructed instance must own a real ``threading.Lock``.
+    lock_field = next(
+        (f for f in VideoFrameProvider.__dataclass_fields__.values() if f.name == "_lock"),
+        None,
+    )
+    assert lock_field is not None
+    assert lock_field.default_factory is threading.Lock
+
+
+@pytest.fixture
+def sample_video(tmp_path: Path) -> Path:
+    """A 3 s 10 fps test-pattern mp4, written with ffmpeg."""
+    if shutil.which("ffmpeg") is None:
+        pytest.skip("ffmpeg not available")
+    out = tmp_path / "sample.mp4"
+    subprocess.run(
+        [
+            "ffmpeg",
+            "-y",
+            "-f",
+            "lavfi",
+            "-i",
+            "testsrc=duration=3:size=160x120:rate=10",
+            "-pix_fmt",
+            "yuv420p",
+            str(out),
+        ],
+        check=True,
+        capture_output=True,
+    )
+    return out
+
+
+def _provider_for_video(tmp_path: Path, video: Path, monkeypatch) -> VideoFrameProvider:
+    """A provider whose single camera resolves to ``video`` via fake metadata."""
+    fake = _FakeMeta(video_keys=["observation.images.cam"], image_keys=[], video_path=video)
+    import lerobot.datasets.dataset_metadata as meta_mod
+
+    monkeypatch.setattr(meta_mod, "LeRobotDatasetMetadata", lambda *a, **k: fake, raising=True)
+    return VideoFrameProvider(root=tmp_path, tolerance_s=0.2)
+
+
+def test_decode_returns_one_uint8_frame_per_timestamp(
+    sample_video: Path, tmp_path: Path, monkeypatch
+) -> None:
+    """``_decode`` routes through ``decode_video_frames`` (torchcodec when
+    available, PyAV otherwise) — no subprocess fallback.
+    """
+    provider = _provider_for_video(tmp_path, sample_video, monkeypatch)
+    timestamps = [0.0, 1.0, 2.5]
+    frames = provider._decode(0, timestamps, "observation.images.cam")
+
+    assert len(frames) == len(timestamps)
+    for frame in frames:
+        assert isinstance(frame, torch.Tensor)
+        assert frame.dtype == torch.uint8
+        assert frame.shape == (3, 120, 160)
+
+
+def test_frames_at_snaps_mid_frame_grid_to_real_frames(
+    sample_video: Path, tmp_path: Path, monkeypatch
+) -> None:
+    """Uniform sampling grids land mid-frame; ``frames_at`` must snap them to
+    real frame timestamps before decoding.
+
+    Regression: ``decode_video_frames`` rejects queries farther than
+    ``tolerance_s`` (default 10 ms) from a decodable frame, so un-snapped
+    mid-frame queries raised ``FrameTimestampError`` wholesale and the plan
+    module silently lost its contact sheets for most episodes.
+    """
+    from types import SimpleNamespace
+
+    fake = _FakeMeta(video_keys=["observation.images.cam"], image_keys=[], video_path=sample_video)
+    import lerobot.datasets.dataset_metadata as meta_mod
+
+    monkeypatch.setattr(meta_mod, "LeRobotDatasetMetadata", lambda *a, **k: fake, raising=True)
+    provider = VideoFrameProvider(root=tmp_path)  # default 10 ms tolerance
+    # 10 fps fixture -> frames at 0.0, 0.1, ...; queries sit mid-frame.
+    record = SimpleNamespace(episode_index=0, frame_timestamps=[i / 10 for i in range(30)])
+
+    frames = provider.frames_at(record, [0.149, 1.234, 2.04], camera_key="observation.images.cam")
+
+    assert len(frames) == 3
+    for frame in frames:
+        assert isinstance(frame, torch.Tensor)
+        assert frame.shape == (3, 120, 160)
+
+
+def test_decode_returns_empty_list_on_missing_file(tmp_path: Path, monkeypatch) -> None:
+    """A missing video is a recoverable no-frames condition, never a crash."""
+    provider = _provider_for_video(tmp_path, tmp_path / "does_not_exist.mp4", monkeypatch)
+    assert provider._decode(0, [0.0], "observation.images.cam") == []
+
+
+def test_episode_clip_path_trims_via_reencode_video(tmp_path: Path, monkeypatch) -> None:
+    """Clip extraction delegates to ``video_utils.reencode_video`` with the
+    episode's ``[from_timestamp, to_timestamp)`` trim window — no subprocess.
+    """
+    from types import SimpleNamespace
+
+    import lerobot.annotations.steerable_pipeline.frames as frames_mod
+
+    src = tmp_path / "src.mp4"
+    src.write_bytes(b"src")
+    fake = _FakeMeta(video_keys=["observation.images.cam"], image_keys=[], video_path=src)
+    fake.episodes[0]["videos/observation.images.cam/from_timestamp"] = 1.5
+    fake.episodes[0]["videos/observation.images.cam/to_timestamp"] = 4.0
+    import lerobot.datasets.dataset_metadata as meta_mod
+
+    monkeypatch.setattr(meta_mod, "LeRobotDatasetMetadata", lambda *a, **k: fake, raising=True)
+
+    captured = {}
+
+    def fake_reencode(
+        input_video_path,
+        output_video_path,
+        video_encoder=None,
+        overwrite=False,
+        start_time_s=None,
+        end_time_s=None,
+    ):
+        captured.update(
+            src=Path(input_video_path),
+            encoder=video_encoder,
+            start_time_s=start_time_s,
+            end_time_s=end_time_s,
+        )
+        Path(output_video_path).write_bytes(b"clip")
+
+    monkeypatch.setattr(frames_mod, "reencode_video", fake_reencode, raising=True)
+    provider = VideoFrameProvider(root=tmp_path)
+    record = SimpleNamespace(episode_index=0, frame_timestamps=[0.0, 1.0])
+
+    out = provider.episode_clip_path(record, tmp_path / "clips")
+
+    assert out == tmp_path / "clips" / "ep_000000.mp4"
+    assert captured["src"] == src
+    assert captured["start_time_s"] == 1.5
+    assert captured["end_time_s"] == 4.0
+    # H.264 so the clip is decodable by vllm's libav build (sources are often AV1).
+    assert captured["encoder"].vcodec == "h264"
+
+
+def test_videoframeprovider_serializes_decodes_with_a_lock() -> None:
+    """torchcodec's cached per-file decoder is single-threaded; the provider
+    must own a dedicated lock that ``_decode`` holds around the decoder call.
+    """
+    import threading
+
+    lock_field = VideoFrameProvider.__dataclass_fields__.get("_decode_lock")
+    assert lock_field is not None
+    assert lock_field.default_factory is threading.Lock

tests/annotations/test_modules.py

@@ -0,0 +1,390 @@
+#!/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.
+"""Module 1/2/3 unit tests with stubbed VLMs."""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+import PIL.Image
+import pytest
+
+# ``lerobot.annotations`` imports pull in ``lerobot.datasets`` (-> the HF
+# ``datasets`` library), which only ships under the ``dataset`` extra. Skip
+# this module in tiers without it instead of erroring at import.
+pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
+pytest.importorskip("pandas", reason="pandas is required (install lerobot[dataset])")
+
+from lerobot.annotations.steerable_pipeline.config import (  # noqa: E402
+    InterjectionsConfig,
+    PlanConfig,
+    VqaConfig,
+)
+from lerobot.annotations.steerable_pipeline.modules import (  # noqa: E402
+    GeneralVqaModule,
+    InterjectionsAndSpeechModule,
+    PlanSubtasksMemoryModule,
+)
+from lerobot.annotations.steerable_pipeline.reader import iter_episodes  # noqa: E402
+from lerobot.annotations.steerable_pipeline.staging import EpisodeStaging  # noqa: E402
+from lerobot.annotations.steerable_pipeline.vlm_client import StubVlmClient  # noqa: E402
+
+from ._helpers import make_canned_responder  # noqa: E402
+
+
+@dataclass
+class _StubFrameProvider:
+    """Returns one sentinel object per requested timestamp."""
+
+    # A real (tiny) PIL image so the contact-sheet builder, which resizes and
+    # tiles frames, has something to draw. VQA still passes it through by
+    # identity via ``to_image_blocks``.
+    sentinel: Any = field(default_factory=lambda: PIL.Image.new("RGB", (32, 24)))
+    cameras: tuple[str, ...] = ("observation.images.top",)
+    calls: list[tuple[int, tuple[float, ...], str | None]] = field(default_factory=list)
+    video_calls: list[tuple[int, int, str | None]] = field(default_factory=list)
+
+    @property
+    def camera_keys(self) -> list[str]:
+        return list(self.cameras)
+
+    def frames_at(self, record, timestamps, camera_key=None):
+        self.calls.append((record.episode_index, tuple(timestamps), camera_key))
+        return [self.sentinel] * len(timestamps)
+
+    def video_for_episode(self, record, max_frames, camera_key=None):
+        self.video_calls.append((record.episode_index, max_frames, camera_key))
+        n = min(max_frames, len(record.frame_timestamps))
+        return [self.sentinel] * n
+
+
+def _spy_responder(captured: list[list[dict[str, Any]]], reply: Any):
+    def responder(messages):
+        captured.append(list(messages))
+        return reply
+
+    return StubVlmClient(responder=responder)
+
+
+def test_module1_plan_memory_subtask_smoke(fixture_dataset_root: Path, tmp_path: Path) -> None:
+    vlm = make_canned_responder(
+        {
+            "atomic subtasks": {
+                "subtasks": [
+                    {"text": "grasp the handle of the sponge", "start": 0.0, "end": 0.4},
+                    {"text": "wipe the counter from left to right", "start": 0.4, "end": 0.8},
+                    {"text": "place the sponge into the sink", "start": 0.8, "end": 1.1},
+                ]
+            },
+            "compressed semantic memory": {"memory": "wiped the counter once"},
+        },
+    )
+    module = PlanSubtasksMemoryModule(vlm=vlm, config=PlanConfig())
+    record = next(iter_episodes(fixture_dataset_root))
+    staging = EpisodeStaging(tmp_path / "stage", record.episode_index)
+    module.run_episode(record, staging)
+    rows = staging.read("plan")
+
+    styles = {r["style"] for r in rows}
+    assert {"subtask", "plan", "memory"}.issubset(styles)
+    # subtask timestamps must be exact frame timestamps
+    frame_set = set(record.frame_timestamps)
+    for row in rows:
+        assert row["timestamp"] in frame_set
+    # one plan row per subtask boundary; the first lands at t0 and each
+    # plan is the deterministic numbered list of still-todo subtasks
+    plan_rows = sorted((r for r in rows if r["style"] == "plan"), key=lambda r: r["timestamp"])
+    subtask_rows = [r for r in rows if r["style"] == "subtask"]
+    assert len(plan_rows) == len(subtask_rows)
+    assert plan_rows[0]["timestamp"] == record.frame_timestamps[0]
+    # the t0 plan enumerates all subtasks; later plans shrink
+    assert plan_rows[0]["content"].startswith("1. ")
+    assert len(plan_rows[0]["content"].splitlines()) == len(subtask_rows)
+    assert len(plan_rows[-1]["content"].splitlines()) == 1
+
+
+def test_module1_emit_memory_false_skips_memory_keeps_subtasks_and_plan(
+    fixture_dataset_root: Path, tmp_path: Path
+) -> None:
+    """``emit_memory=False`` drops ``memory`` rows (and their VLM calls) while
+    leaving subtask + plan generation intact — symmetric to ``emit_plan``."""
+    vlm = make_canned_responder(
+        {
+            "atomic subtasks": {
+                "subtasks": [
+                    {"text": "grasp the handle of the sponge", "start": 0.0, "end": 0.4},
+                    {"text": "wipe the counter from left to right", "start": 0.4, "end": 0.8},
+                    {"text": "place the sponge into the sink", "start": 0.8, "end": 1.1},
+                ]
+            },
+            "compressed semantic memory": {"memory": "wiped the counter once"},
+        },
+    )
+    module = PlanSubtasksMemoryModule(vlm=vlm, config=PlanConfig(emit_memory=False))
+    record = next(iter_episodes(fixture_dataset_root))
+    staging = EpisodeStaging(tmp_path / "stage", record.episode_index)
+    module.run_episode(record, staging)
+    rows = staging.read("plan")
+
+    styles = {r["style"] for r in rows}
+    assert "memory" not in styles
+    assert {"subtask", "plan"}.issubset(styles)
+
+
+def test_module2_at_t0_emits_speech_only_no_interjection(fixture_dataset_root: Path, tmp_path: Path) -> None:
+    vlm = make_canned_responder(
+        {"acknowledgement the robot": {"text": "Sure, on it."}},
+    )
+    module = InterjectionsAndSpeechModule(
+        vlm=vlm,
+        config=InterjectionsConfig(max_interjections_per_episode=0),
+    )
+    record = next(iter_episodes(fixture_dataset_root))
+    staging = EpisodeStaging(tmp_path / "stage", record.episode_index)
+    module.run_episode(record, staging)
+    rows = staging.read("interjections")
+    assert len(rows) == 1
+    only = rows[0]
+    assert only["role"] == "assistant"
+    assert only["style"] is None
+    assert only["content"] is None
+    assert only["timestamp"] == record.frame_timestamps[0]
+    assert only["tool_calls"][0]["function"]["name"] == "say"
+
+
+def test_module2_mid_episode_emits_paired_interjection_and_speech(
+    fixture_dataset_root: Path, tmp_path: Path
+) -> None:
+    """Module 2 anchors interjections on Module 1's subtask boundaries.
+
+    The executor runs Module 1 first, then Module 2 reads the subtask
+    rows back from the same staging tree (see
+    ``_mid_episode_interjections``). Reproduce that contract here by
+    seeding the staging with two subtask rows so a single ``0 → 1``
+    boundary exists for Module 2 to anchor on.
+    """
+    vlm = make_canned_responder(
+        {
+            "acknowledgement the robot": {"text": "OK."},
+            # Marker matches the distinctive line of
+            # ``interjections_interjection.txt`` ("Write ONE compact
+            # interjection ..."). Keep this in sync with that prompt's
+            # wording — the canned responder matches on substring.
+            "Write ONE compact interjection": {
+                "interjection": "now wipe the counter please",
+                "speech": "On it.",
+            },
+        },
+    )
+    module = InterjectionsAndSpeechModule(
+        vlm=vlm,
+        config=InterjectionsConfig(max_interjections_per_episode=1, interjection_min_t=0.2),
+        seed=7,
+    )
+    record = next(iter_episodes(fixture_dataset_root))
+    staging = EpisodeStaging(tmp_path / "stage", record.episode_index)
+    # Seed Module 1's subtask staging so Module 2 has a boundary to
+    # anchor on (it bails with zero rows when no spans exist — the
+    # production executor guarantees Module 1 ran first).
+    boundary_ts = float(record.frame_timestamps[len(record.frame_timestamps) // 2])
+    staging.write(
+        "plan",
+        [
+            {
+                "role": "assistant",
+                "content": "grasp the sponge",
+                "style": "subtask",
+                "timestamp": float(record.frame_timestamps[0]),
+                "tool_calls": None,
+            },
+            {
+                "role": "assistant",
+                "content": "wipe the counter",
+                "style": "subtask",
+                "timestamp": boundary_ts,
+                "tool_calls": None,
+            },
+        ],
+    )
+    module.run_episode(record, staging)
+    rows = staging.read("interjections")
+
+    interjections = [r for r in rows if r["style"] == "interjection"]
+    speeches = [r for r in rows if r["style"] is None and r["role"] == "assistant"]
+    assert len(interjections) == 1
+    assert len(speeches) >= 2  # initial t=0 + one paired with the interjection
+    inter_t = interjections[0]["timestamp"]
+    assert any(abs(s["timestamp"] - inter_t) < 1e-9 for s in speeches)
+
+
+def test_module3_vqa_unique_per_frame_and_camera(single_episode_root: Path, tmp_path: Path) -> None:
+    payload = {
+        "question": "How many cups?",
+        "answer": {"label": "cup", "count": 2, "note": "white & blue"},
+    }
+    vlm = make_canned_responder({"frame-grounded visual question": payload})
+    module = GeneralVqaModule(
+        vlm=vlm,
+        config=VqaConfig(vqa_emission_hz=1.0, K=3),
+        seed=1,
+        frame_provider=_StubFrameProvider(cameras=("observation.images.top", "observation.images.wrist")),
+    )
+    record = next(iter_episodes(single_episode_root))
+    staging = EpisodeStaging(tmp_path / "stage", record.episode_index)
+    module.run_episode(record, staging)
+    rows = staging.read("vqa")
+    # every vqa row must carry a camera tag and one of the configured cameras
+    for r in rows:
+        assert r["style"] == "vqa"
+        assert r.get("camera") in {"observation.images.top", "observation.images.wrist"}
+    # at most one (vqa, user) and one (vqa, assistant) per (timestamp, camera)
+    user_keys = [(r["timestamp"], r["camera"]) for r in rows if r["role"] == "user" and r["style"] == "vqa"]
+    assistant_keys = [
+        (r["timestamp"], r["camera"]) for r in rows if r["role"] == "assistant" and r["style"] == "vqa"
+    ]
+    assert len(user_keys) == len(set(user_keys))
+    assert len(assistant_keys) == len(set(assistant_keys))
+    # both cameras must be represented
+    assert {c for _, c in user_keys} == {"observation.images.top", "observation.images.wrist"}
+    # every emitted timestamp must be an exact source frame timestamp
+    frame_set = set(record.frame_timestamps)
+    for ts, _ in user_keys + assistant_keys:
+        assert ts in frame_set
+
+
+def test_module1_attaches_contact_sheets_to_subtask_prompt(
+    fixture_dataset_root: Path, tmp_path: Path
+) -> None:
+    """Module 1 sends timestamped contact-sheet image blocks (not a raw video block)."""
+    captured: list[list[dict[str, Any]]] = []
+    payload = {
+        "subtasks": [
+            {"text": "grasp the handle of the sponge", "start": 0.0, "end": 0.5},
+            {"text": "wipe the counter", "start": 0.5, "end": 1.1},
+        ]
+    }
+    memory_payload = {"memory": "wiped once"}
+
+    def responder(messages):
+        captured.append(list(messages))
+        text = ""
+        for m in messages:
+            for block in m.get("content", []):
+                if isinstance(block, dict) and block.get("type") == "text":
+                    text = block.get("text", "")
+        if "compressed semantic memory" in text:
+            return memory_payload
+        return payload
+
+    provider = _StubFrameProvider()
+    module = PlanSubtasksMemoryModule(
+        vlm=StubVlmClient(responder=responder),
+        # Disable the rephrasings sub-prompt so the test's only video-bearing
+        # call is the subtask one — keeps the assertions below focused on
+        # ``_generate_subtasks`` rather than fighting the order of unrelated
+        # text-only Module-1 sub-prompts.
+        config=PlanConfig(frames_per_second=2.0, max_frames_per_prompt=60, n_task_rephrasings=0),
+        frame_provider=provider,
+    )
+    record = next(iter_episodes(fixture_dataset_root))
+    staging = EpisodeStaging(tmp_path / "stage", record.episode_index)
+    module.run_episode(record, staging)
+
+    # Find the call carrying the subtask prompt rather than blindly taking
+    # captured[0] — Module 1 issues several sub-prompts and their order is
+    # not part of the contract.
+    assert captured, "no VLM calls made"
+
+    def _prompt_text(messages):
+        for m in messages:
+            for block in m.get("content", []):
+                if isinstance(block, dict) and block.get("type") == "text":
+                    return block.get("text", "")
+        return ""
+
+    subtask_calls = [m for m in captured if "atomic subtasks" in _prompt_text(m)]
+    assert len(subtask_calls) == 1, "expected exactly one subtask-prompt VLM call"
+    content = subtask_calls[0][0]["content"]
+    video_blocks = [b for b in content if isinstance(b, dict) and b.get("type") == "video"]
+    image_blocks = [b for b in content if isinstance(b, dict) and b.get("type") == "image"]
+    text_blocks = [b for b in content if isinstance(b, dict) and b.get("type") == "text"]
+    assert video_blocks == [], "contact-sheet mode must not emit a raw video block"
+    assert len(image_blocks) >= 1, f"expected >=1 contact-sheet image block, got {content}"
+    assert all(isinstance(b["image"], PIL.Image.Image) for b in image_blocks)
+    assert len(text_blocks) == 1
+    # the prompt is prefixed with the contact-sheet reading instructions
+    assert text_blocks[0]["text"].startswith("CONTACT SHEETS")
+    # frames were decoded for this episode at episode-relative timestamps
+    assert provider.calls and provider.calls[0][0] == record.episode_index
+
+
+def test_module3_attaches_frame_image_block_to_prompt(single_episode_root: Path, tmp_path: Path) -> None:
+    """Each VQA prompt must carry a single image block at the emission frame."""
+    captured: list[list[dict[str, Any]]] = []
+    payload = {
+        "question": "How many cups?",
+        "answer": {"label": "cup", "count": 1},
+    }
+    provider = _StubFrameProvider()
+    module = GeneralVqaModule(
+        vlm=_spy_responder(captured, payload),
+        config=VqaConfig(vqa_emission_hz=1.0, K=1),
+        seed=0,
+        frame_provider=provider,
+    )
+    record = next(iter_episodes(single_episode_root))
+    staging = EpisodeStaging(tmp_path / "stage", record.episode_index)
+    module.run_episode(record, staging)
+
+    assert captured, "no VLM calls made"
+    for messages in captured:
+        content = messages[0]["content"]
+        image_blocks = [b for b in content if isinstance(b, dict) and b.get("type") == "image"]
+        text_blocks = [b for b in content if isinstance(b, dict) and b.get("type") == "text"]
+        assert len(image_blocks) == 1, f"expected 1 image block per VQA prompt, got {content}"
+        assert image_blocks[0]["image"] is provider.sentinel
+        assert len(text_blocks) == 1
+    # provider was called once per emission per camera with the exact emission timestamp
+    for ep_idx, ts_tuple, camera in provider.calls:
+        assert ep_idx == record.episode_index
+        assert len(ts_tuple) == 1
+        assert ts_tuple[0] in record.frame_timestamps
+        assert camera in provider.cameras
+
+
+def test_module3_assistant_content_is_valid_json(single_episode_root: Path, tmp_path: Path) -> None:
+    payload = {
+        "question": "Where is the cup?",
+        "answer": {"detections": [{"label": "cup", "bbox_format": "xyxy", "bbox": [10, 20, 50, 80]}]},
+    }
+    vlm = make_canned_responder({"frame-grounded visual question": payload})
+    module = GeneralVqaModule(
+        vlm=vlm,
+        config=VqaConfig(vqa_emission_hz=1.0, K=2),
+        seed=2,
+        frame_provider=_StubFrameProvider(),
+    )
+    record = next(iter_episodes(single_episode_root))
+    staging = EpisodeStaging(tmp_path / "stage", record.episode_index)
+    module.run_episode(record, staging)
+    rows = staging.read("vqa")
+    for row in rows:
+        if row["role"] == "assistant" and row["style"] == "vqa":
+            decoded = json.loads(row["content"])
+            assert "detections" in decoded

tests/annotations/test_pipeline_recipe_render.py

@@ -0,0 +1,183 @@
+#!/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.
+"""End-to-end smoke: pipeline output → canonical recipe rendering."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+
+# ``pyarrow`` and the ``lerobot.datasets`` chain (-> the HF ``datasets``
+# library) only ship under the ``dataset`` extra. Skip this module in
+# tiers without it instead of erroring at import.
+pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
+pytest.importorskip("pandas", reason="pandas is required (install lerobot[dataset])")
+
+import pyarrow.parquet as pq  # noqa: E402
+
+from lerobot.annotations.steerable_pipeline.config import (  # noqa: E402
+    AnnotationPipelineConfig,
+    InterjectionsConfig,
+    PlanConfig,
+    VqaConfig,
+)
+from lerobot.annotations.steerable_pipeline.executor import Executor  # noqa: E402
+from lerobot.annotations.steerable_pipeline.modules import (  # noqa: E402
+    GeneralVqaModule,
+    InterjectionsAndSpeechModule,
+    PlanSubtasksMemoryModule,
+)
+from lerobot.annotations.steerable_pipeline.validator import StagingValidator  # noqa: E402
+from lerobot.annotations.steerable_pipeline.writer import LanguageColumnsWriter  # noqa: E402
+from lerobot.configs.recipe import MessageTurn, TrainingRecipe  # noqa: E402
+from lerobot.datasets.language_render import render_sample  # noqa: E402
+
+from ._helpers import make_canned_responder  # noqa: E402
+
+
+def _build_style_blend_recipe() -> TrainingRecipe:
+    """Inline blend recipe that consumes every style this pipeline produces.
+
+    The language schema/DSL work used to ship
+    ``src/lerobot/configs/recipes/pi05_hirobot.yaml`` as a canonical
+    example, but that file was dropped during review. The contract this
+    test guards is "the recipe DSL can render non-empty messages from
+    pipeline output", which doesn't require a specific YAML — so we build
+    the equivalent blend in code.
+    """
+    return TrainingRecipe(
+        blend={
+            "low_level_execution": TrainingRecipe(
+                weight=0.35,
+                messages=[
+                    MessageTurn(
+                        role="user",
+                        content="${task}\nPlan: ${plan}\nMemory: ${memory}",
+                        stream="high_level",
+                    ),
+                    MessageTurn(role="assistant", content="${subtask}", stream="low_level", target=True),
+                ],
+            ),
+            "user_interjection_response": TrainingRecipe(
+                weight=0.16,
+                bindings={
+                    "speech": "emitted_at(t, role=assistant, tool_name=say)",
+                    "interjection": "emitted_at(t, style=interjection)",
+                },
+                messages=[
+                    MessageTurn(role="user", content="${task}", stream="high_level"),
+                    MessageTurn(
+                        role="user",
+                        content="${interjection}",
+                        stream="high_level",
+                        if_present="interjection",
+                    ),
+                    MessageTurn(
+                        role="assistant",
+                        content="${plan}",
+                        stream="high_level",
+                        target=True,
+                        if_present="plan",
+                        tool_calls_from="speech",
+                    ),
+                ],
+            ),
+        }
+    )
+
+
+def _build_executor() -> Executor:
+    vlm = make_canned_responder(
+        {
+            "atomic subtasks": {
+                "subtasks": [
+                    {"text": "grasp the bottle", "start": 0.0, "end": 0.5},
+                    {"text": "pour into the cup", "start": 0.5, "end": 1.0},
+                    {"text": "place the bottle down", "start": 1.0, "end": 1.5},
+                ]
+            },
+            "compressed semantic memory": {"memory": "poured once"},
+            "acknowledgement the robot": {"text": "Sure."},
+            "compact interjection": {
+                "interjection": "use less water",
+                "speech": "Using less water.",
+            },
+            "frame-grounded visual question": {
+                "question": "How many cups?",
+                "answer": {"label": "cup", "count": 1},
+            },
+        },
+    )
+    config = AnnotationPipelineConfig(
+        plan=PlanConfig(),
+        interjections=InterjectionsConfig(max_interjections_per_episode=1, interjection_min_t=0.5),
+        vqa=VqaConfig(vqa_emission_hz=1.0, K=2),
+    )
+    return Executor(
+        config=config,
+        plan=PlanSubtasksMemoryModule(vlm=vlm, config=config.plan),
+        interjections=InterjectionsAndSpeechModule(vlm=vlm, config=config.interjections, seed=config.seed),
+        vqa=GeneralVqaModule(vlm=vlm, config=config.vqa, seed=config.seed),
+        writer=LanguageColumnsWriter(),
+        validator=StagingValidator(),
+    )
+
+
+def test_canonical_recipe_renders_nonempty_from_pipeline_output(
+    single_episode_root: Path,
+) -> None:
+    executor = _build_executor()
+    summary = executor.run(single_episode_root)
+    # validator may emit warnings but no errors for the synthetic fixture
+    assert summary.validation_report.ok, summary.validation_report.summary()
+
+    table = pq.read_table(single_episode_root / "data" / "chunk-000" / "file-000.parquet")
+    persistent_lists = table.column("language_persistent").to_pylist()
+    events_lists = table.column("language_events").to_pylist()
+    timestamps = table.column("timestamp").to_pylist()
+
+    recipe = _build_style_blend_recipe()
+
+    rendered_any = False
+    for ts, persistent, events in zip(timestamps, persistent_lists, events_lists, strict=True):
+        result = render_sample(
+            recipe=recipe,
+            persistent=persistent,
+            events=events,
+            t=float(ts),
+            sample_idx=0,
+            dataset_ctx={"task": "Pour water from the bottle into the cup."},
+        )
+        if result is None:
+            continue
+        if result["messages"]:
+            rendered_any = True
+            assert result["target_message_indices"]
+            break
+    assert rendered_any, "recipe rendered no messages from pipeline output"
+
+    # Sanity: speech atom appears in events column intact
+    flat_events = [r for ev in events_lists for r in ev]
+    speech_rows = [r for r in flat_events if r.get("style") is None and r.get("role") == "assistant"]
+    assert speech_rows
+    say = speech_rows[0]["tool_calls"][0]
+    assert say["function"]["name"] == "say"
+    assert isinstance(say["function"]["arguments"]["text"], str)
+    # The pipeline does not write a ``tools`` column — the say schema lives
+    # as a constant (``SAY_TOOL_SCHEMA``) so the language row struct is the
+    # single source of truth for the v3.1 schema.
+    assert "tools" not in table.column_names

tests/annotations/test_validator.py

@@ -0,0 +1,133 @@
+#!/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.
+"""Validator behavior tests."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import pytest
+
+# ``lerobot.annotations`` imports pull in ``lerobot.datasets`` (-> the HF
+# ``datasets`` library), which only ships under the ``dataset`` extra. Skip
+# this module in tiers without it instead of erroring at import.
+pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
+pytest.importorskip("pandas", reason="pandas is required (install lerobot[dataset])")
+
+from lerobot.annotations.steerable_pipeline.reader import iter_episodes  # noqa: E402
+from lerobot.annotations.steerable_pipeline.staging import EpisodeStaging  # noqa: E402
+from lerobot.annotations.steerable_pipeline.validator import StagingValidator  # noqa: E402
+from lerobot.annotations.steerable_pipeline.writer import speech_atom  # noqa: E402
+
+
+def _validate(root: Path, staging_dir: Path):
+    records = list(iter_episodes(root))
+    return StagingValidator().validate(records, staging_dir)
+
+
+def test_validator_catches_misaligned_timestamps(fixture_dataset_root: Path, tmp_path: Path) -> None:
+    staging_dir = tmp_path / "stage"
+    EpisodeStaging(staging_dir, 0).write(
+        "vqa",
+        [
+            {
+                "role": "assistant",
+                "content": json.dumps({"label": "cup", "count": 2}, sort_keys=True),
+                "style": "vqa",
+                "timestamp": 9.999,  # not on any 10 fps frame
+                "tool_calls": None,
+            }
+        ],
+    )
+    report = _validate(fixture_dataset_root, staging_dir)
+    assert not report.ok
+    assert any("does not match any source frame timestamp" in e for e in report.errors)
+
+
+def test_validator_catches_orphan_speech(fixture_dataset_root: Path, tmp_path: Path) -> None:
+    staging_dir = tmp_path / "stage"
+    EpisodeStaging(staging_dir, 0).write(
+        "interjections",
+        [
+            speech_atom(0.0, "Got it."),
+            # interjection at 0.3s with NO paired speech
+            {
+                "role": "user",
+                "content": "skip it",
+                "style": "interjection",
+                "timestamp": 0.3,
+                "tool_calls": None,
+            },
+        ],
+    )
+    report = _validate(fixture_dataset_root, staging_dir)
+    assert not report.ok
+    assert any("paired speech" in e for e in report.errors)
+
+
+def test_validator_catches_inconsistent_plan_memory(fixture_dataset_root: Path, tmp_path: Path) -> None:
+    staging_dir = tmp_path / "stage"
+    EpisodeStaging(staging_dir, 0).write(
+        "plan",
+        [
+            {
+                "role": "assistant",
+                "content": "1. do x",
+                "style": "plan",
+                "timestamp": 0.0,
+                "tool_calls": None,
+            },
+            {
+                "role": "assistant",
+                "content": "do x",
+                "style": "subtask",
+                "timestamp": 0.0,
+                "tool_calls": None,
+            },
+        ],
+    )
+    EpisodeStaging(staging_dir, 0).write(
+        "interjections",
+        [
+            speech_atom(0.0, "Got it."),
+            speech_atom(0.4, "Replanning."),
+            {
+                "role": "user",
+                "content": "replan",
+                "style": "interjection",
+                "timestamp": 0.4,
+                "tool_calls": None,
+            },
+        ],
+    )
+    report = _validate(fixture_dataset_root, staging_dir)
+    # missing co-timestamped plan refresh at 0.4s → error
+    assert not report.ok
+    assert any("co-timestamped plan update" in e for e in report.errors)
+
+
+def test_validator_catches_wrong_column(fixture_dataset_root: Path, tmp_path: Path) -> None:
+    staging_dir = tmp_path / "stage"
+    EpisodeStaging(staging_dir, 0).write(
+        "plan",
+        [
+            {"role": "user", "content": "where?", "style": "vqa", "timestamp": 0.0, "tool_calls": None},
+        ],
+    )
+    report = _validate(fixture_dataset_root, staging_dir)
+    assert not report.ok
+    assert any("plan emitted style 'vqa'" in e or "must be persistent" in e for e in report.errors)

tests/annotations/test_vlm_client.py

@@ -0,0 +1,41 @@
+#!/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.
+"""Unit tests for ``vlm_client`` helpers."""
+
+from __future__ import annotations
+
+import pytest
+
+pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
+
+from lerobot.annotations.steerable_pipeline.vlm_client import _bind_serve_port  # noqa: E402
+
+
+def test_bind_serve_port_substitutes_placeholder() -> None:
+    # The {port} placeholder is replaced everywhere it appears, regardless of
+    # parallel vs single server — the bug was the single-server path passing
+    # it through unsubstituted.
+    cmd = "vllm serve M --max-model-len 32768 --port {port}"
+    assert _bind_serve_port(cmd, 8000) == "vllm serve M --max-model-len 32768 --port 8000"
+
+
+def test_bind_serve_port_appends_when_missing() -> None:
+    assert _bind_serve_port("vllm serve M", 8001) == "vllm serve M --port 8001"
+
+
+def test_bind_serve_port_leaves_explicit_port_untouched() -> None:
+    cmd = "vllm serve M --port 9000"
+    assert _bind_serve_port(cmd, 8000) == cmd

tests/annotations/test_writer.py

@@ -0,0 +1,430 @@
+#!/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.
+"""Writer correctness tests."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import pytest
+
+# ``pyarrow`` and the ``lerobot.annotations`` -> ``lerobot.datasets`` chain
+# (-> the HF ``datasets`` library) only ship under the ``dataset`` extra.
+# Skip this module in tiers without it instead of erroring at import.
+pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
+pytest.importorskip("pandas", reason="pandas is required (install lerobot[dataset])")
+
+import pandas as pd  # noqa: E402
+import pyarrow.parquet as pq  # noqa: E402
+
+from lerobot.annotations.steerable_pipeline.reader import iter_episodes  # noqa: E402
+from lerobot.annotations.steerable_pipeline.staging import EpisodeStaging  # noqa: E402
+from lerobot.annotations.steerable_pipeline.writer import (  # noqa: E402
+    LanguageColumnsWriter,
+    speech_atom,
+)
+
+
+def _stage_episode(
+    staging_dir: Path,
+    episode_index: int,
+    *,
+    plan: list[dict] | None = None,
+    interjections: list[dict] | None = None,
+    vqa: list[dict] | None = None,
+) -> None:
+    staging = EpisodeStaging(staging_dir, episode_index)
+    if plan is not None:
+        staging.write("plan", plan)
+    if interjections is not None:
+        staging.write("interjections", interjections)
+    if vqa is not None:
+        staging.write("vqa", vqa)
+
+
+def test_writer_persistence_identity(fixture_dataset_root: Path, tmp_path: Path) -> None:
+    """Every frame in an episode has a byte-identical persistent list."""
+    staging_dir = tmp_path / "stage"
+    _stage_episode(
+        staging_dir,
+        0,
+        plan=[
+            {
+                "role": "assistant",
+                "content": "grasp the sponge",
+                "style": "subtask",
+                "timestamp": 0.0,
+                "tool_calls": None,
+            },
+            {
+                "role": "assistant",
+                "content": "1. wipe\n2. dry",
+                "style": "plan",
+                "timestamp": 0.0,
+                "tool_calls": None,
+            },
+            {
+                "role": "assistant",
+                "content": "wiped the counter",
+                "style": "memory",
+                "timestamp": 0.5,
+                "tool_calls": None,
+            },
+        ],
+    )
+    records = list(iter_episodes(fixture_dataset_root))
+    LanguageColumnsWriter().write_all(records, staging_dir, fixture_dataset_root)
+
+    table = pq.read_table(fixture_dataset_root / "data" / "chunk-000" / "file-000.parquet")
+    persistent = table.column("language_persistent").to_pylist()
+    first = persistent[0]
+    assert first  # non-empty
+    for row in persistent:
+        assert row == first, "persistent slice must be byte-identical across all frames"
+
+
+def test_writer_events_exact_timestamp(fixture_dataset_root: Path, tmp_path: Path) -> None:
+    staging_dir = tmp_path / "stage"
+    _stage_episode(
+        staging_dir,
+        0,
+        interjections=[
+            speech_atom(0.0, "Got it."),
+            {
+                "role": "user",
+                "content": "skip the dishes",
+                "style": "interjection",
+                "timestamp": 0.5,
+                "tool_calls": None,
+            },
+            speech_atom(0.5, "Skipping the dishes."),
+        ],
+    )
+    records = list(iter_episodes(fixture_dataset_root))
+    LanguageColumnsWriter().write_all(records, staging_dir, fixture_dataset_root)
+
+    table = pq.read_table(fixture_dataset_root / "data" / "chunk-000" / "file-000.parquet")
+    timestamps = table.column("timestamp").to_pylist()
+    events = table.column("language_events").to_pylist()
+    for ts, ev in zip(timestamps, events, strict=True):
+        if abs(ts - 0.0) < 1e-9:
+            assert any(r["role"] == "assistant" and r.get("style") is None for r in ev), ev
+        elif abs(ts - 0.5) < 1e-9:
+            assert any(r.get("style") == "interjection" for r in ev), ev
+            assert any(r.get("style") is None for r in ev), ev
+        else:
+            assert ev == []
+
+
+def test_writer_column_routing(fixture_dataset_root: Path, tmp_path: Path) -> None:
+    staging_dir = tmp_path / "stage"
+    _stage_episode(
+        staging_dir,
+        0,
+        plan=[
+            {
+                "role": "assistant",
+                "content": "do X",
+                "style": "subtask",
+                "timestamp": 0.0,
+                "tool_calls": None,
+            },
+            {
+                "role": "assistant",
+                "content": "1. do X",
+                "style": "plan",
+                "timestamp": 0.0,
+                "tool_calls": None,
+            },
+            {
+                "role": "assistant",
+                "content": "did X",
+                "style": "memory",
+                "timestamp": 0.3,
+                "tool_calls": None,
+            },
+        ],
+        interjections=[
+            speech_atom(0.0, "OK"),
+            {
+                "role": "user",
+                "content": "wait",
+                "style": "interjection",
+                "timestamp": 0.2,
+                "tool_calls": None,
+            },
+            speech_atom(0.2, "Waiting"),
+        ],
+        vqa=[
+            {
+                "role": "user",
+                "content": "where is the cup?",
+                "style": "vqa",
+                "timestamp": 0.4,
+                "camera": "observation.images.front",
+                "tool_calls": None,
+            },
+            {
+                "role": "assistant",
+                "content": json.dumps(
+                    {"detections": [{"label": "cup", "bbox_format": "xyxy", "bbox": [1, 2, 3, 4]}]},
+                    sort_keys=True,
+                ),
+                "style": "vqa",
+                "timestamp": 0.4,
+                "camera": "observation.images.front",
+                "tool_calls": None,
+            },
+        ],
+    )
+    records = list(iter_episodes(fixture_dataset_root))
+    LanguageColumnsWriter().write_all(records, staging_dir, fixture_dataset_root)
+    table = pq.read_table(fixture_dataset_root / "data" / "chunk-000" / "file-000.parquet")
+
+    persistent = table.column("language_persistent").to_pylist()[0]
+    persistent_styles = {r["style"] for r in persistent}
+    assert persistent_styles == {"subtask", "plan", "memory"}
+
+    all_events = [r for ev in table.column("language_events").to_pylist() for r in ev]
+    event_styles = {r.get("style") for r in all_events}
+    assert event_styles == {None, "interjection", "vqa"}
+
+
+def test_writer_drops_subtask_index_idempotent(fixture_dataset_root: Path, tmp_path: Path) -> None:
+    staging_dir = tmp_path / "stage"
+    _stage_episode(
+        staging_dir,
+        0,
+        plan=[
+            {
+                "role": "assistant",
+                "content": "do X",
+                "style": "subtask",
+                "timestamp": 0.0,
+                "tool_calls": None,
+            },
+        ],
+    )
+    records = list(iter_episodes(fixture_dataset_root))
+    writer = LanguageColumnsWriter()
+    writer.write_all(records, staging_dir, fixture_dataset_root)
+
+    path = fixture_dataset_root / "data" / "chunk-000" / "file-000.parquet"
+    table_a = pq.read_table(path)
+    assert "subtask_index" not in table_a.column_names
+    assert "language_persistent" in table_a.column_names
+    assert "language_events" in table_a.column_names
+    # The writer no longer emits a dataset-level ``tools`` column; the
+    # ``say`` tool schema lives as a code constant (``SAY_TOOL_SCHEMA``)
+    # so the parquet stays small and the pipeline doesn't extend the schema.
+    assert "tools" not in table_a.column_names
+
+    # second pass — must produce identical bytes for the language columns
+    records_again = list(iter_episodes(fixture_dataset_root))
+    writer.write_all(records_again, staging_dir, fixture_dataset_root)
+    table_b = pq.read_table(path)
+    assert (
+        table_a.column("language_persistent").to_pylist() == table_b.column("language_persistent").to_pylist()
+    )
+    assert table_a.column("language_events").to_pylist() == table_b.column("language_events").to_pylist()
+
+
+def test_writer_normalize_rejects_misrouted_persistent_style() -> None:
+    """``_normalize_persistent_row`` must reject any non-persistent style."""
+    from lerobot.annotations.steerable_pipeline.writer import _normalize_persistent_row
+
+    with pytest.raises(ValueError, match="non-persistent style"):
+        _normalize_persistent_row(
+            {"role": "assistant", "content": "oops", "style": "vqa", "timestamp": 0.0, "tool_calls": None}
+        )
+
+
+def test_writer_normalize_rejects_misrouted_event_style() -> None:
+    """``_normalize_event_row`` must reject any persistent style."""
+    from lerobot.annotations.steerable_pipeline.writer import _normalize_event_row
+
+    with pytest.raises(ValueError):
+        _normalize_event_row({"role": "assistant", "content": "oops", "style": "subtask", "tool_calls": None})
+
+
+def test_say_tool_schema_constant_is_well_formed() -> None:
+    """``SAY_TOOL_SCHEMA`` (and ``DEFAULT_TOOLS``) replace the parquet
+    ``tools`` column — chat-template consumers import them directly.
+    """
+    from lerobot.annotations.steerable_pipeline.writer import (
+        DEFAULT_TOOLS,
+        SAY_TOOL_SCHEMA,
+    )
+
+    assert DEFAULT_TOOLS == [SAY_TOOL_SCHEMA]
+    assert SAY_TOOL_SCHEMA["function"]["name"] == "say"
+    params = SAY_TOOL_SCHEMA["function"]["parameters"]
+    assert params["properties"]["text"]["type"] == "string"
+    assert params["required"] == ["text"]
+
+
+def test_writer_does_not_add_tools_column(fixture_dataset_root: Path, tmp_path: Path) -> None:
+    """Re-running on a parquet that already has a legacy ``tools`` column
+    must drop it cleanly so reruns converge to the v3.1 schema.
+    """
+    staging_dir = tmp_path / "stage"
+    _stage_episode(
+        staging_dir,
+        0,
+        plan=[
+            {"role": "assistant", "content": "x", "style": "subtask", "timestamp": 0.0, "tool_calls": None}
+        ],
+    )
+    records = list(iter_episodes(fixture_dataset_root))
+    LanguageColumnsWriter().write_all(records, staging_dir, fixture_dataset_root)
+    table = pq.read_table(fixture_dataset_root / "data" / "chunk-000" / "file-000.parquet")
+    assert "tools" not in table.column_names
+
+
+def test_annotation_metadata_sync_allows_non_streaming_load(
+    fixture_dataset_root: Path, tmp_path: Path
+) -> None:
+    """Annotated parquet columns must be declared in ``meta/info.json``.
+
+    ``LeRobotDataset`` loads non-streaming datasets by casting parquet
+    against metadata-derived HF features. If the annotation writer adds
+    language columns but metadata stays stale, that cast fails with a column
+    mismatch.
+    """
+    from lerobot.annotations.steerable_pipeline.executor import Executor
+    from lerobot.datasets.feature_utils import get_hf_features_from_features
+    from lerobot.datasets.io_utils import load_info, load_nested_dataset
+    from lerobot.datasets.language import LANGUAGE_EVENTS, LANGUAGE_PERSISTENT, language_feature_info
+
+    info_path = fixture_dataset_root / "meta" / "info.json"
+    info = json.loads(info_path.read_text())
+    info["features"] = {
+        "episode_index": {"dtype": "int64", "shape": (1,), "names": None},
+        "frame_index": {"dtype": "int64", "shape": (1,), "names": None},
+        "timestamp": {"dtype": "float32", "shape": (1,), "names": None},
+        "task_index": {"dtype": "int64", "shape": (1,), "names": None},
+    }
+    info_path.write_text(json.dumps(info, indent=2))
+
+    staging_dir = tmp_path / "stage"
+    _stage_episode(
+        staging_dir,
+        0,
+        plan=[
+            {"role": "assistant", "content": "do X", "style": "subtask", "timestamp": 0.0, "tool_calls": None}
+        ],
+    )
+    records = list(iter_episodes(fixture_dataset_root))
+    LanguageColumnsWriter().write_all(records, staging_dir, fixture_dataset_root)
+
+    Executor._ensure_annotation_metadata_in_info(fixture_dataset_root)
+
+    synced = load_info(fixture_dataset_root)
+    for key, feature in language_feature_info().items():
+        assert synced["features"][key] == feature
+
+    hf_features = get_hf_features_from_features(synced["features"])
+    dataset = load_nested_dataset(fixture_dataset_root / "data", features=hf_features)
+
+    assert LANGUAGE_PERSISTENT in dataset.column_names
+    assert LANGUAGE_EVENTS in dataset.column_names
+    assert len(dataset) == 24
+
+
+def _build_packed_dataset(root: Path, episode_lengths: list[int], *, fps: int = 10) -> Path:
+    """Pack several episodes into a single shard (vs build_annotation_dataset's one-per-file),
+    so the writer's rewrite must re-emit one row group per episode instead of collapsing them."""
+    from lerobot.datasets.io_utils import write_tasks
+    from lerobot.utils.io_utils import write_json
+
+    data_dir = root / "data" / "chunk-000"
+    data_dir.mkdir(parents=True, exist_ok=True)
+
+    episode_index, frame_index, timestamp, task_index, subtask_index = [], [], [], [], []
+    for ep, length in enumerate(episode_lengths):
+        episode_index += [ep] * length
+        frame_index += list(range(length))
+        timestamp += [round(i / fps, 6) for i in range(length)]
+        task_index += [0] * length
+        subtask_index += [0] * length  # legacy column the writer must drop
+    pd.DataFrame(
+        {
+            "episode_index": episode_index,
+            "frame_index": frame_index,
+            "timestamp": timestamp,
+            "task_index": task_index,
+            "subtask_index": subtask_index,
+        }
+    ).to_parquet(data_dir / "file-000.parquet", index=False)
+
+    tasks_df = pd.DataFrame({"task_index": [0]}, index=pd.Index(["do the thing"], name="task"))
+    write_tasks(tasks_df, root)
+    write_json(
+        {"codebase_version": "v3.1", "fps": fps, "features": {}, "total_episodes": len(episode_lengths)},
+        root / "meta" / "info.json",
+    )
+    return root
+
+
+def test_writer_one_row_group_per_episode(tmp_path: Path) -> None:
+    """Rewriting a packed shard must keep one row group per episode, not collapse
+    every episode into a single giant row group."""
+    episode_lengths = [4, 6, 5]  # unequal lengths, all in one shard
+    root = _build_packed_dataset(tmp_path / "ds", episode_lengths)
+    shard = root / "data" / "chunk-000" / "file-000.parquet"
+    assert pq.ParquetFile(shard).metadata.num_row_groups == 1, "fixture should start collapsed"
+
+    staging_dir = tmp_path / "stage"
+    for ep in range(len(episode_lengths)):
+        _stage_episode(
+            staging_dir,
+            ep,
+            plan=[
+                {
+                    "role": "assistant",
+                    "content": f"subtask for ep {ep}",
+                    "style": "subtask",
+                    "timestamp": 0.0,
+                    "tool_calls": None,
+                }
+            ],
+        )
+
+    records = list(iter_episodes(root))
+    LanguageColumnsWriter().write_all(records, staging_dir, root)
+
+    # One row group per episode, with row counts matching the episode lengths.
+    md = pq.ParquetFile(shard).metadata
+    assert md.num_row_groups == len(episode_lengths)
+    assert [md.row_group(i).num_rows for i in range(md.num_row_groups)] == episode_lengths
+    # Language columns are still present after the per-episode rewrite.
+    table = pq.read_table(shard)
+    assert "language_persistent" in table.column_names
+    assert "language_events" in table.column_names
+
+
+def test_speech_atom_shape_matches_plan_spec() -> None:
+    atom = speech_atom(2.5, "I'm cleaning up!")
+    assert atom["role"] == "assistant"
+    assert atom["style"] is None
+    assert atom["content"] is None
+    assert atom["timestamp"] == 2.5
+    assert isinstance(atom["tool_calls"], list)
+    call = atom["tool_calls"][0]
+    assert call["type"] == "function"
+    assert call["function"]["name"] == "say"
+    assert call["function"]["arguments"]["text"] == "I'm cleaning up!"

tests/datasets/test_aggregate.py

@@ -29,7 +29,30 @@ from lerobot.configs import VIDEO_ENCODER_INFO_KEYS
 from lerobot.datasets.aggregate import aggregate_datasets
 from lerobot.datasets.feature_utils import features_equal_for_merge
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from tests.fixtures.constants import DUMMY_REPO_ID
+from tests.fixtures.constants import (
+    DUMMY_CAMERA_FEATURES_WITH_DEPTH,
+    DUMMY_REPO_ID,
+)
+
+
+def assert_data_shards_one_row_group_per_episode(root):
+    """Every aggregated DATA shard must have exactly one parquet row group per episode."""
+    import pyarrow.parquet as pq
+
+    shards = sorted((root / "data").rglob("*.parquet"))
+    assert shards, f"no data shards found under {root}/data"
+    n_episodes = 0
+    for shard in shards:
+        pf = pq.ParquetFile(shard)
+        episodes = pf.read(columns=["episode_index"]).column("episode_index").to_pylist()
+        assert pf.metadata.num_row_groups == len(set(episodes)), shard
+        for i in range(pf.metadata.num_row_groups):
+            rg_episodes = set(
+                pf.read_row_group(i, columns=["episode_index"]).column("episode_index").to_pylist()
+            )
+            assert len(rg_episodes) == 1, f"{shard} row group {i} spans episodes {rg_episodes}"
+        n_episodes += len(set(episodes))
+    return n_episodes
 
 
 def assert_episode_and_frame_counts(aggr_ds, expected_episodes, expected_frames):
@@ -191,6 +214,26 @@ def assert_dataset_iteration_works(aggr_ds):
         pass
 
 
+def assert_depth_keys_preserved(aggr_ds, ds_0, ds_1):
+    """Test that depth keys are correctly preserved after aggregation.
+
+    Ensures that the ``is_depth_map`` marker on visual features survives
+    aggregation, so that downstream consumers (e.g. the dataset reader's
+    depth decoding path) keep working on the merged dataset.
+    """
+    expected_depth_keys = set(ds_0.meta.depth_keys)
+    assert expected_depth_keys == set(ds_1.meta.depth_keys), (
+        "Source datasets disagree on depth_keys; test setup is inconsistent"
+    )
+    actual_depth_keys = set(aggr_ds.meta.depth_keys)
+    assert actual_depth_keys == expected_depth_keys, (
+        f"Expected depth_keys {expected_depth_keys}, got {actual_depth_keys}"
+    )
+    for key in expected_depth_keys:
+        info = aggr_ds.meta.info.features[key].get("info") or {}
+        assert info.get("is_depth_map") is True, f"Depth marker lost on feature {key!r} after aggregation"
+
+
 def assert_video_timestamps_within_bounds(aggr_ds):
     """Test that all video timestamps are within valid bounds for their respective video files.
 
@@ -240,7 +283,11 @@ def assert_video_timestamps_within_bounds(aggr_ds):
 
 
 def test_aggregate_datasets(tmp_path, lerobot_dataset_factory):
-    """Test basic aggregation functionality with standard parameters."""
+    """Test basic aggregation functionality with standard parameters.
+
+    Source datasets include both RGB and depth video features so the same
+    aggregation flow is exercised on the ``is_depth_map`` branch.
+    """
     ds_0_num_frames = 400
     ds_1_num_frames = 800
     ds_0_num_episodes = 10
@@ -252,14 +299,21 @@ def test_aggregate_datasets(tmp_path, lerobot_dataset_factory):
         repo_id=f"{DUMMY_REPO_ID}_0",
         total_episodes=ds_0_num_episodes,
         total_frames=ds_0_num_frames,
+        camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH,
     )
     ds_1 = lerobot_dataset_factory(
         root=tmp_path / "test_1",
         repo_id=f"{DUMMY_REPO_ID}_1",
         total_episodes=ds_1_num_episodes,
         total_frames=ds_1_num_frames,
+        camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH,
     )
 
+    # Confirm depth was actually wired into the source datasets so the
+    # rest of the assertions exercise the depth aggregation path.
+    assert len(ds_0.meta.depth_keys) > 0, "ds_0 should expose at least one depth key"
+    assert len(ds_1.meta.depth_keys) > 0, "ds_1 should expose at least one depth key"
+
     aggregate_datasets(
         repo_ids=[ds_0.repo_id, ds_1.repo_id],
         roots=[ds_0.root, ds_1.root],
@@ -286,9 +340,56 @@ def test_aggregate_datasets(tmp_path, lerobot_dataset_factory):
     assert_episode_indices_updated_correctly(aggr_ds, ds_0, ds_1)
     assert_video_frames_integrity(aggr_ds, ds_0, ds_1)
     assert_video_timestamps_within_bounds(aggr_ds)
+    assert_depth_keys_preserved(aggr_ds, ds_0, ds_1)
     assert_dataset_iteration_works(aggr_ds)
 
 
+def test_aggregate_datasets_without_concatenation(tmp_path, lerobot_dataset_factory):
+    """With concatenation disabled, each source file is kept as its own destination file."""
+    ds_0 = lerobot_dataset_factory(
+        root=tmp_path / "no_stitch_0",
+        repo_id=f"{DUMMY_REPO_ID}_no_stitch_0",
+        total_episodes=3,
+        total_frames=60,
+    )
+    ds_1 = lerobot_dataset_factory(
+        root=tmp_path / "no_stitch_1",
+        repo_id=f"{DUMMY_REPO_ID}_no_stitch_1",
+        total_episodes=4,
+        total_frames=80,
+    )
+
+    aggr_root = tmp_path / "no_stitch_aggr"
+    aggregate_datasets(
+        repo_ids=[ds_0.repo_id, ds_1.repo_id],
+        roots=[ds_0.root, ds_1.root],
+        aggr_repo_id=f"{DUMMY_REPO_ID}_no_stitch_aggr",
+        aggr_root=aggr_root,
+        concatenate_videos=False,
+        concatenate_data=False,
+    )
+
+    with (
+        patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
+        patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
+    ):
+        mock_get_safe_version.return_value = "v3.0"
+        mock_snapshot_download.return_value = str(aggr_root)
+        aggr_ds = LeRobotDataset(f"{DUMMY_REPO_ID}_no_stitch_aggr", root=aggr_root)
+
+    assert_episode_and_frame_counts(
+        aggr_ds, ds_0.num_episodes + ds_1.num_episodes, ds_0.num_frames + ds_1.num_frames
+    )
+    assert_dataset_iteration_works(aggr_ds)
+    assert_video_timestamps_within_bounds(aggr_ds)
+
+    # Two single-file sources stay as two files each, instead of being packed together.
+    assert len(list((aggr_root / "data").rglob("*.parquet"))) == 2
+    assert aggr_ds.meta.video_keys, "Test fixture should produce at least one video feature"
+    for key in aggr_ds.meta.video_keys:
+        assert len(list((aggr_root / "videos" / key).rglob("*.mp4"))) == 2
+
+
 @pytest.mark.parametrize("mutation", ["mismatched_value", "missing_key"])
 def test_aggregate_incomplete_video_encoder_info_warns_and_nuls_encoders(
     tmp_path, lerobot_dataset_factory, caplog, mutation
@@ -357,7 +458,11 @@ def test_aggregate_incomplete_video_encoder_info_warns_and_nuls_encoders(
 
 
 def test_aggregate_with_low_threshold(tmp_path, lerobot_dataset_factory):
-    """Test aggregation with small file size limits to force file rotation/sharding."""
+    """Test aggregation with small file size limits to force file rotation/sharding.
+
+    Depth video features are included to verify that file rotation/concat
+    correctly handles depth-marked features alongside regular RGB ones.
+    """
     ds_0_num_episodes = ds_1_num_episodes = 10
     ds_0_num_frames = ds_1_num_frames = 400
 
@@ -366,14 +471,19 @@ def test_aggregate_with_low_threshold(tmp_path, lerobot_dataset_factory):
         repo_id=f"{DUMMY_REPO_ID}_small_0",
         total_episodes=ds_0_num_episodes,
         total_frames=ds_0_num_frames,
+        camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH,
     )
     ds_1 = lerobot_dataset_factory(
         root=tmp_path / "small_1",
         repo_id=f"{DUMMY_REPO_ID}_small_1",
         total_episodes=ds_1_num_episodes,
         total_frames=ds_1_num_frames,
+        camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH,
     )
 
+    assert len(ds_0.meta.depth_keys) > 0, "ds_0 should expose at least one depth key"
+    assert len(ds_1.meta.depth_keys) > 0, "ds_1 should expose at least one depth key"
+
     # Use the new configurable parameters to force file rotation
     aggregate_datasets(
         repo_ids=[ds_0.repo_id, ds_1.repo_id],
@@ -404,6 +514,7 @@ def test_aggregate_with_low_threshold(tmp_path, lerobot_dataset_factory):
     assert_episode_indices_updated_correctly(aggr_ds, ds_0, ds_1)
     assert_video_frames_integrity(aggr_ds, ds_0, ds_1)
     assert_video_timestamps_within_bounds(aggr_ds)
+    assert_depth_keys_preserved(aggr_ds, ds_0, ds_1)
     assert_dataset_iteration_works(aggr_ds)
 
     # Check that multiple files were actually created due to small size limits
@@ -423,7 +534,8 @@ def test_video_timestamps_regression(tmp_path, lerobot_dataset_factory):
     """Regression test for video timestamp bug when merging datasets.
 
     This test specifically checks that video timestamps are correctly calculated
-    and accumulated when merging multiple datasets.
+    and accumulated when merging multiple datasets. Depth video features are
+    included so depth timestamps are also covered by the regression.
     """
     datasets = []
     for i in range(3):
@@ -432,9 +544,13 @@ def test_video_timestamps_regression(tmp_path, lerobot_dataset_factory):
             repo_id=f"{DUMMY_REPO_ID}_regression_{i}",
             total_episodes=2,
             total_frames=100,
+            camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH,
         )
         datasets.append(ds)
 
+    for i, ds in enumerate(datasets):
+        assert len(ds.meta.depth_keys) > 0, f"Dataset {i} should expose at least one depth key"
+
     aggregate_datasets(
         repo_ids=[ds.repo_id for ds in datasets],
         roots=[ds.root for ds in datasets],
@@ -451,12 +567,21 @@ def test_video_timestamps_regression(tmp_path, lerobot_dataset_factory):
         aggr_ds = LeRobotDataset(f"{DUMMY_REPO_ID}_regression_aggr", root=tmp_path / "regression_aggr")
 
     assert_video_timestamps_within_bounds(aggr_ds)
+    # Depth keys must survive the merge for the regression to cover the
+    # ``is_depth_map`` decoding branch.
+    assert set(aggr_ds.meta.depth_keys) == set(datasets[0].meta.depth_keys)
 
+    depth_keys = set(aggr_ds.meta.depth_keys)
     for i in range(len(aggr_ds)):
         item = aggr_ds[i]
         for key in aggr_ds.meta.video_keys:
             assert key in item, f"Video key {key} missing from item {i}"
-            assert item[key].shape[0] == 3, f"Expected 3 channels for video key {key}"
+            # Depth frames are single-channel (1, H, W) after dequantization;
+            # standard RGB frames keep the 3-channel layout.
+            expected_channels = 1 if key in depth_keys else 3
+            assert item[key].shape[0] == expected_channels, (
+                f"Expected {expected_channels} channels for video key {key}, got {item[key].shape}"
+            )
 
 
 def assert_image_schema_preserved(aggr_ds):
@@ -520,6 +645,41 @@ def assert_image_frames_integrity(aggr_ds, ds_0, ds_1):
             )
 
 
+@pytest.mark.parametrize("use_videos", [True, False], ids=["video", "image"])
+def test_aggregate_one_row_group_per_episode(tmp_path, lerobot_dataset_factory, use_videos):
+    """Aggregated DATA shards keep one row group per episode (not one collapsed group).
+
+    Covers both the non-image (``df.to_parquet``) and image
+    (``to_parquet_with_hf_images``) write branches, including the merge-into-
+    existing-file branch via a low file-size threshold that forces packing.
+    """
+    ds_0 = lerobot_dataset_factory(
+        root=tmp_path / "rg_0",
+        repo_id=f"{DUMMY_REPO_ID}_rg_0",
+        total_episodes=3,
+        total_frames=60,
+        use_videos=use_videos,
+    )
+    ds_1 = lerobot_dataset_factory(
+        root=tmp_path / "rg_1",
+        repo_id=f"{DUMMY_REPO_ID}_rg_1",
+        total_episodes=4,
+        total_frames=80,
+        use_videos=use_videos,
+    )
+
+    aggr_root = tmp_path / "rg_aggr"
+    aggregate_datasets(
+        repo_ids=[ds_0.repo_id, ds_1.repo_id],
+        roots=[ds_0.root, ds_1.root],
+        aggr_repo_id=f"{DUMMY_REPO_ID}_rg_aggr",
+        aggr_root=aggr_root,
+    )
+
+    n_episodes = assert_data_shards_one_row_group_per_episode(aggr_root)
+    assert n_episodes == ds_0.num_episodes + ds_1.num_episodes
+
+
 def test_aggregate_image_datasets(tmp_path, lerobot_dataset_factory):
     """Test aggregation of image-based datasets preserves HuggingFace Image schema.
 
@@ -538,25 +698,31 @@ def test_aggregate_image_datasets(tmp_path, lerobot_dataset_factory):
     ds_0_num_episodes = 2
     ds_1_num_episodes = 3
 
-    # Create two image-based datasets (use_videos=False)
+    # Create two image-based datasets (use_videos=False) with a mix of RGB
+    # and depth-marked cameras so the depth path is exercised in image mode.
     ds_0 = lerobot_dataset_factory(
         root=tmp_path / "image_0",
         repo_id=f"{DUMMY_REPO_ID}_image_0",
         total_episodes=ds_0_num_episodes,
         total_frames=ds_0_num_frames,
-        use_videos=False,  # Image-based dataset
+        use_videos=False,
+        camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH,
     )
     ds_1 = lerobot_dataset_factory(
         root=tmp_path / "image_1",
         repo_id=f"{DUMMY_REPO_ID}_image_1",
         total_episodes=ds_1_num_episodes,
         total_frames=ds_1_num_frames,
-        use_videos=False,  # Image-based dataset
+        use_videos=False,
+        camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH,
     )
 
     # Verify source datasets have image keys
     assert len(ds_0.meta.image_keys) > 0, "ds_0 should have image keys"
     assert len(ds_1.meta.image_keys) > 0, "ds_1 should have image keys"
+    # And that the depth marker actually made it onto an image feature.
+    assert len(ds_0.meta.depth_keys) > 0, "ds_0 should expose at least one depth key"
+    assert len(ds_1.meta.depth_keys) > 0, "ds_1 should expose at least one depth key"
 
     # Aggregate the datasets
     aggregate_datasets(
@@ -591,6 +757,7 @@ def test_aggregate_image_datasets(tmp_path, lerobot_dataset_factory):
     # Image-specific assertions
     assert_image_schema_preserved(aggr_ds)
     assert_image_frames_integrity(aggr_ds, ds_0, ds_1)
+    assert_depth_keys_preserved(aggr_ds, ds_0, ds_1)
 
     # Verify images can be accessed and have correct shape
     sample_item = aggr_ds[0]

tests/datasets/test_compute_stats.py

@@ -35,7 +35,11 @@ from lerobot.utils.constants import OBS_IMAGE, OBS_STATE
 
 
 def mock_load_image_as_numpy(path, dtype, channel_first):
-    return np.ones((3, 32, 32), dtype=dtype) if channel_first else np.ones((32, 32, 3), dtype=dtype)
+    is_depth = "depth" in str(path)
+    channels = 1 if is_depth else 3
+    out_dtype = np.uint16 if is_depth else dtype
+    arr = np.arange(channels * 32 * 32, dtype=out_dtype).reshape(channels, 32, 32)
+    return arr if channel_first else arr.transpose(1, 2, 0)
 
 
 @pytest.fixture
@@ -83,6 +87,29 @@ def test_get_feature_stats_images():
     assert stats["min"].shape == stats["max"].shape == stats["mean"].shape == stats["std"].shape
 
 
+def test_get_feature_stats_uint8_images_preserves_std():
+    data = np.array(
+        [
+            [
+                [[0, 64], [128, 255]],
+                [[255, 128], [64, 0]],
+                [[32, 96], [160, 224]],
+            ],
+            [
+                [[16, 80], [144, 240]],
+                [[240, 144], [80, 16]],
+                [[48, 112], [176, 208]],
+            ],
+        ],
+        dtype=np.uint8,
+    )
+
+    stats = get_feature_stats(data, axis=(0, 2, 3), keepdims=True)
+
+    expected_std = data.transpose(0, 2, 3, 1).reshape(-1, 3).std(axis=0).reshape(1, 3, 1, 1)
+    np.testing.assert_allclose(stats["std"], expected_std)
+
+
 def test_get_feature_stats_axis_0_keepdims(sample_array):
     expected = {
         "min": np.array([[1, 2, 3]]),
@@ -145,22 +172,33 @@ def test_get_feature_stats_single_value():
 
 
 def test_compute_episode_stats():
+    depth_key = "observation.images.depth"
     episode_data = {
         OBS_IMAGE: [f"image_{i}.jpg" for i in range(100)],
+        depth_key: [f"depth_{i}.tiff" for i in range(100)],
         OBS_STATE: np.random.rand(100, 10),
     }
     features = {
         OBS_IMAGE: {"dtype": "image"},
+        depth_key: {"dtype": "image", "info": {"is_depth_map": True}},
         OBS_STATE: {"dtype": "numeric"},
     }
 
     with patch("lerobot.datasets.compute_stats.load_image_as_numpy", side_effect=mock_load_image_as_numpy):
         stats = compute_episode_stats(episode_data, features)
 
-    assert OBS_IMAGE in stats and OBS_STATE in stats
+    assert OBS_IMAGE in stats and depth_key in stats and OBS_STATE in stats
     assert stats[OBS_IMAGE]["count"].item() == 100
+    assert stats[depth_key]["count"].item() == 100
     assert stats[OBS_STATE]["count"].item() == 100
     assert stats[OBS_IMAGE]["mean"].shape == (3, 1, 1)
+    assert stats[depth_key]["mean"].shape == (1, 1, 1)
+    # Depth keeps raw values: max far exceeds 255, proving no /255 and no uint8 downcast.
+    assert stats[depth_key]["min"].item() == 0.0
+    assert stats[depth_key]["max"].item() == 1023.0
+    # RGB is normalized to [0, 1].
+    np.testing.assert_allclose(stats[OBS_IMAGE]["min"], 0.0)
+    np.testing.assert_allclose(stats[OBS_IMAGE]["max"], 1.0)
 
 
 def test_assert_type_and_shape_valid():
@@ -595,25 +633,31 @@ def test_compute_episode_stats_with_custom_quantiles():
 def test_compute_episode_stats_with_image_data():
     """Test quantile computation with image features."""
     image_paths = [f"image_{i}.jpg" for i in range(50)]
+    depth_paths = [f"depth_{i}.tiff" for i in range(50)]
     episode_data = {
         "observation.image": image_paths,
+        "observation.images.depth": depth_paths,
         "action": np.random.normal(0, 1, (50, 5)),
     }
     features = {
         "observation.image": {"dtype": "image"},
+        "observation.images.depth": {"dtype": "image", "info": {"is_depth_map": True}},
         "action": {"dtype": "float32", "shape": (5,)},
     }
 
     with patch("lerobot.datasets.compute_stats.load_image_as_numpy", side_effect=mock_load_image_as_numpy):
         stats = compute_episode_stats(episode_data, features)
 
-    # Image quantiles should be normalized and have correct shape
-    assert "q01" in stats["observation.image"]
-    assert "q50" in stats["observation.image"]
-    assert "q99" in stats["observation.image"]
-    assert stats["observation.image"]["q01"].shape == (3, 1, 1)
-    assert stats["observation.image"]["q50"].shape == (3, 1, 1)
-    assert stats["observation.image"]["q99"].shape == (3, 1, 1)
+    # RGB image quantiles should be normalized and per-channel.
+    for q in ("q01", "q50", "q99"):
+        assert stats["observation.image"][q].shape == (3, 1, 1)
+
+    # Depth quantiles are single-channel and kept in raw (un-normalized) units.
+    for q in ("q01", "q50", "q99"):
+        assert stats["observation.images.depth"][q].shape == (1, 1, 1)
+    # Depth max stays in raw units (not /255, not uint8-capped); RGB is normalized.
+    assert stats["observation.images.depth"]["max"].item() == 1023.0
+    np.testing.assert_allclose(stats["observation.image"]["max"], 1.0)
 
     # Action quantiles should have correct shape
     assert stats["action"]["q01"].shape == (5,)

tests/datasets/test_dataset_metadata.py

@@ -59,11 +59,13 @@ def _make_dummy_stats(features: dict) -> dict:
     stats = {}
     for key, ft in features.items():
         if ft["dtype"] in ("image", "video"):
+            channels = ft["shape"][-1]
+            stat_shape = (channels, 1, 1)
             stats[key] = {
-                "max": np.ones((3, 1, 1), dtype=np.float32),
-                "mean": np.full((3, 1, 1), 0.5, dtype=np.float32),
-                "min": np.zeros((3, 1, 1), dtype=np.float32),
-                "std": np.full((3, 1, 1), 0.25, dtype=np.float32),
+                "max": np.ones(stat_shape, dtype=np.float32),
+                "mean": np.full(stat_shape, 0.5, dtype=np.float32),
+                "min": np.zeros(stat_shape, dtype=np.float32),
+                "std": np.full(stat_shape, 0.25, dtype=np.float32),
                 "count": np.array([5]),
             }
         elif ft["dtype"] in ("float32", "float64", "int64"):
@@ -142,6 +144,45 @@ def test_create_without_videos_has_no_video_path(tmp_path):
     assert meta.video_keys == []
 
 
+@pytest.mark.parametrize(
+    ("marker_field", "marker_key"),
+    [
+        ("info", "is_depth_map"),
+        ("info", "video.is_depth_map"),
+        ("video_info", "video.is_depth_map"),
+    ],
+    ids=["info.is_depth_map", "info.video.is_depth_map_legacy", "video_info.video.is_depth_map_legacy"],
+)
+def test_depth_keys_property_filters_by_marker(tmp_path, marker_field, marker_key):
+    """``depth_keys`` recognises the canonical and the two legacy marker variants."""
+    depth_feature = {
+        "dtype": "video",
+        "shape": (64, 96, 1),
+        "names": ["height", "width", "channels"],
+        marker_field: {marker_key: True},
+    }
+    features = {
+        **VIDEO_FEATURES,
+        "observation.images.laptop_depth": depth_feature,
+    }
+    meta = LeRobotDatasetMetadata.create(
+        repo_id="test/depth_keys",
+        fps=DEFAULT_FPS,
+        features=features,
+        root=tmp_path / f"depth_keys_{marker_field}_{marker_key.replace('.', '_')}",
+    )
+
+    assert set(meta.video_keys) == {"observation.images.laptop", "observation.images.laptop_depth"}
+    assert meta.depth_keys == ["observation.images.laptop_depth"]
+
+
+def test_depth_keys_empty_when_no_marker(tmp_path):
+    meta = LeRobotDatasetMetadata.create(
+        repo_id="test/no_depth", fps=DEFAULT_FPS, features=VIDEO_FEATURES, root=tmp_path / "no_depth"
+    )
+    assert meta.depth_keys == []
+
+
 def test_create_raises_on_existing_directory(tmp_path):
     """create() raises if root directory already exists."""
     root = tmp_path / "existing"

tests/datasets/test_dataset_tools.py

@@ -24,7 +24,7 @@ import torch
 pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
 
 
-from lerobot.configs import VideoEncoderConfig
+from lerobot.configs import DepthEncoderConfig, RGBEncoderConfig
 from lerobot.datasets.dataset_tools import (
     add_features,
     convert_image_to_video_dataset,
@@ -37,7 +37,9 @@ from lerobot.datasets.dataset_tools import (
     split_dataset,
 )
 from lerobot.datasets.io_utils import load_info
-from tests.datasets.test_video_encoding import _add_frames, require_h264, require_libsvtav1
+from tests.datasets.test_video_encoding import require_h264, require_hevc, require_libsvtav1
+from tests.fixtures.constants import DUMMY_DEPTH_FEATURES, DUMMY_DEPTH_KEY
+from tests.fixtures.dataset_factories import add_frames
 
 
 @pytest.fixture
@@ -1251,7 +1253,7 @@ def test_convert_image_to_video_dataset(tmp_path):
             dataset=source_dataset,
             output_dir=output_dir,
             repo_id="lerobot/pusht_video",
-            camera_encoder=VideoEncoderConfig(
+            rgb_encoder=RGBEncoderConfig(
                 vcodec="libsvtav1",
                 pix_fmt="yuv420p",
                 g=2,
@@ -1332,9 +1334,131 @@ def test_convert_image_to_video_dataset_subset_episodes(tmp_path):
             shutil.rmtree(output_dir)
 
 
+@require_libsvtav1
+@require_hevc
+def test_convert_image_to_video_dataset_depth(tmp_path, empty_lerobot_dataset_factory):
+    """Depth image features convert to depth videos using the depth encoder.
+
+    Mirrors :func:`test_convert_image_to_video_dataset` but with a small local
+    image dataset that mixes an RGB camera with a depth camera, so the
+    ``depth_keys`` → ``depth_encoder`` routing and ``is_depth_map`` preservation
+    are exercised end-to-end.
+    """
+    features = {
+        "action": {"dtype": "float32", "shape": (2,), "names": ["a", "b"]},
+        "observation.images.cam": {
+            "dtype": "image",
+            "shape": (64, 96, 3),
+            "names": ["height", "width", "channels"],
+        },
+        "observation.images.depth": {
+            "dtype": "image",
+            "shape": (64, 96, 1),
+            "names": ["height", "width", "channels"],
+            "info": {"is_depth_map": True},
+        },
+    }
+    source_dataset = empty_lerobot_dataset_factory(
+        root=tmp_path / "img_ds",
+        features=features,
+        use_videos=False,
+    )
+
+    add_frames(source_dataset, num_frames=4)
+    source_dataset.save_episode()
+    source_dataset.finalize()
+
+    # Source is an image dataset with the depth marker on the depth camera.
+    assert len(source_dataset.meta.video_keys) == 0
+    assert "observation.images.depth" in source_dataset.meta.depth_keys
+
+    output_dir = tmp_path / "video_ds"
+    with (
+        patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_get_safe_version,
+        patch("lerobot.datasets.dataset_metadata.snapshot_download") as mock_snapshot_download,
+    ):
+        mock_get_safe_version.return_value = "v3.0"
+        mock_snapshot_download.return_value = str(output_dir)
+
+        # Use non-default quantization params so the persisted metadata must
+        # come from the depth encoder (not RGB encoder defaults).
+        depth_encoder = DepthEncoderConfig(
+            vcodec="hevc",
+            pix_fmt="gray12le",
+            g=2,
+            crf=30,
+            depth_min=0.05,
+            depth_max=8.0,
+            shift=2.0,
+            use_log=False,
+        )
+        video_dataset = convert_image_to_video_dataset(
+            dataset=source_dataset,
+            output_dir=output_dir,
+            repo_id="dummy/depth_video",
+            rgb_encoder=RGBEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30),
+            depth_encoder=depth_encoder,
+            num_workers=1,
+        )
+
+    # Both cameras are now videos, and the depth marker survived the conversion.
+    assert "observation.images.cam" in video_dataset.meta.video_keys
+    assert "observation.images.depth" in video_dataset.meta.video_keys
+    assert "observation.images.depth" in video_dataset.meta.depth_keys
+    assert "observation.images.cam" not in video_dataset.meta.depth_keys
+
+    depth_path = video_dataset.root / video_dataset.meta.get_video_file_path(0, "observation.images.depth")
+    assert depth_path.exists(), f"Depth video file should exist: {depth_path}"
+
+    # The persisted depth-video metadata must carry the depth quantization params
+    # from the depth encoder (so frames dequantize correctly on read), and the RGB
+    # camera must not be marked as a depth map.
+    persisted_info = load_info(video_dataset.root)
+    depth_info = persisted_info.features["observation.images.depth"]["info"]
+    assert depth_info["is_depth_map"] is True
+    assert DepthEncoderConfig.from_video_info(depth_info) == depth_encoder
+
+    cam_info = persisted_info.features["observation.images.cam"]["info"]
+    assert cam_info.get("is_depth_map") is False
+    assert "video.codec" in cam_info
+
+
 # ─── reencode_dataset ─────────────────────────────────────────────────
 
 
+@require_hevc
+def test_reencode_dataset_depth_uses_depth_encoder(tmp_path, empty_lerobot_dataset_factory):
+    """Depth videos are re-encoded with the depth encoder and keep their depth metadata.
+
+    Depth-focused companion to :func:`test_reencode_dataset_multi_key_multiprocessing`.
+    """
+    initial_cfg = DepthEncoderConfig(vcodec="hevc", pix_fmt="gray12le", g=2, crf=30)
+    dataset = empty_lerobot_dataset_factory(
+        root=tmp_path / "ds",
+        features=DUMMY_DEPTH_FEATURES,
+        use_videos=True,
+        depth_encoder=initial_cfg,
+    )
+
+    add_frames(dataset, num_frames=4)
+    dataset.save_episode()
+    dataset.finalize()
+
+    assert DUMMY_DEPTH_KEY in dataset.meta.depth_keys
+
+    target_cfg = DepthEncoderConfig(vcodec="hevc", pix_fmt="gray12le", g=6, crf=23)
+    result = reencode_dataset(dataset, depth_encoder=target_cfg, num_workers=0)
+
+    assert result is dataset
+
+    persisted_info = load_info(dataset.root)
+    depth_info = persisted_info.features[DUMMY_DEPTH_KEY].get("info", {})
+    # Re-encode applied the new codec parameters to the depth video ...
+    assert DepthEncoderConfig.from_video_info(depth_info) == target_cfg
+    # ... while preserving the depth marker.
+    assert depth_info["is_depth_map"] is True
+
+
 @require_libsvtav1
 @require_h264
 def test_reencode_dataset_multi_key_multiprocessing(
@@ -1342,29 +1466,29 @@ def test_reencode_dataset_multi_key_multiprocessing(
 ):
     """Re-encode a two-camera dataset with num_workers=2 and verify metadata refresh."""
     features = features_factory(use_videos=True)
-    initial_cfg = VideoEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12)
+    initial_cfg = RGBEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12)
     dataset = empty_lerobot_dataset_factory(
         root=tmp_path / "ds",
         features=features,
         use_videos=True,
-        camera_encoder=initial_cfg,
+        rgb_encoder=initial_cfg,
     )
 
-    _add_frames(dataset, num_frames=4)
+    add_frames(dataset, num_frames=4)
     dataset.save_episode()
-    _add_frames(dataset, num_frames=4)
+    add_frames(dataset, num_frames=4)
     dataset.save_episode()
     dataset.finalize()
 
     assert len(dataset.meta.video_keys) == 2
 
-    target_cfg = VideoEncoderConfig(vcodec="h264", g=6, crf=23, pix_fmt="yuv420p")
+    target_cfg = RGBEncoderConfig(vcodec="h264", g=6, crf=23, pix_fmt="yuv420p")
 
-    result = reencode_dataset(dataset, camera_encoder=target_cfg, num_workers=2)
+    result = reencode_dataset(dataset, rgb_encoder=target_cfg, num_workers=2)
 
     assert result is dataset
 
     persisted_info = load_info(dataset.root)
     for vk in dataset.meta.video_keys:
-        persisted_encoder = VideoEncoderConfig.from_video_info(persisted_info.features[vk].get("info", {}))
+        persisted_encoder = RGBEncoderConfig.from_video_info(persisted_info.features[vk].get("info", {}))
         assert persisted_encoder == target_cfg

tests/datasets/test_dataset_writer.py

@@ -53,8 +53,8 @@ def _make_frame(features: dict, task: str = "Dummy task") -> dict:
 # ── Existing encode_video_worker tests ───────────────────────────────
 
 
-def test_encode_video_worker_forwards_camera_encoder(tmp_path):
-    """_encode_video_worker forwards camera_encoder to encode_video_frames."""
+def test_encode_video_worker_forwards_video_encoder(tmp_path):
+    """_encode_video_worker forwards video_encoder to encode_video_frames."""
     video_key = "observation.images.laptop"
     fpath = DEFAULT_IMAGE_PATH.format(image_key=video_key, episode_index=0, frame_index=0)
     img_dir = tmp_path / Path(fpath).parent
@@ -74,16 +74,16 @@ def test_encode_video_worker_forwards_camera_encoder(tmp_path):
             0,
             tmp_path,
             fps=30,
-            camera_encoder=VideoEncoderConfig(vcodec="h264", preset=None),
+            video_encoder=VideoEncoderConfig(vcodec="h264", preset=None),
             encoder_threads=4,
         )
 
-    assert captured_kwargs["camera_encoder"].vcodec == "h264"
+    assert captured_kwargs["video_encoder"].vcodec == "h264"
     assert captured_kwargs["encoder_threads"] == 4
 
 
-def test_encode_video_worker_default_camera_encoder(tmp_path):
-    """_encode_video_worker passes None camera_encoder which encode_video_frames defaults."""
+def test_encode_video_worker_default_video_encoder(tmp_path):
+    """_encode_video_worker passes None video_encoder which encode_video_frames defaults."""
     video_key = "observation.images.laptop"
     fpath = DEFAULT_IMAGE_PATH.format(image_key=video_key, episode_index=0, frame_index=0)
     img_dir = tmp_path / Path(fpath).parent
@@ -100,7 +100,7 @@ def test_encode_video_worker_default_camera_encoder(tmp_path):
     with patch("lerobot.datasets.dataset_writer.encode_video_frames", side_effect=mock_encode):
         _encode_video_worker(video_key, 0, tmp_path, fps=30)
 
-    assert captured_kwargs["camera_encoder"] is None
+    assert captured_kwargs["video_encoder"] is None
     assert captured_kwargs["encoder_threads"] is None
 
 

tests/datasets/test_datasets.py

@@ -51,7 +51,7 @@ from lerobot.robots import make_robot_from_config
 from lerobot.transforms import ImageTransforms, ImageTransformsConfig
 from lerobot.utils.constants import ACTION, DONE, OBS_IMAGES, OBS_STATE, OBS_STR, REWARD
 from lerobot.utils.feature_utils import hw_to_dataset_features
-from tests.fixtures.constants import DUMMY_CHW, DUMMY_HWC, DUMMY_REPO_ID
+from tests.fixtures.constants import DUMMY_CHW, DUMMY_HWC, DUMMY_MOTOR_FEATURES, DUMMY_REPO_ID
 from tests.mocks.mock_robot import MockRobotConfig
 from tests.utils import require_x86_64_kernel
 
@@ -133,6 +133,21 @@ def test_dataset_feature_with_forward_slash_raises_error():
         )
 
 
+def test_create_does_not_mutate_input_features(tmp_path, empty_lerobot_dataset_factory):
+    # ``create`` must deep-copy features so a dataset built from another's features stays independent.
+    dataset = empty_lerobot_dataset_factory(
+        root=tmp_path / "ds1", features=DUMMY_MOTOR_FEATURES, use_videos=False
+    )
+    dataset_copy = empty_lerobot_dataset_factory(
+        root=tmp_path / "ds2", features=dataset.meta.features, use_videos=False
+    )
+
+    original_shape = dataset.meta.info.features["state"]["shape"]
+    dataset_copy.meta.info.features["state"]["shape"] = (999,)
+
+    assert dataset.meta.info.features["state"]["shape"] == original_shape
+
+
 def test_add_frame_missing_task(tmp_path, empty_lerobot_dataset_factory):
     features = {"state": {"dtype": "float32", "shape": (1,), "names": None}}
     dataset = empty_lerobot_dataset_factory(root=tmp_path / "test", features=features)
@@ -1519,6 +1534,10 @@ def test_valid_video_codecs_constant():
     assert "auto" in VALID_VIDEO_CODECS
     assert "h264_videotoolbox" in VALID_VIDEO_CODECS
     assert "h264_nvenc" in VALID_VIDEO_CODECS
+    assert "h264_vaapi" in VALID_VIDEO_CODECS
+    assert "h264_qsv" in VALID_VIDEO_CODECS
+    assert "hevc_videotoolbox" in VALID_VIDEO_CODECS
+    assert "hevc_nvenc" in VALID_VIDEO_CODECS
     assert len(VALID_VIDEO_CODECS) == 10
 
 

tests/datasets/test_depth.py

@@ -0,0 +1,247 @@
+"""Tests for the depth-integration feature.
+
+Covers:
+- ``depth_utils`` quantize/dequantize round-trips and backend agreement.
+- Image-writer support for single-channel depth.
+- Hardware-feature → depth flag routing.
+- Feature-to-file-format routing through the dataset writer.
+
+Depth metadata detection on ``LeRobotDatasetMetadata.depth_keys`` lives in
+``test_dataset_metadata.py``. Depth video encoding/decoding lives in
+``test_video_encoding.py``.
+"""
+
+from pathlib import Path
+
+import pytest
+
+pytest.importorskip("av", reason="av is required (install lerobot[dataset])")
+
+import av
+import numpy as np
+import PIL.Image
+import torch
+
+from lerobot.configs import DepthEncoderConfig
+from lerobot.configs.video import (
+    DEFAULT_DEPTH_MAX,
+    DEFAULT_DEPTH_MIN,
+    DEPTH_METER_UNIT,
+    DEPTH_MILLIMETER_UNIT,
+    DEPTH_QMAX,
+)
+from lerobot.datasets.depth_utils import dequantize_depth, quantize_depth
+from lerobot.datasets.image_writer import image_array_to_pil_image, write_image
+from tests.fixtures.constants import (
+    DEFAULT_FPS,
+    DUMMY_CAMERA_FEATURES,
+    DUMMY_CAMERA_FEATURES_WITH_DEPTH,
+    DUMMY_CHW,
+    DUMMY_DEPTH_CAMERA_FEATURES,
+    DUMMY_REPO_ID,
+)
+from tests.fixtures.dataset_factories import add_frames
+
+_, H, W = DUMMY_CHW
+
+
+def _depth_metres_ramp() -> np.ndarray:
+    """Linearly-spaced float32 depth in metres covering the default range."""
+    return np.linspace(DEFAULT_DEPTH_MIN, DEFAULT_DEPTH_MAX, H * W, dtype=np.float32).reshape(H, W)
+
+
+# ── 1. Quantize / dequantize round-trips ──────────────────────────────
+
+
+class TestQuantizeDequantize:
+    """Numerical contract of ``quantize_depth`` / ``dequantize_depth``."""
+
+    @pytest.mark.parametrize("use_log", [False, True])
+    @pytest.mark.parametrize("output_unit", [DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT])
+    @pytest.mark.parametrize("output_channel_last", [False, True])
+    def test_roundtrip(self, use_log, output_unit, output_channel_last):
+        """quantize → dequantize recovers depth; layout and unit are honored."""
+        depth = _depth_metres_ramp()
+        quantized = quantize_depth(depth, use_log=use_log, video_backend=None)
+        recovered = dequantize_depth(
+            quantized,
+            use_log=use_log,
+            output_unit=output_unit,
+            output_tensor=False,
+            output_channel_last=output_channel_last,
+        )
+
+        expected_shape = (H, W, 1) if output_channel_last else (1, H, W)
+        assert recovered.shape == expected_shape
+
+        recovered_m = recovered.astype(np.float32)
+        if output_unit == DEPTH_MILLIMETER_UNIT:
+            recovered_m = recovered_m / 1000.0
+        recovered_2d = recovered_m[..., 0] if output_channel_last else recovered_m[0]
+
+        if use_log:
+            # Log mode: tighter near-range error than far-range (the whole point).
+            near = depth < 1.0
+            far = depth > 8.0
+            err_near = np.abs(recovered_2d[near] - depth[near])
+            err_far = np.abs(recovered_2d[far] - depth[far])
+            assert err_near.mean() < err_far.mean()
+        else:
+            # Linear mode: bounded by quant step + 1 mm of unit-conversion rounding.
+            tol = (DEFAULT_DEPTH_MAX - DEFAULT_DEPTH_MIN) / DEPTH_QMAX + 1e-3
+            np.testing.assert_allclose(recovered_2d, depth, atol=tol)
+
+    @pytest.mark.parametrize("use_log", [False, True])
+    @pytest.mark.parametrize("output_unit", [DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT])
+    def test_numpy_torch_agree(self, use_log, output_unit):
+        """Batched torch path produces the same values as the numpy path."""
+        batch_size = 3
+        per_frame = np.linspace(0, DEPTH_QMAX, H * W, dtype=np.uint16).reshape(H, W)
+        batch_np = np.broadcast_to(per_frame[None, None, ...], (batch_size, 1, H, W)).copy()
+        batch_t = torch.from_numpy(batch_np.astype(np.int32))  # torch.uint16 support is patchy.
+
+        ref = dequantize_depth(batch_np, use_log=use_log, output_unit=output_unit, output_tensor=False)
+        out = dequantize_depth(batch_t, use_log=use_log, output_unit=output_unit, output_tensor=True)
+
+        assert isinstance(out, torch.Tensor)
+        assert out.shape == (batch_size, 1, H, W)
+        # ``m``: float32 noise (~10 µm in log mode, after ``exp``) — still 200× below the ~2 mm quant step.
+        # ``mm`` + tensor stays in float32 (no uint16 round-trip), so allow 1 mm slop.
+        atol = 1e-5 if output_unit == DEPTH_METER_UNIT else 1.0
+        np.testing.assert_allclose(out.cpu().numpy().astype(np.float64), ref.astype(np.float64), atol=atol)
+
+    @pytest.mark.parametrize(
+        "input_shape,output_shape",
+        [
+            ((H, W), (1, H, W)),
+            ((1, H, W), (1, H, W)),
+            ((H, W, 1), (1, H, W)),
+            ((3, 1, H, W), (3, 1, H, W)),
+            ((3, H, W, 1), (3, 1, H, W)),
+        ],
+    )
+    def test_input_layouts_accepted(self, input_shape, output_shape):
+        """All documented input layouts decode to the channel-first default."""
+        quantized = np.full(input_shape, DEPTH_QMAX // 2, dtype=np.uint16)
+        out = dequantize_depth(quantized, output_unit=DEPTH_METER_UNIT, output_tensor=False)
+        assert out.shape == output_shape
+
+    def test_pyav_frame_roundtrip(self):
+        """quantize → av.VideoFrame → dequantize works."""
+        depth = _depth_metres_ramp()
+        frame = quantize_depth(depth, use_log=False, video_backend="pyav")
+        assert isinstance(frame, av.VideoFrame)
+
+        recovered = dequantize_depth(frame, use_log=False, output_unit=DEPTH_METER_UNIT, output_tensor=False)
+        assert recovered.shape == (1, H, W)
+        tol = (DEFAULT_DEPTH_MAX - DEFAULT_DEPTH_MIN) / DEPTH_QMAX + 1e-3
+        np.testing.assert_allclose(recovered[0], depth, atol=tol)
+
+    def test_invalid_log_params_raises(self):
+        with pytest.raises(ValueError, match=r"depth_min \+ shift must be positive"):
+            quantize_depth(_depth_metres_ramp(), depth_min=1.0, shift=-2.0, use_log=True, video_backend=None)
+
+
+# ── 2. Image writer depth support ─────────────────────────────────────
+
+
+class TestImageWriterDepth:
+    """``image_array_to_pil_image`` and ``write_image`` for depth maps."""
+
+    @pytest.mark.parametrize("dtype,expected_mode", [(np.uint16, "I;16"), (np.float32, "F")])
+    @pytest.mark.parametrize("shape", [(H, W), (H, W, 1), (1, H, W)])
+    def test_pil_depth_modes_and_squeeze(self, dtype, expected_mode, shape):
+        """Single-channel depth converts to PIL with the right mode and (W, H) size."""
+        arr = np.zeros(shape, dtype=dtype)
+        img = image_array_to_pil_image(arr)
+        assert img.mode == expected_mode
+        assert img.size == (W, H)
+
+    def test_write_image_tiff_roundtrip(self, tmp_path):
+        """uint16 depth round-trips through .tiff."""
+        arr = np.arange(H * W, dtype=np.uint16).reshape(H, W)
+        fpath = tmp_path / "depth.tiff"
+        write_image(arr, fpath)
+        with PIL.Image.open(fpath) as loaded:
+            recovered = np.array(loaded)
+        np.testing.assert_array_equal(recovered, arr)
+
+
+# ── 3. Hardware-feature → depth flag ──────────────────────────────────
+
+
+class TestHwToDatasetFeaturesDepth:
+    """``hw_to_dataset_features`` flags single-channel cameras as depth."""
+
+    @pytest.mark.parametrize("channels,is_depth", [(1, True), (3, False)])
+    def test_depth_marker_by_channels(self, channels, is_depth):
+        from lerobot.utils.feature_utils import hw_to_dataset_features
+
+        features = hw_to_dataset_features({"cam": (480, 640, channels)}, prefix="observation")
+        assert features["observation.images.cam"]["info"]["is_depth_map"] is is_depth
+
+    def test_invalid_channel_count_raises(self):
+        from lerobot.utils.feature_utils import hw_to_dataset_features
+
+        with pytest.raises(ValueError, match="Expected a 3-tuple"):
+            hw_to_dataset_features({"cam": (480, 640, 2)}, prefix="observation")
+
+
+# ── 4. Feature-to-file-format routing ────────────────────────────────
+
+
+# Keys derived from DUMMY_CAMERA_FEATURES_WITH_DEPTH; pick one RGB and the depth camera.
+RGB_KEY = next(iter(DUMMY_CAMERA_FEATURES))
+DEPTH_KEY = next(iter(DUMMY_DEPTH_CAMERA_FEATURES))
+
+
+class TestFeatureFileRouting:
+    """Depth vs RGB features route to the correct file format."""
+
+    NUM_FRAMES = 5
+
+    def test_image_mode_depth_tiff_rgb_png(self, tmp_path, features_factory):
+        """Without video encoding: depth → .tiff, RGB → .png."""
+        from lerobot.datasets.lerobot_dataset import LeRobotDataset
+
+        features = features_factory(camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH, use_videos=False)
+        dataset = LeRobotDataset.create(
+            repo_id=DUMMY_REPO_ID,
+            fps=DEFAULT_FPS,
+            features=features,
+            root=tmp_path / "ds",
+            use_videos=False,
+        )
+
+        add_frames(dataset, num_frames=self.NUM_FRAMES)
+
+        buf = dataset.writer.episode_buffer
+        assert all(Path(p).suffix == ".tiff" for p in buf[DEPTH_KEY])
+        assert all(Path(p).suffix == ".png" for p in buf[RGB_KEY])
+
+        dataset.save_episode()
+        dataset.finalize()
+
+    def test_video_mode_depth_uses_depth_encoder(self, tmp_path, features_factory):
+        """With streaming video encoding: depth → DepthEncoderConfig, RGB does not."""
+        from lerobot.datasets.lerobot_dataset import LeRobotDataset
+
+        features = features_factory(camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH, use_videos=True)
+        dataset = LeRobotDataset.create(
+            repo_id=DUMMY_REPO_ID,
+            fps=DEFAULT_FPS,
+            features=features,
+            root=tmp_path / "ds",
+            use_videos=True,
+            streaming_encoding=True,
+        )
+
+        add_frames(dataset, num_frames=self.NUM_FRAMES)
+
+        encoder = dataset.writer._streaming_encoder
+        assert encoder is not None
+        assert isinstance(encoder._threads[DEPTH_KEY].video_encoder, DepthEncoderConfig)
+        assert not isinstance(encoder._threads[RGB_KEY].video_encoder, DepthEncoderConfig)
+
+        dataset.save_episode()
+        dataset.finalize()

tests/datasets/test_image_writer.py

@@ -94,7 +94,7 @@ def test_image_array_to_pil_image_pytorch_format(img_array_factory):
 
 def test_image_array_to_pil_image_single_channel(img_array_factory):
     img_array = img_array_factory(channels=1)
-    with pytest.raises(NotImplementedError):
+    with pytest.raises(ValueError, match="Unsupported single-channel image dtype"):
         image_array_to_pil_image(img_array)
 
 
@@ -344,7 +344,7 @@ def test_with_different_image_formats(tmp_path, img_array_factory):
     writer = AsyncImageWriter()
     try:
         image_array = img_array_factory()
-        formats = ["png", "jpeg", "bmp"]
+        formats = ["png", "tiff", "tif"]
         for fmt in formats:
             fpath = tmp_path / f"test_image.{fmt}"
             write_image(image_array, fpath)

tests/datasets/test_sampler.py

@@ -114,6 +114,19 @@ def test_shuffle():
     assert set(sampler) == {0, 1, 2, 3, 4, 5}
 
 
+def test_shuffle_is_reproducible_across_instances():
+    # The order is a pure function of (seed, epoch), so two fresh samplers (e.g. two ranks)
+    # produce the same permutation without any generator synchronization.
+    sampler_a = EpisodeAwareSampler([0], [6], shuffle=True, seed=42)
+    sampler_b = EpisodeAwareSampler([0], [6], shuffle=True, seed=42)
+    epoch_0 = list(sampler_a)
+    assert list(sampler_b) == epoch_0
+    # Desyncing the global RNG must not affect the permutation.
+    sampler_c = EpisodeAwareSampler([0], [6], shuffle=True, seed=42)
+    torch.randperm(1000)  # consume global RNG, as rank-asymmetric code (e.g. eval) would
+    assert list(sampler_c) == epoch_0
+
+
 def test_negative_drop_first_frames_raises():
     with pytest.raises(ValueError, match="drop_n_first_frames must be >= 0"):
         EpisodeAwareSampler([0], [10], drop_n_first_frames=-1)
@@ -137,3 +150,87 @@ def test_partial_episode_drop_warns(caplog):
     # Episode 0 is skipped (1 frame, drop 1), Episode 1 keeps frames 2-5
     assert sampler.indices == [2, 3, 4, 5]
     assert "Episode 0" in caplog.text
+
+
+# --- seeded (seed, epoch) shuffling, resume, and state ---
+
+from lerobot.datasets.sampler import compute_sampler_state  # noqa: E402
+
+EPISODE_BOUNDS = ([0, 2, 3], [2, 3, 6])  # episodes of 2, 1 and 3 frames
+
+
+@pytest.mark.parametrize("num_frames", [1, 2, 3, 37, 64, 100])
+def test_deterministic_sampler_shuffle_is_permutation(num_frames):
+    for seed in (0, 1, 1234):
+        sampler = EpisodeAwareSampler([0], [num_frames], shuffle=True, seed=seed)
+        assert sorted(sampler) == list(range(num_frames))
+
+
+def test_deterministic_sampler_epochs_reproduce_and_differ():
+    sampler_a = EpisodeAwareSampler([0], [100], shuffle=True, seed=42)
+    sampler_b = EpisodeAwareSampler([0], [100], shuffle=True, seed=42)
+    epoch_0 = list(sampler_a)
+    assert list(sampler_b) == epoch_0  # same (seed, epoch) -> same order on any process
+    epoch_1 = list(sampler_a)  # __iter__ auto-advances the epoch
+    assert epoch_1 != epoch_0
+    assert sorted(epoch_1) == sorted(epoch_0)
+    sampler_a.set_epoch(0)
+    assert list(sampler_a) == epoch_0
+    assert list(EpisodeAwareSampler([0], [100], shuffle=True, seed=7)) != epoch_0
+
+
+def test_deterministic_sampler_resume_mid_epoch():
+    reference = EpisodeAwareSampler(*EPISODE_BOUNDS, shuffle=True, seed=42)
+    epoch_0 = list(reference)
+    epoch_1 = list(reference)
+    for start in (0, 1, 4, len(epoch_0)):
+        resumed = EpisodeAwareSampler(*EPISODE_BOUNDS, shuffle=True, seed=42)
+        resumed.load_state_dict({"epoch": 0, "start_index": start})
+        assert list(resumed) == epoch_0[start:]
+        # the resumed sampler continues into the same epoch 1 as the uninterrupted one
+        assert list(resumed) == epoch_1
+
+
+def test_deterministic_sampler_construction_stores_only_boundaries():
+    # Construction is O(num_episodes), not O(num_frames): a million-frame single episode
+    # instantiates from just its boundaries without materializing a per-frame index list.
+    num_frames = 1_000_000
+    sampler = EpisodeAwareSampler([0], [num_frames], shuffle=True, seed=0)
+    assert len(sampler) == num_frames
+    assert sampler._starts.shape == (1,) and sampler._cum_lengths.shape == (1,)
+
+
+def test_deterministic_sampler_resume_is_exact_at_scale():
+    # Seeded randperm makes resume sample-exact at non-trivial sizes: regenerating the epoch's
+    # permutation and slicing from the saved offset reproduces the remaining order exactly.
+    num_frames = 100_000
+    reference = EpisodeAwareSampler([0], [num_frames], shuffle=True, seed=0)
+    epoch_0 = list(reference)
+    assert sorted(epoch_0) == list(range(num_frames))
+    start = num_frames - 5
+    resumed = EpisodeAwareSampler([0], [num_frames], shuffle=True, seed=0)
+    resumed.load_state_dict({"epoch": 0, "start_index": start})
+    assert list(resumed) == epoch_0[start:]
+
+
+def test_compute_sampler_state():
+    # 100 frames, batch 10, 2 ranks -> 10 underlying batches, 5 per rank per epoch.
+    assert compute_sampler_state(step=0, num_frames=100, batch_size=10, num_processes=2) == {
+        "epoch": 0,
+        "start_index": 0,
+    }
+    # step 7 -> epoch 1, 2 per-rank batches in = 2 * 10 * 2 = 40 samples in
+    assert compute_sampler_state(step=7, num_frames=100, batch_size=10, num_processes=2) == {
+        "epoch": 1,
+        "start_index": 40,
+    }
+    # uneven epoch: 95 frames -> 10 underlying batches (last short), still 5 per rank
+    assert compute_sampler_state(step=12, num_frames=95, batch_size=10, num_processes=2) == {
+        "epoch": 2,
+        "start_index": 40,
+    }
+    # uneven sharding: 105 frames -> 11 underlying batches, 6 per rank (even_batches pads)
+    assert compute_sampler_state(step=11, num_frames=105, batch_size=10, num_processes=2) == {
+        "epoch": 1,
+        "start_index": 100,
+    }

tests/datasets/test_streaming_video_encoder.py

@@ -26,7 +26,7 @@ pytest.importorskip("av", reason="av is required (install lerobot[dataset])")
 
 import av  # noqa: E402
 
-from lerobot.configs import VideoEncoderConfig
+from lerobot.configs import RGBEncoderConfig
 from lerobot.datasets.pyav_utils import get_codec
 from lerobot.datasets.video_utils import (
     StreamingVideoEncoder,
@@ -57,13 +57,11 @@ class TestCameraEncoderThread:
         result_queue: queue.Queue = queue.Queue(maxsize=1)
         stop_event = threading.Event()
 
-        enc_cfg = VideoEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13)
+        enc_cfg = RGBEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13)
         encoder_thread = _CameraEncoderThread(
             video_path=video_path,
             fps=fps,
-            vcodec=enc_cfg.vcodec,
-            pix_fmt=enc_cfg.pix_fmt,
-            codec_options=enc_cfg.get_codec_options(as_strings=True),
+            video_encoder=enc_cfg,
             frame_queue=frame_queue,
             result_queue=result_queue,
             stop_event=stop_event,
@@ -108,13 +106,11 @@ class TestCameraEncoderThread:
         result_queue: queue.Queue = queue.Queue(maxsize=1)
         stop_event = threading.Event()
 
-        enc_cfg = VideoEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13)
+        enc_cfg = RGBEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13)
         encoder_thread = _CameraEncoderThread(
             video_path=video_path,
             fps=fps,
-            vcodec=enc_cfg.vcodec,
-            pix_fmt=enc_cfg.pix_fmt,
-            codec_options=enc_cfg.get_codec_options(as_strings=True),
+            video_encoder=enc_cfg,
             frame_queue=frame_queue,
             result_queue=result_queue,
             stop_event=stop_event,
@@ -142,13 +138,11 @@ class TestCameraEncoderThread:
         result_queue: queue.Queue = queue.Queue(maxsize=1)
         stop_event = threading.Event()
 
-        enc_cfg = VideoEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13)
+        enc_cfg = RGBEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13)
         encoder_thread = _CameraEncoderThread(
             video_path=video_path,
             fps=fps,
-            vcodec=enc_cfg.vcodec,
-            pix_fmt=enc_cfg.pix_fmt,
-            codec_options=enc_cfg.get_codec_options(as_strings=True),
+            video_encoder=enc_cfg,
             frame_queue=frame_queue,
             result_queue=result_queue,
             stop_event=stop_event,
@@ -171,15 +165,15 @@ class TestCameraEncoderThread:
 
 class TestStreamingVideoEncoder:
     def _make_encoder_config(self, **kwargs):
-        """Helper to build a VideoEncoderConfig."""
-        return VideoEncoderConfig(**kwargs)
+        """Helper to build an RGBEncoderConfig."""
+        return RGBEncoderConfig(**kwargs)
 
     def test_single_camera_episode(self, tmp_path):
         """Test encoding a single camera episode."""
         video_keys = [f"{OBS_IMAGES}.laptop"]
         encoder = StreamingVideoEncoder(
             fps=30,
-            camera_encoder=self._make_encoder_config(
+            rgb_encoder=self._make_encoder_config(
                 vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13
             ),
         )
@@ -211,7 +205,7 @@ class TestStreamingVideoEncoder:
         video_keys = [f"{OBS_IMAGES}.laptop", f"{OBS_IMAGES}.phone"]
         encoder = StreamingVideoEncoder(
             fps=30,
-            camera_encoder=self._make_encoder_config(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30),
+            rgb_encoder=self._make_encoder_config(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30),
         )
         encoder.start_episode(video_keys, tmp_path)
 
@@ -237,7 +231,7 @@ class TestStreamingVideoEncoder:
         video_keys = [f"{OBS_IMAGES}.cam"]
         encoder = StreamingVideoEncoder(
             fps=30,
-            camera_encoder=self._make_encoder_config(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30),
+            rgb_encoder=self._make_encoder_config(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30),
         )
 
         for ep in range(3):
@@ -263,7 +257,7 @@ class TestStreamingVideoEncoder:
         video_keys = [f"{OBS_IMAGES}.cam"]
         encoder = StreamingVideoEncoder(
             fps=30,
-            camera_encoder=self._make_encoder_config(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30),
+            rgb_encoder=self._make_encoder_config(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30),
         )
 
         encoder.start_episode(video_keys, tmp_path)
@@ -309,7 +303,7 @@ class TestStreamingVideoEncoder:
         video_keys = [f"{OBS_IMAGES}.cam"]
         encoder = StreamingVideoEncoder(
             fps=30,
-            camera_encoder=self._make_encoder_config(
+            rgb_encoder=self._make_encoder_config(
                 vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13
             ),
         )
@@ -346,7 +340,7 @@ class TestStreamingVideoEncoder:
         video_keys = [f"{OBS_IMAGES}.cam1", f"{OBS_IMAGES}.cam2"]
         encoder = StreamingVideoEncoder(
             fps=30,
-            camera_encoder=self._make_encoder_config(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30),
+            rgb_encoder=self._make_encoder_config(vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30),
         )
         encoder.start_episode(video_keys, tmp_path)
 
@@ -375,7 +369,7 @@ class TestStreamingVideoEncoder:
     def test_encoder_threads_passed_to_thread(self, tmp_path):
         """Test that encoder_threads is stored and passed through to encoder threads."""
         video_keys = [f"{OBS_IMAGES}.cam"]
-        cfg = VideoEncoderConfig(
+        cfg = RGBEncoderConfig(
             vcodec="libsvtav1",
             pix_fmt="yuv420p",
             g=2,
@@ -383,7 +377,7 @@ class TestStreamingVideoEncoder:
         )
         encoder = StreamingVideoEncoder(
             fps=30,
-            camera_encoder=cfg,
+            rgb_encoder=cfg,
             encoder_threads=2,
         )
         assert encoder._encoder_threads == 2
@@ -391,7 +385,8 @@ class TestStreamingVideoEncoder:
 
         # Verify codec options include thread tuning for libsvtav1 (lp=…)
         thread = encoder._threads[f"{OBS_IMAGES}.cam"]
-        assert "svtav1-params" in thread.codec_options or "threads" in thread.codec_options
+        codec_opts = thread.video_encoder.get_codec_options(encoder_threads=thread.encoder_threads)
+        assert "svtav1-params" in codec_opts or "threads" in codec_opts
 
         # Feed some frames and finish to ensure it works end-to-end
         num_frames = 10
@@ -422,7 +417,7 @@ class TestStreamingVideoEncoder:
         video_keys = [f"{OBS_IMAGES}.cam"]
         encoder = StreamingVideoEncoder(
             fps=30,
-            camera_encoder=self._make_encoder_config(
+            rgb_encoder=self._make_encoder_config(
                 vcodec="libsvtav1", pix_fmt="yuv420p", g=2, crf=30, preset=13
             ),
             queue_maxsize=1,

tests/datasets/test_video_encoding.py

@@ -26,7 +26,7 @@ pytest.importorskip("av", reason="av is required (install lerobot[dataset])")
 
 import av  # noqa: E402
 
-from lerobot.configs import VALID_VIDEO_CODECS, VideoEncoderConfig
+from lerobot.configs import VALID_VIDEO_CODECS, DepthEncoderConfig, RGBEncoderConfig, VideoEncoderConfig
 from lerobot.datasets.image_writer import write_image
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.datasets.pyav_utils import get_codec
@@ -37,7 +37,15 @@ from lerobot.datasets.video_utils import (
     get_video_info,
     reencode_video,
 )
-from tests.fixtures.constants import DUMMY_VIDEO_INFO
+from tests.fixtures.constants import (
+    DUMMY_DEPTH_FEATURES,
+    DUMMY_DEPTH_KEY,
+    DUMMY_DEPTH_VIDEO_INFO_FULL,
+    DUMMY_VIDEO_FEATURES,
+    DUMMY_VIDEO_INFO,
+    DUMMY_VIDEO_KEY,
+)
+from tests.fixtures.dataset_factories import add_frames
 
 
 # Per-codec skip markers — validation tests only fire when the codec is available
@@ -48,19 +56,74 @@ def _require_encoder(vcodec: str) -> pytest.MarkDecorator:
 
 require_libsvtav1 = _require_encoder("libsvtav1")
 require_h264 = _require_encoder("h264")
+require_hevc = _require_encoder("hevc")
 require_videotoolbox = _require_encoder("h264_videotoolbox")
 require_nvenc = _require_encoder("h264_nvenc")
 require_vaapi = _require_encoder("h264_vaapi")
 require_qsv = _require_encoder("h264_qsv")
 
 
-# ─── VideoEncoderConfig / codec options ──────────────────────────────
+TEST_ARTIFACTS_DIR = Path(__file__).parent.parent / "artifacts" / "encoded_videos"
+
+
+def _write_color_frames(imgs_dir: Path, num_frames: int = 4, height: int = 64, width: int = 96) -> None:
+    imgs_dir.mkdir(parents=True, exist_ok=True)
+    for i in range(num_frames):
+        arr = np.random.randint(0, 256, (height, width, 3), dtype=np.uint8)
+        write_image(arr, imgs_dir / f"frame-{i:06d}.png")
+
+
+def _write_depth_frames(imgs_dir: Path, num_frames: int = 4, height: int = 64, width: int = 96) -> None:
+    """Write synthetic uint16 depth TIFFs (millimetres) for depth encoder tests.
+
+    Uses a smooth linear ramp + per-frame offset (not white noise) so HEVC Main 12
+    on ``gray12le`` compresses well. Values span ~100 mm to 10 m, covering most
+    of the default ``[DEPTH_MIN, DEPTH_MAX]`` metres range after
+    ``quantize_depth(input_unit="auto"="mm")``.
+    """
+    imgs_dir.mkdir(parents=True, exist_ok=True)
+    base = np.linspace(100.0, 10_000.0, height * width, dtype=np.float32).reshape(height, width)
+    for i in range(num_frames):
+        arr = (base + 50.0 * i).clip(0, 65535).astype(np.uint16)
+        write_image(arr, imgs_dir / f"frame-{i:06d}.tiff")
+
+
+def _encode_video(
+    path: Path,
+    num_frames: int = 4,
+    fps: int = 30,
+    cfg: VideoEncoderConfig | None = None,
+    depth: bool = False,
+) -> Path:
+    """Write synthetic frames to a temp dir and encode them to ``path``.
+
+    ``depth=False`` writes uint8 RGB PNG noise and encodes with ``cfg``
+    (defaulting to the library default). ``depth=True`` writes synthetic uint16
+    depth TIFFs and encodes with ``cfg`` or a default :class:`DepthEncoderConfig`
+    (HEVC Main 12 / ``gray12le``).
+    """
+    imgs_dir = path.parent / f"imgs_{path.stem}"
+    if depth:
+        _write_depth_frames(imgs_dir, num_frames=num_frames)
+        cfg = cfg or DepthEncoderConfig()
+    else:
+        _write_color_frames(imgs_dir, num_frames=num_frames)
+    encode_video_frames(imgs_dir, path, fps=fps, video_encoder=cfg, overwrite=True)
+    return path
+
+
+def _read_feature_info(dataset: LeRobotDataset, key: str = DUMMY_VIDEO_KEY) -> dict:
+    info = json.loads((dataset.root / INFO_PATH).read_text())
+    return info["features"][key]["info"]
+
+
+# ─── RGBEncoderConfig / codec options ──────────────────────────────
 
 
 class TestCodecOptions:
     @require_libsvtav1
     def test_libsvtav1_defaults(self):
-        cfg = VideoEncoderConfig()
+        cfg = RGBEncoderConfig()
         opts = cfg.get_codec_options()
         assert opts["g"] == 2
         assert opts["crf"] == 30
@@ -68,12 +131,12 @@ class TestCodecOptions:
 
     @require_libsvtav1
     def test_libsvtav1_custom_preset(self):
-        cfg = VideoEncoderConfig(preset=8)
+        cfg = RGBEncoderConfig(preset=8)
         assert cfg.get_codec_options()["preset"] == 8
 
     @require_h264
     def test_h264_options(self):
-        cfg = VideoEncoderConfig(vcodec="h264", g=10, crf=23, preset=None)
+        cfg = RGBEncoderConfig(vcodec="h264", g=10, crf=23, preset=None)
         opts = cfg.get_codec_options()
         assert opts["g"] == 10
         assert opts["crf"] == 23
@@ -81,120 +144,120 @@ class TestCodecOptions:
 
     @require_videotoolbox
     def test_videotoolbox_options(self):
-        cfg = VideoEncoderConfig(vcodec="h264_videotoolbox", g=2, crf=30, preset=None)
+        cfg = RGBEncoderConfig(vcodec="h264_videotoolbox", g=2, crf=30, preset=None)
         opts = cfg.get_codec_options()
         assert opts["g"] == 2
         assert opts["q:v"] == 40
         assert "crf" not in opts
 
-    @_require_encoder("h264_nvenc")
+    @require_nvenc
     def test_nvenc_options(self):
-        cfg = VideoEncoderConfig(vcodec="h264_nvenc", g=2, crf=25, preset=None)
+        cfg = RGBEncoderConfig(vcodec="h264_nvenc", g=2, crf=25, preset=None)
         opts = cfg.get_codec_options()
         assert opts["rc"] == 0
         assert opts["qp"] == 25
         assert "crf" not in opts
         assert opts["g"] == 2
 
-    @_require_encoder("h264_vaapi")
+    @require_vaapi
     def test_vaapi_options(self):
-        cfg = VideoEncoderConfig(vcodec="h264_vaapi", crf=28, preset=None)
+        cfg = RGBEncoderConfig(vcodec="h264_vaapi", crf=28, preset=None)
         assert cfg.get_codec_options()["qp"] == 28
 
-    @_require_encoder("h264_qsv")
+    @require_qsv
     def test_qsv_options(self):
-        cfg = VideoEncoderConfig(vcodec="h264_qsv", crf=25, preset=None)
+        cfg = RGBEncoderConfig(vcodec="h264_qsv", crf=25, preset=None)
         assert cfg.get_codec_options()["global_quality"] == 25
 
     @require_h264
     def test_no_g_no_crf(self):
-        cfg = VideoEncoderConfig(vcodec="h264", g=None, crf=None, preset=None)
+        cfg = RGBEncoderConfig(vcodec="h264", g=None, crf=None, preset=None)
         opts = cfg.get_codec_options()
         assert "g" not in opts
         assert "crf" not in opts
 
     @require_libsvtav1
     def test_encoder_threads_libsvtav1(self):
-        cfg = VideoEncoderConfig(fast_decode=0)
+        cfg = RGBEncoderConfig(fast_decode=0)
         opts = cfg.get_codec_options(encoder_threads=4)
         assert "lp=4" in opts.get("svtav1-params", "")
 
     @require_h264
     def test_encoder_threads_h264(self):
-        cfg = VideoEncoderConfig(vcodec="h264", preset=None)
+        cfg = RGBEncoderConfig(vcodec="h264", preset=None)
         assert cfg.get_codec_options(encoder_threads=2)["threads"] == 2
 
     @require_libsvtav1
     def test_fast_decode_libsvtav1(self):
-        cfg = VideoEncoderConfig(fast_decode=1)
+        cfg = RGBEncoderConfig(fast_decode=1)
         opts = cfg.get_codec_options()
         assert "fast-decode=1" in opts.get("svtav1-params", "")
 
     @require_libsvtav1
     def test_libsvtav1_fast_decode_clamped_to_svt_range(self):
         """Out-of-range fast_decode is clamped to [0, 2] in svtav1-params (SVT-AV1 FastDecode)."""
-        cfg = VideoEncoderConfig(fast_decode=100)
+        cfg = RGBEncoderConfig(fast_decode=100)
         assert "fast-decode=2" in cfg.get_codec_options().get("svtav1-params", "")
-        cfg_neg = VideoEncoderConfig(fast_decode=-5)
+        cfg_neg = RGBEncoderConfig(fast_decode=-5)
         assert "fast-decode=0" in cfg_neg.get_codec_options().get("svtav1-params", "")
 
     @require_h264
     def test_fast_decode_h264(self):
-        cfg = VideoEncoderConfig(vcodec="h264", fast_decode=1, preset=None)
+        cfg = RGBEncoderConfig(vcodec="h264", fast_decode=1, preset=None)
         assert cfg.get_codec_options()["tune"] == "fastdecode"
 
     @require_libsvtav1
     def test_pix_fmt_unsupported_raises(self):
         """Passing an unsupported pix_fmt is a hard error."""
         with pytest.raises(ValueError, match="pix_fmt"):
-            VideoEncoderConfig(pix_fmt="yuv444p")  # libsvtav1 only supports yuv420p variants
+            RGBEncoderConfig(pix_fmt="yuv444p")  # libsvtav1 only supports yuv420p variants
 
     @require_libsvtav1
     @require_h264
     def test_preset_default_behaviour(self):
         """Empty constructor picks preset=12 (libsvtav1 path); other codecs stay None."""
-        assert VideoEncoderConfig().preset == 12
-        assert VideoEncoderConfig(vcodec="libsvtav1").preset == 12
-        assert VideoEncoderConfig(vcodec="h264").preset is None
-        assert VideoEncoderConfig(vcodec="h264", preset=None).preset is None
+        assert RGBEncoderConfig().preset == 12
+        assert RGBEncoderConfig(vcodec="libsvtav1").preset == 12
+        assert RGBEncoderConfig(vcodec="h264").preset is None
+        assert RGBEncoderConfig(vcodec="h264", preset=None).preset is None
 
     @require_h264
     def test_preset_string_on_h264(self):
         """h264 accepts string presets and forwards them to FFmpeg."""
-        cfg = VideoEncoderConfig(vcodec="h264", preset="slow")
+        cfg = RGBEncoderConfig(vcodec="h264", preset="slow")
         assert cfg.get_codec_options()["preset"] == "slow"
 
     @require_videotoolbox
     def test_preset_on_videotoolbox_not_set(self):
         """videotoolbox has no preset option at all."""
-        cfg = VideoEncoderConfig(vcodec="h264_videotoolbox", preset="slow")
+        cfg = RGBEncoderConfig(vcodec="h264_videotoolbox", preset="slow")
         assert "preset" not in cfg.get_codec_options()
 
     @require_libsvtav1
     def test_libsvtav1_preset_out_of_range_raises(self):
         """libsvtav1 preset must sit in [-2, 13] as exposed by PyAV."""
         with pytest.raises(ValueError, match="out of range"):
-            VideoEncoderConfig(vcodec="libsvtav1", preset=100)
+            RGBEncoderConfig(vcodec="libsvtav1", preset=100)
         with pytest.raises(ValueError, match="out of range"):
-            VideoEncoderConfig(vcodec="libsvtav1", preset=-3)
+            RGBEncoderConfig(vcodec="libsvtav1", preset=-3)
 
     @require_libsvtav1
     def test_libsvtav1_crf_out_of_range_raises(self):
         """libsvtav1 crf must sit in [0, 63]."""
         with pytest.raises(ValueError, match="crf.*out of range"):
-            VideoEncoderConfig(vcodec="libsvtav1", crf=64)
+            RGBEncoderConfig(vcodec="libsvtav1", crf=64)
 
     @require_libsvtav1
     def test_libsvtav1_crf_rejects_python_float(self):
         """libsvtav1 exposes ``crf`` as an INT AVOption; Python float must not pass validation."""
         with pytest.raises(ValueError, match="float values are not allowed"):
-            VideoEncoderConfig(vcodec="libsvtav1", crf=2.5)
+            RGBEncoderConfig(vcodec="libsvtav1", crf=2.5)
 
     @require_libsvtav1
     def test_libsvtav1_extra_crf_rejects_fractional_string(self):
         """INT options reject fractional values even when supplied only via ``extra_options``."""
         with pytest.raises(ValueError, match="float values are not allowed"):
-            VideoEncoderConfig(
+            RGBEncoderConfig(
                 vcodec="libsvtav1",
                 crf=None,
                 extra_options={"crf": "2.5"},
@@ -203,7 +266,7 @@ class TestCodecOptions:
     @require_libsvtav1
     def test_libsvtav1_extra_crf_rejects_float(self):
         with pytest.raises(ValueError, match="float values are not allowed"):
-            VideoEncoderConfig(
+            RGBEncoderConfig(
                 vcodec="libsvtav1",
                 crf=None,
                 extra_options={"crf": 2.5},
@@ -212,13 +275,13 @@ class TestCodecOptions:
     @require_h264
     def test_h264_crf_accepts_float_and_int(self):
         """x264 exposes crf as a FLOAT option, so both int and float are accepted."""
-        assert VideoEncoderConfig(vcodec="h264", crf=23).get_codec_options()["crf"] == 23
-        assert VideoEncoderConfig(vcodec="h264", crf=23.5).get_codec_options()["crf"] == 23.5
+        assert RGBEncoderConfig(vcodec="h264", crf=23).get_codec_options()["crf"] == 23
+        assert RGBEncoderConfig(vcodec="h264", crf=23.5).get_codec_options()["crf"] == 23.5
 
     @require_libsvtav1
     def test_validate_is_rerunnable(self):
         """After mutating a field, validate() re-checks and surfaces new issues."""
-        cfg = VideoEncoderConfig(vcodec="libsvtav1")
+        cfg = RGBEncoderConfig(vcodec="libsvtav1")
         cfg.preset = 100  # now out of range
         with pytest.raises(ValueError, match="out of range"):
             cfg.validate()
@@ -227,58 +290,58 @@ class TestCodecOptions:
 class TestExtraOptions:
     @require_libsvtav1
     def test_default_is_empty_dict(self):
-        cfg = VideoEncoderConfig()
+        cfg = RGBEncoderConfig()
         assert cfg.extra_options == {}
 
     @require_libsvtav1
     def test_unknown_key_passes_through(self):
         """Keys not published as AVOptions are forwarded to FFmpeg."""
-        cfg = VideoEncoderConfig(extra_options={"totally_made_up_option": "value"})
+        cfg = RGBEncoderConfig(extra_options={"totally_made_up_option": "value"})
         assert cfg.extra_options == {"totally_made_up_option": "value"}
 
     @require_libsvtav1
     def test_numeric_value_in_range_ok(self):
         """libsvtav1 exposes ``qp`` as INT in [0, 63]."""
-        cfg = VideoEncoderConfig(extra_options={"qp": 30})
+        cfg = RGBEncoderConfig(extra_options={"qp": 30})
         assert cfg.extra_options == {"qp": 30}
 
     @require_libsvtav1
     def test_numeric_out_of_range_raises(self):
         with pytest.raises(ValueError, match=r"qp=.*out of range"):
-            VideoEncoderConfig(extra_options={"qp": 999})
+            RGBEncoderConfig(extra_options={"qp": 999})
 
     @require_libsvtav1
     def test_numeric_string_accepted_in_range(self):
         """Numeric strings are accepted for numeric options (mirrors FFmpeg)."""
-        cfg = VideoEncoderConfig(extra_options={"qp": "18"})
+        cfg = RGBEncoderConfig(extra_options={"qp": "18"})
         assert cfg.extra_options == {"qp": "18"}
 
     @require_libsvtav1
     def test_numeric_string_out_of_range_raises(self):
         with pytest.raises(ValueError, match=r"qp=.*out of range"):
-            VideoEncoderConfig(extra_options={"qp": "999"})
+            RGBEncoderConfig(extra_options={"qp": "999"})
 
     @require_libsvtav1
     def test_non_numeric_string_on_numeric_option_raises(self):
         with pytest.raises(ValueError, match=r"qp=.*not numeric"):
-            VideoEncoderConfig(extra_options={"qp": "medium"})
+            RGBEncoderConfig(extra_options={"qp": "medium"})
 
     @require_libsvtav1
     def test_bool_on_numeric_option_raises(self):
         """``bool`` is explicitly rejected for numeric options."""
         with pytest.raises(ValueError, match=r"qp=.*not numeric"):
-            VideoEncoderConfig(extra_options={"qp": True})
+            RGBEncoderConfig(extra_options={"qp": True})
 
     @require_h264
     def test_string_option_passes_through_unchecked(self):
         """String-typed AVOptions are NOT enum-checked (too many accept freeform)."""
-        cfg = VideoEncoderConfig(vcodec="h264", preset=None, extra_options={"tune": "some-future-tune"})
+        cfg = RGBEncoderConfig(vcodec="h264", preset=None, extra_options={"tune": "some-future-tune"})
         assert cfg.extra_options == {"tune": "some-future-tune"}
 
     @require_libsvtav1
     def test_merged_into_codec_options_and_stringified(self):
         """Typed merge by default; ``as_strings=True`` matches FFmpeg option dict."""
-        cfg = VideoEncoderConfig(extra_options={"qp": 20})
+        cfg = RGBEncoderConfig(extra_options={"qp": 20})
         opts = cfg.get_codec_options()
         assert opts["qp"] == 20
         assert isinstance(opts["qp"], int)
@@ -287,25 +350,25 @@ class TestExtraOptions:
     @require_libsvtav1
     def test_structured_fields_win_on_collision(self):
         """A colliding extra_options key is discarded; the structured field wins."""
-        cfg = VideoEncoderConfig(crf=30, extra_options={"crf": 18})
+        cfg = RGBEncoderConfig(crf=30, extra_options={"crf": 18})
         assert cfg.get_codec_options()["crf"] == 30
 
 
 class TestEncoderDetection:
     @require_h264
     def test_explicit_codec_kept_when_available(self):
-        cfg = VideoEncoderConfig(vcodec="h264")
+        cfg = RGBEncoderConfig(vcodec="h264")
         assert cfg.vcodec == "h264"
 
     @require_videotoolbox
     def test_auto_picks_videotoolbox_when_available(self):
         """``h264_videotoolbox`` sits at the top of ``HW_VIDEO_CODECS`` so it wins when present."""
-        cfg = VideoEncoderConfig(vcodec="auto")
+        cfg = RGBEncoderConfig(vcodec="auto")
         assert cfg.vcodec == "h264_videotoolbox"
 
     def test_invalid_codec_raises(self):
         with pytest.raises(ValueError, match="Invalid vcodec"):
-            VideoEncoderConfig(vcodec="not_a_real_codec")
+            RGBEncoderConfig(vcodec="not_a_real_codec")
 
     def test_hw_encoder_names_listed_as_valid(self):
         assert "auto" in VALID_VIDEO_CODECS
@@ -313,59 +376,6 @@ class TestEncoderDetection:
         assert "h264_nvenc" in VALID_VIDEO_CODECS
 
 
-TEST_ARTIFACTS_DIR = Path(__file__).parent.parent / "artifacts" / "encoded_videos"
-
-# Default video feature set used by persistence tests.
-VIDEO_FEATURES = {
-    "observation.images.cam": {
-        "dtype": "video",
-        "shape": (64, 96, 3),
-        "names": ["height", "width", "channels"],
-    },
-    "action": {"dtype": "float32", "shape": (2,), "names": ["a", "b"]},
-}
-VIDEO_KEY = "observation.images.cam"
-
-
-def _write_frames(imgs_dir: Path, num_frames: int = 4, height: int = 64, width: int = 96) -> None:
-    imgs_dir.mkdir(parents=True, exist_ok=True)
-    for i in range(num_frames):
-        arr = np.random.randint(0, 256, (height, width, 3), dtype=np.uint8)
-        write_image(arr, imgs_dir / f"frame-{i:06d}.png")
-
-
-def _encode_video(
-    path: Path, num_frames: int = 4, fps: int = 30, cfg: VideoEncoderConfig | None = None
-) -> Path:
-    imgs_dir = path.parent / f"imgs_{path.stem}"
-    _write_frames(imgs_dir, num_frames=num_frames)
-    encode_video_frames(imgs_dir, path, fps=fps, camera_encoder=cfg, overwrite=True)
-    return path
-
-
-def _read_feature_info(dataset: LeRobotDataset) -> dict:
-    info = json.loads((dataset.root / INFO_PATH).read_text())
-    return info["features"][VIDEO_KEY]["info"]
-
-
-def _add_frames(dataset: LeRobotDataset, num_frames: int, video_keys: list[str] | None = None) -> None:
-    from lerobot.utils.constants import DEFAULT_FEATURES
-
-    if video_keys is None:
-        video_keys = dataset.meta.video_keys
-    for _ in range(num_frames):
-        frame: dict = {"task": "test"}
-        for key, ft in dataset.meta.features.items():
-            if key in DEFAULT_FEATURES:
-                continue
-            shape = ft["shape"]
-            if key in video_keys:
-                frame[key] = np.random.randint(0, 256, shape, dtype=np.uint8)
-            else:
-                frame[key] = np.zeros(shape, dtype=np.float32)
-        dataset.add_frame(frame)
-
-
 class TestGetVideoInfo:
     def test_returns_all_stream_fields(self):
         info = get_video_info(TEST_ARTIFACTS_DIR / "clip_4frames.mp4")
@@ -375,7 +385,7 @@ class TestGetVideoInfo:
         assert info["video.pix_fmt"] == "yuv420p"
         assert info["video.fps"] == 30
         assert info["video.channels"] == 3
-        assert info["video.is_depth_map"] is False
+        assert info["is_depth_map"] is False
         assert info["has_audio"] is False
         assert "video.g" not in info
         assert "video.crf" not in info
@@ -383,9 +393,9 @@ class TestGetVideoInfo:
 
     @require_libsvtav1
     def test_merges_encoder_config_as_video_prefixed_entries(self):
-        cfg = VideoEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12)
+        cfg = RGBEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12)
 
-        info = get_video_info(TEST_ARTIFACTS_DIR / "clip_4frames.mp4", camera_encoder=cfg)
+        info = get_video_info(TEST_ARTIFACTS_DIR / "clip_4frames.mp4", video_encoder=cfg)
 
         assert info["video.g"] == 2
         assert info["video.crf"] == 30
@@ -396,13 +406,18 @@ class TestGetVideoInfo:
 
     @require_libsvtav1
     def test_stream_derived_keys_take_precedence_over_config(self):
-        cfg = VideoEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p")
+        cfg = RGBEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p")
 
-        info = get_video_info(TEST_ARTIFACTS_DIR / "clip_4frames.mp4", camera_encoder=cfg)
+        info = get_video_info(TEST_ARTIFACTS_DIR / "clip_4frames.mp4", video_encoder=cfg)
 
         assert info["video.codec"]  # populated from stream, not from config's vcodec
         assert info["video.pix_fmt"] == "yuv420p"
 
+    def test_depth_encoder_config_sets_is_depth_map_true(self):
+        """A ``DepthEncoderConfig`` causes ``get_video_info`` to mark the stream as depth."""
+        info = get_video_info(TEST_ARTIFACTS_DIR / "clip_4frames.mp4", video_encoder=DepthEncoderConfig())
+        assert info["is_depth_map"] is True
+
 
 class TestEncodeVideoFrames:
     @require_libsvtav1
@@ -434,7 +449,7 @@ class TestEncodeVideoFrames:
 
     def test_overwrite_false_skips_existing_file(self, tmp_path):
         imgs_dir = tmp_path / "imgs"
-        _write_frames(imgs_dir)
+        _write_color_frames(imgs_dir)
         video_path = tmp_path / "out.mp4"
         sentinel = b"pre-existing content"
         video_path.write_bytes(sentinel)
@@ -446,7 +461,7 @@ class TestEncodeVideoFrames:
     @require_libsvtav1
     def test_overwrite_true_replaces_existing_file(self, tmp_path):
         imgs_dir = tmp_path / "imgs"
-        _write_frames(imgs_dir)
+        _write_color_frames(imgs_dir)
         video_path = tmp_path / "out.mp4"
         video_path.write_bytes(b"stale content")
 
@@ -458,10 +473,10 @@ class TestEncodeVideoFrames:
     @require_libsvtav1
     def test_custom_encoder_config_fields_stored_in_info(self, tmp_path):
         """All stream-derived and encoder config fields are present after encoding."""
-        cfg = VideoEncoderConfig(vcodec="libsvtav1", g=4, crf=25, preset=10)
+        cfg = RGBEncoderConfig(vcodec="libsvtav1", g=4, crf=25, preset=10)
         video_path = _encode_video(tmp_path / "out.mp4", num_frames=4, fps=30, cfg=cfg)
 
-        info = get_video_info(video_path, camera_encoder=cfg)
+        info = get_video_info(video_path, video_encoder=cfg)
 
         # Stream-derived
         assert info["video.height"] == 64
@@ -470,7 +485,7 @@ class TestEncodeVideoFrames:
         assert info["video.codec"] == "av1"
         assert info["video.pix_fmt"] == "yuv420p"
         assert info["video.fps"] == 30
-        assert info["video.is_depth_map"] is False
+        assert info["is_depth_map"] is False
         assert info["has_audio"] is False
         # Encoder config
         assert info["video.g"] == 4
@@ -487,15 +502,15 @@ class TestReencodeVideo:
     def test_reencode_video(self, tmp_path):
         src = TEST_ARTIFACTS_DIR / "clip_4frames.mp4"
         out = tmp_path / "reencoded.mp4"
-        cfg = VideoEncoderConfig(vcodec="h264", g=6, crf=23, pix_fmt="yuv444p")
-        reencode_video(src, out, camera_encoder=cfg, overwrite=True)
+        cfg = RGBEncoderConfig(vcodec="h264", g=6, crf=23, pix_fmt="yuv444p")
+        reencode_video(src, out, video_encoder=cfg, overwrite=True)
 
         assert out.exists()
         with av.open(str(out)) as container:
             n_frames = sum(1 for _ in container.decode(video=0))
         assert n_frames == 4
 
-        info = get_video_info(out, camera_encoder=cfg)
+        info = get_video_info(out, video_encoder=cfg)
         assert info["video.codec"] == "h264"
         assert info["video.pix_fmt"] == "yuv444p"
         assert info["video.height"] == 64
@@ -504,6 +519,19 @@ class TestReencodeVideo:
         assert info["video.g"] == 6
         assert info["video.crf"] == 23
 
+    @require_h264
+    def test_reencode_video_trim_window(self, tmp_path):
+        src = TEST_ARTIFACTS_DIR / "clip_6frames.mp4"
+        out = tmp_path / "trim_window.mp4"
+        cfg = RGBEncoderConfig(vcodec="h264")
+        reencode_video(src, out, video_encoder=cfg, start_time_s=0.05, end_time_s=0.12, overwrite=True)
+
+        with av.open(str(out)) as container:
+            frames = list(container.decode(video=0))
+        # Only the frames at 0.067 and 0.1 s fall inside [0.05, 0.12).
+        assert len(frames) == 2
+        assert frames[0].time == pytest.approx(0.0, abs=1e-3)
+
 
 class TestConcatenateVideoFiles:
     def test_two_clips_frame_count(self, tmp_path):
@@ -565,12 +593,12 @@ class TestEncoderConfigPersistence:
 
     @require_libsvtav1
     def test_first_episode_save_persists_encoder_config(self, tmp_path, empty_lerobot_dataset_factory):
-        cfg = VideoEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12)
+        cfg = RGBEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12)
         dataset = empty_lerobot_dataset_factory(
-            root=tmp_path / "ds", features=VIDEO_FEATURES, use_videos=True, camera_encoder=cfg
+            root=tmp_path / "ds", features=DUMMY_VIDEO_FEATURES, use_videos=True, rgb_encoder=cfg
         )
 
-        _add_frames(dataset, num_frames=4)
+        add_frames(dataset, num_frames=4)
         dataset.save_episode()
         dataset.finalize()
 
@@ -588,16 +616,16 @@ class TestEncoderConfigPersistence:
 
     @require_libsvtav1
     def test_second_episode_does_not_overwrite_encoder_fields(self, tmp_path, empty_lerobot_dataset_factory):
-        cfg = VideoEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12)
+        cfg = RGBEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12)
         dataset = empty_lerobot_dataset_factory(
-            root=tmp_path / "ds", features=VIDEO_FEATURES, use_videos=True, camera_encoder=cfg
+            root=tmp_path / "ds", features=DUMMY_VIDEO_FEATURES, use_videos=True, rgb_encoder=cfg
         )
 
-        _add_frames(dataset, num_frames=4)
+        add_frames(dataset, num_frames=4)
         dataset.save_episode()
         first_info = dict(_read_feature_info(dataset))
 
-        _add_frames(dataset, num_frames=4)
+        add_frames(dataset, num_frames=4)
         dataset.save_episode()
         dataset.finalize()
 
@@ -605,13 +633,13 @@ class TestEncoderConfigPersistence:
 
 
 class TestFromVideoInfo:
-    """``VideoEncoderConfig.from_video_info`` reconstructs an encoder config
+    """``RGBEncoderConfig.from_video_info`` reconstructs an encoder config
     from the ``video.*`` keys persisted in a dataset's ``info.json``.
     """
 
     @require_libsvtav1
     def test_reconstructs_from_dummy_video_info(self):
-        cfg = VideoEncoderConfig.from_video_info(DUMMY_VIDEO_INFO)
+        cfg = RGBEncoderConfig.from_video_info(DUMMY_VIDEO_INFO)
 
         # Canonical stream codec ``"av1"`` is aliased to the encoder name.
         assert cfg.vcodec == "libsvtav1"
@@ -623,4 +651,220 @@ class TestFromVideoInfo:
         assert cfg.video_backend == DUMMY_VIDEO_INFO["video.video_backend"]
         # ``{}`` placeholder (typical after a merge with disagreeing sources)
         # must not leak into the reconstructed config.
-        assert cfg.extra_options == VideoEncoderConfig().extra_options
+        assert cfg.extra_options == RGBEncoderConfig().extra_options
+
+
+# ─── Depth-specific encoding tests ────────────────────────────────────
+
+
+class TestEncodeDepthVideoFrames:
+    """Depth mirror of :class:`TestEncodeVideoFrames`.
+
+    Exercises ``encode_video_frames`` end-to-end through
+    :class:`DepthEncoderConfig` (HEVC Main 12 / ``gray12le``) on synthetic
+    uint16 depth TIFFs.
+    """
+
+    @require_hevc
+    def test_produces_readable_file(self, tmp_path):
+        video_path = _encode_video(tmp_path / "out.mp4", depth=True)
+
+        assert video_path.exists()
+        info = get_video_info(video_path, video_encoder=DepthEncoderConfig())
+        assert info["video.height"] == 64
+        assert info["video.width"] == 96
+        assert info["video.codec"] == "hevc"
+        assert info["video.pix_fmt"] == "gray12le"
+        assert info["video.channels"] == 1
+        assert info["is_depth_map"] is True
+
+    @require_hevc
+    def test_frame_count_and_duration_match_input(self, tmp_path):
+        num_frames = 10
+        fps = 30
+        video_path = _encode_video(tmp_path / "out.mp4", num_frames=num_frames, fps=fps, depth=True)
+
+        with av.open(str(video_path)) as container:
+            stream = container.streams.video[0]
+            actual_frames = sum(1 for _ in container.decode(stream))
+            duration = (
+                float(stream.duration * stream.time_base)
+                if stream.duration is not None
+                else float(container.duration / av.time_base)
+            )
+
+        assert actual_frames == num_frames
+        assert abs(duration - num_frames / fps) < 0.1
+
+    def test_overwrite_false_skips_existing_file(self, tmp_path):
+        """Codec-agnostic: file-system semantics must hold even without an HEVC encoder."""
+        imgs_dir = tmp_path / "imgs"
+        _write_depth_frames(imgs_dir)
+        video_path = tmp_path / "out.mp4"
+        sentinel = b"pre-existing depth content"
+        video_path.write_bytes(sentinel)
+
+        encode_video_frames(imgs_dir, video_path, fps=30, video_encoder=DepthEncoderConfig(), overwrite=False)
+
+        assert video_path.read_bytes() == sentinel
+
+    @require_hevc
+    def test_overwrite_true_replaces_existing_file(self, tmp_path):
+        imgs_dir = tmp_path / "imgs"
+        _write_depth_frames(imgs_dir)
+        video_path = tmp_path / "out.mp4"
+        video_path.write_bytes(b"stale content")
+
+        encode_video_frames(imgs_dir, video_path, fps=30, video_encoder=DepthEncoderConfig(), overwrite=True)
+
+        info = get_video_info(video_path, video_encoder=DepthEncoderConfig())
+        assert info["video.height"] == 64
+        assert info["video.pix_fmt"] == "gray12le"
+        assert info["is_depth_map"] is True
+
+    @require_hevc
+    def test_custom_encoder_config_fields_stored_in_info(self, tmp_path):
+        """All stream-derived and depth-encoder config fields are present after encoding."""
+        cfg = DepthEncoderConfig(
+            vcodec="hevc",
+            pix_fmt="gray12le",
+            g=4,
+            crf=25,
+            extra_options={},
+            depth_min=0.05,
+            depth_max=8.0,
+            shift=2.5,
+            use_log=False,
+        )
+        video_path = _encode_video(tmp_path / "out.mp4", num_frames=4, fps=30, cfg=cfg, depth=True)
+
+        info = get_video_info(video_path, video_encoder=cfg)
+
+        # Stream-derived
+        assert info["video.height"] == 64
+        assert info["video.width"] == 96
+        assert info["video.channels"] == 1
+        assert info["video.codec"] == "hevc"
+        assert info["video.pix_fmt"] == "gray12le"
+        assert info["video.fps"] == 30
+        assert info["is_depth_map"] is True
+        assert info["has_audio"] is False
+        # Base encoder config
+        assert info["video.g"] == 4
+        assert info["video.crf"] == 25
+        assert info["video.fast_decode"] == 0
+        assert info["video.video_backend"] == "pyav"
+        assert info["video.extra_options"] == {}
+        # Depth-specific tuning
+        assert info["video.depth_min"] == 0.05
+        assert info["video.depth_max"] == 8.0
+        assert info["video.shift"] == 2.5
+        assert info["video.use_log"] is False
+
+
+class TestDepthEncoderConfigPersistence:
+    """Depth mirror of :class:`TestEncoderConfigPersistence`.
+
+    ``DepthEncoderConfig`` must be stored as ``video.<field>`` entries
+    (including the depth-specific ``depth_min`` / ``depth_max`` / ``shift`` /
+    ``use_log``) under ``info["features"][<depth_key>]["info"]`` when the
+    first episode is saved.
+    """
+
+    @require_hevc
+    def test_first_episode_save_persists_depth_encoder_config(self, tmp_path, empty_lerobot_dataset_factory):
+        cfg = DepthEncoderConfig(
+            vcodec="hevc",
+            pix_fmt="gray12le",
+            g=2,
+            crf=30,
+            extra_options={},
+            depth_min=0.05,
+            depth_max=8.0,
+            shift=2.5,
+            use_log=False,
+        )
+        dataset = empty_lerobot_dataset_factory(
+            root=tmp_path / "ds", features=DUMMY_DEPTH_FEATURES, use_videos=True, depth_encoder=cfg
+        )
+
+        add_frames(dataset, num_frames=4)
+        dataset.save_episode()
+        dataset.finalize()
+
+        info = _read_feature_info(dataset, key=DUMMY_DEPTH_KEY)
+
+        # Stream-derived
+        assert info["video.height"] == 64
+        assert info["video.width"] == 96
+        assert info["video.fps"] == 30
+        assert info["video.codec"] == "hevc"
+        assert info["video.pix_fmt"] == "gray12le"
+        assert info["is_depth_map"] is True
+        # Base encoder config
+        assert info["video.g"] == 2
+        assert info["video.crf"] == 30
+        assert info["video.fast_decode"] == 0
+        assert info["video.video_backend"] == "pyav"
+        assert info["video.extra_options"] == {}
+        # Depth-specific tuning
+        assert info["video.depth_min"] == 0.05
+        assert info["video.depth_max"] == 8.0
+        assert info["video.shift"] == 2.5
+        assert info["video.use_log"] is False
+
+    @require_hevc
+    def test_second_episode_does_not_overwrite_depth_encoder_fields(
+        self, tmp_path, empty_lerobot_dataset_factory
+    ):
+        cfg = DepthEncoderConfig(
+            vcodec="hevc",
+            pix_fmt="gray12le",
+            g=2,
+            crf=30,
+            depth_min=0.05,
+            depth_max=8.0,
+            shift=2.5,
+            use_log=False,
+        )
+        dataset = empty_lerobot_dataset_factory(
+            root=tmp_path / "ds", features=DUMMY_DEPTH_FEATURES, use_videos=True, depth_encoder=cfg
+        )
+
+        add_frames(dataset, num_frames=4)
+        dataset.save_episode()
+        first_info = dict(_read_feature_info(dataset, key=DUMMY_DEPTH_KEY))
+
+        add_frames(dataset, num_frames=4)
+        dataset.save_episode()
+        dataset.finalize()
+
+        assert _read_feature_info(dataset, key=DUMMY_DEPTH_KEY) == first_info
+
+
+class TestDepthFromVideoInfo:
+    """``DepthEncoderConfig.from_video_info`` reconstructs a depth encoder
+    config from the ``video.*`` keys persisted in a dataset's ``info.json``.
+
+    Depth mirror of :class:`TestFromVideoInfo`.
+    """
+
+    @require_hevc
+    def test_reconstructs_from_dummy_depth_video_info(self):
+        cfg = DepthEncoderConfig.from_video_info(DUMMY_DEPTH_VIDEO_INFO_FULL)
+
+        # No alias for ``"hevc"``; the canonical stream codec is reused as-is.
+        assert cfg.vcodec == "hevc"
+        assert cfg.pix_fmt == DUMMY_DEPTH_VIDEO_INFO_FULL["video.pix_fmt"]
+        assert cfg.g == DUMMY_DEPTH_VIDEO_INFO_FULL["video.g"]
+        assert cfg.crf == DUMMY_DEPTH_VIDEO_INFO_FULL["video.crf"]
+        assert cfg.fast_decode == DUMMY_DEPTH_VIDEO_INFO_FULL["video.fast_decode"]
+        assert cfg.video_backend == DUMMY_DEPTH_VIDEO_INFO_FULL["video.video_backend"]
+        # ``{}`` placeholder (typical after a merge with disagreeing sources)
+        # must not leak into the reconstructed config.
+        assert cfg.extra_options == DepthEncoderConfig().extra_options
+        # Depth-specific tuning round-trips through ``info.json``.
+        assert cfg.depth_min == DUMMY_DEPTH_VIDEO_INFO_FULL["video.depth_min"]
+        assert cfg.depth_max == DUMMY_DEPTH_VIDEO_INFO_FULL["video.depth_max"]
+        assert cfg.shift == DUMMY_DEPTH_VIDEO_INFO_FULL["video.shift"]
+        assert cfg.use_log == DUMMY_DEPTH_VIDEO_INFO_FULL["video.use_log"]

tests/fixtures/constants.py

@@ -39,12 +39,56 @@ DUMMY_VIDEO_INFO = {
     "video.crf": 30,
     "video.preset": 12,
     "video.fast_decode": 0,
-    "video.is_depth_map": False,
+    "is_depth_map": False,
     "has_audio": False,
 }
 DUMMY_CAMERA_FEATURES = {
     "laptop": {"shape": (64, 96, 3), "names": ["height", "width", "channels"], "info": DUMMY_VIDEO_INFO},
     "phone": {"shape": (64, 96, 3), "names": ["height", "width", "channels"], "info": DUMMY_VIDEO_INFO},
 }
+DUMMY_DEPTH_VIDEO_INFO = {
+    **DUMMY_VIDEO_INFO,
+    "is_depth_map": True,
+}
+DUMMY_DEPTH_VIDEO_INFO_FULL = {
+    **{k: v for k, v in DUMMY_VIDEO_INFO.items() if k != "video.preset"},
+    "video.codec": "hevc",
+    "video.pix_fmt": "gray12le",
+    "is_depth_map": True,
+    "video.depth_min": 0.05,
+    "video.depth_max": 8.0,
+    "video.shift": 2.5,
+    "video.use_log": True,
+}
+DUMMY_DEPTH_CAMERA_FEATURES = {
+    "laptop_depth": {
+        "shape": (64, 96, 1),
+        "names": ["height", "width", "channels"],
+        "info": DUMMY_DEPTH_VIDEO_INFO,
+    },
+}
+DUMMY_CAMERA_FEATURES_WITH_DEPTH = {**DUMMY_CAMERA_FEATURES, **DUMMY_DEPTH_CAMERA_FEATURES}
 DUMMY_CHW = (3, 96, 128)
 DUMMY_HWC = (96, 128, 3)
+
+# Default video feature set used by video-encoding persistence tests.
+DUMMY_VIDEO_FEATURES = {
+    "observation.images.cam": {
+        "dtype": "video",
+        "shape": (64, 96, 3),
+        "names": ["height", "width", "channels"],
+    },
+    "action": {"dtype": "float32", "shape": (2,), "names": ["a", "b"]},
+}
+DUMMY_VIDEO_KEY = "observation.images.cam"
+
+DUMMY_DEPTH_FEATURES = {
+    "observation.images.depth": {
+        "dtype": "video",
+        "shape": (64, 96, 1),
+        "names": ["height", "width", "channels"],
+        "info": {"is_depth_map": True},
+    },
+    "action": {"dtype": "float32", "shape": (2,), "names": ["a", "b"]},
+}
+DUMMY_DEPTH_KEY = "observation.images.depth"

tests/fixtures/dataset_factories.py

@@ -49,6 +49,39 @@ from tests.fixtures.constants import (
 )
 
 
+def add_frames(dataset: LeRobotDataset, num_frames: int) -> None:
+    """Append ``num_frames`` synthetic frames to ``dataset``.
+
+    Generates per-feature payloads from ``dataset.meta``: uint16 depth ramps for
+    keys in ``dataset.meta.depth_keys``, uint8 random noise for video/image keys,
+    and float32 zeros for everything else. ``DEFAULT_FEATURES`` (timestamp,
+    frame_index, ...) are auto-populated by ``add_frame`` and skipped here.
+    """
+    video_keys = dataset.meta.video_keys
+    depth_keys = dataset.meta.depth_keys
+    # Smooth gradient base reused per (H, W) to keep depth frames cheap to
+    # encode (HEVC Main 12 hates white noise).
+    _depth_base_cache: dict[tuple[int, int], np.ndarray] = {}
+    for i in range(num_frames):
+        frame: dict = {"task": "test"}
+        for key, ft in dataset.meta.features.items():
+            if key in DEFAULT_FEATURES:
+                continue
+            shape = ft["shape"]
+            if key in depth_keys:
+                h, w, _ = shape
+                base = _depth_base_cache.setdefault(
+                    (h, w),
+                    np.linspace(100.0, 10_000.0, h * w, dtype=np.float32).reshape(h, w, 1),
+                )
+                frame[key] = (base + 50.0 * i).clip(0, 65535).astype(np.uint16)
+            elif key in video_keys:
+                frame[key] = np.random.randint(0, 256, shape, dtype=np.uint8)
+            else:
+                frame[key] = np.zeros(shape, dtype=np.float32)
+        dataset.add_frame(frame)
+
+
 class LeRobotDatasetFactory(Protocol):
     def __call__(self, *args, **kwargs) -> LeRobotDataset: ...
 
@@ -485,10 +518,14 @@ def lerobot_dataset_factory(
         hf_dataset: datasets.Dataset | None = None,
         data_files_size_in_mb: float = DEFAULT_DATA_FILE_SIZE_IN_MB,
         chunks_size: int = DEFAULT_CHUNK_SIZE,
+        camera_features: dict | None = None,
         **kwargs,
     ) -> LeRobotDataset:
         # Instantiate objects
         if info is None:
+            info_kwargs = {}
+            if camera_features is not None:
+                info_kwargs["camera_features"] = camera_features
             info = info_factory(
                 total_episodes=total_episodes,
                 total_frames=total_frames,
@@ -496,6 +533,7 @@ def lerobot_dataset_factory(
                 use_videos=use_videos,
                 data_files_size_in_mb=data_files_size_in_mb,
                 chunks_size=chunks_size,
+                **info_kwargs,
             )
         if stats is None:
             stats = stats_factory(features=info.features)
@@ -552,3 +590,64 @@ def lerobot_dataset_factory(
 @pytest.fixture(scope="session")
 def empty_lerobot_dataset_factory() -> LeRobotDatasetFactory:
     return partial(LeRobotDataset.create, repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS)
+
+
+def build_annotation_dataset(
+    root: Path,
+    episode_specs: list[tuple[int, int, str]],
+    *,
+    fps: int = 10,
+) -> Path:
+    """Build a minimal LeRobot-shaped dataset on disk for annotation tests.
+
+    ``episode_specs`` is a list of ``(episode_index, num_frames, task_text)``.
+    Each episode is written to its own
+    ``data/chunk-000/file-{ep:03d}.parquet`` so the writer's per-shard
+    rewrite path is exercised. The dataset carries the minimum
+    ``meta/tasks.parquet`` + ``meta/info.json`` the reader / executor need;
+    it has no videos, so the modules fall back to text-only prompts.
+
+    Shared by the annotation-pipeline pytest fixtures (``tests/annotations/
+    conftest.py``) and the opt-in E2E smoke run so the fixture shape lives
+    in exactly one place.
+    """
+    from lerobot.datasets.io_utils import write_tasks
+    from lerobot.utils.io_utils import write_json
+
+    data_dir = root / "data" / "chunk-000"
+    data_dir.mkdir(parents=True, exist_ok=True)
+
+    tasks: dict[int, str] = {}
+    for episode_index, num_frames, task_text in episode_specs:
+        if task_text not in tasks.values():
+            tasks[len(tasks)] = task_text
+        task_index = next(k for k, v in tasks.items() if v == task_text)
+        frame = pd.DataFrame(
+            {
+                "episode_index": [episode_index] * num_frames,
+                "frame_index": list(range(num_frames)),
+                "timestamp": [round(i / fps, 6) for i in range(num_frames)],
+                "task_index": [task_index] * num_frames,
+                "subtask_index": [0] * num_frames,  # legacy column the writer must drop
+            }
+        )
+        frame.to_parquet(data_dir / f"file-{episode_index:03d}.parquet", index=False)
+
+    # Canonical tasks frame: indexed by task string with a ``task_index``
+    # column, matching what ``lerobot.datasets.io_utils.load_tasks`` expects.
+    tasks_df = pd.DataFrame(
+        {"task_index": list(tasks.keys())},
+        index=pd.Index(list(tasks.values()), name="task"),
+    )
+    write_tasks(tasks_df, root)
+
+    write_json(
+        {
+            "codebase_version": "v3.1",
+            "fps": fps,
+            "features": {},
+            "total_episodes": len(episode_specs),
+        },
+        root / "meta" / "info.json",
+    )
+    return root

tests/optim/test_optimizers.py

@@ -20,6 +20,7 @@ from lerobot.optim.optimizers import (
     MultiAdamConfig,
     SGDConfig,
     load_optimizer_state,
+    load_optimizer_state_dict,
     save_optimizer_state,
 )
 from lerobot.utils.constants import (
@@ -65,6 +66,44 @@ def test_save_and_load_optimizer_state(model_params, optimizer, tmp_path):
     torch.testing.assert_close(optimizer.state_dict(), loaded_optimizer.state_dict())
 
 
+def test_save_and_load_fsdp_optimizer_state_dict_roundtrip(tmp_path):
+    """The FSDP full optimizer state dict is keyed by parameter FQNs (dotted strings), not the
+    integer indices of the single-GPU path. Verify it survives the safetensors save -> read
+    round-trip used by the FSDP save/resume path (save_optimizer_state(optim_state_dict=...) then
+    load_optimizer_state_dict), which the flatten/unflatten "/" separator must not corrupt."""
+    full_osd = {
+        "state": {
+            "model.layers.0.weight": {
+                "step": torch.tensor(3.0),
+                "exp_avg": torch.randn(4, 4),
+                "exp_avg_sq": torch.randn(4, 4),
+            },
+            "model.layers.0.bias": {
+                "step": torch.tensor(3.0),
+                "exp_avg": torch.randn(4),
+                "exp_avg_sq": torch.randn(4),
+            },
+        },
+        "param_groups": [
+            {"lr": 1e-4, "betas": [0.9, 0.999], "eps": 1e-8, "weight_decay": 0.0, "params": [0, 1]}
+        ],
+    }
+
+    save_optimizer_state(
+        torch.optim.Adam([torch.nn.Parameter(torch.randn(1))]), tmp_path, optim_state_dict=full_osd
+    )
+    assert (tmp_path / OPTIMIZER_STATE).is_file()
+    assert (tmp_path / OPTIMIZER_PARAM_GROUPS).is_file()
+
+    loaded = load_optimizer_state_dict(tmp_path)
+    # FQN keys must be preserved verbatim (not int-cast, not split on their dots).
+    assert set(loaded["state"].keys()) == set(full_osd["state"].keys())
+    for fqn, sub in full_osd["state"].items():
+        for k, v in sub.items():
+            torch.testing.assert_close(loaded["state"][fqn][k], v)
+    assert loaded["param_groups"] == full_osd["param_groups"]
+
+
 @pytest.fixture
 def base_params_dict():
     return {

tests/policies/molmoact2/test_molmoact2.py

@@ -1,5 +1,3 @@
-#!/usr/bin/env python
-
 # Copyright 2026 The Allen Institute for Artificial Intelligence and The HuggingFace Inc. team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -35,16 +33,16 @@ pytest.importorskip("scipy")
 from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature
 from lerobot.policies import get_policy_class, make_policy_config
 from lerobot.policies.molmoact2 import (
-    configuration_molmoact2 as molmoact2_config,
     modeling_molmoact2 as molmoact2_modeling,
     processor_molmoact2 as molmoact2_processor,
 )
-from lerobot.policies.molmoact2.configuration_molmoact2 import (
-    MolmoAct2Config,
-    MolmoAct2CosineDecayWithWarmupSchedulerConfig,
-    infer_molmoact2_max_sequence_length,
+from lerobot.policies.molmoact2.configuration_molmoact2 import MolmoAct2Config
+from lerobot.policies.molmoact2.modeling_molmoact2 import (
+    MolmoAct2Policy,
+    _apply_action_chunk_padding_mask,
+    _apply_action_dim_padding_mask,
+    _combine_rollout_seeds,
 )
-from lerobot.policies.molmoact2.modeling_molmoact2 import MolmoAct2Policy
 from lerobot.policies.molmoact2.processor_molmoact2 import (
     MolmoAct2ClampNormalizedProcessorStep,
     MolmoAct2MaskedNormalizerProcessorStep,
@@ -53,6 +51,7 @@ from lerobot.policies.molmoact2.processor_molmoact2 import (
     _add_gripper_masks_to_stats,
     _build_discrete_state_string,
     _normalize_question_text,
+    infer_molmoact2_max_sequence_length,
     make_molmoact2_pre_post_processors,
 )
 from lerobot.policies.rtc.configuration_rtc import RTCConfig
@@ -71,34 +70,38 @@ def test_molmoact2_policy_registration():
     assert cfg.per_episode_seed is False
     assert cfg.eval_seed is None
     assert cfg.normalize_language is True
-    assert cfg.get_scheduler_preset().num_decay_steps is None
+    assert cfg.get_scheduler_preset().num_decay_steps == 100_000
     assert cfg.action_delta_indices == list(range(cfg.chunk_size))
     assert get_policy_class("molmoact2") is MolmoAct2Policy
 
 
 def test_molmoact2_checkpoint_download_ignores_remote_python(monkeypatch):
+    import huggingface_hub
+
     download_kwargs = {}
 
     def fake_snapshot_download(**kwargs):
         download_kwargs.update(kwargs)
         return "/tmp/downloaded-molmoact2"
 
-    monkeypatch.setattr(molmoact2_config, "snapshot_download", fake_snapshot_download)
+    monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download)
 
-    checkpoint_location = molmoact2_config._resolve_checkpoint_location("allenai/MolmoAct2")
+    checkpoint_location = molmoact2_modeling._resolve_checkpoint_location("allenai/MolmoAct2")
 
     assert checkpoint_location == "/tmp/downloaded-molmoact2"
     assert download_kwargs["ignore_patterns"] == ["*.py", "*.pyc", "__pycache__/*"]
 
 
-def test_molmoact2_scheduler_decay_steps_auto_match_training_steps():
+def test_molmoact2_scheduler_auto_scales_to_training_steps():
+    from lerobot.optim import CosineDecayWithWarmupSchedulerConfig
+
     param = torch.nn.Parameter(torch.ones(()))
     optimizer = torch.optim.AdamW([param], lr=0.001)
-    config = MolmoAct2CosineDecayWithWarmupSchedulerConfig(
+    config = CosineDecayWithWarmupSchedulerConfig(
         peak_lr=0.01,
         decay_lr=0.001,
         num_warmup_steps=10,
-        num_decay_steps=None,
+        num_decay_steps=100_000,
     )
 
     scheduler = config.build(optimizer, num_training_steps=100)
@@ -123,9 +126,7 @@ def test_molmoact2_rollout_generator_uses_eval_seed_per_task():
         batch_size=3,
         device=torch.device("cpu"),
     )
-    expected_first = torch.Generator().manual_seed(
-        MolmoAct2Policy._combine_rollout_seeds(first_seed=1000, batch_size=3)
-    )
+    expected_first = torch.Generator().manual_seed(_combine_rollout_seeds(first_seed=1000, batch_size=3))
     assert torch.allclose(torch.rand(4, generator=first), torch.rand(4, generator=expected_first))
 
     policy.reset()
@@ -134,9 +135,7 @@ def test_molmoact2_rollout_generator_uses_eval_seed_per_task():
         batch_size=3,
         device=torch.device("cpu"),
     )
-    expected_second = torch.Generator().manual_seed(
-        MolmoAct2Policy._combine_rollout_seeds(first_seed=1003, batch_size=3)
-    )
+    expected_second = torch.Generator().manual_seed(_combine_rollout_seeds(first_seed=1003, batch_size=3))
     assert torch.allclose(torch.rand(4, generator=second), torch.rand(4, generator=expected_second))
 
     policy.reset()
@@ -145,9 +144,7 @@ def test_molmoact2_rollout_generator_uses_eval_seed_per_task():
         batch_size=3,
         device=torch.device("cpu"),
     )
-    expected_new_task = torch.Generator().manual_seed(
-        MolmoAct2Policy._combine_rollout_seeds(first_seed=1000, batch_size=3)
-    )
+    expected_new_task = torch.Generator().manual_seed(_combine_rollout_seeds(first_seed=1000, batch_size=3))
     assert torch.allclose(torch.rand(4, generator=new_task), torch.rand(4, generator=expected_new_task))
 
 
@@ -537,36 +534,26 @@ def test_train_action_expert_only_requires_continuous_action_mode():
 
 
 def test_molmoact2_sequence_length_is_inferred_from_fixed_token_budget():
-    cfg = MolmoAct2Config(
-        action_mode="both",
-        chunk_size=10,
-        n_action_steps=10,
-        image_keys=["observation.images.image", "observation.images.wrist_image"],
-        input_features={OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(8,))},
-        output_features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(7,))},
-    )
-
-    assert cfg.max_sequence_length is None
-    assert cfg.inferred_max_sequence_length() == 640
-    assert cfg.inferred_max_sequence_length(include_discrete_action=False) == 576
     assert (
         infer_molmoact2_max_sequence_length(
-            num_images=2,
-            state_dim=8,
-            action_dim=7,
-            action_horizon=30,
-            include_discrete_action=True,
+            num_images=2, state_dim=8, action_dim=7, action_horizon=10, include_discrete_action=True
+        )
+        == 640
+    )
+    assert (
+        infer_molmoact2_max_sequence_length(
+            num_images=2, state_dim=8, action_dim=7, action_horizon=10, include_discrete_action=False
+        )
+        == 576
+    )
+    assert (
+        infer_molmoact2_max_sequence_length(
+            num_images=2, state_dim=8, action_dim=7, action_horizon=30, include_discrete_action=True
         )
         == 768
     )
 
 
-def test_molmoact2_sequence_length_override_is_preserved():
-    cfg = MolmoAct2Config(max_sequence_length=1024)
-
-    assert cfg.inferred_max_sequence_length(num_images=2, state_dim=8, action_dim=7) == 1024
-
-
 def test_train_action_expert_only_freezes_non_action_expert_params():
     class DummyBackbone(torch.nn.Module):
         def __init__(self):
@@ -963,7 +950,7 @@ def test_action_dim_padding_loss_reduces_like_old_trainer():
         ]
     )
 
-    reduced = MolmoAct2Policy._apply_action_dim_padding_mask(loss, action_dim_is_pad)
+    reduced = _apply_action_dim_padding_mask(loss, action_dim_is_pad)
 
     expected = torch.stack(
         [
@@ -979,7 +966,7 @@ def test_action_chunk_padding_keeps_old_mean_denominator():
     loss = torch.ones(1, 2, 4, 3)
     action_horizon_is_pad = torch.tensor([[False, False, True, True]])
 
-    masked = MolmoAct2Policy._apply_action_chunk_padding_mask(loss, action_horizon_is_pad)
+    masked = _apply_action_chunk_padding_mask(loss, action_horizon_is_pad)
 
     assert masked.mean().item() == 0.5
 

tests/policies/test_policies.py

@@ -23,6 +23,7 @@ import torch
 
 pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
 
+from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
 from packaging import version
 from safetensors.torch import load_file
 
@@ -300,6 +301,29 @@ def test_save_and_load_pretrained(dummy_dataset_metadata, tmp_path, policy_name:
     torch.testing.assert_close(list(policy.parameters()), list(loaded_policy.parameters()), rtol=0, atol=0)
 
 
+def test_save_pretrained_with_state_dict(dummy_dataset_metadata, tmp_path):
+    """Exercise the FSDP checkpoint path: save_pretrained with a pre-gathered state_dict."""
+    policy_cls = get_policy_class("act")
+    policy_cfg = make_policy_config("act")
+    features = dataset_to_policy_features(dummy_dataset_metadata.features)
+    policy_cfg.output_features = {key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION}
+    policy_cfg.input_features = {
+        key: ft for key, ft in features.items() if key not in policy_cfg.output_features
+    }
+    policy = policy_cls(policy_cfg)
+    policy.to(policy_cfg.device)
+
+    save_dir = tmp_path / "fsdp_state_dict"
+    policy.save_pretrained(save_dir, state_dict=policy.state_dict())
+
+    # A single, unsharded safetensors file (no sharded set + index).
+    assert (save_dir / SAFETENSORS_SINGLE_FILE).is_file()
+    assert not (save_dir / f"{SAFETENSORS_SINGLE_FILE}.index.json").exists()
+
+    loaded_policy = policy_cls.from_pretrained(save_dir, config=policy_cfg)
+    torch.testing.assert_close(list(policy.parameters()), list(loaded_policy.parameters()), rtol=0, atol=0)
+
+
 @pytest.mark.parametrize("multikey", [True, False])
 def test_multikey_construction(multikey: bool):
     """

tests/processor/test_pipeline.py

@@ -2370,14 +2370,32 @@ def test_aggregate_images_when_use_videos_false():
     out = aggregate_pipeline_dataset_features(
         pipeline=rp,
         initial_features={PipelineFeatureType.ACTION: {}, PipelineFeatureType.OBSERVATION: initial},
-        use_videos=False,  # expect "image" dtype
+        use_videos=False,  # images kept, stored as "image" dtype
         patterns=None,
     )
 
     key = f"{OBS_IMAGES}.back"
     key_front = f"{OBS_IMAGES}.front"
-    assert key not in out
-    assert key_front not in out
+    assert key in out
+    assert key_front in out
+    assert out[key]["dtype"] == "image"
+    assert out[key_front]["dtype"] == "image"
+    assert out[key]["shape"] == initial["back"]
+
+
+def test_aggregate_images_excluded():
+    rp = DataProcessorPipeline([AddObservationStateFeatures(add_front_image=True)])
+    initial = {"back": (480, 640, 3)}
+
+    out = aggregate_pipeline_dataset_features(
+        pipeline=rp,
+        initial_features={PipelineFeatureType.ACTION: {}, PipelineFeatureType.OBSERVATION: initial},
+        exclude_images=True,
+        patterns=None,
+    )
+
+    assert f"{OBS_IMAGES}.back" not in out
+    assert f"{OBS_IMAGES}.front" not in out
 
 
 def test_aggregate_images_when_use_videos_true():

tests/scripts/test_edit_dataset_parsing.py

@@ -27,6 +27,7 @@ from lerobot.scripts.lerobot_edit_dataset import (
     MergeConfig,
     ModifyTasksConfig,
     OperationConfig,
+    ReencodeVideosConfig,
     RemoveFeatureConfig,
     SplitConfig,
     _validate_config,
@@ -66,6 +67,20 @@ class TestOperationTypeParsing:
         with pytest.raises(ValueError, match="--new_repo_id is required for merge"):
             _validate_config(cfg)
 
+    @pytest.mark.parametrize("flag", ["concatenate_videos", "concatenate_data"])
+    def test_merge_concatenate_flag_defaults_true(self, flag):
+        cfg = parse_cfg(["--new_repo_id", "test/merged", "--operation.type", "merge"])
+        assert isinstance(cfg.operation, MergeConfig)
+        assert getattr(cfg.operation, flag) is True
+
+    @pytest.mark.parametrize("flag", ["concatenate_videos", "concatenate_data"])
+    def test_merge_concatenate_flag_can_be_disabled(self, flag):
+        cfg = parse_cfg(
+            ["--new_repo_id", "test/merged", "--operation.type", "merge", f"--operation.{flag}", "false"]
+        )
+        assert isinstance(cfg.operation, MergeConfig)
+        assert getattr(cfg.operation, flag) is False
+
     def test_non_merge_requires_repo_id(self):
         cfg = parse_cfg(["--operation.type", "delete_episodes"])
         with pytest.raises(ValueError, match="--repo_id is required for delete_episodes"):
@@ -89,3 +104,47 @@ class TestOperationTypeParsing:
         )
         resolved_name = OperationConfig.get_choice_name(type(cfg.operation))
         assert resolved_name == type_name
+
+
+class TestDepthEncoderParsing:
+    """Test that the depth encoder is exposed and parsed for video operations."""
+
+    def test_reencode_has_default_depth_encoder(self):
+        cfg = parse_cfg(["--repo_id", "test/repo", "--operation.type", "reencode_videos"])
+        assert isinstance(cfg.operation, ReencodeVideosConfig)
+        # A depth encoder is configured by default so depth videos are re-encoded too.
+        assert cfg.operation.depth_encoder is not None
+        assert hasattr(cfg.operation.depth_encoder, "depth_min")
+
+    def test_reencode_parses_depth_encoder_overrides(self):
+        cfg = parse_cfg(
+            [
+                "--repo_id",
+                "test/repo",
+                "--operation.type",
+                "reencode_videos",
+                "--operation.depth_encoder.extra_options",
+                '{"x265-params": "lossless=1"}',
+                "--operation.depth_encoder.depth_max",
+                "12.0",
+                "--operation.depth_encoder.use_log",
+                "false",
+            ]
+        )
+        assert cfg.operation.depth_encoder.extra_options == {"x265-params": "lossless=1"}
+        assert cfg.operation.depth_encoder.depth_max == 12.0
+        assert cfg.operation.depth_encoder.use_log is False
+
+    def test_convert_image_to_video_parses_depth_encoder_overrides(self):
+        cfg = parse_cfg(
+            [
+                "--repo_id",
+                "test/repo",
+                "--operation.type",
+                "convert_image_to_video",
+                "--operation.depth_encoder.depth_min",
+                "0.05",
+            ]
+        )
+        assert isinstance(cfg.operation, ConvertImageToVideoConfig)
+        assert cfg.operation.depth_encoder.depth_min == 0.05

tests/scripts/test_lerobot_annotate.py

@@ -0,0 +1,86 @@
+#!/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 json
+from types import SimpleNamespace
+
+import pytest
+
+# ``lerobot.scripts.lerobot_annotate`` (and the ``_push_to_hub`` path it
+# exercises) imports ``lerobot.datasets``, which only ships under the
+# ``dataset`` extra. Skip in tiers without it instead of erroring.
+pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
+
+
+def test_push_to_hub_tags_uploaded_dataset_revision(tmp_path, monkeypatch):
+    from lerobot.scripts.lerobot_annotate import _push_to_hub
+
+    root = tmp_path / "dataset"
+    (root / "meta").mkdir(parents=True)
+    (root / "meta" / "info.json").write_text(json.dumps({"codebase_version": "v3.0"}))
+
+    calls = {}
+
+    class FakeHfApi:
+        def create_repo(self, **kwargs):
+            calls["create_repo"] = kwargs
+
+        def upload_folder(self, **kwargs):
+            calls["upload_folder"] = kwargs
+            return SimpleNamespace(oid="abc123")
+
+        def delete_tag(self, repo_id, **kwargs):
+            import requests
+            from huggingface_hub.errors import RevisionNotFoundError
+
+            calls["delete_tag"] = {"repo_id": repo_id, **kwargs}
+            # Simulate the common case: no stale tag to delete.
+            raise RevisionNotFoundError("no such tag", response=requests.Response())
+
+        def create_tag(self, **kwargs):
+            calls["create_tag"] = kwargs
+
+    monkeypatch.setattr("huggingface_hub.HfApi", FakeHfApi)
+
+    cfg = SimpleNamespace(
+        repo_id="source/dataset",
+        new_repo_id="annotated/dataset",
+        push_private=True,
+        push_commit_message=None,
+    )
+
+    _push_to_hub(root, cfg)
+
+    assert calls["create_repo"] == {
+        "repo_id": "annotated/dataset",
+        "repo_type": "dataset",
+        "private": True,
+        "exist_ok": True,
+    }
+    assert calls["upload_folder"]["repo_id"] == "annotated/dataset"
+    # A stale tag (e.g. from a previous annotation run) is deleted first so
+    # the new tag always points at the upload we just made.
+    assert calls["delete_tag"] == {
+        "repo_id": "annotated/dataset",
+        "tag": "v3.0",
+        "repo_type": "dataset",
+    }
+    assert calls["create_tag"] == {
+        "repo_id": "annotated/dataset",
+        "tag": "v3.0",
+        "repo_type": "dataset",
+        "revision": "abc123",
+    }

tests/teleoperators/test_rebot_102_leader.py

@@ -18,7 +18,7 @@ from unittest.mock import MagicMock, patch
 
 import pytest
 
-from lerobot.teleoperators.bi_rebot_102_leader import BiRebotArm102Leader, BiRebotArm102LeaderConfig
+from lerobot.teleoperators.bi_rebot_102_leader import BiRebot102Leader, BiRebot102LeaderConfig
 from lerobot.teleoperators.rebot_102_leader import (
     RebotArm102Leader,
     RebotArm102LeaderConfig,
@@ -91,11 +91,11 @@ def test_send_feedback_not_implemented(leader):
 
 def test_bimanual_prefixes_features():
     with patch(f"{_MODULE}.require_package", lambda *a, **kw: None):
-        cfg = BiRebotArm102LeaderConfig(
+        cfg = BiRebot102LeaderConfig(
             left_arm_config=RebotArm102LeaderConfig(port="/dev/null0"),
             right_arm_config=RebotArm102LeaderConfig(port="/dev/null1"),
         )
-        teleop = BiRebotArm102Leader(cfg)
+        teleop = BiRebot102Leader(cfg)
     assert any(k.startswith("left_") for k in teleop.action_features)
     assert any(k.startswith("right_") for k in teleop.action_features)
     assert "left_gripper.pos" in teleop.action_features

tests/training/test_multi_gpu.py

@@ -58,7 +58,46 @@ def download_dataset(repo_id, episodes):
     print(f"Dataset {repo_id} downloaded successfully")
 
 
-def run_accelerate_training(config_args, num_processes=4, temp_dir=None):
+def _write_multi_gpu_config(f, num_processes):
+    f.write("compute_environment: LOCAL_MACHINE\n")
+    f.write("distributed_type: MULTI_GPU\n")
+    f.write("mixed_precision: 'no'\n")
+    f.write(f"num_processes: {num_processes}\n")
+    f.write("use_cpu: false\n")
+    f.write("gpu_ids: all\n")
+    f.write("downcast_bf16: 'no'\n")
+    f.write("machine_rank: 0\n")
+    f.write("main_training_function: main\n")
+    f.write("num_machines: 1\n")
+    f.write("rdzv_backend: static\n")
+    f.write("same_network: true\n")
+
+
+def _write_fsdp_config(f, num_processes):
+    # FSDP1 with FULL_SHARD (ZeRO-3-equivalent) and FULL_STATE_DICT, matching
+    # docs/source/multi_gpu_training.mdx. ACT's repeated transformer blocks are the wrap units;
+    # fsdp_use_orig_params is required because LeRobot builds the optimizer before prepare().
+    f.write("compute_environment: LOCAL_MACHINE\n")
+    f.write("distributed_type: FSDP\n")
+    f.write("mixed_precision: 'no'\n")
+    f.write(f"num_processes: {num_processes}\n")
+    f.write("use_cpu: false\n")
+    f.write("gpu_ids: all\n")
+    f.write("machine_rank: 0\n")
+    f.write("main_training_function: main\n")
+    f.write("num_machines: 1\n")
+    f.write("rdzv_backend: static\n")
+    f.write("same_network: true\n")
+    f.write("fsdp_config:\n")
+    f.write("  fsdp_version: 1\n")
+    f.write("  fsdp_sharding_strategy: FULL_SHARD\n")
+    f.write("  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP\n")
+    f.write("  fsdp_transformer_layer_cls_to_wrap: ACTEncoderLayer,ACTDecoderLayer\n")
+    f.write("  fsdp_use_orig_params: true\n")
+    f.write("  fsdp_state_dict_type: FULL_STATE_DICT\n")
+
+
+def run_accelerate_training(config_args, num_processes=4, temp_dir=None, distributed_type="MULTI_GPU"):
     """
     Helper function to run training with accelerate launch.
 
@@ -66,6 +105,7 @@ def run_accelerate_training(config_args, num_processes=4, temp_dir=None):
         config_args: List of config arguments to pass to lerobot_train.py
         num_processes: Number of processes (GPUs) to use
         temp_dir: Temporary directory for outputs
+        distributed_type: "MULTI_GPU" (DDP) or "FSDP" — selects the generated accelerate config.
 
     Returns:
         subprocess.CompletedProcess result
@@ -75,18 +115,10 @@ def run_accelerate_training(config_args, num_processes=4, temp_dir=None):
 
     # Write YAML config
     with open(config_path, "w") as f:
-        f.write("compute_environment: LOCAL_MACHINE\n")
-        f.write("distributed_type: MULTI_GPU\n")
-        f.write("mixed_precision: 'no'\n")
-        f.write(f"num_processes: {num_processes}\n")
-        f.write("use_cpu: false\n")
-        f.write("gpu_ids: all\n")
-        f.write("downcast_bf16: 'no'\n")
-        f.write("machine_rank: 0\n")
-        f.write("main_training_function: main\n")
-        f.write("num_machines: 1\n")
-        f.write("rdzv_backend: static\n")
-        f.write("same_network: true\n")
+        if distributed_type == "FSDP":
+            _write_fsdp_config(f, num_processes)
+        else:
+            _write_multi_gpu_config(f, num_processes)
 
     cmd = [
         "accelerate",
@@ -134,7 +166,7 @@ class TestMultiGPUTraining:
                 f"--output_dir={output_dir}",
                 "--batch_size=4",
                 "--steps=10",
-                "--eval_freq=-1",
+                "--env_eval_freq=-1",
                 "--log_freq=5",
                 "--save_freq=10",
                 "--seed=42",
@@ -177,7 +209,7 @@ class TestMultiGPUTraining:
                 f"--output_dir={output_dir}",
                 "--batch_size=4",
                 "--steps=20",
-                "--eval_freq=-1",
+                "--env_eval_freq=-1",
                 "--log_freq=5",
                 "--save_freq=10",
                 "--seed=42",
@@ -211,3 +243,66 @@ class TestMultiGPUTraining:
                 # Verify optimizer state exists
                 optimizer_state = training_state_dir / "optimizer_state.safetensors"
                 assert optimizer_state.exists(), f"No optimizer state in checkpoint {checkpoint_dir}"
+
+    def test_fsdp_optimizer_save_and_resume(self):
+        """
+        Test that FSDP saves the (gathered) optimizer state and can resume from it.
+
+        Trains a few steps under FSDP, verifies the gathered optimizer state is written next to the
+        rest of the training state, then resumes from the checkpoint for more steps and checks it
+        completes without shape/key errors in the FSDP optimizer load path.
+        """
+        # Pre-download dataset to avoid race conditions
+        download_dataset("lerobot/pusht", episodes=[0])
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            output_dir = Path(temp_dir) / "outputs"
+
+            config_args = [
+                "--dataset.repo_id=lerobot/pusht",
+                "--dataset.episodes=[0]",
+                "--policy.type=act",
+                "--policy.device=cuda",
+                "--policy.push_to_hub=false",
+                f"--output_dir={output_dir}",
+                "--batch_size=4",
+                "--steps=10",
+                "--env_eval_freq=-1",
+                "--log_freq=5",
+                "--save_freq=10",
+                "--seed=42",
+                "--num_workers=0",
+            ]
+
+            result = run_accelerate_training(
+                config_args, num_processes=2, temp_dir=temp_dir, distributed_type="FSDP"
+            )
+            assert result.returncode == 0, (
+                f"FSDP training failed:\nSTDOUT:\n{result.stdout}\n\nSTDERR:\n{result.stderr}"
+            )
+
+            # The gathered optimizer state must be written under FSDP (proves the save collective ran),
+            # in the same safetensors format as single-GPU training.
+            training_state_dir = output_dir / "checkpoints" / "last" / "training_state"
+            optimizer_state = training_state_dir / "optimizer_state.safetensors"
+            optimizer_param_groups = training_state_dir / "optimizer_param_groups.json"
+            assert optimizer_state.exists(), f"FSDP optimizer state not saved in {training_state_dir}"
+            assert optimizer_param_groups.exists(), (
+                f"FSDP optimizer param groups not saved in {training_state_dir}"
+            )
+
+            # Resume from the checkpoint for more steps. A successful run proves load_fsdp_optimizer
+            # accepts the saved state and reshards it without shape/key errors.
+            resume_config = output_dir / "checkpoints" / "last" / "pretrained_model" / "train_config.json"
+            resume_args = [
+                f"--config_path={resume_config}",
+                "--resume=true",
+                "--steps=20",
+            ]
+            resume_result = run_accelerate_training(
+                resume_args, num_processes=2, temp_dir=temp_dir, distributed_type="FSDP"
+            )
+            assert resume_result.returncode == 0, (
+                f"FSDP resume failed:\nSTDOUT:\n{resume_result.stdout}\n\nSTDERR:\n{resume_result.stderr}"
+            )
+            assert "End of training" in resume_result.stdout or "End of training" in resume_result.stderr

tests/utils/test_keyboard_input.py

@@ -0,0 +1,228 @@
+#!/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.
+
+"""Unit tests for the display-independent keyboard input helpers.
+
+These cover the parts most likely to regress: the environment-detection decision
+table (the heart of the Wayland/headless fix), the macOS trust probe, the control
+mapping, the terminal escape-sequence parsing, and backend selection. They require
+neither ``pynput`` nor a real terminal.
+"""
+
+import io
+import platform
+import sys
+
+import pytest
+
+import lerobot.utils.keyboard_input as ki
+from lerobot.utils.keyboard_input import (
+    TerminalKeyListener,
+    apply_recording_control,
+    create_key_listener,
+    init_keyboard_listener,
+    is_headless,
+    is_wayland,
+    pynput_can_capture,
+    pynput_listener_is_trusted,
+)
+
+
+@pytest.fixture(autouse=True)
+def _clear_detection_caches():
+    """The detection helpers are ``@cache``-decorated; clear around each test."""
+    for fn in (is_headless, is_wayland, pynput_can_capture):
+        fn.cache_clear()
+    yield
+    for fn in (is_headless, is_wayland, pynput_can_capture):
+        fn.cache_clear()
+
+
+def _set_platform(monkeypatch, name):
+    monkeypatch.setattr(platform, "system", lambda: name)
+
+
+def _set_tty(monkeypatch, is_tty):
+    stdin = io.StringIO("")
+    stdin.isatty = lambda: is_tty
+    monkeypatch.setattr(sys, "stdin", stdin)
+
+
+# --- Environment detection (the core of the fix) ---------------------------
+@pytest.mark.parametrize(
+    ("system", "env", "expected"),
+    [
+        ("Linux", {}, True),  # no display server
+        ("Linux", {"DISPLAY": ":0"}, False),  # X11
+        ("Linux", {"WAYLAND_DISPLAY": "wayland-0"}, False),  # Wayland
+        ("Darwin", {}, False),  # display always assumed present
+    ],
+)
+def test_is_headless(monkeypatch, system, env, expected):
+    _set_platform(monkeypatch, system)
+    monkeypatch.delenv("DISPLAY", raising=False)
+    monkeypatch.delenv("WAYLAND_DISPLAY", raising=False)
+    for key, value in env.items():
+        monkeypatch.setenv(key, value)
+    assert is_headless() is expected
+
+
+@pytest.mark.parametrize(
+    ("env", "expected"),
+    [
+        ({"XDG_SESSION_TYPE": "wayland"}, True),
+        ({"WAYLAND_DISPLAY": "wayland-0"}, True),
+        ({"XDG_SESSION_TYPE": "x11"}, False),
+        ({}, False),
+    ],
+)
+def test_is_wayland(monkeypatch, env, expected):
+    monkeypatch.delenv("XDG_SESSION_TYPE", raising=False)
+    monkeypatch.delenv("WAYLAND_DISPLAY", raising=False)
+    for key, value in env.items():
+        monkeypatch.setenv(key, value)
+    assert is_wayland() is expected
+
+
+@pytest.mark.parametrize(
+    ("system", "env", "pynput_available", "expected"),
+    [
+        ("Linux", {"DISPLAY": ":0"}, True, True),  # X11
+        ("Linux", {"DISPLAY": ":0", "WAYLAND_DISPLAY": "wayland-0"}, True, False),  # Wayland
+        ("Linux", {}, True, False),  # headless
+        ("Darwin", {}, True, True),
+        ("Linux", {"DISPLAY": ":0"}, False, False),  # pynput not installed
+    ],
+)
+def test_pynput_can_capture(monkeypatch, system, env, pynput_available, expected):
+    _set_platform(monkeypatch, system)
+    monkeypatch.setattr(ki, "_pynput_available", pynput_available)
+    for var in ("DISPLAY", "WAYLAND_DISPLAY", "XDG_SESSION_TYPE"):
+        monkeypatch.delenv(var, raising=False)
+    for key, value in env.items():
+        monkeypatch.setenv(key, value)
+    assert pynput_can_capture() is expected
+
+
+# --- macOS trust probe ------------------------------------------------------
+class _FakeListener:
+    def __init__(self, is_trusted):
+        self.IS_TRUSTED = is_trusted
+
+
+def test_pynput_listener_is_trusted(monkeypatch):
+    _set_platform(monkeypatch, "Linux")
+    assert pynput_listener_is_trusted(_FakeListener(False)) is True  # non-macOS: always assumed ok
+    _set_platform(monkeypatch, "Darwin")
+    assert pynput_listener_is_trusted(_FakeListener(False), timeout_s=0.05) is False
+
+
+# --- Control mapping --------------------------------------------------------
+def test_apply_recording_control():
+    events = {"exit_early": False, "rerecord_episode": False, "stop_recording": False}
+    apply_recording_control("left", events)
+    assert events == {"exit_early": True, "rerecord_episode": True, "stop_recording": False}
+    apply_recording_control("esc", events)
+    assert events["stop_recording"] is True
+    apply_recording_control("up", events)  # unknown control -> no-op (no error)
+
+
+# --- Terminal escape-sequence parsing (the tricky bit) ----------------------
+def _drive(listener, byte_seq):
+    """Run the listener's read loop over a scripted list of bytes (no real terminal)."""
+    script = list(byte_seq)
+
+    def fake_read(timeout):
+        if script:
+            return script.pop(0)
+        listener._running = False
+        return None
+
+    listener._read_char = fake_read
+    listener._running = True
+    listener._run()
+
+
+@pytest.mark.parametrize(
+    ("byte_seq", "expected"),
+    [
+        (["\x1b", "[", "C"], ["right"]),  # CSI arrow
+        (["\x1b", "O", "D"], ["left"]),  # SS3 arrow (e.g. over SSH/tmux)
+        (["\x1b"], ["esc"]),  # bare ESC
+        (["\x1b", "[", "A"], ["up"]),  # decoded even though the record handler ignores it
+        (["n"], ["n"]),  # letter passthrough
+    ],
+)
+def test_terminal_parsing(byte_seq, expected):
+    collected = []
+    _drive(TerminalKeyListener(collected.append), byte_seq)
+    assert collected == expected
+
+
+# --- Backend selection ------------------------------------------------------
+def test_init_selects_terminal_when_pynput_cannot_capture(monkeypatch):
+    monkeypatch.setattr(ki, "pynput_can_capture", lambda: False)
+    _set_tty(monkeypatch, is_tty=True)
+    monkeypatch.setattr(TerminalKeyListener, "start", lambda self: None)  # avoid touching termios
+    listener, _ = init_keyboard_listener()
+    assert isinstance(listener, TerminalKeyListener)
+
+
+def test_init_returns_none_without_tty(monkeypatch):
+    monkeypatch.setattr(ki, "pynput_can_capture", lambda: False)
+    _set_tty(monkeypatch, is_tty=False)
+    listener, _ = init_keyboard_listener()
+    assert listener is None
+
+
+@pytest.mark.parametrize(
+    ("key", "flag"),
+    [("right", "exit_early"), ("r", "rerecord_episode"), ("q", "stop_recording")],
+)
+def test_init_terminal_key_routing(monkeypatch, key, flag):
+    """Arrows and their letter equivalents drive the same events (terminal backend)."""
+    monkeypatch.setattr(ki, "pynput_can_capture", lambda: False)
+    _set_tty(monkeypatch, is_tty=True)
+    monkeypatch.setattr(TerminalKeyListener, "start", lambda self: None)
+    listener, events = init_keyboard_listener()
+    listener._on_key(key)
+    assert events[flag] is True
+
+
+# --- Shared factory + pynput key resolver -----------------------------------
+def test_resolve_pynput_key_char_fallback():
+    """Unmapped keys fall back to ``.char`` (and yield None when there is none)."""
+    assert ki._resolve_pynput_key(type("K", (), {"char": "s"})()) == "s"
+    assert ki._resolve_pynput_key(type("K", (), {"char": None})()) is None
+    assert ki._resolve_pynput_key(type("K", (), {"char": ""})()) is None  # empty char -> no key
+
+
+def test_create_key_listener_routes_to_dispatch(monkeypatch):
+    """The terminal backend forwards canonical key names straight to ``dispatch``."""
+    monkeypatch.setattr(ki, "pynput_can_capture", lambda: False)
+    _set_tty(monkeypatch, is_tty=True)
+    monkeypatch.setattr(TerminalKeyListener, "start", lambda self: None)
+    seen = []
+    listener = create_key_listener(seen.append, controls_help="save='s'")
+    assert isinstance(listener, TerminalKeyListener)
+    listener._on_key("space")
+    assert seen == ["space"]
+
+
+def test_create_key_listener_none_without_tty(monkeypatch):
+    monkeypatch.setattr(ki, "pynput_can_capture", lambda: False)
+    _set_tty(monkeypatch, is_tty=False)
+    assert create_key_listener(lambda name: None) is None

tests/utils/test_logging_utils.py

@@ -15,6 +15,7 @@
 # limitations under the License.
 
 import pytest
+import torch
 
 from lerobot.utils.logging_utils import AverageMeter, MetricsTracker
 
@@ -25,8 +26,16 @@ def mock_metrics():
 
 
 class MockAccelerator:
-    def __init__(self, num_processes: int):
+    def __init__(self, num_processes: int, reduce_fn=None):
         self.num_processes = num_processes
+        self.device = torch.device("cpu")
+        self._reduce_fn = reduce_fn
+
+    def reduce(self, tensor, reduction="mean"):
+        # In single-process tests we just want a deterministic stand-in for accelerate's reduce.
+        if self._reduce_fn is not None:
+            return self._reduce_fn(tensor, reduction)
+        return tensor
 
 
 def test_average_meter_initialization():
@@ -157,3 +166,70 @@ def test_metrics_tracker_reset_averages(mock_metrics):
     tracker.reset_averages()
     assert tracker.loss.avg == 0.0
     assert tracker.accuracy.avg == 0.0
+
+
+def test_average_meter_invalid_reduction():
+    with pytest.raises(ValueError):
+        AverageMeter("loss", reduction="median")
+
+
+def test_average_meter_reduction_stored():
+    meter = AverageMeter("updt_s", reduction="max")
+    assert meter.reduction == "max"
+
+
+def test_metrics_tracker_reduce_across_ranks_no_accelerator():
+    metrics = {"update_s": AverageMeter("update_s", reduction="max")}
+    tracker = MetricsTracker(batch_size=32, num_frames=1000, num_episodes=50, metrics=metrics)
+    tracker.update_s = 0.5
+    tracker.reduce_across_ranks()  # no-op without accelerator
+    assert tracker.update_s.avg == 0.5
+
+
+def test_metrics_tracker_reduce_across_ranks_single_process():
+    metrics = {"update_s": AverageMeter("update_s", reduction="max")}
+    tracker = MetricsTracker(
+        batch_size=32,
+        num_frames=1000,
+        num_episodes=50,
+        metrics=metrics,
+        accelerator=MockAccelerator(num_processes=1),
+    )
+    tracker.update_s = 0.5
+    tracker.reduce_across_ranks()  # no-op when world size is 1
+    assert tracker.update_s.avg == 0.5
+
+
+def test_metrics_tracker_reduce_across_ranks_invokes_reduce():
+    captured = {}
+
+    def fake_reduce(tensor, reduction):
+        captured["reduction"] = reduction
+        captured["values"] = tensor.clone()
+        # Pretend the slowest rank reported 0.9 instead of this rank's 0.4.
+        return torch.tensor([0.9], dtype=tensor.dtype, device=tensor.device)
+
+    metrics = {
+        "loss": AverageMeter("loss"),  # reduction="none" -> not touched
+        "update_s": AverageMeter("update_s", reduction="max"),
+    }
+    tracker = MetricsTracker(
+        batch_size=32,
+        num_frames=1000,
+        num_episodes=50,
+        metrics=metrics,
+        accelerator=MockAccelerator(num_processes=4, reduce_fn=fake_reduce),
+    )
+    tracker.loss = 1.0
+    tracker.update_s = 0.4
+    tracker.reduce_across_ranks()
+
+    assert captured["reduction"] == "max"
+    assert torch.allclose(captured["values"], torch.tensor([0.4]))
+    assert tracker.update_s.avg == pytest.approx(0.9)
+    # Metrics without a reduction stay untouched.
+    assert tracker.loss.avg == 1.0
+    # Invariant: avg == sum / count must hold after reduce, so subsequent .update() calls
+    # accumulate against the cluster view rather than the stale per-rank sum.
+    meter = tracker.update_s
+    assert meter.sum / meter.count == pytest.approx(meter.avg)

tests/utils/test_train_utils.py

@@ -20,6 +20,8 @@ from unittest.mock import Mock, patch
 from lerobot.common.train_utils import (
     get_step_checkpoint_dir,
     get_step_identifier,
+    load_training_batch_size,
+    load_training_num_processes,
     load_training_state,
     load_training_step,
     save_checkpoint,
@@ -63,6 +65,28 @@ def test_load_training_step(tmp_path):
     assert loaded_step == step
 
 
+def test_save_training_state_records_num_processes(tmp_path, optimizer, scheduler):
+    save_training_state(tmp_path, 10, optimizer, scheduler, num_processes=4)
+    assert load_training_num_processes(tmp_path) == 4
+
+
+def test_load_training_num_processes_absent_returns_none(tmp_path, optimizer, scheduler):
+    # Checkpoints written before the world size was recorded must still load (back-compat).
+    save_training_state(tmp_path, 10, optimizer, scheduler)
+    assert load_training_num_processes(tmp_path) is None
+
+
+def test_save_training_state_records_batch_size(tmp_path, optimizer, scheduler):
+    save_training_state(tmp_path, 10, optimizer, scheduler, batch_size=32)
+    assert load_training_batch_size(tmp_path) == 32
+
+
+def test_load_training_batch_size_absent_returns_none(tmp_path, optimizer, scheduler):
+    # Checkpoints written before the batch size was recorded must still load (back-compat).
+    save_training_state(tmp_path, 10, optimizer, scheduler)
+    assert load_training_batch_size(tmp_path) is None
+
+
 def test_update_last_checkpoint(tmp_path):
     checkpoint = tmp_path / "0005"
     checkpoint.mkdir()
@@ -112,3 +136,18 @@ def test_save_load_training_state(tmp_path, optimizer, scheduler):
     assert loaded_step == 10
     assert loaded_optimizer is optimizer
     assert loaded_scheduler is scheduler
+
+
+def test_load_training_state_skip_optimizer(tmp_path, optimizer, scheduler):
+    # FSDP loads optimizer separately (after accelerator.prepare)
+    # load_training_state(load_optimizer=False) must restore step + scheduler but leave the
+    # optimizer untouched and never touch the on-disk optimizer state.
+    save_training_state(tmp_path, 10, optimizer, scheduler)
+    with patch("lerobot.common.train_utils.load_optimizer_state") as mock_load_optimizer_state:
+        loaded_step, loaded_optimizer, loaded_scheduler = load_training_state(
+            tmp_path, optimizer, scheduler, load_optimizer=False
+        )
+    mock_load_optimizer_state.assert_not_called()
+    assert loaded_step == 10
+    assert loaded_optimizer is optimizer
+    assert loaded_scheduler is scheduler

tests/utils/test_visualization_utils.py

@@ -48,6 +48,11 @@ def mock_rerun(monkeypatch):
 
         def compress(self, *a, **k):
             return self
+    
+    class DummyDepthImage:
+        def __init__(self, arr, colormap=None):
+            self.arr = arr
+            self.colormap = colormap
 
     def dummy_log(key, obj=None, **kwargs):
         # Accept either positional `obj` or keyword `entity` and record remaining kwargs.
@@ -76,6 +81,8 @@ def mock_rerun(monkeypatch):
         __spec__=SimpleNamespace(name="rerun", submodule_search_locations=None),
         Scalars=DummyScalar,
         Image=DummyImage,
+        DepthImage=DummyDepthImage,
+        components=SimpleNamespace(Colormap=SimpleNamespace(Viridis="viridis")),
         log=dummy_log,
         send_blueprint=dummy_send_blueprint,
         init=lambda *a, **k: None,
@@ -272,7 +279,7 @@ def test_log_rerun_data_kwargs_only(mock_rerun):
     assert float(temp.value) == pytest.approx(10.0)
 
     img = _obj_for(calls, "observation.gray")
-    assert type(img).__name__ == "DummyImage"
+    assert type(img).__name__ == "DummyDepthImage"  # single-channel -> DepthImage
     assert img.arr.shape == (8, 8, 1)  # remains HWC
     assert _kwargs_for(calls, "observation.gray").get("static", False) is True