From e2804c9fbddc540d1a9610b23f126a085c87a8cf Mon Sep 17 00:00:00 2001 From: "dongmao.zhang" Date: Thu, 16 Jul 2026 14:11:33 -0700 Subject: [PATCH] test(datasets): support the multi-file video layout in dataset fixtures 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. --- tests/datasets/test_streaming.py | 153 ++++++++++++++++++++++++++++ tests/fixtures/dataset_factories.py | 84 +++++++++++---- tests/fixtures/hub.py | 10 +- 3 files changed, 228 insertions(+), 19 deletions(-) diff --git a/tests/datasets/test_streaming.py b/tests/datasets/test_streaming.py index 099f373ca..b40d831e9 100644 --- a/tests/datasets/test_streaming.py +++ b/tests/datasets/test_streaming.py @@ -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}") diff --git a/tests/fixtures/dataset_factories.py b/tests/fixtures/dataset_factories.py index 5c0b0f524..fddbd5038 100644 --- a/tests/fixtures/dataset_factories.py +++ b/tests/fixtures/dataset_factories.py @@ -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,21 +406,27 @@ 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(): - # 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): - 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" - pil_img.save(path) + 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 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-{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.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) - shutil.rmtree(tmp_dir) + 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) + shutil.rmtree(tmp_dir) return _create_video_directory @@ -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( diff --git a/tests/fixtures/hub.py b/tests/fixtures/hub.py index 2f521c766..55b750295 100644 --- a/tests/fixtures/hub.py +++ b/tests/fixtures/hub.py @@ -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//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)