From 2ded9ba7837048fa8364a0bcf16ef5977958e984 Mon Sep 17 00:00:00 2001 From: Khalil Meftah Date: Mon, 22 Jun 2026 15:08:05 +0200 Subject: [PATCH] feat(rollout): add episode success labeling to DAgger strategy --- src/lerobot/rollout/configs.py | 2 + src/lerobot/rollout/strategies/dagger.py | 27 ++++++- tests/test_rollout.py | 97 ++++++++++++++++++++++++ 3 files changed, 125 insertions(+), 1 deletion(-) diff --git a/src/lerobot/rollout/configs.py b/src/lerobot/rollout/configs.py index bd8deb37e..c5907c2bf 100644 --- a/src/lerobot/rollout/configs.py +++ b/src/lerobot/rollout/configs.py @@ -121,6 +121,8 @@ class DAggerPedalConfig: pause_resume: str = "KEY_A" correction: str = "KEY_B" upload: str = "KEY_C" + success: str = "KEY_D" + failure: str = "KEY_E" @RolloutStrategyConfig.register_subclass("episodic") diff --git a/src/lerobot/rollout/strategies/dagger.py b/src/lerobot/rollout/strategies/dagger.py index f41fd832c..ef7c48c94 100644 --- a/src/lerobot/rollout/strategies/dagger.py +++ b/src/lerobot/rollout/strategies/dagger.py @@ -118,6 +118,9 @@ class DAggerEvents: # Episode success labeling self._episode_success: bool | None = None + # Episode success labeling + self._episode_success: bool | None = None + # -- Thread-safe phase access ------------------------------------------ @property @@ -181,6 +184,23 @@ class DAggerEvents: self._episode_success = None return result + def mark_success(self) -> None: + """Mark the current episode as successful (called from input threads).""" + with self._lock: + self._episode_success = True + + def mark_failure(self) -> None: + """Mark the current episode as failed (called from input threads).""" + with self._lock: + self._episode_success = False + + def consume_episode_success(self) -> bool | None: + """Consume and reset the episode success label. Returns None if unlabeled.""" + with self._lock: + result = self._episode_success + self._episode_success = None + return result + # --------------------------------------------------------------------------- # Input device handlers @@ -243,6 +263,12 @@ def _init_dagger_pedal(events: DAggerEvents, cfg: DAggerPedalConfig): events.request_transition(code_to_event[code]) if code == cfg.upload: events.upload_requested.set() + if code == cfg.success: + events.mark_success() + logger.info("Episode marked as SUCCESS (pedal)") + if code == cfg.failure: + events.mark_failure() + logger.info("Episode marked as FAILURE (pedal)") logger.info("Initializing DAgger foot pedal listener (device=%s)", cfg.device_path) return start_pedal_listener(on_press, device_path=cfg.device_path) @@ -365,7 +391,6 @@ class DAggerStrategy(RolloutStrategy): return label = self._events.consume_episode_success() - logger.info("_stamp_episode_success: label=%s, buffer_len=%d", label, len(success_buf)) if label: success_buf[-1] = np.array([True], dtype=bool) diff --git a/tests/test_rollout.py b/tests/test_rollout.py index 85a29ff4c..cc1c48b87 100644 --- a/tests/test_rollout.py +++ b/tests/test_rollout.py @@ -338,6 +338,103 @@ def test_dagger_events_reset(): assert not events.upload_requested.is_set() +def test_dagger_mark_success(): + """mark_success sets the episode label to True.""" + from lerobot.rollout.strategies import DAggerEvents + + events = DAggerEvents() + assert events.consume_episode_success() is None + + events.mark_success() + assert events.consume_episode_success() is True + # Consuming clears the label + assert events.consume_episode_success() is None + + +def test_dagger_mark_failure(): + """mark_failure sets the episode label to False.""" + from lerobot.rollout.strategies import DAggerEvents + + events = DAggerEvents() + events.mark_failure() + assert events.consume_episode_success() is False + + +def test_dagger_success_overrides_failure(): + """Last label wins — success after failure overrides.""" + from lerobot.rollout.strategies import DAggerEvents + + events = DAggerEvents() + events.mark_failure() + events.mark_success() + assert events.consume_episode_success() is True + + +def test_dagger_reset_clears_success_label(): + """reset() clears any pending episode success label.""" + from lerobot.rollout.strategies import DAggerEvents + + events = DAggerEvents() + events.mark_success() + events.reset() + assert events.consume_episode_success() is None + + +def test_stamp_episode_success_labels_terminal_frame(): + """_stamp_episode_success sets last frame's next.success to True.""" + import numpy as np + + from lerobot.rollout.strategies.dagger import DAggerStrategy + + strategy = DAggerStrategy.__new__(DAggerStrategy) + strategy.config = MagicMock() + + from lerobot.rollout.strategies import DAggerEvents + + strategy._events = DAggerEvents() + strategy._events.mark_success() + + dataset = MagicMock() + dataset.writer.episode_buffer = { + "next.success": [ + np.array([False], dtype=bool), + np.array([False], dtype=bool), + np.array([False], dtype=bool), + ], + } + + strategy._stamp_episode_success(dataset) + + assert dataset.writer.episode_buffer["next.success"][-1].item() is True + assert dataset.writer.episode_buffer["next.success"][0].item() is False + + +def test_stamp_episode_success_no_label_stays_false(): + """Without a label, all frames remain False.""" + import numpy as np + + from lerobot.rollout.strategies.dagger import DAggerStrategy + + strategy = DAggerStrategy.__new__(DAggerStrategy) + strategy.config = MagicMock() + + from lerobot.rollout.strategies import DAggerEvents + + strategy._events = DAggerEvents() + + dataset = MagicMock() + dataset.writer.episode_buffer = { + "next.success": [ + np.array([False], dtype=bool), + np.array([False], dtype=bool), + ], + } + + strategy._stamp_episode_success(dataset) + + assert all(v.item() is False for v in dataset.writer.episode_buffer["next.success"]) + + # --------------------------------------------------------------------------- # Context dataclass # ---------------------------------------------------------------------------