diff --git a/agentdescent/__init__.py b/agentdescent/__init__.py index 8b590aba..03109d0e 100644 --- a/agentdescent/__init__.py +++ b/agentdescent/__init__.py @@ -50,7 +50,7 @@ check_dependency_graph, check_invariants, ) -from .fusion import KeepContradictions, ReflectiveFusion, reflective_merge +from .fusion import KeepContradictions, ReflectiveFusion, LearnableFusionPolicy, reflective_merge from .sampling import DifficultyWeighted, ReplaySampler, RoundRobin, TaskSampler from .replay import ReplayAwareProposal from .selection import ( @@ -267,6 +267,7 @@ "TaskSampler", "ReflectiveFusion", "KeepContradictions", + "LearnableFusionPolicy", "DefaultConflict", "DefaultFusion", "reflective_merge", diff --git a/agentdescent/fusion.py b/agentdescent/fusion.py index 6ed2a277..f44ffe45 100644 --- a/agentdescent/fusion.py +++ b/agentdescent/fusion.py @@ -56,13 +56,16 @@ from __future__ import annotations +import math +import random from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple +from .defaults import DefaultFusion from .evolvable import Diff, Evolvable, combined_rationale from .policies import FusionTrial __all__ = ["MERGE_PROMPT", "KeepContradictions", "ReflectiveFusion", - "reflective_merge"] + "LearnableFusionPolicy", "reflective_merge"] #: What the model is asked, and it is asked for a **union of deltas** rather than @@ -335,3 +338,242 @@ def reflective_merge(complete, **kwargs) -> Dict[str, Any]: """ return {"fusion": ReflectiveFusion(complete, **kwargs), "conflict": KeepContradictions()} + + +class LearnableFusionPolicy: + """A fusion policy that learns when merging wins -- and stops paying for + the ranking it cannot change. + + The tournament (``DefaultFusion(tournament=True)``) answers "does merging + beat the best single?" by scoring every candidate on the cheap layer, every + merge. The answer lands in a :class:`~agentdescent.policies.FusionTrial` + and is then thrown away. This policy learns the answer instead: an online + logistic over **diff-only features** (zero verifier calls) is trained on + the tournament's own trials, and once it is confident that the union is the + right answer for a merge shape, it returns the union **without buying the + cheap sweep**. When the prediction is uncertain -- or says the tournament + would change the outcome -- it runs the tournament and observes a fresh + label. + + This is the allocator's idea (#204) applied to the fusion decision: value- + directed spending of the ranking cost, learned from the run's own history, + with the same safety. Three guarantees: + + * **The default is the default.** With no trials the policy is + :class:`DefaultFusion(tournament=True)` verbatim -- every merge is + ranked, every trial recorded. Learning changes nothing until it has data. + * **Counterfactual verification.** ``epsilon`` of the confident skips are + still run through the tournament, so a wrong prediction is observed + rather than reinforced (the allocator's counterfactual exploration). + * **A rollback guard.** A counterfactual merge the model was *confident* + about skipping but the tournament ranked a single instead breaks the + trust, and after ``max_wrong_skips`` such events (default one) skipping + is retired permanently -- a rule that learned to be wrong is disabled, + not patched. The cost of a false alarm is merely losing the saving; the + cost of a wrong skip is committing a worse merge. + + Commits are unaffected: the skip returns the same union the non-tournament + default would have committed, and the acceptance gate still scores it on + the full held-out set before anything commits. + """ + + def __init__( + self, + inner: Optional[Any] = None, + *, + tournament: bool = True, + seed: int = 0, + min_trials: int = 20, + uncertainty_threshold: float = 0.15, + epsilon: float = 0.1, + max_wrong_skips: int = 1, + lr: float = 0.2, + reg: float = 0.05, + max_logit: float = 3.0, + ) -> None: + if not (0.0 < uncertainty_threshold < 1.0): + raise ValueError("uncertainty_threshold must be in (0, 1)") + if not (0.0 <= epsilon <= 1.0): + raise ValueError("epsilon must be in [0, 1]") + if max_wrong_skips < 1: + raise ValueError("max_wrong_skips must be >= 1") + if not (isinstance(min_trials, int) and min_trials >= 1): + raise ValueError("min_trials must be an int >= 1") + for name, value in (("lr", lr), ("reg", reg), ("max_logit", max_logit)): + if not (isinstance(value, (int, float)) and math.isfinite(value) + and value > 0.0): + raise ValueError(f"{name} must be a finite number > 0") + self.inner = inner or DefaultFusion(tournament=tournament) + #: The engine's budget governor flips this when tokens run low + #: (``aggregator.fusion_policy.tournament = should``); turning it off + #: degrades to the union path, no labels collected. + self.tournament = tournament + #: The engine reads ``aggregator.fusion_policy.trials`` for + #: ``EvolutionResult.fusion_stats``. These are the inner tournament's + #: real measured trials; a skipped merge has none by construction. + self.trials = self.inner.trials + self.min_trials = int(min_trials) + self.threshold = float(uncertainty_threshold) + self.epsilon = float(epsilon) + self.max_wrong_skips = int(max_wrong_skips) + self._rng = random.Random(seed) + self._model = _FusionModel(n_features=5, lr=lr, reg=reg, + max_logit=max_logit) + self._union_policy = DefaultFusion(tournament=False) + #: The budget-degraded union path records unranked trials, and those + #: must land on the **same** list the engine reads -- otherwise + #: ``fusion_stats`` goes silent exactly when the governor flipped the + #: tournament off, which is the one time it should still speak. + self._union_policy.trials = self.inner.trials + self._skip_enabled = True + #: Consecutive counterfactual merges the model would have skipped but + #: the tournament ranked a single instead (a wrong skip in a row). + self._wrong_skips = 0 + #: Confident skips that avoided the cheap sweep. + self.skipped = 0 + #: Merges actually ranked (the labels). + self.tournament_runs = 0 + + # -- engine hooks --------------------------------------------------------- + + def bind(self, verifier: Any) -> None: + """Forward the verifier to the inner tournament and the union path.""" + self.inner.bind(verifier) + self._union_policy.bind(verifier) + + def configure(self, config: Any) -> None: + pass + + # -- the fusion decision -------------------------------------------------- + + def select(self, artifact: Evolvable, + diffs: List[Diff]) -> Tuple[Diff, Evolvable, bool]: + """Learn or predict: skip the ranking sweep when the union is + confidently predicted to win (a learned logistic, not a proof).""" + if not self.tournament: + return self._union_policy.select(artifact, diffs) + union, _reason = self._union_of(diffs) + if (self._skip_enabled and self._model.n_obs >= self.min_trials + and union is not None + and self._rng.random() >= self.epsilon): + p = self._model.predict(self._features(diffs, union)) + if p >= 1.0 - self.threshold: + self.skipped += 1 + return self._union_policy.select(artifact, diffs) + return self._run_tournament(artifact, diffs) + + def _run_tournament(self, artifact: Evolvable, + diffs: List[Diff]) -> Tuple[Diff, Evolvable, bool]: + """Rank the candidates (the label-producing path), then observe.""" + self.tournament_runs += 1 + result = self.inner.select(artifact, diffs) + fused = bool(result[2]) + union, _reason = self._union_of(diffs) + self._observe(self._features(diffs, union), fused) + return result + + def _observe(self, features: Tuple[float, ...], fused: bool) -> None: + # Snapshot the **pre-update** prediction first. The rollback question + # is "would the rule that was in force have skipped this merge?" -- and + # a negative label immediately drags the logistic down, often below the + # confidence cut in one step, so scoring after the update misses exactly + # the confident-and-wrong counterfactual the guard exists for. + predicted_skip = self._model.predict(features) >= 1.0 - self.threshold + self._model.update(features, fused) + if self._model.n_obs < self.min_trials: + return + # Skipping is right iff the tournament also picked the union + # (``fused``). The rollback guard is deliberately unforgiving: a + # counterfactual merge the model was *confident* about skipping but the + # tournament ranked a single instead breaks the trust, and after + # ``max_wrong_skips`` such events (default one) skipping is retired + # permanently. The cost of a false alarm is merely losing the saving; + # the cost of a wrong skip is committing a worse merge. Conservative is + # correct here. + if predicted_skip and not fused: + self._wrong_skips += 1 + if self._wrong_skips >= self.max_wrong_skips: + self._skip_enabled = False + else: + self._wrong_skips = 0 + + # -- features (diff-only, zero verifier calls) ----------------------------- + + def _union_of(self, diffs: List[Diff]) -> Tuple[Optional[Diff], str]: + """The fused candidate, via the inner policy's own rule. + + Deliberately delegated rather than copied: ``DefaultFusion._union_of`` + is the single source of the three-way "nothing-to-fuse" rule, and a + byte-for-byte copy here would silently drift if that rule ever changed + again. The inner is the tournament policy the learner wraps; it must + have a ``_union_of`` (``DefaultFusion`` does).""" + return self.inner._union_of(diffs) + + @staticmethod + def _features(diffs: List[Diff], union: Optional[Diff]) -> Tuple[float, ...]: + """The pre-fusion signal the learner decides on: merge shape, size and + key overlap, all computable from the diffs alone.""" + import math as _math + + keys = [set(d.ops) for d in diffs] + union_keys = set().union(*keys) if keys else set() + inter = set.intersection(*keys) if keys else set() + overlap = len(inter) / len(union_keys) if union_keys else 0.0 + total_ops = sum(len(d.ops) for d in diffs) + max_ops = max((len(d.ops) for d in diffs), default=0) + return ( + _math.log1p(len(diffs)), + overlap, + _math.log1p(total_ops), + _math.log1p(max_ops), + 1.0 if union is not None else 0.0, + ) + + def status(self) -> dict: + # The skip fraction's denominator is *all* merges decided (skipped or + # ranked) -- "skips / tournament_runs" would climb past 1.0 on a + # warm skip-heavy run and read as an occupancy instead of a fraction. + decided = self.skipped + self.tournament_runs + return { + "skipped": self.skipped, + "tournament_runs": self.tournament_runs, + "model_observations": self._model.n_obs, + "skip_enabled": self._skip_enabled, + "skips_per_merge": (self.skipped / decided + if decided else 0.0), + } + + +class _FusionModel: + """A tiny online logistic over the fusion features. + + Deliberately a private twin of the allocator's + :class:`~agentdescent.allocator.LearningValueModel` rather than a reuse: + same idea (momentum learning rate, L2, logit clip), different feature count + and weight layout (bias at ``[-1]`` here), and sharing the allocator's + hard-coded three-feature class would force this module's five features into + a layout that does not fit. None of this is a research contribution -- it + is the smallest thing that learns without diverging on one noisy run.""" + + def __init__(self, n_features: int, lr: float = 0.2, reg: float = 0.05, + max_logit: float = 3.0) -> None: + self.n_features = int(n_features) + self.lr = lr + self.reg = reg + self.max_logit = max_logit + self._w = [0.0] * (n_features + 1) + self.n_obs = 0 + + def predict(self, features: Sequence[float]) -> float: + logit = self._w[-1] + sum( + w * x for w, x in zip(self._w[:-1], features)) + logit = max(-self.max_logit, min(self.max_logit, logit)) + return 1.0 / (1.0 + math.exp(-logit)) + + def update(self, features: Sequence[float], positive: bool) -> None: + p = self.predict(features) + err = (1.0 if positive else 0.0) - p + self._w[-1] += self.lr * err - self.lr * self.reg * self._w[-1] + for i, x in enumerate(features): + self._w[i] += self.lr * err * x - self.lr * self.reg * self._w[i] + self.n_obs += 1 diff --git a/changelog.d/298-learnable-fusion.md b/changelog.d/298-learnable-fusion.md new file mode 100644 index 00000000..feb54910 --- /dev/null +++ b/changelog.d/298-learnable-fusion.md @@ -0,0 +1,15 @@ +Added +LearnableFusionPolicy (#298): a `FusionPolicy` that learns when merging wins. +An online logistic over diff-only features (candidate count, key overlap, diff +sizes) is trained on the tournament's own `FusionTrial`s, and once confident +that the union is the right answer for a merge shape it returns the union +**without buying the cheap ranking sweep**; uncertain merges still run the +tournament and observe a fresh label. The rollback guard scores the **pre- +update** prediction -- the rule actually in force -- and retires skipping +permanently after a confident-and-wrong counterfactual (default +`max_wrong_skips=1`). The budget-governor's `tournament=False` degrades to the +plain union and keeps recording unranked trials on the list `fusion_stats` +reads. With no trials the policy is `DefaultFusion(tournament=True)` verbatim. +Commits are unaffected: a skip returns the same union the non-tournament +default would have committed, and the acceptance gate still scores it on the +full held-out set. diff --git a/docs/api.md b/docs/api.md index b99e04d1..15f0e6ab 100644 --- a/docs/api.md +++ b/docs/api.md @@ -13,7 +13,7 @@ means the parameter has none. Each section links to the page that explains *why* the module is shaped the way it is; this page is the *what*. -342 public names across 56 modules. +343 public names across 56 modules. --- @@ -2877,6 +2877,31 @@ Combine competing values for the same key, when a dict update cannot.  ยท&n A conflict policy that leaves contradicting diffs for fusion to resolve. +### `LearnableFusionPolicy(...)` + +A fusion policy that learns when merging wins -- and stops paying for the ranking it cannot change. + +```python +LearnableFusionPolicy( + inner: Optional[Any] = None, + *, + tournament: bool = True, + seed: int = 0, + min_trials: int = 20, + uncertainty_threshold: float = 0.15, + epsilon: float = 0.1, + max_wrong_skips: int = 1, + lr: float = 0.2, + reg: float = 0.05, + max_logit: float = 3.0 +) -> None +``` + +| method | what it does | +|---|---| +| `bind(verifier: Any) -> None` | Forward the verifier to the inner tournament and the union path. | +| `select(artifact: Evolvable, diffs: List[Diff]) -> Tuple[Diff, Evolvable, bool]` | Learn or predict: skip the ranking sweep when the union is confidently predicted to win (a learned logistic, not a proof). | + ### `ReflectiveFusion(...)` Combine contradicting diffs by asking a model to synthesise their values. diff --git a/tests/test_learnable_fusion.py b/tests/test_learnable_fusion.py new file mode 100644 index 00000000..e3accae0 --- /dev/null +++ b/tests/test_learnable_fusion.py @@ -0,0 +1,202 @@ +"""The learned fusion decision has to be right in the places a learner is +dangerous and a merge is load-bearing. + +A policy that learns "when does merging win" from the tournament's own trials +must keep the default as the default (no data -> tournament verbatim), must +only skip the ranking sweep when the prediction is confident, must observe its +mistakes through counterfactual exploration, and must retire itself when it +starts being wrong. Every fixture here pins one of those. +""" + +from agentdescent import Diff, EvolvingArtifact, SingleSlot +from agentdescent.fusion import LearnableFusionPolicy + + +def _setup(): + base = EvolvingArtifact("artifact", {"k": "v"}, strategy=SingleSlot()) + d1 = Diff(diff_id="1", target="artifact", ops={"a": "x"}) + d2 = Diff(diff_id="2", target="artifact", ops={"b": "y"}) + return base, [d1, d2] + + +class _ScoreVerifier: + """A cheap layer whose score is a pure function of the candidate's state.""" + + def __init__(self, fn): + self.fn = fn + + def cheap_eval(self, artifact): + return self.fn(artifact) + + +def _union_wins(artifact): + return 0.5 + 0.1 * len(artifact.state) # union has the most keys + + +def _single_wins(artifact): + return 0.9 - 0.1 * len(artifact.state) # singles have fewer keys + + +def _tie(artifact): + return 0.6 # everyone equal + + +def test_default_is_the_tournament_verbatim(): + base, diffs = _setup() + policy = LearnableFusionPolicy(seed=0) + policy.bind(_ScoreVerifier(_union_wins)) + out = policy.select(base, diffs) + assert out[2] is True # the union was ranked and won + assert policy.tournament_runs == 1 + assert policy.skipped == 0 + assert len(policy.trials) == 1 # a real measured trial + + +def test_features_are_diff_only_and_zero_verifier_calls(): + base, diffs = _setup() + policy = LearnableFusionPolicy(seed=0) + union, _ = policy._union_of(diffs) + f = LearnableFusionPolicy._features(diffs, union) + assert len(f) == 5 + assert f[0] > 0 # log1p(n_candidates) + assert 0.0 <= f[1] <= 1.0 # overlap fraction + assert f[4] == 1.0 # a union exists + none, _ = policy._union_of(diffs[:1]) + f2 = LearnableFusionPolicy._features(diffs[:1], none) + assert f2[4] == 0.0 + + +def test_the_policy_learns_to_skip_the_ranking_sweep(): + base, diffs = _setup() + policy = LearnableFusionPolicy(seed=0, epsilon=0.0, min_trials=20) + policy.bind(_ScoreVerifier(_union_wins)) + for _ in range(60): + policy.select(base, diffs) + assert policy.skipped >= 1 # once warm, the union wins -> skip + assert policy.status()["skips_per_merge"] > 0 + # a skipped merge returns the union without ranking it + out = policy.select(base, diffs) + assert out[2] is True + assert set(out[0].ops) == {"a", "b"} + + +def test_the_policy_does_not_skip_when_the_union_loses(): + base, diffs = _setup() + policy = LearnableFusionPolicy(seed=0, epsilon=0.0, min_trials=20) + policy.bind(_ScoreVerifier(_single_wins)) + for _ in range(60): + policy.select(base, diffs) + assert policy.skipped == 0 # a single wins -> never confident + assert policy.status()["skip_enabled"] is True + + +def test_a_tie_is_read_conservatively_and_never_skipped(): + base, diffs = _setup() + policy = LearnableFusionPolicy(seed=0, epsilon=0.0, min_trials=20) + policy.bind(_ScoreVerifier(_tie)) + for _ in range(60): + policy.select(base, diffs) + assert policy.skipped == 0 + # the conservative reading: a fusion that merely ties does not win + assert all(t.winner != "fused" for t in policy.trials) + + +def test_the_rollback_guard_retires_a_wrong_learner(): + # Regime 1: the union always wins (the model learns to skip). Regime 2: a + # single starts winning. The counterfactual exploration proves the model + # confident-and-wrong, and the guard retires skipping permanently. + base, diffs = _setup() + mode = {"single_wins": False} + + def two_regime(artifact): + return _single_wins(artifact) if mode["single_wins"] else _union_wins(artifact) + + policy = LearnableFusionPolicy(seed=0, epsilon=1.0, min_trials=15) + policy.bind(_ScoreVerifier(two_regime)) + for _ in range(25): + policy.select(base, diffs) # warm-up on union-wins + assert policy.status()["skip_enabled"] is True + mode["single_wins"] = True + for _ in range(10): + policy.select(base, diffs) # regime flips; guard must fire + assert policy.status()["skip_enabled"] is False + + +def test_turning_the_tournament_off_degrades_to_the_union_path(): + base, diffs = _setup() + policy = LearnableFusionPolicy(seed=0, tournament=True) + policy.bind(_ScoreVerifier(_union_wins)) + policy.tournament = False # the budget governor's flip + out = policy.select(base, diffs) + assert out[2] is True # the plain union + assert policy.skipped == 0 + assert policy.tournament_runs == 0 + # the unranked trial still lands on the list the engine reads: a budget + # governor that turned the tournament off must not make fusion_stats silent + assert len(policy.trials) == 1 + assert policy.trials[0].ranked is False + + +def test_the_guard_scores_the_pre_update_prediction(): + # The rollback question is "would the rule *in force* have skipped this + # merge?". A negative label drags the logistic down -- often below the + # confidence cut in one step -- so scoring after the update misses the + # confident-and-wrong counterfactual. Inject barely-confident weights + # (p ~ 0.86 pre-update, ~0.84 after one negative label) and prove the + # guard still fires. + base, diffs = _setup() + policy = LearnableFusionPolicy(seed=0, min_trials=5) + union, _ = policy._union_of(diffs) + features = policy._features(diffs, union) + # bias = logit for p = 0.86: ln(0.86/0.14) ~ 1.815, all feature weights 0 + import math as _m + policy._model._w = [0.0, 0.0, 0.0, 0.0, _m.log(0.86 / 0.14)] + assert abs(policy._model.predict(features) - 0.86) < 0.02 # barely confident + policy._model.n_obs = policy.min_trials # past warm-up + policy._observe(features, False) # a wrong skip + assert policy.status()["skip_enabled"] is False + + +def test_bind_forwards_the_verifier_to_the_inner_tournament(): + base, diffs = _setup() + policy = LearnableFusionPolicy(seed=0) + verifier = _ScoreVerifier(_union_wins) + policy.bind(verifier) + assert policy.inner.verifier is verifier + assert policy._union_policy.verifier is verifier + + +def test_bad_knobs_are_refused(): + import pytest + with pytest.raises(ValueError): + LearnableFusionPolicy(uncertainty_threshold=0.0) + with pytest.raises(ValueError): + LearnableFusionPolicy(epsilon=1.5) + + +def test_the_policy_runs_end_to_end_through_evolve(tmp_path): + # Installed as the engine's fusion slot, the learner ranks every merge + # (collecting labels) and reports fusion_stats -- the pipeline the + # tournament already provides, with the learner in the seat. + from agentdescent import Policies, evolve + from agentdescent.evolution import SingleSlot, Task + + tasks = [Task(id=f"t{i}", prompt=f"item {i}") for i in range(8)] + + def run(rendered, task): + return "answer" + (" 2026" if "year" in rendered else "") + + def reward(task, output): + return 1.0 if "2026" in output else 0.0 + + policy = LearnableFusionPolicy(seed=0) + result = evolve( + tasks, reward, run=run, + propose=lambda r, t, o, re: "always state the year", + strategy=SingleSlot(initial_value="seed"), rounds=2, n_workers=2, + repo_path=str(tmp_path / "repo"), + policies=Policies(fusion=policy)) + assert result.outcomes().get("committed", 0) >= 1 + assert policy.tournament_runs >= 1 + stats = result.fusion_stats() + assert stats.trials >= 1