admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-08-08 17:39:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

test(datasets): support the multi-file video layout in dataset fixtures

Browse Source 
Add `episodes_per_video_file` to `episodes_factory`: episode `i` goes to
`file_index = i // episodes_per_video_file` and `from_timestamp` restarts at 0 on
each rollover. `create_videos` encodes one .mp4 per file, and
`mock_snapshot_download` lists and creates every video file rather than only
file-000. Without the option the fixtures emit the same single-file layout as
before.

Add streaming tests over that layout, on the plain and the delta path.
This commit is contained in:
dongmao.zhang
2026-07-16 14:11:33 -07:00
committed by CarolinePascal
parent e0226b23c8
commit e2804c9fbd
3 changed files with 228 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/datasets/test_streaming.py
+153
View File
@@ -24,11 +24,18 @@ pytest.importorskip("datasets", reason="datasets is required (install lerobot[da
import lerobot.datasets.streaming_dataset as streaming_dataset_module
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset
from lerobot.datasets.utils import safe_shard
from lerobot.utils.constants import ACTION
from tests.fixtures.constants import DUMMY_REPO_ID
# A dataset whose videos roll over into a second file, as v3.0 does past
# DEFAULT_VIDEO_FILE_SIZE_IN_MB: episodes 4-7 live in file-001, whose timeline restarts at 0.
MULTI_FILE_EPISODES = 8
MULTI_FILE_FRAMES = 200
MULTI_FILE_EPISODES_PER_VIDEO_FILE = 4
def get_frames_expected_order(streaming_ds: StreamingLeRobotDataset) -> list[int]:
"""Replicates the shuffling logic of StreamingLeRobotDataset to get the expected order of indices."""
@@ -111,6 +118,59 @@ def test_streaming_dataset_forwards_hub_token_only_for_remote_data(tmp_path, mon
assert not hasattr(dataset, "_token")
def assert_videos_roll_over(ds: LeRobotDataset) -> None:
"""Guard the fixture: these tests are only meaningful if episodes live past ``file-000``.
If ``episodes_per_video_file`` ever stops splitting the videos, the rollover tests below
would still pass while silently covering nothing.
"""
for key in ds.meta.video_keys:
episodes = [ds.meta.episodes[ep_idx] for ep_idx in range(ds.meta.total_episodes)]
file_indices = {ep[f"videos/{key}/file_index"] for ep in episodes}
assert len(file_indices) > 1, f"{key} is not split across video files (file_index: {file_indices})"
# The property under test: a later file's timeline starts back at 0, so an episode's
# `from_timestamp` is no longer its global position in the dataset.
assert any(
ep[f"videos/{key}/file_index"] > 0 and ep[f"videos/{key}/from_timestamp"] == 0.0
for ep in episodes
), f"No episode of {key} restarts a video file's timeline at 0"
def assert_frame_matches(streaming_frame: dict, target_frame: dict, ds: LeRobotDataset, context: str) -> None:
"""Assert a streamed frame equals the same frame read by the non-streaming reader."""
assert set(streaming_frame.keys()) == set(target_frame.keys()), (
f"Keys differ between streaming frame and target one ({context}). "
f"Differ at: {set(streaming_frame.keys()) ^ set(target_frame.keys())}"
)
mismatched = []
for key in streaming_frame:
left, right = streaming_frame[key], target_frame[key]
if isinstance(left, str):
check = left == right
elif isinstance(left, float):
check = left == right.item() # right is a torch.Tensor
elif isinstance(left, torch.Tensor):
if key not in ds.meta.camera_keys and "is_pad" not in key and f"{key}_is_pad" in streaming_frame:
# comparing frames only on non-padded regions. Padding is applied to last-valid broadcasting
left = left[~streaming_frame[f"{key}_is_pad"]]
right = right[~target_frame[f"{key}_is_pad"]]
check = left.shape == right.shape and torch.allclose(left, right)
else:
check = left == right
if not check:
mismatched.append(key)
assert not mismatched, f"Streaming and target frames differ on {mismatched} ({context})"
def test_single_frame_consistency(tmp_path, lerobot_dataset_factory):
"""Test if are correctly accessed"""
ds_num_frames = 400
@@ -645,3 +705,96 @@ def test_bucket_root_caches_metadata_without_switching_to_local_streaming(tmp_pa
def test_invalid_repo_type_fails_before_io():
with pytest.raises(ValueError, match="repo_type must be 'dataset' or 'bucket'"):
StreamingLeRobotDataset(DUMMY_REPO_ID, repo_type="space")
def test_single_frame_consistency_across_video_files(tmp_path, lerobot_dataset_factory):
"""Streaming a dataset whose videos span several files must decode each frame from its own file.
Regression test for decoding at a *global* timestamp (`index / fps`). That position only
exists while the whole dataset fits in one .mp4; once v3.0 rolls the video over, every
episode in a later file asked for a frame past the end of the file being read
(`IndexError: Invalid frame index=... must be less than ...`).
"""
buffer_size = 100
local_path = tmp_path / "test"
repo_id = f"{DUMMY_REPO_ID}-video-rollover"
ds = lerobot_dataset_factory(
root=local_path,
repo_id=repo_id,
total_episodes=MULTI_FILE_EPISODES,
total_frames=MULTI_FILE_FRAMES,
episodes_per_video_file=MULTI_FILE_EPISODES_PER_VIDEO_FILE,
)
assert_videos_roll_over(ds)
streaming_ds = iter(
StreamingLeRobotDataset(
repo_id=repo_id,
root=local_path,
buffer_size=buffer_size,
shuffle=False,
)
)
for _ in range(MULTI_FILE_FRAMES):
streaming_frame = next(streaming_ds)
frame_idx = streaming_frame["index"]
assert_frame_matches(streaming_frame, ds[frame_idx], ds, context=f"frame_idx: {frame_idx}")
@pytest.mark.parametrize(
"state_deltas, action_deltas",
[
([-1, -0.5, -0.20, 0], [0, 1, 2, 3]),
([-2, -1, -0.5, 0], [-1.5, -1, -0.5, -0.20, -0.10, 0]),
],
)
def test_frames_with_delta_consistency_across_video_files(
tmp_path, lerobot_dataset_factory, state_deltas, action_deltas
):
"""Same rollover, on the delta path.
Here the old global timestamp failed silently rather than raising: the query was clamped to
the episode's `to_timestamp`, so every frame decoded the episode's *last* frame — a frozen
video paired with advancing state/action.
"""
buffer_size = 100
seed = 42
local_path = tmp_path / "test"
repo_id = f"{DUMMY_REPO_ID}-video-rollover-deltas"
camera_key = "phone"
delta_timestamps = {
camera_key: state_deltas,
"state": state_deltas,
ACTION: action_deltas,
}
ds = lerobot_dataset_factory(
root=local_path,
repo_id=repo_id,
total_episodes=MULTI_FILE_EPISODES,
total_frames=MULTI_FILE_FRAMES,
episodes_per_video_file=MULTI_FILE_EPISODES_PER_VIDEO_FILE,
delta_timestamps=delta_timestamps,
)
assert_videos_roll_over(ds)
streaming_ds = iter(
StreamingLeRobotDataset(
repo_id=repo_id,
root=local_path,
buffer_size=buffer_size,
seed=seed,
shuffle=False,
delta_timestamps=delta_timestamps,
)
)
for i in range(MULTI_FILE_FRAMES):
streaming_frame = next(streaming_ds)
frame_idx = streaming_frame["index"]
assert_frame_matches(streaming_frame, ds[frame_idx], ds, context=f"i: {i}, frame_idx: {frame_idx}")
tests/fixtures/dataset_factories.py
Vendored
+56 -6
View File
@@ -268,7 +268,17 @@ def episodes_factory(tasks_factory, stats_factory):
video_keys: list[str] | None = None,
tasks: pd.DataFrame | None = None,
multi_task: bool = False,
episodes_per_video_file: int | None = None,
):
"""Build episode metadata.
``episodes_per_video_file`` splits the video keys across several files, as v3.0 does
once a video grows past ``DEFAULT_VIDEO_FILE_SIZE_IN_MB``: episode ``i`` lands in
``file_index = i // episodes_per_video_file`` and each file's timeline restarts at 0,
so ``from_timestamp`` is relative to the file the episode lives in — not to the
dataset. Left ``None``, everything goes to ``file-000`` with a cumulative
``from_timestamp`` (the single-file layout, where the two happen to coincide).
"""
if total_episodes <= 0 or total_frames <= 0:
raise ValueError("num_episodes and total_length must be positive integers.")
if total_frames < total_episodes:
@@ -310,8 +320,16 @@ def episodes_factory(tasks_factory, stats_factory):
d[stats_key] = []
num_frames = 0
# Frames written to the current video file. Resets on every file rollover, since each
# .mp4 carries its own timeline.
num_frames_in_video_file = 0
video_file_index = 0
remaining_tasks = list(tasks.index)
for ep_idx in range(total_episodes):
if episodes_per_video_file is not None and ep_idx // episodes_per_video_file != video_file_index:
video_file_index = ep_idx // episodes_per_video_file
num_frames_in_video_file = 0
num_tasks_in_episode = random.randint(1, min(3, num_tasks_available)) if multi_task else 1
tasks_to_sample = remaining_tasks if len(remaining_tasks) > 0 else list(tasks.index)
episode_tasks = random.sample(tasks_to_sample, min(num_tasks_in_episode, len(tasks_to_sample)))
@@ -333,21 +351,44 @@ def episodes_factory(tasks_factory, stats_factory):
if video_keys is not None:
for video_key in video_keys:
d[f"videos/{video_key}/chunk_index"].append(0)
d[f"videos/{video_key}/file_index"].append(0)
d[f"videos/{video_key}/from_timestamp"].append(num_frames / fps)
d[f"videos/{video_key}/to_timestamp"].append((num_frames + lengths[ep_idx]) / fps)
d[f"videos/{video_key}/file_index"].append(video_file_index)
d[f"videos/{video_key}/from_timestamp"].append(num_frames_in_video_file / fps)
d[f"videos/{video_key}/to_timestamp"].append(
(num_frames_in_video_file + lengths[ep_idx]) / fps
)
# Add stats columns like "stats/action/max"
for stats_key, stats in flatten_dict({"stats": stats_factory(features)}).items():
d[stats_key].append(stats)
num_frames += lengths[ep_idx]
num_frames_in_video_file += lengths[ep_idx]
return Dataset.from_dict(d)
return _create_episodes
def video_file_frames(
episodes: datasets.Dataset | None, video_key: str, total_frames: int
) -> dict[int, list[int]]:
"""Map each of ``video_key``'s files to the global frame indices it holds, in file order.
``episodes`` is the source of truth for how a video key is split across files. Without it
(or without the videos columns), the whole key is one file — the pre-v3.0 layout.
"""
if episodes is None or f"videos/{video_key}/file_index" not in episodes.column_names:
return {0: list(range(total_frames))}
frames_per_file: dict[int, list[int]] = {}
for ep in episodes:
file_index = ep[f"videos/{video_key}/file_index"]
frames = range(ep["dataset_from_index"], ep["dataset_to_index"])
frames_per_file.setdefault(file_index, []).extend(frames)
return frames_per_file
@pytest.fixture(scope="session")
def create_videos(info_factory, img_array_factory):
def _create_video_directory(
@@ -356,6 +397,7 @@ def create_videos(info_factory, img_array_factory):
total_episodes: int = 3,
total_frames: int = 150,
total_tasks: int = 1,
episodes: datasets.Dataset | None = None,
):
if info is None:
info = info_factory(
@@ -364,17 +406,23 @@ def create_videos(info_factory, img_array_factory):
video_feats = {key: feats for key, feats in info.features.items() if feats["dtype"] == "video"}
for key, ft in video_feats.items():
for file_index, frame_indices in video_file_frames(episodes, key, info.total_frames).items():
# create and save images with identifiable content
tmp_dir = root / "tmp_images"
tmp_dir.mkdir(parents=True, exist_ok=True)
for frame_index in range(info.total_frames):
for position, frame_index in enumerate(frame_indices):
# Content stays keyed on the *global* frame index so a frame remains
# identifiable across files, but its position in the file is what the
# file's timeline addresses.
content = f"{key}-{frame_index}"
img = img_array_factory(height=ft["shape"][0], width=ft["shape"][1], content=content)
pil_img = PIL.Image.fromarray(img)
path = tmp_dir / f"frame-{frame_index:06d}.png"
path = tmp_dir / f"frame-{position:06d}.png"
pil_img.save(path)
video_path = root / DEFAULT_VIDEO_PATH.format(video_key=key, chunk_index=0, file_index=0)
video_path = root / DEFAULT_VIDEO_PATH.format(
video_key=key, chunk_index=0, file_index=file_index
)
video_path.parent.mkdir(parents=True, exist_ok=True)
# Use the global fps from info, not video-specific fps which might not exist
encode_video_frames(tmp_dir, video_path, fps=info.fps)
@@ -520,6 +568,7 @@ def lerobot_dataset_factory(
data_files_size_in_mb: float = DEFAULT_DATA_FILE_SIZE_IN_MB,
chunks_size: int = DEFAULT_CHUNK_SIZE,
camera_features: dict | None = None,
episodes_per_video_file: int | None = None,
**kwargs,
) -> LeRobotDataset:
# Instantiate objects
@@ -557,6 +606,7 @@ def lerobot_dataset_factory(
video_keys=video_keys,
tasks=tasks,
multi_task=multi_task,
episodes_per_video_file=episodes_per_video_file,
)
if hf_dataset is None:
hf_dataset = hf_dataset_factory(
tests/fixtures/hub.py
Vendored
+8 -2
View File
@@ -29,6 +29,7 @@ from lerobot.datasets.utils import (
STATS_PATH,
)
from tests.fixtures.constants import LEROBOT_TEST_DIR
from tests.fixtures.dataset_factories import video_file_frames
@pytest.fixture(scope="session")
@@ -99,7 +100,12 @@ def mock_snapshot_download_factory(
video_keys = [key for key, feats in info.features.items() if feats["dtype"] == "video"]
for key in video_keys:
all_files.append(DEFAULT_VIDEO_PATH.format(video_key=key, chunk_index=0, file_index=0))
# A video key spans one file per `videos/<key>/file_index` in the episodes
# metadata — several of them once v3.0 rolls a video over.
for file_index in video_file_frames(episodes, key, info.total_frames):
all_files.append(
DEFAULT_VIDEO_PATH.format(video_key=key, chunk_index=0, file_index=file_index)
)
allowed_files = filter_repo_objects(
all_files, allow_patterns=allow_patterns, ignore_patterns=ignore_patterns
@@ -138,7 +144,7 @@ def mock_snapshot_download_factory(
if request_data:
create_hf_dataset(local_dir, hf_dataset, data_files_size_in_mb, chunks_size)
if request_videos:
create_videos(root=local_dir, info=info)
create_videos(root=local_dir, info=info, episodes=episodes)
return str(local_dir)