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: