From bf5148c6e832933d51da22edcb5809456cc46192 Mon Sep 17 00:00:00 2001
From: Maxime Ellerbach <maxime.ellerbach@huggingface.co>
Date: Mon, 27 Apr 2026 17:25:37 +0200
Subject: [PATCH] fix(rollout): properly propagating video_files_size_in_mb to
 lerobot_dataset (#3470)

---
 src/lerobot/datasets/lerobot_dataset.py |  4 ++++
 src/lerobot/rollout/configs.py          |  6 +++---
 src/lerobot/rollout/context.py          |  2 ++
 src/lerobot/rollout/ring_buffer.py      |  2 +-
 tests/datasets/test_lerobot_dataset.py  | 12 ++++++++++++
 5 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/src/lerobot/datasets/lerobot_dataset.py b/src/lerobot/datasets/lerobot_dataset.py
index 644ce14db..b6ab0f5f0 100644
--- a/src/lerobot/datasets/lerobot_dataset.py
+++ b/src/lerobot/datasets/lerobot_dataset.py
@@ -630,6 +630,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
         streaming_encoding: bool = False,
         encoder_queue_maxsize: int = 30,
         encoder_threads: int | None = None,
+        video_files_size_in_mb: int | None = None,
+        data_files_size_in_mb: int | None = None,
     ) -> "LeRobotDataset":
         """Create a new LeRobotDataset from scratch for recording data.
 
@@ -677,6 +679,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
             root=root,
             use_videos=use_videos,
             metadata_buffer_size=metadata_buffer_size,
+            video_files_size_in_mb=video_files_size_in_mb,
+            data_files_size_in_mb=data_files_size_in_mb,
         )
         obj.repo_id = obj.meta.repo_id
         obj._requested_root = obj.meta.root
diff --git a/src/lerobot/rollout/configs.py b/src/lerobot/rollout/configs.py
index e877ce253..9d019c887 100644
--- a/src/lerobot/rollout/configs.py
+++ b/src/lerobot/rollout/configs.py
@@ -75,7 +75,7 @@ class SentryStrategyConfig(RolloutStrategyConfig):
     # Target video file size in MB for episode rotation.  Episodes are
     # saved once the estimated video duration would exceed this limit.
     # Defaults to DEFAULT_VIDEO_FILE_SIZE_IN_MB when set to None.
-    target_video_file_size_mb: float | None = None
+    target_video_file_size_mb: int | None = None
 
 
 @RolloutStrategyConfig.register_subclass("highlight")
@@ -90,7 +90,7 @@ class HighlightStrategyConfig(RolloutStrategyConfig):
     """
 
     ring_buffer_seconds: float = 10.0
-    ring_buffer_max_memory_mb: float = 1024.0
+    ring_buffer_max_memory_mb: int = 1024
     save_key: str = "s"
     push_key: str = "h"
 
@@ -150,7 +150,7 @@ class DAggerStrategyConfig(RolloutStrategyConfig):
     upload_every_n_episodes: int = 5
     # Target video file size in MB for episode rotation (record_autonomous
     # mode only).  Defaults to DEFAULT_VIDEO_FILE_SIZE_IN_MB when None.
-    target_video_file_size_mb: float | None = None
+    target_video_file_size_mb: int | None = None
     input_device: str = "keyboard"
     keyboard: DAggerKeyboardConfig = field(default_factory=DAggerKeyboardConfig)
     pedal: DAggerPedalConfig = field(default_factory=DAggerPedalConfig)
diff --git a/src/lerobot/rollout/context.py b/src/lerobot/rollout/context.py
index f07ed7e2f..382909e6c 100644
--- a/src/lerobot/rollout/context.py
+++ b/src/lerobot/rollout/context.py
@@ -355,6 +355,7 @@ def build_rollout_context(
                     "Use --dataset.repo_id=<user>/rollout_<name> for policy deployment datasets."
                 )
             cfg.dataset.stamp_repo_id()
+            target_video_mb = getattr(cfg.strategy, "target_video_file_size_mb", None)
             dataset = LeRobotDataset.create(
                 cfg.dataset.repo_id,
                 cfg.dataset.fps,
@@ -370,6 +371,7 @@ def build_rollout_context(
                 streaming_encoding=cfg.dataset.streaming_encoding,
                 encoder_queue_maxsize=cfg.dataset.encoder_queue_maxsize,
                 encoder_threads=cfg.dataset.encoder_threads,
+                video_files_size_in_mb=target_video_mb,
             )
 
     if dataset is not None:
diff --git a/src/lerobot/rollout/ring_buffer.py b/src/lerobot/rollout/ring_buffer.py
index 6b362869a..2c0a06301 100644
--- a/src/lerobot/rollout/ring_buffer.py
+++ b/src/lerobot/rollout/ring_buffer.py
@@ -47,7 +47,7 @@ class RolloutRingBuffer:
         count.
     """
 
-    def __init__(self, max_seconds: float = 30.0, max_memory_mb: float = 2048.0, fps: float = 30.0) -> None:
+    def __init__(self, max_seconds: float = 30.0, max_memory_mb: int = 2048, fps: float = 30.0) -> None:
         self._max_frames = int(max_seconds * fps)
         self._max_bytes = int(max_memory_mb * 1024 * 1024)
         self._buffer: deque[dict] = deque(maxlen=self._max_frames)
diff --git a/tests/datasets/test_lerobot_dataset.py b/tests/datasets/test_lerobot_dataset.py
index 49efa84d9..26406dea2 100644
--- a/tests/datasets/test_lerobot_dataset.py
+++ b/tests/datasets/test_lerobot_dataset.py
@@ -416,6 +416,18 @@ def test_create_initial_counts_zero(tmp_path):
     assert dataset.num_frames == 0
 
 
+def test_create_propagates_video_files_size_in_mb(tmp_path):
+    """video_files_size_in_mb passed to create() is reflected in the dataset metadata."""
+    dataset = LeRobotDataset.create(
+        repo_id=DUMMY_REPO_ID,
+        fps=DEFAULT_FPS,
+        features=SIMPLE_FEATURES,
+        root=tmp_path / "ds",
+        video_files_size_in_mb=42.0,
+    )
+    assert dataset.meta.video_files_size_in_mb == 42.0
+
+
 def test_add_frame_works_in_write_mode(tmp_path):
     """add_frame() succeeds on a dataset created via create()."""
     dataset = LeRobotDataset.create(