From a51682b266458eaeab136c2207cae2d7cfc8fbdc Mon Sep 17 00:00:00 2001
From: Antoine <antoine.dandigne@gmail.com>
Date: Fri, 17 Oct 2025 15:18:21 +0200
Subject: [PATCH] Optimized episode cache verification (#2166)

Signed-off-by: Antoine <antoine.dandigne@gmail.com>
Co-authored-by: Michel Aractingi <michel.aractingi@huggingface.co>
---
 src/lerobot/datasets/lerobot_dataset.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/lerobot/datasets/lerobot_dataset.py b/src/lerobot/datasets/lerobot_dataset.py
index ae142c1e..ec419d55 100644
--- a/src/lerobot/datasets/lerobot_dataset.py
+++ b/src/lerobot/datasets/lerobot_dataset.py
@@ -842,7 +842,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
         # Get available episode indices from cached dataset
         available_episodes = {
             ep_idx.item() if isinstance(ep_idx, torch.Tensor) else ep_idx
-            for ep_idx in self.hf_dataset["episode_index"]
+            for ep_idx in self.hf_dataset.unique("episode_index")
         }
 
         # Determine requested episodes