From 64473524393fb70c336c50a9e782eacd9fbe0c76 Mon Sep 17 00:00:00 2001
From: Michel Aractingi <michel.aractingi@huggingface.co>
Date: Wed, 30 Jul 2025 00:32:28 +0200
Subject: [PATCH] added a check for comparing cached episodes in order to
 trigger a new download if the requested episodes dont match the cached ones

---
 src/lerobot/datasets/lerobot_dataset.py | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/src/lerobot/datasets/lerobot_dataset.py b/src/lerobot/datasets/lerobot_dataset.py
index 74d5a7fd2..b059c35ed 100644
--- a/src/lerobot/datasets/lerobot_dataset.py
+++ b/src/lerobot/datasets/lerobot_dataset.py
@@ -551,6 +551,9 @@ class LeRobotDataset(torch.utils.data.Dataset):
             if force_cache_sync:
                 raise FileNotFoundError
             self.hf_dataset = self.load_hf_dataset()
+            # Check if cached dataset contains all requested episodes
+            if not self._check_cached_episodes_sufficient():
+                raise FileNotFoundError("Cached dataset doesn't contain all requested episodes")
         except (AssertionError, FileNotFoundError, NotADirectoryError):
             self.revision = get_safe_version(self.repo_id, self.revision)
             self.download(download_videos)
@@ -666,6 +669,25 @@ class LeRobotDataset(torch.utils.data.Dataset):
         hf_dataset.set_transform(hf_transform_to_torch)
         return hf_dataset
 
+    def _check_cached_episodes_sufficient(self) -> bool:
+        """Check if the cached dataset contains all requested episodes."""
+        if self.hf_dataset is None or len(self.hf_dataset) == 0:
+            return False
+
+        # Get available episode indices from cached dataset
+        available_episodes = set(self.hf_dataset["episode_index"])
+
+        # Determine requested episodes
+        if self.episodes is None:
+            # Requesting all episodes - check if we have all episodes from metadata
+            requested_episodes = set(range(self.meta.total_episodes))
+        else:
+            # Requesting specific episodes
+            requested_episodes = set(self.episodes)
+
+        # Check if all requested episodes are available in cached data
+        return requested_episodes.issubset(available_episodes)
+
     def create_hf_dataset(self) -> datasets.Dataset:
         features = get_hf_features_from_features(self.features)
         ft_dict = {col: [] for col in features}