From 1f1e8774473d7c11e9ae541bd5aa8a316983074a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christian=20Sch=C3=BCller?= Date: Mon, 10 Aug 2026 11:54:20 +0200 Subject: [PATCH 1/3] fix(experimentalist): retain the baseline when no candidate beats it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Finalization picked the winner with `select_diverse_survivors(..., k=1)`. That agent is prompted to always include a candidate created in the current round, so at k=1 the rule consumes the only slot and the round-0 baseline can never be returned. Every run therefore shipped a diff, including runs where no candidate measured better than doing nothing. Winner selection is not the same question as survivor selection. Picking a set to carry forward wants diversity and fresh blood; picking a winner wants the best measured result, and no judgement call. Finalization now uses `select_winner_node`: the Pareto front by the configured objectives, with the lowest round winning ties. A candidate that genuinely improves dominates its ancestor and removes it from the front, so the age preference only ever decides ties and never blocks a real gain. Mid-loop survivor selection is unchanged and still uses the diversity selector. Objective-metric eligibility (`has_metric_dimensions`) and `pareto_objectives` ranking are both preserved. The previous behaviour was correct only by accident: `pareto_front` preserves input order and agent-0 is inserted first, so ties fell to the baseline. The rule is now stated explicitly rather than relying on dict ordering. This path had no test coverage, which is how the regression shipped. Adds `test_winner_selection.py` pinning the tie case, order-independence, real improvements, regressions, multi-objective trades, and minimized objectives. Signed-off-by: Christian Schüller --- .../experimentalist/components/loop.py | 34 +++-- .../experimentalist/test_winner_selection.py | 142 ++++++++++++++++++ 2 files changed, 165 insertions(+), 11 deletions(-) create mode 100644 plugins/nemo-experimentalist/tests/experimentalist/test_winner_selection.py diff --git a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/loop.py b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/loop.py index 4c7c7eafb6..14b3e210ff 100644 --- a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/loop.py +++ b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/loop.py @@ -52,8 +52,10 @@ restore_heldout_splits, ) from nemo_experimentalist_plugin.experimentalist.components.models import ( + EvolutionNode, EvolutionTree, OptimizationType, + pareto_front, pareto_sort, ) from nemo_experimentalist_plugin.experimentalist.components.proposer import Improvement, Proposer @@ -129,6 +131,25 @@ def _coerce_optimization_type(optimization_type: str | None) -> OptimizationType return None +def select_winner_node( + eligible: list[EvolutionNode], + objective_function: list[MetricTarget], +) -> EvolutionNode | None: + """Return the run's winner: non-dominated, oldest candidate wins ties. + + Deliberately *not* ``select_diverse_survivors``. Picking a winner and picking a + set to carry into the next round are different questions. The selector is told to + always include a candidate created this round, so at ``k=1`` that rule consumes the + only slot and the round-0 baseline can never be returned -- meaning a run always + ships a diff, even one that measured no better than doing nothing. + + Preferring the lowest round only decides ties: a candidate that genuinely improves + dominates its ancestor and removes it from the front before this is consulted. + """ + front = pareto_front(eligible, lambda node: pareto_objectives(node.val_reward, objective_function)) + return min(front, key=lambda node: node.round) if front else None + + def _with_insight_objective( config: EvolutionaryOptimizerConfig, metric_keys: tuple[str, ...] ) -> EvolutionaryOptimizerConfig: @@ -1786,17 +1807,8 @@ async def _finalize( # Only survivors that actually have a validation reward are eligible winners. scored = [n for n in evolution_tree.nodes.values() if n.is_survivor and n.val_reward] eligible = [node for node in scored if has_metric_dimensions(node.val_reward, self.config.objective_function)] - ranked_nodes = pareto_sort( - eligible, - lambda node: pareto_objectives(node.val_reward, self.config.objective_function), - ) - finalists = await self.select_diverse_survivors( - [node.candidate for node in ranked_nodes], - 1, - self.config.objective_function, - self.config.regression_metrics, - ) - best_id = finalists[0].label if finalists else None + best = select_winner_node(eligible, self.config.objective_function) + best_id = best.label if best else None restore_heldout_splits(self.working_dir) diff --git a/plugins/nemo-experimentalist/tests/experimentalist/test_winner_selection.py b/plugins/nemo-experimentalist/tests/experimentalist/test_winner_selection.py new file mode 100644 index 0000000000..7422ba6c0f --- /dev/null +++ b/plugins/nemo-experimentalist/tests/experimentalist/test_winner_selection.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Pin how the run's winner is chosen at finalization. + +This had no coverage, and a regression shipped through the gap: finalization was +routed through ``select_diverse_survivors``, whose prompt requires including a +candidate created in the current round. At ``k=1`` that rule takes the only slot, +so the round-0 baseline could never win and every run shipped a diff -- including +runs where nothing measured better than doing nothing. + +The rule these tests encode: the winner is non-dominated, and among candidates +nothing dominates, the oldest wins. A real improvement dominates its ancestor and +removes it from the front, so the age preference only ever decides ties. +""" + +from __future__ import annotations + +from nemo_experimentalist_plugin.config import MetricTarget +from nemo_experimentalist_plugin.entities import Candidate, RewardRecord +from nemo_experimentalist_plugin.experimentalist.components.loop import select_winner_node +from nemo_experimentalist_plugin.experimentalist.components.models import EvolutionNode + +MAXIMIZE = [MetricTarget(name="reward", direction="maximize")] + + +def _node(label: str, round_: int, **metrics: float) -> EvolutionNode: + return EvolutionNode( + candidate=Candidate( + run_id="run-1", + label=label, + round=round_, + optimization="baseline" if round_ == 0 else "some change", + rewards={"validation": RewardRecord(metrics=dict(metrics))}, + ) + ) + + +def test_baseline_wins_a_tie() -> None: + """The regression this file exists for: equal score means the diff bought nothing.""" + winner = select_winner_node( + [_node("agent-0", 0, reward=0.5), _node("agent-1", 1, reward=0.5)], + MAXIMIZE, + ) + assert winner is not None + assert winner.label == "agent-0" + + +def test_baseline_wins_a_tie_regardless_of_input_order() -> None: + """The old behaviour was correct only because agent-0 happened to be inserted first.""" + winner = select_winner_node( + [_node("agent-1", 1, reward=0.5), _node("agent-0", 0, reward=0.5)], + MAXIMIZE, + ) + assert winner is not None + assert winner.label == "agent-0" + + +def test_a_real_improvement_wins() -> None: + """Age must never block a candidate that actually scored better.""" + winner = select_winner_node( + [_node("agent-0", 0, reward=0.5), _node("agent-1", 1, reward=0.7)], + MAXIMIZE, + ) + assert winner is not None + assert winner.label == "agent-1" + + +def test_a_regression_loses() -> None: + winner = select_winner_node( + [_node("agent-0", 0, reward=0.5), _node("agent-1", 1, reward=0.3)], + MAXIMIZE, + ) + assert winner is not None + assert winner.label == "agent-0" + + +def test_oldest_ancestor_wins_among_tied_descendants() -> None: + """Not just the baseline: an ancestor carried forward unchanged also outranks a tie.""" + winner = select_winner_node( + [ + _node("agent-3", 3, reward=0.8), + _node("agent-1", 1, reward=0.8), + _node("agent-2", 2, reward=0.8), + ], + MAXIMIZE, + ) + assert winner is not None + assert winner.label == "agent-1" + + +def test_incomparable_multi_objective_candidate_does_not_displace_the_baseline() -> None: + """Better on one objective and worse on another is not an improvement, it is a trade.""" + objectives = [ + MetricTarget(name="reward", direction="maximize"), + MetricTarget(name="shape_ok", direction="maximize"), + ] + winner = select_winner_node( + [ + _node("agent-0", 0, reward=0.5, shape_ok=0.9), + _node("agent-1", 1, reward=0.7, shape_ok=0.4), + ], + objectives, + ) + assert winner is not None + assert winner.label == "agent-0" + + +def test_dominating_on_every_objective_wins() -> None: + objectives = [ + MetricTarget(name="reward", direction="maximize"), + MetricTarget(name="shape_ok", direction="maximize"), + ] + winner = select_winner_node( + [ + _node("agent-0", 0, reward=0.5, shape_ok=0.5), + _node("agent-1", 1, reward=0.7, shape_ok=0.9), + ], + objectives, + ) + assert winner is not None + assert winner.label == "agent-1" + + +def test_minimized_objective_direction_is_respected() -> None: + """Lower is better here, so the baseline's higher cost must lose.""" + winner = select_winner_node( + [_node("agent-0", 0, cost=0.9), _node("agent-1", 1, cost=0.2)], + [MetricTarget(name="cost", direction="minimize")], + ) + assert winner is not None + assert winner.label == "agent-1" + + +def test_no_eligible_candidates_returns_none() -> None: + assert select_winner_node([], MAXIMIZE) is None + + +def test_single_candidate_is_the_winner() -> None: + winner = select_winner_node([_node("agent-0", 0, reward=0.4)], MAXIMIZE) + assert winner is not None + assert winner.label == "agent-0" From be2f2c38a4776750818d1efc8e952a8b3e596c6f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christian=20Sch=C3=BCller?= Date: Mon, 10 Aug 2026 12:27:21 +0200 Subject: [PATCH 2/3] fix(experimentalist): bar a winner that regresses a protected metric MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-up. Three gaps in the first pass. `regression_metrics` was dropped. The old finalization call passed it to the selector; ranking on the objectives alone let a candidate buy an objective gain with a regression the config says must not happen. Candidates that worsen a regression target against the baseline are now filtered out before ranking. The baseline is exempt from its own comparison, so a round where every candidate regresses keeps the baseline rather than shipping the least-bad regression. A target absent from either side is skipped rather than treated as a regression: the evaluator not reporting a dimension is not evidence of harm. Tie-breaking relied on input order. `min` returns the first minimum and `pareto_front` preserves input order, so two same-round candidates that nothing dominates could swap winners if the caller reordered its input -- the same implicit-ordering dependence this change set out to remove, one level down. Ties now fall to creation order, recovered from the run-scoped label's numeric suffix so `agent-9` precedes `agent-10`. The docstring described how the previous behaviour broke. That belongs in history, not in an interface contract, so it now states what the function guarantees and the reasoning stays in the commit that made the change. Adds nine tests: regression in both directions, an objective gain that cannot pay for one, multiple objectives alongside a protected metric, every candidate regressing, a metric missing from one side, reversed input at equal round, and numeric-versus-lexicographic creation order. Signed-off-by: Christian Schüller --- .../experimentalist/components/loop.py | 59 +++++++-- .../experimentalist/test_winner_selection.py | 117 +++++++++++++++++- 2 files changed, 164 insertions(+), 12 deletions(-) diff --git a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/loop.py b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/loop.py index 14b3e210ff..aadec439d1 100644 --- a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/loop.py +++ b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/loop.py @@ -131,23 +131,60 @@ def _coerce_optimization_type(optimization_type: str | None) -> OptimizationType return None +def _creation_order(node: EvolutionNode) -> tuple[int, str]: + """Order candidates by creation, independently of how the caller ordered them. + + Labels are run-scoped and assigned in sequence (``agent-0``, ``agent-1``, ...), so + the numeric suffix recovers creation order. The label itself is the final component + so the key stays total if a label ever stops matching that shape. + """ + _, _, suffix = node.label.rpartition("-") + return (int(suffix) if suffix.isdigit() else 0, node.label) + + +def _regresses( + node: EvolutionNode, + baseline: EvolutionNode, + regression_metrics: list[MetricTarget], +) -> bool: + """Whether *node* is worse than *baseline* on any metric that must not regress. + + A target missing from either side is skipped: absent evidence is not a regression. + """ + if node is baseline: + return False + for target in regression_metrics: + before = baseline.val_reward.get(target.name) + after = node.val_reward.get(target.name) + if before is None or after is None: + continue + if after < before if target.direction == "maximize" else after > before: + return True + return False + + def select_winner_node( eligible: list[EvolutionNode], objective_function: list[MetricTarget], + regression_metrics: list[MetricTarget] | None = None, ) -> EvolutionNode | None: - """Return the run's winner: non-dominated, oldest candidate wins ties. + """Return the run's winner, or None when there is nothing eligible. - Deliberately *not* ``select_diverse_survivors``. Picking a winner and picking a - set to carry into the next round are different questions. The selector is told to - always include a candidate created this round, so at ``k=1`` that rule consumes the - only slot and the round-0 baseline can never be returned -- meaning a run always - ships a diff, even one that measured no better than doing nothing. + Candidates that worsen a ``regression_metrics`` target against the baseline are + dropped before ranking, so a gain on the objectives cannot pay for a regression. + The winner is then taken from the Pareto front over ``objective_function``, with + ties going to the lowest round and then to creation order. - Preferring the lowest round only decides ties: a candidate that genuinely improves - dominates its ancestor and removes it from the front before this is consulted. + The baseline is the oldest eligible candidate and is never dropped, so a run whose + candidates all regress -- or all merely tie -- keeps it rather than shipping a diff + that bought nothing. """ - front = pareto_front(eligible, lambda node: pareto_objectives(node.val_reward, objective_function)) - return min(front, key=lambda node: node.round) if front else None + if not eligible: + return None + baseline = min(eligible, key=lambda node: (node.round, _creation_order(node))) + kept = [node for node in eligible if not _regresses(node, baseline, regression_metrics or [])] + front = pareto_front(kept, lambda node: pareto_objectives(node.val_reward, objective_function)) + return min(front, key=lambda node: (node.round, _creation_order(node))) if front else None def _with_insight_objective( @@ -1807,7 +1844,7 @@ async def _finalize( # Only survivors that actually have a validation reward are eligible winners. scored = [n for n in evolution_tree.nodes.values() if n.is_survivor and n.val_reward] eligible = [node for node in scored if has_metric_dimensions(node.val_reward, self.config.objective_function)] - best = select_winner_node(eligible, self.config.objective_function) + best = select_winner_node(eligible, self.config.objective_function, self.config.regression_metrics) best_id = best.label if best else None restore_heldout_splits(self.working_dir) diff --git a/plugins/nemo-experimentalist/tests/experimentalist/test_winner_selection.py b/plugins/nemo-experimentalist/tests/experimentalist/test_winner_selection.py index 7422ba6c0f..13b1cea389 100644 --- a/plugins/nemo-experimentalist/tests/experimentalist/test_winner_selection.py +++ b/plugins/nemo-experimentalist/tests/experimentalist/test_winner_selection.py @@ -9,7 +9,8 @@ so the round-0 baseline could never win and every run shipped a diff -- including runs where nothing measured better than doing nothing. -The rule these tests encode: the winner is non-dominated, and among candidates +The rule these tests encode: a candidate that regresses a protected metric is +dropped, the winner is non-dominated on the objectives, and among candidates nothing dominates, the oldest wins. A real improvement dominates its ancestor and removes it from the front, so the age preference only ever decides ties. """ @@ -22,6 +23,12 @@ from nemo_experimentalist_plugin.experimentalist.components.models import EvolutionNode MAXIMIZE = [MetricTarget(name="reward", direction="maximize")] +TWO_OBJECTIVES = [ + MetricTarget(name="reward", direction="maximize"), + MetricTarget(name="shape_ok", direction="maximize"), +] +LATENCY_MUST_NOT_RISE = [MetricTarget(name="latency", direction="minimize")] +SHAPE_MUST_NOT_FALL = [MetricTarget(name="shape_ok", direction="maximize")] def _node(label: str, round_: int, **metrics: float) -> EvolutionNode: @@ -140,3 +147,111 @@ def test_single_candidate_is_the_winner() -> None: winner = select_winner_node([_node("agent-0", 0, reward=0.4)], MAXIMIZE) assert winner is not None assert winner.label == "agent-0" + + +def test_same_round_ties_are_stable_under_input_order() -> None: + """Two same-round candidates nothing dominates must not depend on caller ordering.""" + forward = select_winner_node( + [_node("agent-1", 1, reward=0.7), _node("agent-2", 1, reward=0.7)], + MAXIMIZE, + ) + reversed_ = select_winner_node( + [_node("agent-2", 1, reward=0.7), _node("agent-1", 1, reward=0.7)], + MAXIMIZE, + ) + assert forward is not None and reversed_ is not None + assert forward.label == reversed_.label == "agent-1" + + +def test_creation_order_is_numeric_not_lexicographic() -> None: + """`agent-9` was created before `agent-10`; a string sort would invert that.""" + winner = select_winner_node( + [_node("agent-10", 2, reward=0.6), _node("agent-9", 2, reward=0.6)], + MAXIMIZE, + ) + assert winner is not None + assert winner.label == "agent-9" + + +def test_objective_gain_cannot_pay_for_a_regression() -> None: + """A protected metric that worsens disqualifies the candidate however good the objective.""" + winner = select_winner_node( + [ + _node("agent-0", 0, reward=0.5, latency=1.0), + _node("agent-1", 1, reward=0.9, latency=2.5), + ], + MAXIMIZE, + LATENCY_MUST_NOT_RISE, + ) + assert winner is not None + assert winner.label == "agent-0" + + +def test_improvement_that_holds_the_protected_metric_still_wins() -> None: + winner = select_winner_node( + [ + _node("agent-0", 0, reward=0.5, latency=1.0), + _node("agent-1", 1, reward=0.9, latency=1.0), + ], + MAXIMIZE, + LATENCY_MUST_NOT_RISE, + ) + assert winner is not None + assert winner.label == "agent-1" + + +def test_maximized_regression_metric_uses_the_opposite_direction() -> None: + """`shape_ok` must not *fall*; a candidate that improves reward but drops it is out.""" + winner = select_winner_node( + [ + _node("agent-0", 0, reward=0.5, shape_ok=1.0), + _node("agent-1", 1, reward=0.9, shape_ok=0.6), + ], + MAXIMIZE, + SHAPE_MUST_NOT_FALL, + ) + assert winner is not None + assert winner.label == "agent-0" + + +def test_multiple_objectives_and_a_regression_metric_together() -> None: + """agent-1 dominates on both objectives but regresses latency; agent-2 is the honest win.""" + winner = select_winner_node( + [ + _node("agent-0", 0, reward=0.5, shape_ok=0.5, latency=1.0), + _node("agent-1", 1, reward=0.9, shape_ok=0.9, latency=3.0), + _node("agent-2", 1, reward=0.7, shape_ok=0.7, latency=0.8), + ], + TWO_OBJECTIVES, + LATENCY_MUST_NOT_RISE, + ) + assert winner is not None + assert winner.label == "agent-2" + + +def test_every_candidate_regressing_falls_back_to_the_baseline() -> None: + winner = select_winner_node( + [ + _node("agent-0", 0, reward=0.5, latency=1.0), + _node("agent-1", 1, reward=0.9, latency=2.0), + _node("agent-2", 1, reward=0.8, latency=1.5), + ], + MAXIMIZE, + LATENCY_MUST_NOT_RISE, + ) + assert winner is not None + assert winner.label == "agent-0" + + +def test_a_regression_metric_missing_from_a_candidate_is_not_a_regression() -> None: + """Absent evidence must not disqualify: the evaluator simply did not report it.""" + winner = select_winner_node( + [ + _node("agent-0", 0, reward=0.5, latency=1.0), + _node("agent-1", 1, reward=0.9), + ], + MAXIMIZE, + LATENCY_MUST_NOT_RISE, + ) + assert winner is not None + assert winner.label == "agent-1" From e959ca41d53658f8a1cb00ec3e5d1d7e634e9ce8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christian=20Sch=C3=BCller?= Date: Mon, 10 Aug 2026 12:37:05 +0200 Subject: [PATCH 3/3] fix(experimentalist): keep the baseline eligible when it is killed mid-loop MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-up. Baseline retention still had a second way to fail. The candidate pool passed to survivor selection starts as the baseline and stays in it every round -- `candidates = list(evolution_tree.survivors(0))`, then `candidates = survivors + new_candidates`. Everything the selector does not keep gets a `killed_round`, and that selector is the one told to always keep a candidate created this round. So the baseline can be killed mid-loop. `_finalize` filtered on `is_survivor`, so a killed baseline never reached winner selection. Two consequences, both silent: a run whose baseline was killed could only choose among diffs, and the regression comparison rehomed onto the oldest surviving candidate, measuring "does not regress" against a candidate that had itself already regressed. Adds `finalization_pool`, which re-admits the round-0 node regardless of `is_survivor` and returns it separately as the regression reference. Survivor selection and finalization answer different questions -- "what is worth carrying forward" versus "is any of this better than shipping nothing" -- and only the second requires the baseline to remain on the table. `select_winner_node` now takes that baseline explicitly rather than inferring it from the oldest eligible node, which is only the same thing when the baseline survived. A baseline lacking the configured objective dimensions cannot be ranked, so it is not used as a reference and selection falls back rather than failing. Adds five tests over an EvolutionTree with a killed baseline, including that the regression reference stays the real baseline rather than the oldest survivor. Signed-off-by: Christian Schüller --- .../experimentalist/components/loop.py | 60 +++++++++--- .../experimentalist/test_winner_selection.py | 95 ++++++++++++++++++- 2 files changed, 142 insertions(+), 13 deletions(-) diff --git a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/loop.py b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/loop.py index aadec439d1..c363581eda 100644 --- a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/loop.py +++ b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/loop.py @@ -163,27 +163,60 @@ def _regresses( return False +def finalization_pool( + evolution_tree: EvolutionTree, + objective_function: list[MetricTarget], +) -> tuple[list[EvolutionNode], EvolutionNode | None]: + """Return the nodes that may win, and the round-0 baseline they are measured against. + + Survivor selection runs mid-loop and can kill the baseline: it is told to always + keep a candidate created this round, and the baseline never is one. Finalization + asks a different question -- "is any of this better than shipping nothing?" -- so + the baseline is re-admitted here regardless of ``is_survivor``. Without that, a run + whose baseline was killed mid-loop can only choose among diffs, and the regression + comparison silently rehomes onto the oldest surviving candidate. + + The baseline is returned separately so it is used as the regression reference even + if it is not itself the winner. It is None when the run has no round-0 result, or + when that result lacks the configured objective dimensions. + """ + scored = [node for node in evolution_tree.nodes.values() if node.is_survivor and node.val_reward] + baseline = next( + (node for node in evolution_tree.nodes.values() if node.round == 0 and node.val_reward), + None, + ) + if baseline is not None and all(node is not baseline for node in scored): + scored.insert(0, baseline) + eligible = [node for node in scored if has_metric_dimensions(node.val_reward, objective_function)] + if baseline is not None and all(node is not baseline for node in eligible): + baseline = None + return eligible, baseline + + def select_winner_node( eligible: list[EvolutionNode], objective_function: list[MetricTarget], regression_metrics: list[MetricTarget] | None = None, + baseline: EvolutionNode | None = None, ) -> EvolutionNode | None: """Return the run's winner, or None when there is nothing eligible. - Candidates that worsen a ``regression_metrics`` target against the baseline are + Candidates that worsen a ``regression_metrics`` target against *baseline* are dropped before ranking, so a gain on the objectives cannot pay for a regression. The winner is then taken from the Pareto front over ``objective_function``, with ties going to the lowest round and then to creation order. - The baseline is the oldest eligible candidate and is never dropped, so a run whose - candidates all regress -- or all merely tie -- keeps it rather than shipping a diff - that bought nothing. + *baseline* is the reference for regression and is never dropped by it, so a run + whose candidates all regress -- or all merely tie -- keeps the baseline rather than + shipping a diff that bought nothing. It defaults to the oldest eligible node, which + is only correct when the true baseline is still present; callers that can tell + should pass it explicitly (see ``finalization_pool``). """ if not eligible: return None - baseline = min(eligible, key=lambda node: (node.round, _creation_order(node))) - kept = [node for node in eligible if not _regresses(node, baseline, regression_metrics or [])] - front = pareto_front(kept, lambda node: pareto_objectives(node.val_reward, objective_function)) + reference = baseline if baseline is not None else min(eligible, key=lambda n: (n.round, _creation_order(n))) + kept = [node for node in eligible if not _regresses(node, reference, regression_metrics or [])] + front = pareto_front(kept or eligible, lambda node: pareto_objectives(node.val_reward, objective_function)) return min(front, key=lambda node: (node.round, _creation_order(node))) if front else None @@ -1841,10 +1874,15 @@ async def _finalize( agent_name: str, ) -> Candidate | None: """Select the winner, copy to workspace root, write final report.""" - # Only survivors that actually have a validation reward are eligible winners. - scored = [n for n in evolution_tree.nodes.values() if n.is_survivor and n.val_reward] - eligible = [node for node in scored if has_metric_dimensions(node.val_reward, self.config.objective_function)] - best = select_winner_node(eligible, self.config.objective_function, self.config.regression_metrics) + # Survivors with a validation reward, plus the baseline even if it was killed + # mid-loop -- finalization must still be able to decide to ship nothing. + eligible, baseline = finalization_pool(evolution_tree, self.config.objective_function) + best = select_winner_node( + eligible, + self.config.objective_function, + self.config.regression_metrics, + baseline, + ) best_id = best.label if best else None restore_heldout_splits(self.working_dir) diff --git a/plugins/nemo-experimentalist/tests/experimentalist/test_winner_selection.py b/plugins/nemo-experimentalist/tests/experimentalist/test_winner_selection.py index 13b1cea389..d02657dbe4 100644 --- a/plugins/nemo-experimentalist/tests/experimentalist/test_winner_selection.py +++ b/plugins/nemo-experimentalist/tests/experimentalist/test_winner_selection.py @@ -19,8 +19,8 @@ from nemo_experimentalist_plugin.config import MetricTarget from nemo_experimentalist_plugin.entities import Candidate, RewardRecord -from nemo_experimentalist_plugin.experimentalist.components.loop import select_winner_node -from nemo_experimentalist_plugin.experimentalist.components.models import EvolutionNode +from nemo_experimentalist_plugin.experimentalist.components.loop import finalization_pool, select_winner_node +from nemo_experimentalist_plugin.experimentalist.components.models import EvolutionNode, EvolutionTree MAXIMIZE = [MetricTarget(name="reward", direction="maximize")] TWO_OBJECTIVES = [ @@ -255,3 +255,94 @@ def test_a_regression_metric_missing_from_a_candidate_is_not_a_regression() -> N ) assert winner is not None assert winner.label == "agent-1" + + +def _tree(*nodes: EvolutionNode) -> EvolutionTree: + tree = EvolutionTree() + for node in nodes: + tree.add(node.candidate) + return tree + + +def _killed(node: EvolutionNode, round_: int) -> EvolutionNode: + node.candidate.killed_round = round_ + return node + + +def test_a_baseline_killed_mid_loop_can_still_win() -> None: + """Survivor selection may kill agent-0; deciding to ship nothing must survive that. + + `_finalize` filters on `is_survivor`, so without re-admitting the baseline a run + whose baseline was killed could only choose among diffs -- reintroducing the very + bug this module exists to prevent, one level down. + """ + eligible, baseline = finalization_pool( + _tree( + _killed(_node("agent-0", 0, reward=0.5), 1), + _node("agent-1", 1, reward=0.5), + ), + MAXIMIZE, + ) + assert baseline is not None and baseline.label == "agent-0" + assert "agent-0" in {node.label for node in eligible} + + winner = select_winner_node(eligible, MAXIMIZE, [], baseline) + assert winner is not None + assert winner.label == "agent-0" + + +def test_regression_is_measured_against_the_real_baseline_not_the_oldest_survivor() -> None: + """A killed baseline must still anchor the regression comparison.""" + eligible, baseline = finalization_pool( + _tree( + _killed(_node("agent-0", 0, reward=0.5, latency=1.0), 1), + _node("agent-1", 1, reward=0.7, latency=2.0), + _node("agent-2", 2, reward=0.9, latency=2.5), + ), + MAXIMIZE, + ) + assert baseline is not None and baseline.label == "agent-0" + + # Against agent-1 (the oldest *survivor*) agent-2 looks fine on latency at 2.5 > 2.0 + # only by a little; against the real baseline at 1.0 both candidates regress. + winner = select_winner_node(eligible, MAXIMIZE, LATENCY_MUST_NOT_RISE, baseline) + assert winner is not None + assert winner.label == "agent-0" + + +def test_surviving_candidates_are_all_eligible() -> None: + eligible, baseline = finalization_pool( + _tree( + _node("agent-0", 0, reward=0.5), + _node("agent-1", 1, reward=0.7), + _killed(_node("agent-2", 1, reward=0.9), 1), + ), + MAXIMIZE, + ) + assert baseline is not None and baseline.label == "agent-0" + assert {node.label for node in eligible} == {"agent-0", "agent-1"} + + +def test_a_baseline_missing_the_objective_dimensions_is_not_used_as_reference() -> None: + """It cannot be ranked, so it cannot anchor; the selector falls back rather than crash.""" + eligible, baseline = finalization_pool( + _tree( + _node("agent-0", 0, unrelated=1.0), + _node("agent-1", 1, reward=0.7), + ), + MAXIMIZE, + ) + assert baseline is None + assert {node.label for node in eligible} == {"agent-1"} + + winner = select_winner_node(eligible, MAXIMIZE, [], baseline) + assert winner is not None + assert winner.label == "agent-1" + + +def test_a_run_with_no_baseline_result_still_finalizes() -> None: + eligible, baseline = finalization_pool(_tree(_node("agent-1", 1, reward=0.7)), MAXIMIZE) + assert baseline is None + winner = select_winner_node(eligible, MAXIMIZE, [], baseline) + assert winner is not None + assert winner.label == "agent-1"