From a99943ca26c849197c43c37dc7597c6e052cd5dc Mon Sep 17 00:00:00 2001 From: Shu Jiuhe <134381798+hello3x3@users.noreply.github.com> Date: Sun, 19 Apr 2026 01:28:50 +0800 Subject: [PATCH] Improve loading performance in `_absolute_to_relative_idx` when remapping indices (#3279) --- src/lerobot/datasets/dataset_reader.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/lerobot/datasets/dataset_reader.py b/src/lerobot/datasets/dataset_reader.py index fc7ce36ce..718b33b12 100644 --- a/src/lerobot/datasets/dataset_reader.py +++ b/src/lerobot/datasets/dataset_reader.py @@ -105,10 +105,8 @@ class DatasetReader: """Build absolute-to-relative index mapping from loaded hf_dataset.""" self._absolute_to_relative_idx = None if self.episodes is not None and self.hf_dataset 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"]) - } + indices = self.hf_dataset.data.column("index").to_numpy() + self._absolute_to_relative_idx = dict(zip(indices.tolist(), range(len(indices)), strict=True)) @property def num_frames(self) -> int: