diff --git a/skillopt_sleep/consolidate.py b/skillopt_sleep/consolidate.py index 827d0395..754c492c 100644 --- a/skillopt_sleep/consolidate.py +++ b/skillopt_sleep/consolidate.py @@ -244,6 +244,8 @@ def _gate_apply(doc: str, edits: List[EditRecord], which: str) -> str: edits = contrastive_reflect( backend, sets, cand_skill, cand_memory, edit_budget=edit_budget, target="skill", + gate_metric=gate_metric, + gate_mixed_weight=gate_mixed_weight, ) # fall back to single-shot reflect if contrast yielded nothing if not edits: diff --git a/skillopt_sleep/rollout.py b/skillopt_sleep/rollout.py index 8dc2c95c..f889333b 100644 --- a/skillopt_sleep/rollout.py +++ b/skillopt_sleep/rollout.py @@ -19,6 +19,7 @@ from typing import List, Optional, Tuple from skillopt_sleep.backend import Backend, _extract_json +from skillopt_sleep.gate import select_gate_score from skillopt_sleep.replay import replay_one from skillopt_sleep.types import EditRecord, ReplayResult, TaskRecord @@ -97,37 +98,55 @@ def contrastive_reflect( *, edit_budget: int = 4, target: str = "skill", + gate_metric: str = "hard", + gate_mixed_weight: float = 0.5, ) -> List[EditRecord]: """Distill a rule from the contrast between good and bad attempts. - We pick tasks with the highest score *spread* (some attempts passed, some - failed) — those are the most informative — and show the optimizer a - high-scoring vs a low-scoring attempt of each, asking what general rule makes - the good behavior reliable. + We pick tasks with the highest score *spread* under the same objective the + validation gate uses. This matters for soft and mixed gates: attempts may + have identical binary outcomes but materially different partial-credit + scores. The default remains hard-score selection for direct callers that do + not provide gate settings. """ - informative = [rs for rs in rollout_sets if rs.spread > 0 and rs.best and rs.worst] - informative.sort(key=lambda rs: rs.spread, reverse=True) + informative: List[Tuple[float, RolloutSet, ReplayResult, ReplayResult, float, float]] = [] + for rs in rollout_sets: + if len(rs.attempts) < 2: + continue + scored = [ + (select_gate_score(r.hard, r.soft, gate_metric, gate_mixed_weight), r) + for r in rs.attempts + ] + best_score, best = max(scored, key=lambda pair: pair[0]) + worst_score, worst = min(scored, key=lambda pair: pair[0]) + spread = best_score - worst_score + if spread > 0: + informative.append((spread, rs, best, worst, best_score, worst_score)) + informative.sort(key=lambda item: item[0], reverse=True) informative = informative[:6] if not informative: return [] blocks = [] - for rs in informative: + for _spread, rs, best, worst, best_score, worst_score in informative: blocks.append( f"## Task: {rs.task.intent[:160]}\n" - f"- GOOD attempt (score {rs.best.hard:.1f}): {rs.best.response[:200]}\n" - f"- BAD attempt (score {rs.worst.hard:.1f}): {rs.worst.response[:200]}\n" - f" (bad failed: {rs.worst.fail_reason[:100]})" + f"- GOOD attempt ({gate_metric} score {best_score:.3f}; " + f"hard {best.hard:.3f}, soft {best.soft:.3f}): {best.response[:200]}\n" + f"- BAD attempt ({gate_metric} score {worst_score:.3f}; " + f"hard {worst.hard:.3f}, soft {worst.soft:.3f}): {worst.response[:200]}\n" + f" (bad failed: {worst.fail_reason[:100]})" ) # the output contract the proposed rules must not violate (same guardrail the # single-shot reflect uses — prevents harness-violating rules like "return VBA" # or "ask the user for the range" on SpreadsheetBench). from skillopt_sleep.backend import _task_guardrail - guard = _task_guardrail([(rs.task, rs.best) for rs in informative]) + guard = _task_guardrail([(rs.task, best) for _, rs, best, _, _, _ in informative]) prompt = ( "You are SkillOpt's optimizer doing CONTRASTIVE reflection. For each task " - "below the agent was run multiple times; some attempts succeeded and some " - "failed. Identify what the GOOD attempts did that the BAD ones did not, " + "below the agent was run multiple times; some attempts scored better than " + "others under the gate objective. Identify what the GOOD attempts did that " + "the BAD ones did not, " f"and propose at most {edit_budget} SHORT, GENERAL, reusable rules for the " f"{target} that would make the good behavior reliable every time. Quote " "concrete thresholds/formats verbatim; do not paraphrase vaguely. " diff --git a/tests/test_sleep_dream_metric.py b/tests/test_sleep_dream_metric.py new file mode 100644 index 00000000..48a72926 --- /dev/null +++ b/tests/test_sleep_dream_metric.py @@ -0,0 +1,104 @@ +"""Gate-aligned scoring for SkillOpt-Sleep contrastive dream rollouts.""" +from __future__ import annotations + +from unittest import mock + +from skillopt_sleep.backend import Backend, MockBackend +from skillopt_sleep.consolidate import consolidate +from skillopt_sleep.rollout import RolloutSet, contrastive_reflect +from skillopt_sleep.types import ReplayResult, TaskRecord + + +class RecordingBackend(Backend): + name = "recording" + + def __init__(self) -> None: + super().__init__() + self.prompts = [] + + def _call(self, prompt: str, *, max_tokens: int = 1024) -> str: + self.prompts.append(prompt) + return '[{"op":"add","content":"prefer the stronger attempt"}]' + + +def _soft_spread() -> RolloutSet: + task = TaskRecord(id="t1", project="/p", intent="produce the best answer") + return RolloutSet( + task=task, + attempts=[ + ReplayResult(id="t1", hard=1.0, soft=0.2, response="weak partial answer"), + ReplayResult(id="t1", hard=1.0, soft=0.9, response="strong complete answer"), + ], + ) + + +def test_hard_metric_preserves_no_contrast_for_equal_hard_scores(): + backend = RecordingBackend() + + edits = contrastive_reflect(backend, [_soft_spread()], "skill", "") + + assert edits == [] + assert backend.prompts == [] + + +def test_soft_metric_learns_from_partial_credit_spread(): + backend = RecordingBackend() + + edits = contrastive_reflect( + backend, + [_soft_spread()], + "skill", + "", + gate_metric="soft", + ) + + assert len(edits) == 1 + assert "strong complete answer" in backend.prompts[0] + assert "weak partial answer" in backend.prompts[0] + assert "soft score 0.900" in backend.prompts[0] + assert "soft score 0.200" in backend.prompts[0] + + +def test_mixed_metric_uses_the_configured_weight(): + backend = RecordingBackend() + + edits = contrastive_reflect( + backend, + [_soft_spread()], + "skill", + "", + gate_metric="mixed", + gate_mixed_weight=0.25, + ) + + assert len(edits) == 1 + assert "mixed score 0.975" in backend.prompts[0] + assert "mixed score 0.800" in backend.prompts[0] + + +def test_consolidate_passes_gate_objective_to_dream_reflection(): + tasks = [ + TaskRecord(id="train", project="/p", intent="train", split="train"), + TaskRecord(id="val", project="/p", intent="validate", split="val"), + ] + rollout = RolloutSet( + task=tasks[0], + attempts=[ReplayResult(id="train", hard=0.0, soft=0.1)], + ) + + with mock.patch("skillopt_sleep.rollout.multi_rollout", return_value=rollout), mock.patch( + "skillopt_sleep.rollout.contrastive_reflect", return_value=[] + ) as reflect: + consolidate( + MockBackend(), + tasks, + "# skill\n", + "", + rollouts_k=2, + gate_metric="soft", + gate_mixed_weight=0.7, + evolve_memory=False, + ) + + assert reflect.call_args.kwargs["gate_metric"] == "soft" + assert reflect.call_args.kwargs["gate_mixed_weight"] == 0.7