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..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 @@ -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,95 @@ 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 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 *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. + + *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 + 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 + + def _with_insight_objective( config: EvolutionaryOptimizerConfig, metric_keys: tuple[str, ...] ) -> EvolutionaryOptimizerConfig: @@ -1783,20 +1874,16 @@ 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)] - ranked_nodes = pareto_sort( + # 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, - 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, + baseline, ) - best_id = finalists[0].label if finalists else None + 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..d02657dbe4 --- /dev/null +++ b/plugins/nemo-experimentalist/tests/experimentalist/test_winner_selection.py @@ -0,0 +1,348 @@ +# 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: 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. +""" + +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 finalization_pool, select_winner_node +from nemo_experimentalist_plugin.experimentalist.components.models import EvolutionNode, EvolutionTree + +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: + 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" + + +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" + + +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"