fix(annotation): remove dropout when doing annotation

This commit is contained in:
Khalil Meftah
2026-07-11 11:09:24 +02:00
parent 5fccaf0477
commit 235e88c743
3 changed files with 8 additions and 34 deletions
@@ -176,7 +176,7 @@ class AdvantageConfig:
enabled: bool = True
# Constant advantage label for all frames (e.g. "positive" for SFT iteration 0).
# Skips VF inference, dropout still applies for CFG.
# Skips VF inference.
constant_value: str | None = None
# Trained value function checkpoint (local path or Hub repo ID).
@@ -195,9 +195,6 @@ class AdvantageConfig:
# Actions with advantage > ε_ get I_t = True (positive).
threshold_percentile: float = 0.3
# Fraction of frames to randomly omit advantage labels (enables CFG).
dropout_rate: float = 0.3
# Force I_t = True for frames marked as human interventions.
force_positive_on_intervention: bool = True
@@ -210,9 +207,6 @@ class AdvantageConfig:
# Batch size for value function inference.
batch_size: int = 32
# Random seed for dropout reproducibility.
seed: int = 1729
@dataclass
class AnnotationPipelineConfig:
@@ -262,13 +262,8 @@ class AdvantageModule:
threshold = self._compute_threshold(advantages, intervention_mask)
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
if (
self.config.force_positive_on_intervention
and intervention_mask[t]
@@ -303,15 +298,11 @@ class AdvantageModule:
)
def _run_constant_mode(self, record: EpisodeRecord, staging: EpisodeStaging) -> None:
"""Emit a fixed advantage value for every frame (with dropout for CFG)."""
"""Emit a fixed advantage value for every frame."""
num_frames = len(record.frame_timestamps)
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
rows.append(
{
"role": "user",
@@ -325,12 +316,11 @@ class AdvantageModule:
staging.write("advantage", rows)
logger.debug(
"Episode %d: %d/%d frames labeled constant '%s' (dropout=%.2f)",
"Episode %d: %d/%d frames labeled constant '%s'",
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:
+5 -15
View File
@@ -111,7 +111,6 @@ def test_binarization_with_mock_values(staging: EpisodeStaging):
cfg = AdvantageConfig(
value_function_path="/fake/vf",
dropout_rate=0.0,
threshold_percentile=0.5,
)
module = AdvantageModule(config=cfg)
@@ -146,7 +145,6 @@ def test_intervention_frames_forced_positive(staging: EpisodeStaging):
cfg = AdvantageConfig(
value_function_path="/fake/vf",
dropout_rate=0.0,
force_positive_on_intervention=True,
)
module = AdvantageModule(config=cfg)
@@ -163,16 +161,13 @@ def test_intervention_frames_forced_positive(staging: EpisodeStaging):
assert rows[2]["content"] == "positive"
def test_dropout_reduces_output_rows(staging: EpisodeStaging):
"""Non-zero dropout rate omits some frames."""
def test_all_frames_labeled(staging: EpisodeStaging):
"""Every frame gets an advantage label (no annotation-level dropout)."""
num_frames = 100
mc_returns = np.linspace(-0.9, -0.1, num_frames).astype(np.float32)
mock_values = np.full(num_frames, -0.5, dtype=np.float32)
cfg = AdvantageConfig(
value_function_path="/fake/vf",
dropout_rate=0.3,
)
cfg = AdvantageConfig(value_function_path="/fake/vf")
module = AdvantageModule(config=cfg)
record = _make_record(num_frames=num_frames, mc_returns=mc_returns)
@@ -183,8 +178,7 @@ def test_dropout_reduces_output_rows(staging: EpisodeStaging):
module.run_episode(record, staging)
rows = staging.read("advantage")
# With 30% dropout on 100 frames, expect ~70 rows (with some variance)
assert 50 < len(rows) < 90
assert len(rows) == num_frames
def test_staged_row_format(staging: EpisodeStaging):
@@ -193,10 +187,7 @@ def test_staged_row_format(staging: EpisodeStaging):
mc_returns = np.array([-0.5, -0.4, -0.3, -0.2, -0.1], dtype=np.float32)
mock_values = np.full(5, -0.3, dtype=np.float32)
cfg = AdvantageConfig(
value_function_path="/fake/vf",
dropout_rate=0.0,
)
cfg = AdvantageConfig(value_function_path="/fake/vf")
module = AdvantageModule(config=cfg)
record = _make_record(num_frames=num_frames, mc_returns=mc_returns)
@@ -225,7 +216,6 @@ def test_n_step_advantage():
cfg = AdvantageConfig(
value_function_path="/fake/vf",
n_step=3,
dropout_rate=0.0,
)
module = AdvantageModule(config=cfg)
record = _make_record(num_frames=num_frames, mc_returns=mc_returns)