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

add absolute_to_reative_idx for remapping indicies when a subset of data is loaded (#2490)

Browse Source
This commit is contained in:
Michel Aractingi
2025-11-20 14:05:31 +01:00
committed by GitHub
parent 6e86a69dcd
commit 0f551df8f4
1 changed files with 23 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/datasets/lerobot_dataset.py
+23 -3
View File
@@ -712,6 +712,15 @@ class LeRobotDataset(torch.utils.data.Dataset):
self.download(download_videos) self.download(download_videos)
self.hf_dataset = self.load_hf_dataset() self.hf_dataset = self.load_hf_dataset()
# Create mapping from absolute indices to relative indices when only a subset of the episodes are loaded
# Build a mapping: absolute_index -> relative_index_in_filtered_dataset
self._absolute_to_relative_idx = None
if self.episodes is not None:
self._absolute_to_relative_idx = {
abs_idx.item() if isinstance(abs_idx, torch.Tensor) else abs_idx: rel_idx
for rel_idx, abs_idx in enumerate(self.hf_dataset["index"])
}
# Setup delta_indices # Setup delta_indices
if self.delta_timestamps is not None: if self.delta_timestamps is not None:
check_delta_timestamps(self.delta_timestamps, self.fps, self.tolerance_s) check_delta_timestamps(self.delta_timestamps, self.fps, self.tolerance_s)
@@ -930,7 +939,11 @@ class LeRobotDataset(torch.utils.data.Dataset):
query_timestamps = {} query_timestamps = {}
for key in self.meta.video_keys: for key in self.meta.video_keys:
if query_indices is not None and key in query_indices: if query_indices is not None and key in query_indices:
timestamps = self.hf_dataset[query_indices[key]]["timestamp"] if self._absolute_to_relative_idx is not None:
relative_indices = [self._absolute_to_relative_idx[idx] for idx in query_indices[key]]
timestamps = self.hf_dataset[relative_indices]["timestamp"]
else:
timestamps = self.hf_dataset[query_indices[key]]["timestamp"]
query_timestamps[key] = torch.stack(timestamps).tolist() query_timestamps[key] = torch.stack(timestamps).tolist()
else: else:
query_timestamps[key] = [current_ts] query_timestamps[key] = [current_ts]
@@ -953,10 +966,16 @@ class LeRobotDataset(torch.utils.data.Dataset):
for key, q_idx in query_indices.items(): for key, q_idx in query_indices.items():
if key in self.meta.video_keys: if key in self.meta.video_keys:
continue continue
# Map absolute indices to relative indices if needed
relative_indices = (
q_idx
if self._absolute_to_relative_idx is None
else [self._absolute_to_relative_idx[idx] for idx in q_idx]
)
try: try:
result[key] = torch.stack(self.hf_dataset[key][q_idx]) result[key] = torch.stack(self.hf_dataset[key][relative_indices])
except (KeyError, TypeError, IndexError): except (KeyError, TypeError, IndexError):
result[key] = torch.stack(self.hf_dataset[q_idx][key]) result[key] = torch.stack(self.hf_dataset[relative_indices][key])
return result return result
def _query_videos(self, query_timestamps: dict[str, list[float]], ep_idx: int) -> dict[str, torch.Tensor]: def _query_videos(self, query_timestamps: dict[str, list[float]], ep_idx: int) -> dict[str, torch.Tensor]:
@@ -1496,6 +1515,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
obj.image_transforms = None obj.image_transforms = None
obj.delta_timestamps = None obj.delta_timestamps = None
obj.delta_indices = None obj.delta_indices = None
obj._absolute_to_relative_idx = None
obj.video_backend = video_backend if video_backend is not None else get_safe_default_codec() obj.video_backend = video_backend if video_backend is not None else get_safe_default_codec()
obj.writer = None obj.writer = None
obj.latest_episode = None obj.latest_episode = None