mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-21 17:01:53 +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
|
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 = ""
|
value_function_path: str = ""
|
||||||
|
|
||||||
# Device to run the value function on.
|
# Device to run the value function on.
|
||||||
@@ -205,6 +210,9 @@ class AdvantageConfig:
|
|||||||
# Batch size for value function inference.
|
# Batch size for value function inference.
|
||||||
batch_size: int = 32
|
batch_size: int = 32
|
||||||
|
|
||||||
|
# Random seed for dropout reproducibility.
|
||||||
|
seed: int = 1729
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
class AnnotationPipelineConfig:
|
class AnnotationPipelineConfig:
|
||||||
|
|||||||
@@ -206,8 +206,12 @@ class AdvantageModule:
|
|||||||
|
|
||||||
def run_episode(self, record: EpisodeRecord, staging: EpisodeStaging) -> None:
|
def run_episode(self, record: EpisodeRecord, staging: EpisodeStaging) -> None:
|
||||||
"""Score one episode and write advantage rows to staging."""
|
"""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:
|
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
|
return
|
||||||
|
|
||||||
advantages, intervention_mask = self.compute_advantages_for_episode(record)
|
advantages, intervention_mask = self.compute_advantages_for_episode(record)
|
||||||
@@ -215,7 +219,7 @@ class AdvantageModule:
|
|||||||
|
|
||||||
threshold = self._compute_threshold(advantages, intervention_mask)
|
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]] = []
|
rows: list[dict[str, Any]] = []
|
||||||
for t in range(num_frames):
|
for t in range(num_frames):
|
||||||
@@ -255,6 +259,39 @@ class AdvantageModule:
|
|||||||
sum(1 for r in rows if r["content"] == "negative"),
|
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:
|
def _compute_threshold(self, advantages: np.ndarray, intervention_mask: np.ndarray) -> float:
|
||||||
"""Compute the binarization threshold as the configured percentile of advantages."""
|
"""Compute the binarization threshold as the configured percentile of advantages."""
|
||||||
non_intervention = advantages[~intervention_mask] if intervention_mask.any() else advantages
|
non_intervention = advantages[~intervention_mask] if intervention_mask.any() else advantages
|
||||||
|
|||||||
Reference in New Issue
Block a user