mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-06 09:37:06 +00:00
feat(annotate): add constant advantage labeling for RECAP SFT phase
- Add constant_value and seed fields to AdvantageConfig - Implement _run_constant_mode in AdvantageModule with CFG dropout - Use deterministic seeding (config.seed + episode_index) for reproducibility
This commit is contained in:
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user