From ab87fd97646f323d5267b029b2c2f2181e9da886 Mon Sep 17 00:00:00 2001
From: MihaiAnca13 <mihai.anca13@gmail.com>
Date: Mon, 27 Jul 2026 12:44:32 +0100
Subject: [PATCH] fix(datasets): clear video frame staging on episode reset
 (#3683)

* fix video frame staging cleanup on episode reset

* linting

---------

Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
---
 src/lerobot/datasets/dataset_writer.py | 37 ++++++++++++++++----------
 tests/datasets/test_dataset_writer.py  | 32 ++++++++++++++++++++++
 2 files changed, 55 insertions(+), 14 deletions(-)

diff --git a/src/lerobot/datasets/dataset_writer.py b/src/lerobot/datasets/dataset_writer.py
index a6049312f..5693e9bb2 100644
--- a/src/lerobot/datasets/dataset_writer.py
+++ b/src/lerobot/datasets/dataset_writer.py
@@ -172,6 +172,23 @@ class DatasetWriter:
     def _get_image_file_dir(self, episode_index: int, image_key: str) -> Path:
         return self._get_image_file_path(episode_index, image_key, frame_index=0).parent
 
+    def _get_episode_buffer_index(self) -> int:
+        episode_index = self.episode_buffer["episode_index"]
+        # episode_index is `int` when freshly created, but becomes `np.ndarray` after
+        # save_episode() mutates the buffer. Handle both types here.
+        if isinstance(episode_index, np.ndarray):
+            episode_index = episode_index.item() if episode_index.size == 1 else episode_index[0]
+        return int(episode_index)
+
+    def _delete_camera_frame_dirs(self, camera_keys: list[str]) -> None:
+        if self.image_writer is not None:
+            self._wait_image_writer()
+        episode_index = self._get_episode_buffer_index()
+        for camera_key in camera_keys:
+            img_dir = self._get_image_file_dir(episode_index, camera_key)
+            if img_dir.is_dir():
+                shutil.rmtree(img_dir)
+
     def _save_image(
         self, image: torch.Tensor | np.ndarray | PIL.Image.Image, fpath: Path, compress_level: int = 1
     ) -> None:
@@ -369,7 +386,9 @@ class DatasetWriter:
                 self._episodes_since_last_encoding = 0
 
         if episode_data is None:
-            self.clear_episode_buffer(delete_images=len(self._meta.image_keys) > 0)
+            if len(self._meta.image_keys) > 0:
+                self._delete_camera_frame_dirs(self._meta.image_keys)
+            self.episode_buffer = self._create_episode_buffer()
 
     def _batch_save_episode_video(self, start_episode: int, end_episode: int | None = None) -> None:
         """Batch save videos for multiple episodes."""
@@ -561,10 +580,10 @@ class DatasetWriter:
         return metadata
 
     def clear_episode_buffer(self, delete_images: bool = True) -> None:
-        """Discard the current episode buffer and optionally delete temp images.
+        """Discard the current episode buffer and optionally delete temp camera frames.
 
         Args:
-            delete_images: If ``True``, remove temporary image directories
+            delete_images: If ``True``, remove temporary camera frame directories
                 written for the current episode.
         """
         # Cancel streaming encoder if active
@@ -572,17 +591,7 @@ class DatasetWriter:
             self._streaming_encoder.cancel_episode()
 
         if delete_images:
-            if self.image_writer is not None:
-                self._wait_image_writer()
-            episode_index = self.episode_buffer["episode_index"]
-            # episode_index is `int` when freshly created, but becomes `np.ndarray` after
-            # save_episode() mutates the buffer. Handle both types here.
-            if isinstance(episode_index, np.ndarray):
-                episode_index = episode_index.item() if episode_index.size == 1 else episode_index[0]
-            for cam_key in self._meta.image_keys:
-                img_dir = self._get_image_file_dir(episode_index, cam_key)
-                if img_dir.is_dir():
-                    shutil.rmtree(img_dir)
+            self._delete_camera_frame_dirs(self._meta.camera_keys)
 
         self.episode_buffer = self._create_episode_buffer()
 
diff --git a/tests/datasets/test_dataset_writer.py b/tests/datasets/test_dataset_writer.py
index 17785ad74..bd51a2687 100644
--- a/tests/datasets/test_dataset_writer.py
+++ b/tests/datasets/test_dataset_writer.py
@@ -204,6 +204,38 @@ def test_clear_resets_buffer(tmp_path):
     assert dataset.writer.episode_buffer["size"] == 0
 
 
+def test_clear_removes_video_frame_staging_dir(tmp_path):
+    """clear_episode_buffer() removes PNG staging dirs for video features."""
+    video_key = "observation.images.cam"
+    features = {
+        video_key: {
+            "dtype": "video",
+            "shape": (64, 96, 3),
+            "names": ["height", "width", "channels"],
+        },
+        "action": {"dtype": "float32", "shape": (2,), "names": None},
+    }
+    dataset = LeRobotDataset.create(
+        repo_id=DUMMY_REPO_ID,
+        fps=DEFAULT_FPS,
+        features=features,
+        root=tmp_path / "ds",
+        use_videos=True,
+    )
+
+    dataset.add_frame(_make_frame(features))
+    video_staging_dir = (
+        dataset.root
+        / Path(DEFAULT_IMAGE_PATH.format(image_key=video_key, episode_index=0, frame_index=0)).parent
+    )
+    assert video_staging_dir.is_dir()
+
+    dataset.clear_episode_buffer()
+
+    assert dataset.writer.episode_buffer["size"] == 0
+    assert not video_staging_dir.exists()
+
+
 def test_finalize_is_idempotent(tmp_path):
     """Calling finalize() twice does not raise."""
     dataset = LeRobotDataset.create(