diff --git a/src/lerobot/datasets/lerobot_dataset.py b/src/lerobot/datasets/lerobot_dataset.py index aba86efe3..77b3032df 100644 --- a/src/lerobot/datasets/lerobot_dataset.py +++ b/src/lerobot/datasets/lerobot_dataset.py @@ -478,18 +478,19 @@ class LeRobotDataset(torch.utils.data.Dataset): """Return the number of frames in the selected episodes.""" return self.num_frames - def __getitem__(self, idx) -> dict: - """Return a single frame by index, with all transforms applied. + def __getitem__(self, idx: int | slice) -> dict | list[dict]: + """Return one frame or a slice of frames, with all transforms applied. Loads the frame from the underlying HF dataset, expands delta-timestamp windows, decodes video frames, and applies image transforms. Delegates - the core logic to :meth:`DatasetReader.get_item`. + the core logic to :class:`DatasetReader`. Args: - idx: Index into the (possibly episode-filtered) dataset. + idx: Integer index or slice into the possibly episode-filtered dataset. Returns: - Dict mapping feature names to their tensor values for this frame. + A frame dictionary for an integer index, or a list of frame + dictionaries for a slice. Raises: RuntimeError: If the dataset is currently being recorded and @@ -499,6 +500,9 @@ class LeRobotDataset(torch.utils.data.Dataset): raise RuntimeError( "Cannot read from a dataset that is being recorded. Call finalize() first, then access items." ) + if isinstance(idx, slice): + return [self[item_idx] for item_idx in range(*idx.indices(len(self)))] + reader = self._ensure_reader() if reader.hf_dataset is None: # One-shot load after finalize() diff --git a/tests/datasets/test_datasets.py b/tests/datasets/test_datasets.py index 97f799d9f..d76ed2e1d 100644 --- a/tests/datasets/test_datasets.py +++ b/tests/datasets/test_datasets.py @@ -114,6 +114,20 @@ def test_dataset_initialization(tmp_path, lerobot_dataset_factory): assert dataset.num_frames == len(dataset) +def test_dataset_slice(tmp_path, lerobot_dataset_factory): + dataset = lerobot_dataset_factory( + root=tmp_path / "test", total_episodes=3, total_frames=30, use_videos=False + ) + + assert len(dataset[:5]) == 5 + assert len(dataset[::2]) == (len(dataset) + 1) // 2 + assert [item["index"].item() for item in dataset[4::-1]] == [4, 3, 2, 1, 0] + assert [item["index"].item() for item in dataset[-3:]] == list(range(len(dataset) - 3, len(dataset))) + assert dataset[len(dataset) :] == [] + assert isinstance(dataset[0], dict) + assert dataset[:1][0].keys() == dataset[0].keys() + + # TODO(rcadene, aliberts): do not run LeRobotDataset.create, instead refactor LeRobotDatasetMetadata.create # and test the small resulting function that validates the features def test_dataset_feature_with_forward_slash_raises_error(): @@ -1741,6 +1755,38 @@ def test_delta_timestamps_query_returns_correct_values(tmp_path, empty_lerobot_d assert is_pad == [True, False], f"Expected [True, False], got {is_pad}" +def test_dataset_slice_with_delta_timestamps(tmp_path, empty_lerobot_dataset_factory): + features = { + "observation.state": {"dtype": "float32", "shape": (1,), "names": ["x"]}, + } + dataset = empty_lerobot_dataset_factory( + root=tmp_path / "test_slice_delta", features=features, use_videos=False, fps=10 + ) + + for frame_idx in range(5): + dataset.add_frame( + { + "observation.state": torch.tensor([frame_idx], dtype=torch.float32), + "task": "task_0", + } + ) + dataset.save_episode() + dataset.finalize() + + sliced_dataset = LeRobotDataset( + dataset.repo_id, + root=dataset.root, + delta_timestamps={"observation.state": [-0.1, 0.0]}, + tolerance_s=0.04, + ) + + items = sliced_dataset[:2] + + assert items[0]["observation.state"].tolist() == [0.0, 0.0] + assert items[0]["observation.state_is_pad"].tolist() == [True, False] + assert items[1]["observation.state"].tolist() == [0.0, 1.0] + + def test_episode_filter_filters_dataset(tmp_path, lerobot_dataset_factory): """episode_filter on LeRobotDataset narrows the loaded dataset to matching episodes.""" dataset = lerobot_dataset_factory(root=tmp_path / "test", total_episodes=8, total_frames=200)