From 14c7a25ce47764e02721cf1e9dbf837bd89d4b19 Mon Sep 17 00:00:00 2001
From: Steven Palma <steven.palma@huggingface.co>
Date: Sun, 19 Apr 2026 17:29:08 +0200
Subject: [PATCH] fix(rollout) require dataset in dagger + use duration too

---
 src/lerobot/rollout/__init__.py           |  1 -
 src/lerobot/rollout/configs.py            | 23 ++++++++++++++---------
 src/lerobot/rollout/strategies/dagger.py  |  5 +++++
 src/lerobot/rollout/strategies/factory.py |  2 +-
 4 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/src/lerobot/rollout/__init__.py b/src/lerobot/rollout/__init__.py
index 6a2444762..8ef6c26cd 100644
--- a/src/lerobot/rollout/__init__.py
+++ b/src/lerobot/rollout/__init__.py
@@ -23,7 +23,6 @@ from .configs import (
     DAggerKeyboardConfig,
     DAggerPedalConfig,
     DAggerStrategyConfig,
-    DatasetRecordConfig,
     HighlightStrategyConfig,
     RolloutConfig,
     RolloutStrategyConfig,
diff --git a/src/lerobot/rollout/configs.py b/src/lerobot/rollout/configs.py
index 22ed39031..46f46eae1 100644
--- a/src/lerobot/rollout/configs.py
+++ b/src/lerobot/rollout/configs.py
@@ -216,7 +216,10 @@ class RolloutConfig:
         if isinstance(self.strategy, DAggerStrategyConfig) and self.teleop is None:
             raise ValueError("DAgger strategy requires --teleop.type to be set")
 
-        needs_dataset = isinstance(self.strategy, (SentryStrategyConfig, HighlightStrategyConfig))
+        # TODO(Steven): DAgger shouldn't require a dataset (user may want to just rollout+intervene without recording), but for now we require it to simplify the implementation.
+        needs_dataset = isinstance(
+            self.strategy, (SentryStrategyConfig, HighlightStrategyConfig, DAggerStrategyConfig)
+        )
         if needs_dataset and (self.dataset is None or not self.dataset.repo_id):
             raise ValueError(f"{self.strategy.type} strategy requires --dataset.repo_id to be set")
 
@@ -244,14 +247,16 @@ class RolloutConfig:
             self.dataset.streaming_encoding = True
 
         # DAgger: streaming is mandatory only when the autonomous phase is also recorded.
-        if (
-            isinstance(self.strategy, DAggerStrategyConfig)
-            and self.strategy.record_autonomous
-            and self.dataset is not None
-            and not self.dataset.streaming_encoding
-        ):
-            logger.warning("DAgger with record_autonomous=True forces streaming_encoding=True")
-            self.dataset.streaming_encoding = True
+        if isinstance(self.strategy, DAggerStrategyConfig) and self.dataset is not None:
+            if self.strategy.record_autonomous and not self.dataset.streaming_encoding:
+                logger.warning("DAgger with record_autonomous=True forces streaming_encoding=True")
+                self.dataset.streaming_encoding = True
+            elif not self.strategy.record_autonomous and not self.dataset.streaming_encoding:
+                logger.info(
+                    "Streaming encoding is disabled for DAgger corrections-only mode. "
+                    "Consider enabling it for faster episode saving: "
+                    "--dataset.streaming_encoding=true --dataset.encoder_threads=2"
+                )
 
         # --- Policy loading ---
         if self.robot is None:
diff --git a/src/lerobot/rollout/strategies/dagger.py b/src/lerobot/rollout/strategies/dagger.py
index b346e50b8..ba842844f 100644
--- a/src/lerobot/rollout/strategies/dagger.py
+++ b/src/lerobot/rollout/strategies/dagger.py
@@ -556,6 +556,7 @@ class DAggerStrategy(RolloutStrategy):
         engine.resume()
 
         last_action: dict[str, Any] | None = None
+        start_time = time.perf_counter()
         record_tick = 0
         recorded = 0
         logger.info(
@@ -571,6 +572,10 @@ class DAggerStrategy(RolloutStrategy):
                 ):
                     loop_start = time.perf_counter()
 
+                    if cfg.duration > 0 and (time.perf_counter() - start_time) >= cfg.duration:
+                        logger.info("Duration limit reached (%.0fs)", cfg.duration)
+                        break
+
                     # Process transitions
                     transition = events.consume_transition()
                     if transition is not None:
diff --git a/src/lerobot/rollout/strategies/factory.py b/src/lerobot/rollout/strategies/factory.py
index 0a07d4750..8a9727769 100644
--- a/src/lerobot/rollout/strategies/factory.py
+++ b/src/lerobot/rollout/strategies/factory.py
@@ -25,7 +25,7 @@ from .highlight import HighlightStrategy
 from .sentry import SentryStrategy
 
 if TYPE_CHECKING:
-    from lerobot.rollout import RolloutStrategyConfig
+    from ..configs import RolloutStrategyConfig
 
 
 def create_strategy(config: RolloutStrategyConfig) -> RolloutStrategy: