mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-07 01:51:47 +00:00
Merge remote-tracking branch 'origin/main' into feat/smolvla-on-steerable
# Conflicts: # docs/source/annotation_pipeline.mdx # examples/annotations/run_hf_job.py # pyproject.toml # src/lerobot/annotations/steerable_pipeline/config.py # src/lerobot/annotations/steerable_pipeline/frames.py # src/lerobot/annotations/steerable_pipeline/modules/plan_subtasks_memory.py # src/lerobot/annotations/steerable_pipeline/vlm_client.py # src/lerobot/annotations/steerable_pipeline/writer.py # src/lerobot/datasets/__init__.py # src/lerobot/datasets/sampler.py # src/lerobot/scripts/lerobot_annotate.py # src/lerobot/scripts/lerobot_train.py # tests/annotations/test_frames.py # tests/annotations/test_modules.py # tests/annotations/test_writer.py # tests/datasets/test_sampler.py # tests/scripts/test_lerobot_annotate.py # uv.lock
This commit is contained in:
Pepijn
@@ -38,19 +38,20 @@ import torch
|
||||
|
||||
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
|
||||
|
||||
from lerobot.annotations.steerable_pipeline.frames import ( # noqa: E402
|
||||
VideoFrameProvider,
|
||||
_decode_frames_av,
|
||||
_decode_frames_ffmpeg,
|
||||
)
|
||||
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]) -> None:
|
||||
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._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:
|
||||
@@ -124,15 +125,24 @@ def sample_video(tmp_path: Path) -> Path:
|
||||
return out
|
||||
|
||||
|
||||
def test_decode_frames_av_returns_one_uint8_frame_per_timestamp(sample_video: Path) -> None:
|
||||
"""``_decode_frames_av`` decodes via PyAV directly — no torchcodec/torchvision.
|
||||
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
|
||||
|
||||
This is the always-available fallback: torchcodec is unusable in some
|
||||
containers and lerobot's ``pyav`` backend routes through the removed
|
||||
``torchvision.io.VideoReader``.
|
||||
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 = _decode_frames_av(sample_video, timestamps)
|
||||
frames = provider._decode(0, timestamps, "observation.images.cam")
|
||||
|
||||
assert len(frames) == len(timestamps)
|
||||
for frame in frames:
|
||||
@@ -141,39 +151,96 @@ def test_decode_frames_av_returns_one_uint8_frame_per_timestamp(sample_video: Pa
|
||||
assert frame.shape == (3, 120, 160)
|
||||
|
||||
|
||||
def test_decode_frames_av_picks_nearest_frame(sample_video: Path) -> None:
|
||||
"""Repeated and out-of-order timestamps each resolve to the nearest frame."""
|
||||
frames = _decode_frames_av(sample_video, [2.0, 0.0, 2.0])
|
||||
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
|
||||
assert torch.equal(frames[0], frames[2])
|
||||
assert not torch.equal(frames[0], frames[1])
|
||||
|
||||
|
||||
def test_decode_frames_av_raises_on_missing_file(tmp_path: Path) -> None:
|
||||
"""A missing video surfaces as an exception the caller can fall back on."""
|
||||
with pytest.raises(Exception): # noqa: B017, PT011
|
||||
_decode_frames_av(tmp_path / "does_not_exist.mp4", [0.0])
|
||||
|
||||
|
||||
def test_decode_frames_ffmpeg_returns_one_uint8_frame_per_timestamp(sample_video: Path) -> None:
|
||||
"""``_decode_frames_ffmpeg`` shells out to the ffmpeg CLI — the always-
|
||||
available fallback that decodes AV1 and isolates crashes to a child
|
||||
process.
|
||||
"""
|
||||
timestamps = [0.0, 1.0, 2.5]
|
||||
frames = _decode_frames_ffmpeg(sample_video, timestamps)
|
||||
|
||||
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_decode_frames_ffmpeg_raises_on_missing_file(tmp_path: Path) -> None:
|
||||
"""A missing video raises (non-zero ffmpeg exit), never crashes the job."""
|
||||
if shutil.which("ffmpeg") is None:
|
||||
pytest.skip("ffmpeg not available")
|
||||
with pytest.raises(Exception): # noqa: B017, PT011
|
||||
_decode_frames_ffmpeg(tmp_path / "does_not_exist.mp4", [0.0])
|
||||
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,
|
||||
camera_encoder=None,
|
||||
overwrite=False,
|
||||
start_time_s=None,
|
||||
end_time_s=None,
|
||||
):
|
||||
captured.update(
|
||||
src=Path(input_video_path),
|
||||
encoder=camera_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
|
||||
|
||||
@@ -22,6 +22,7 @@ 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
|
||||
@@ -51,7 +52,10 @@ from ._helpers import make_canned_responder # noqa: E402
|
||||
class _StubFrameProvider:
|
||||
"""Returns one sentinel object per requested timestamp."""
|
||||
|
||||
sentinel: Any = field(default_factory=lambda: object())
|
||||
# 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)
|
||||
@@ -115,6 +119,34 @@ def test_module1_plan_memory_subtask_smoke(fixture_dataset_root: Path, tmp_path:
|
||||
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."}},
|
||||
@@ -236,8 +268,10 @@ def test_module3_vqa_unique_per_frame_and_camera(single_episode_root: Path, tmp_
|
||||
assert ts in frame_set
|
||||
|
||||
|
||||
def test_module1_attaches_video_block_to_subtask_prompt(fixture_dataset_root: Path, tmp_path: Path) -> None:
|
||||
"""Module 1 sends one ``type=video`` block covering the whole episode."""
|
||||
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": [
|
||||
@@ -265,7 +299,7 @@ def test_module1_attaches_video_block_to_subtask_prompt(fixture_dataset_root: Pa
|
||||
# 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(max_video_frames=5, frames_per_second=10.0, n_task_rephrasings=0),
|
||||
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))
|
||||
@@ -290,16 +324,14 @@ def test_module1_attaches_video_block_to_subtask_prompt(fixture_dataset_root: Pa
|
||||
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 len(video_blocks) == 1, f"expected exactly 1 video block, got {content}"
|
||||
assert image_blocks == [], "subtask prompt must not mix image blocks with the video block"
|
||||
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
|
||||
# video block must wrap a list of frames covering the episode
|
||||
assert isinstance(video_blocks[0]["video"], list)
|
||||
assert len(video_blocks[0]["video"]) <= 5
|
||||
# provider is called with target_count = min(duration * fps, max). With
|
||||
# fps=10 on a ~1s episode that requests >max, so max=5 wins.
|
||||
assert provider.video_calls and provider.video_calls[0][0] == record.episode_index
|
||||
assert provider.video_calls[0][1] <= 5
|
||||
# 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:
|
||||
|
||||
@@ -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
|
||||
@@ -28,6 +28,7 @@ import pytest
|
||||
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
|
||||
@@ -344,6 +345,78 @@ def test_annotation_metadata_sync_allows_non_streaming_load(
|
||||
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"
|
||||
|
||||
@@ -32,6 +32,26 @@ from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
||||
from tests.fixtures.constants import 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):
|
||||
"""Test that total number of episodes and frames are correctly aggregated."""
|
||||
assert aggr_ds.num_episodes == expected_episodes, (
|
||||
@@ -289,6 +309,52 @@ def test_aggregate_datasets(tmp_path, lerobot_dataset_factory):
|
||||
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
|
||||
@@ -520,6 +586,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.
|
||||
|
||||
|
||||
@@ -83,6 +83,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]]),
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -25,7 +25,7 @@ from datasets import Dataset # noqa: E402
|
||||
from lerobot.datasets.io_utils import (
|
||||
hf_transform_to_torch,
|
||||
)
|
||||
from lerobot.datasets.sampler import EpisodeAwareSampler, WeightedEpisodeAwareSampler
|
||||
from lerobot.datasets.sampler import EpisodeAwareSampler, WeightedEpisodeAwareSampler, compute_sampler_state
|
||||
|
||||
|
||||
def calculate_episode_data_index(hf_dataset: Dataset) -> dict[str, torch.Tensor]:
|
||||
@@ -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)
|
||||
@@ -183,3 +196,85 @@ def test_weighted_sampler_zero_weights_fall_back_to_uniform():
|
||||
def test_weighted_sampler_rejects_short_weight_vector():
|
||||
with pytest.raises(ValueError, match="frame_weights"):
|
||||
WeightedEpisodeAwareSampler([0], [10], frame_weights=torch.ones(5))
|
||||
|
||||
|
||||
# --- seeded (seed, epoch) shuffling, resume, and state ---
|
||||
|
||||
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,
|
||||
}
|
||||
|
||||
@@ -504,6 +504,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 = VideoEncoderConfig(vcodec="h264")
|
||||
reencode_video(src, out, camera_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):
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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):
|
||||
"""
|
||||
|
||||
@@ -24,6 +24,7 @@ from typing import Any
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from safetensors.torch import load_file
|
||||
|
||||
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
|
||||
|
||||
@@ -174,6 +175,53 @@ class MockStepWithTensorState(ProcessorStep):
|
||||
return features
|
||||
|
||||
|
||||
class MockLazyTensorStateStep(ProcessorStep):
|
||||
"""Mock step whose tensor state is not present in constructor config."""
|
||||
|
||||
def __init__(
|
||||
self, name: str = "lazy_tensor_step", scale: float = 1.0, initial_value: float | None = None
|
||||
):
|
||||
self.name = name
|
||||
self.scale = scale
|
||||
self.tensor_state: torch.Tensor | None = None
|
||||
|
||||
if initial_value is not None:
|
||||
self.tensor_state = torch.tensor([initial_value], dtype=torch.float32)
|
||||
|
||||
def __call__(self, transition: EnvTransition) -> EnvTransition:
|
||||
"""Return the transition unchanged."""
|
||||
return transition
|
||||
|
||||
def get_config(self) -> dict[str, Any]:
|
||||
"""Return constructor config while intentionally omitting tensor state."""
|
||||
return {
|
||||
"name": self.name,
|
||||
"scale": self.scale,
|
||||
}
|
||||
|
||||
def state_dict(self) -> dict[str, torch.Tensor]:
|
||||
"""Return tensor state only after it has been initialized or loaded."""
|
||||
if self.tensor_state is None:
|
||||
return {}
|
||||
|
||||
return {"tensor_state": self.tensor_state}
|
||||
|
||||
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
|
||||
"""Load tensor state."""
|
||||
self.tensor_state = state["tensor_state"].clone()
|
||||
|
||||
def transform_features(
|
||||
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||
"""Return features unchanged."""
|
||||
return features
|
||||
|
||||
|
||||
@ProcessorStepRegistry.register("registered_lazy_tensor_state_step")
|
||||
class RegisteredLazyTensorStateStep(MockLazyTensorStateStep):
|
||||
"""Registered lazy tensor state step for registry-based serialization tests."""
|
||||
|
||||
|
||||
def test_empty_pipeline():
|
||||
"""Test pipeline with no steps."""
|
||||
pipeline = DataProcessorPipeline([], to_transition=identity_transition, to_output=identity_transition)
|
||||
@@ -620,6 +668,178 @@ def test_mixed_json_and_tensor_state():
|
||||
assert torch.allclose(loaded_step.running_mean, step.running_mean)
|
||||
|
||||
|
||||
def test_get_config_matches_saved_json():
|
||||
"""Test that in-memory config matches the config written by save_pretrained."""
|
||||
stateless_step = MockStep(name="stateless")
|
||||
stateful_step = MockLazyTensorStateStep(name="stateful", initial_value=4.0)
|
||||
pipeline = DataProcessorPipeline([stateless_step, stateful_step], name="Memory Pipeline")
|
||||
|
||||
in_memory_config = pipeline.get_config()
|
||||
|
||||
assert pipeline.get_config() == in_memory_config
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
pipeline.save_pretrained(tmp_dir)
|
||||
|
||||
config_path = Path(tmp_dir) / "memory_pipeline.json"
|
||||
with open(config_path) as file_pointer:
|
||||
saved_config = json.load(file_pointer)
|
||||
|
||||
assert in_memory_config == saved_config
|
||||
assert "state_file" not in in_memory_config["steps"][0]
|
||||
assert in_memory_config["steps"][1]["state_file"] == "memory_pipeline_step_1.safetensors"
|
||||
|
||||
|
||||
def test_state_dict_matches_saved_safetensors():
|
||||
"""Test that in-memory state matches the safetensors written by save_pretrained."""
|
||||
stateful_step = MockLazyTensorStateStep(initial_value=7.0)
|
||||
pipeline = DataProcessorPipeline([stateful_step], name="Stateful Pipeline")
|
||||
|
||||
in_memory_state_dict = pipeline.state_dict()
|
||||
state_filename = "stateful_pipeline_step_0.safetensors"
|
||||
state_key = "stateful_pipeline_step_0"
|
||||
|
||||
assert set(in_memory_state_dict) == {state_key}
|
||||
assert set(in_memory_state_dict[state_key]) == {"tensor_state"}
|
||||
|
||||
in_memory_state_dict[state_key]["tensor_state"].add_(1)
|
||||
assert stateful_step.tensor_state is not None
|
||||
assert torch.equal(stateful_step.tensor_state, torch.tensor([7.0]))
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
pipeline.save_pretrained(tmp_dir)
|
||||
saved_state_dict = load_file(Path(tmp_dir) / state_filename)
|
||||
|
||||
torch.testing.assert_close(saved_state_dict["tensor_state"], torch.tensor([7.0]))
|
||||
|
||||
|
||||
def test_save_pretrained_still_writes_expected_serialization_files():
|
||||
"""Test that save_pretrained keeps the existing config and state filenames."""
|
||||
stateful_step = MockLazyTensorStateStep(initial_value=3.0)
|
||||
pipeline = DataProcessorPipeline([stateful_step], name="Policy Preprocessor")
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
pipeline.save_pretrained(tmp_dir)
|
||||
|
||||
save_path = Path(tmp_dir)
|
||||
assert (save_path / "policy_preprocessor.json").exists()
|
||||
assert (save_path / "policy_preprocessor_step_0.safetensors").exists()
|
||||
|
||||
|
||||
def test_from_config_round_trips_stateful_pipeline():
|
||||
"""Test that from_config rebuilds a stateful pipeline from in-memory artifacts."""
|
||||
stateful_step = MockLazyTensorStateStep(name="roundtrip", initial_value=11.0)
|
||||
pipeline = DataProcessorPipeline([stateful_step], name="Roundtrip Pipeline")
|
||||
config = pipeline.get_config()
|
||||
pipeline_state_dict = pipeline.state_dict()
|
||||
|
||||
loaded_pipeline = DataProcessorPipeline.from_config(config, state_dict=pipeline_state_dict)
|
||||
loaded_step = loaded_pipeline.steps[0]
|
||||
|
||||
assert len(loaded_pipeline) == 1
|
||||
assert isinstance(loaded_step, MockLazyTensorStateStep)
|
||||
torch.testing.assert_close(loaded_step.tensor_state, torch.tensor([11.0]))
|
||||
|
||||
|
||||
def test_from_config_round_trips_registered_stateful_pipeline():
|
||||
"""Test that from_config resolves registry steps and loads their named tensor state."""
|
||||
stateful_step = RegisteredLazyTensorStateStep(name="registered", initial_value=29.0)
|
||||
pipeline = DataProcessorPipeline([stateful_step], name="Registry Pipeline")
|
||||
config = pipeline.get_config()
|
||||
pipeline_state_dict = pipeline.state_dict()
|
||||
state_filename = "registry_pipeline_step_0_registered_lazy_tensor_state_step.safetensors"
|
||||
state_key = "registry_pipeline_step_0_registered_lazy_tensor_state_step"
|
||||
|
||||
assert config["steps"][0]["registry_name"] == "registered_lazy_tensor_state_step"
|
||||
assert config["steps"][0]["state_file"] == state_filename
|
||||
assert set(pipeline_state_dict) == {state_key}
|
||||
|
||||
loaded_pipeline = DataProcessorPipeline.from_config(config, state_dict=pipeline_state_dict)
|
||||
loaded_step = loaded_pipeline.steps[0]
|
||||
|
||||
assert isinstance(loaded_step, RegisteredLazyTensorStateStep)
|
||||
assert loaded_step.tensor_state is not None
|
||||
torch.testing.assert_close(loaded_step.tensor_state, torch.tensor([29.0]))
|
||||
|
||||
|
||||
def test_from_config_preserves_state_metadata_for_empty_initial_state():
|
||||
"""Test in-memory loading when rebuilt steps start without tensor state."""
|
||||
stateful_step = MockLazyTensorStateStep(name="lazy", initial_value=13.0)
|
||||
pipeline = DataProcessorPipeline([stateful_step], name="Lazy Pipeline")
|
||||
config = pipeline.get_config()
|
||||
pipeline_state_dict = pipeline.state_dict()
|
||||
|
||||
loaded_pipeline = DataProcessorPipeline.from_config(config)
|
||||
loaded_step = loaded_pipeline.steps[0]
|
||||
|
||||
assert isinstance(loaded_step, MockLazyTensorStateStep)
|
||||
assert loaded_step.state_dict() == {}
|
||||
assert "state_file" not in loaded_pipeline.get_config()["steps"][0]
|
||||
|
||||
loaded_pipeline.load_state_dict(pipeline_state_dict)
|
||||
|
||||
torch.testing.assert_close(loaded_step.tensor_state, torch.tensor([13.0]))
|
||||
|
||||
|
||||
def test_from_config_applies_overrides_before_state_loading():
|
||||
"""Test that constructor overrides and tensor state loading are separate operations."""
|
||||
stateful_step = MockLazyTensorStateStep(name="override", scale=1.0, initial_value=17.0)
|
||||
pipeline = DataProcessorPipeline([stateful_step], name="Override Pipeline")
|
||||
config = pipeline.get_config()
|
||||
pipeline_state_dict = pipeline.state_dict()
|
||||
|
||||
loaded_pipeline = DataProcessorPipeline.from_config(
|
||||
config,
|
||||
state_dict=pipeline_state_dict,
|
||||
overrides={"MockLazyTensorStateStep": {"scale": 5.0}},
|
||||
)
|
||||
loaded_step = loaded_pipeline.steps[0]
|
||||
|
||||
assert isinstance(loaded_step, MockLazyTensorStateStep)
|
||||
assert loaded_step.scale == 5.0
|
||||
torch.testing.assert_close(loaded_step.tensor_state, torch.tensor([17.0]))
|
||||
|
||||
|
||||
def test_load_state_dict_raises_on_missing_expected_state():
|
||||
"""Test loading raises when serialized config expects missing state."""
|
||||
stateful_step = MockLazyTensorStateStep(initial_value=19.0)
|
||||
pipeline = DataProcessorPipeline([stateful_step], name="Missing Pipeline")
|
||||
loaded_pipeline = DataProcessorPipeline.from_config(pipeline.get_config())
|
||||
|
||||
with pytest.raises(KeyError, match="missing_pipeline_step_0"):
|
||||
loaded_pipeline.load_state_dict({})
|
||||
|
||||
|
||||
def test_load_state_dict_raises_on_unexpected_extra_state():
|
||||
"""Test loading raises on unexpected top-level state keys."""
|
||||
pipeline = DataProcessorPipeline([MockStep(name="stateless")], name="Unexpected Pipeline")
|
||||
|
||||
with pytest.raises(KeyError, match="extra"):
|
||||
pipeline.load_state_dict({"extra": {"tensor_state": torch.tensor([1.0])}})
|
||||
|
||||
|
||||
def test_stateless_pipeline_in_memory_serialization_returns_empty_state():
|
||||
"""Test stateless in-memory serialization and loading."""
|
||||
pipeline = DataProcessorPipeline([MockStep(name="stateless")], name="Stateless Pipeline")
|
||||
config = pipeline.get_config()
|
||||
config_without_name = {"steps": config["steps"]}
|
||||
|
||||
assert pipeline.state_dict() == {}
|
||||
assert all("state_file" not in step_entry for step_entry in config["steps"])
|
||||
|
||||
loaded_pipeline = DataProcessorPipeline.from_config(config_without_name, state_dict={})
|
||||
|
||||
assert loaded_pipeline.name == "DataProcessorPipeline"
|
||||
assert loaded_pipeline.state_dict() == {}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("invalid_config", [None, [], "not config"])
|
||||
def test_from_config_rejects_non_dict_config(invalid_config):
|
||||
"""Test from_config reports invalid top-level config values cleanly."""
|
||||
with pytest.raises(ValueError, match="not a valid processor configuration"):
|
||||
DataProcessorPipeline.from_config(invalid_config) # type: ignore[arg-type]
|
||||
|
||||
|
||||
class MockModuleStep(ProcessorStep, nn.Module):
|
||||
"""Mock step that inherits from nn.Module to test state_dict handling of module parameters."""
|
||||
|
||||
@@ -2150,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():
|
||||
|
||||
@@ -66,6 +66,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"):
|
||||
|
||||
@@ -1,5 +1,19 @@
|
||||
#!/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
|
||||
|
||||
@@ -28,6 +42,14 @@ def test_push_to_hub_tags_uploaded_dataset_revision(tmp_path, monkeypatch):
|
||||
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
|
||||
|
||||
@@ -49,10 +71,16 @@ def test_push_to_hub_tags_uploaded_dataset_revision(tmp_path, monkeypatch):
|
||||
"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",
|
||||
"exist_ok": True,
|
||||
"revision": "abc123",
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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",
|
||||
@@ -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",
|
||||
"--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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user