diff --git a/vizier/_src/algorithms/evolution/nsga2.py b/vizier/_src/algorithms/evolution/nsga2.py index 0e3603390..766dd6483 100644 --- a/vizier/_src/algorithms/evolution/nsga2.py +++ b/vizier/_src/algorithms/evolution/nsga2.py @@ -16,8 +16,9 @@ """NSGA-II algorithm: https://ieeexplore.ieee.org/document/996017.""" -from typing import Callable, Optional, Tuple +from typing import Callable, Optional, Sequence, Tuple +from absl import logging import attr import numpy as np from vizier import pyvizier as vz @@ -44,30 +45,62 @@ def pareto_rank(ys: np.ndarray) -> np.ndarray: return np.sum(np.stack(dominated), axis=0) -def crowding_distance(ys: np.ndarray) -> np.ndarray: +def crowding_distance( + ys: np.ndarray, *, extra_tiebreakers: Sequence[np.ndarray] = tuple() +) -> np.ndarray: """Crowding distance. + Reference: + https://medium.com/@rossleecooloh/optimization-algorithm-nsga-ii-and-python-package-deap-fca0be6b2ffc + except that the lower boundary does not get infinity crowding score. + Args: ys: (number of population) x (number of metrics) array. + extra_tiebreakers: A sequence of (number of population) array of floating + numbers. If specified, they are used to break ties when sorting the + population. By default, a random number is always used to break ties + consistently across all metrics. Returns: (number of population) float32 array. Higher numbers mean less crowding and more desirable. """ scores = np.zeros([ys.shape[0]], dtype=np.float32) + + if ys.shape[0] <= 1: + return scores + + rng = np.random.default_rng() + # Use a random number to break ties. But the same random number is used for + # all metrics. + random_tiebreaker = rng.random(ys.shape[0]) + tiebreakers = list(extra_tiebreakers) + [random_tiebreaker] + for m in range(ys.shape[1]): # Sort by the m-th metric. + sid = sorted( + np.arange(ys.shape[0]), + key=lambda i, m=m: (ys[i, m],) + + tuple(tiebreaker[i] for tiebreaker in tiebreakers), + ) + + # Compute the range of the m-th metric. yy = ys[:, m] # Shape: (num_population,) - sid = np.argsort(yy) + yrange = yy[sid[-1]] - yy[sid[0]] + np.finfo(np.float32).eps - # Boundary are assigned infinity. - scores[sid[0]] += np.inf + # Lower boundary is assigned a one-sided score and does not automatically + # get infinity. This is different from the paper. The lower boundary means + # it's dominated by all other points in one dimension. There's no reason to + # favor it over other points. + scores[sid[0]] += (yy[sid[1]] - yy[sid[0]]) / yrange + # Upper boundary is assigned infinity. This point will survive anyways + # because it's pareto-optimal. But in case there are ties, it's useful to + # make only one of them stand out. scores[sid[-1]] += np.inf - # Compute the crowding distance. - yrange = yy[sid[-1]] - yy[sid[0]] + np.finfo(np.float32).eps scores[sid[1:-1]] += (yy[sid[2:]] - yy[sid[:-2]]) / yrange - return scores + # Normalize the score to [0, 1]. + return scores / ys.shape[1] def _constraint_violation(ys: np.ndarray) -> np.ndarray: @@ -163,6 +196,7 @@ def select(self, population: Population) -> Population: # return empty. return population + logging.info('Selecting %s from %s', self._target_size, len(population)) selected = population.empty_like() # Sort by the safety constraint. if selected.cs.shape[1]: @@ -171,19 +205,33 @@ def select(self, population: Population) -> Population: ) selected += population[top] population = population[border] - + logging.info( + 'Selected %s by safety constraints, and will break ties among: %s', + len(selected), + len(population), + ) # Sort by the pareto rank. pareto_ranks = self._ranking_fn(population.ys) top, border = _select_by( pareto_ranks, target=self._target_size - len(selected) ) + considered_pareto_ranks = np.concatenate( + [-np.ones(len(selected)), pareto_ranks[top], pareto_ranks[border]] + ) selected += population[top] population = population[border] - + logging.info( + 'Selected %s by pareto rank, and will break ties among: %s', + len(selected), + len(population), + ) # Sort by the distance. Include the points that are already selected for # the computation. # Flip the sign so it works with ascending sort. - distance = -crowding_distance((selected + population).ys) + distance = -crowding_distance( + (selected + population).ys, + extra_tiebreakers=[considered_pareto_ranks], + ) sids = np.argsort(distance) # Selected points have fewer constraint violations or better pareto rank. # Regardless of the distance, they remain selected. Rank the remainder only. diff --git a/vizier/_src/algorithms/evolution/nsga2_test.py b/vizier/_src/algorithms/evolution/nsga2_test.py index 664c612d4..e69af1b2a 100644 --- a/vizier/_src/algorithms/evolution/nsga2_test.py +++ b/vizier/_src/algorithms/evolution/nsga2_test.py @@ -34,6 +34,27 @@ np.set_printoptions(precision=3) +class CrowdingDistanceTest(absltest.TestCase): + + def test_crowding_distance(self): + ys = np.array([ + [0.5, 4.0], + [0.5, 4.0], + [0.5, 2.0], + [1.0, 2.0], + [1.0, 2.0], + ]) + result = nsga2.crowding_distance( + ys, extra_tiebreakers=[nsga2.pareto_rank(ys)] + ) + np.testing.assert_allclose( + # Only one of the duplicate pareto-optimal points should get + # infinity score. + np.isinf(result).sum(), + 2, + ) + + def nsga2_on_all_types( population_size: int = 50, eviction_limit: Optional[int] = None ) -> templates.CanonicalEvolutionDesigner[nsga2.Population, nsga2.Offspring]: @@ -145,8 +166,9 @@ def test_survival_by_crowding_distance(self): vz.Measurement({'m1': 1.001, 'm2': -1.001, 's1': 2.0, 's2': 0.0}) ) - # 4 safe trials with the same pareto rank. Crowding distance is computed - # among them to break ties. Trial 3 is less "crowded" than Trial 2. + # 4 safe trials with the same pareto rank. Ties are broken by their crowding + # distances, computed *with* the pareto optimal trials. Trial 1 is most + # crowded because of its proximity to trial 0. trial1 = vz.Trial(id=1) trial1.complete( vz.Measurement({'m1': 1.0, 'm2': 0.0, 's1': 2.0, 's2': 0.9}) @@ -166,7 +188,7 @@ def test_survival_by_crowding_distance(self): trials = vza.CompletedTrials([trial0, trial1, trial2, trial3, trial4]) algorithm.update(trials, vza.ActiveTrials()) - self.assertSetEqual(set(algorithm.population.trial_ids), {0, 1, 3, 4}) + self.assertSetEqual(set(algorithm.population.trial_ids), {0, 2, 3, 4}) def test_survival_by_safety(self): algorithm = nsga2_on_all_types(3)