diff --git a/src/lerobot/annotations/steerable_pipeline/config.py b/src/lerobot/annotations/steerable_pipeline/config.py index 1d7ca8af2..f93002609 100644 --- a/src/lerobot/annotations/steerable_pipeline/config.py +++ b/src/lerobot/annotations/steerable_pipeline/config.py @@ -175,7 +175,12 @@ class AdvantageConfig: enabled: bool = True - # Path or Hub repo ID of the trained distributional value function checkpoint. + # Constant advantage label for all frames (e.g. "positive" for SFT iteration 0). + # Skips VF inference, dropout still applies for CFG. + constant_value: str | None = None + + # Trained value function checkpoint (local path or Hub repo ID). + # Ignored when constant_value is set. value_function_path: str = "" # Device to run the value function on. @@ -205,6 +210,9 @@ class AdvantageConfig: # Batch size for value function inference. batch_size: int = 32 + # Random seed for dropout reproducibility. + seed: int = 1729 + @dataclass class AnnotationPipelineConfig: diff --git a/src/lerobot/annotations/steerable_pipeline/modules/advantage.py b/src/lerobot/annotations/steerable_pipeline/modules/advantage.py index 5682ee3ca..55ebcf554 100644 --- a/src/lerobot/annotations/steerable_pipeline/modules/advantage.py +++ b/src/lerobot/annotations/steerable_pipeline/modules/advantage.py @@ -206,8 +206,12 @@ class AdvantageModule: def run_episode(self, record: EpisodeRecord, staging: EpisodeStaging) -> None: """Score one episode and write advantage rows to staging.""" + if self.config.constant_value: + self._run_constant_mode(record, staging) + return + if not self.config.value_function_path: - logger.warning("No value_function_path configured; skipping advantage scoring.") + logger.warning("No value_function_path or constant_value configured; skipping advantage scoring.") return advantages, intervention_mask = self.compute_advantages_for_episode(record) @@ -215,7 +219,7 @@ class AdvantageModule: threshold = self._compute_threshold(advantages, intervention_mask) - rng = np.random.default_rng(seed=hash((record.episode_index, 42)) & 0xFFFFFFFF) + rng = np.random.default_rng(seed=self.config.seed + record.episode_index) rows: list[dict[str, Any]] = [] for t in range(num_frames): @@ -255,6 +259,39 @@ class AdvantageModule: sum(1 for r in rows if r["content"] == "negative"), ) + def _run_constant_mode(self, record: EpisodeRecord, staging: EpisodeStaging) -> None: + """Emit a fixed advantage value for every frame (with dropout for CFG).""" + num_frames = record.num_frames + rng = np.random.default_rng(seed=self.config.seed + record.episode_index) + + rows: list[dict[str, Any]] = [] + for t in range(num_frames): + if rng.random() < self.config.dropout_rate: + continue + + timestamp = float(record.frame_timestamps[t]) if t < len(record.frame_timestamps) else 0.0 + + rows.append( + { + "role": "user", + "content": self.config.constant_value, + "style": "advantage", + "timestamp": timestamp, + "camera": None, + "tool_calls": None, + } + ) + + staging.write("advantage", rows) + logger.debug( + "Episode %d: %d/%d frames labeled constant '%s' (dropout=%.2f)", + record.episode_index, + len(rows), + num_frames, + self.config.constant_value, + self.config.dropout_rate, + ) + def _compute_threshold(self, advantages: np.ndarray, intervention_mask: np.ndarray) -> float: """Compute the binarization threshold as the configured percentile of advantages.""" non_intervention = advantages[~intervention_mask] if intervention_mask.any() else advantages