diff --git a/package/samplers/cmamae/LICENSE b/package/samplers/cmamae/LICENSE new file mode 100644 index 00000000..51c2bdfd --- /dev/null +++ b/package/samplers/cmamae/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2024 Bryon Tjanaka + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/package/samplers/cmamae/README.md b/package/samplers/cmamae/README.md new file mode 100644 index 00000000..03736fa3 --- /dev/null +++ b/package/samplers/cmamae/README.md @@ -0,0 +1,150 @@ +--- +author: Bryon Tjanaka +title: CMA-MAE Sampler +description: This sampler searches for solutions using CMA-MAE, a quality diversity algorihm implemented in pyribs. +tags: [sampler, quality diversity, pyribs] +optuna_versions: [4.0.0] +license: MIT License +--- + +## Abstract + +This package provides a sampler using CMA-MAE as implemented in pyribs. +[CMA-MAE](https://dl.acm.org/doi/abs/10.1145/3583131.3590389) is a quality +diversity algorithm that has demonstrated state-of-the-art performance in a +variety of domains. [Pyribs](https://pyribs.org) is a bare-bones Python library +for quality diversity optimization algorithms. For a primer on CMA-MAE, quality +diversity, and pyribs, we recommend referring to the series of +[pyribs tutorials](https://docs.pyribs.org/en/stable/tutorials.html). + +For simplicity, this implementation provides a default instantiation of CMA-MAE +with a +[GridArchive](https://docs.pyribs.org/en/stable/api/ribs.archives.GridArchive.html) +and +[EvolutionStrategyEmitter](https://docs.pyribs.org/en/stable/api/ribs.emitters.EvolutionStrategyEmitter.html) +with improvement ranking, all wrapped up in a +[Scheduler](https://docs.pyribs.org/en/stable/api/ribs.schedulers.Scheduler.html). +However, it is possible to implement many variations of CMA-MAE and other +quality diversity algorithms using pyribs. + +## Class or Function Names + +- CmaMaeSampler + +## Installation + +```shell +$ pip install ribs +``` + +## Example + +```python +import optuna +import optunahub +from optuna.study import StudyDirection + +module = optunahub.load_module("samplers/cmamae") +CmaMaeSampler = module.CmaMaeSampler + + +def objective(trial: optuna.trial.Trial) -> tuple[float, float, float]: + """Returns an objective followed by two measures.""" + x = trial.suggest_float("x", -10, 10) + y = trial.suggest_float("y", -10, 10) + return x**2 + y**2, x, y + + +if __name__ == "__main__": + sampler = CmaMaeSampler( + param_names=["x", "y"], + archive_dims=[20, 20], + archive_ranges=[(-1, 1), (-1, 1)], + archive_learning_rate=0.1, + archive_threshold_min=-10, + n_emitters=1, + emitter_x0={ + "x": 0, + "y": 0, + }, + emitter_sigma0=0.1, + emitter_batch_size=20, + ) + study = optuna.create_study( + sampler=sampler, + directions=[ + StudyDirection.MINIMIZE, + # The remaining directions are for the measures, which do not have + # an optimization direction. However, we set MINIMIZE as a + # placeholder direction. + StudyDirection.MINIMIZE, + StudyDirection.MINIMIZE, + ], + ) + study.optimize(objective, n_trials=10000) +``` + +## Others + +### Reference + +#### CMA-MAE + +Matthew Fontaine and Stefanos Nikolaidis. 2023. Covariance Matrix Adaptation +MAP-Annealing. In Proceedings of the Genetic and Evolutionary Computation +Conference (GECCO '23). Association for Computing Machinery, New York, NY, USA, +456–465. https://doi.org/10.1145/3583131.3590389 + +#### Pyribs + +Bryon Tjanaka, Matthew C Fontaine, David H Lee, Yulun Zhang, Nivedit Reddy +Balam, Nathaniel Dennler, Sujay S Garlanka, Nikitas Dimitri Klapsis, and +Stefanos Nikolaidis. 2023. Pyribs: A Bare-Bones Python Library for Quality +Diversity Optimization. In Proceedings of the Genetic and Evolutionary +Computation Conference (GECCO '23). Association for Computing Machinery, New +York, NY, USA, 220–229. https://doi.org/10.1145/3583131.3590374 + +### Bibtex + +#### CMA-MAE + +``` +@inproceedings{10.1145/3583131.3590389, + author = {Fontaine, Matthew and Nikolaidis, Stefanos}, + title = {Covariance Matrix Adaptation MAP-Annealing}, + year = {2023}, + isbn = {9798400701191}, + publisher = {Association for Computing Machinery}, + address = {New York, NY, USA}, + url = {https://doi.org/10.1145/3583131.3590389}, + doi = {10.1145/3583131.3590389}, + abstract = {Single-objective optimization algorithms search for the single highest-quality solution with respect to an objective. Quality diversity (QD) optimization algorithms, such as Covariance Matrix Adaptation MAP-Elites (CMA-ME), search for a collection of solutions that are both high-quality with respect to an objective and diverse with respect to specified measure functions. However, CMA-ME suffers from three major limitations highlighted by the QD community: prematurely abandoning the objective in favor of exploration, struggling to explore flat objectives, and having poor performance for low-resolution archives. We propose a new quality diversity algorithm, Covariance Matrix Adaptation MAP-Annealing (CMA-MAE), that addresses all three limitations. We provide theoretical justifications for the new algorithm with respect to each limitation. Our theory informs our experiments, which support the theory and show that CMA-MAE achieves state-of-the-art performance and robustness.}, + booktitle = {Proceedings of the Genetic and Evolutionary Computation Conference}, + pages = {456–465}, + numpages = {10}, + location = {Lisbon, Portugal}, + series = {GECCO '23} +} +``` + +#### Pyribs + +``` +@inproceedings{10.1145/3583131.3590374, + author = {Tjanaka, Bryon and Fontaine, Matthew C and Lee, David H and Zhang, Yulun and Balam, Nivedit Reddy and Dennler, Nathaniel and Garlanka, Sujay S and Klapsis, Nikitas Dimitri and Nikolaidis, Stefanos}, + title = {pyribs: A Bare-Bones Python Library for Quality Diversity Optimization}, + year = {2023}, + isbn = {9798400701191}, + publisher = {Association for Computing Machinery}, + address = {New York, NY, USA}, + url = {https://doi.org/10.1145/3583131.3590374}, + doi = {10.1145/3583131.3590374}, + abstract = {Recent years have seen a rise in the popularity of quality diversity (QD) optimization, a branch of optimization that seeks to find a collection of diverse, high-performing solutions to a given problem. To grow further, we believe the QD community faces two challenges: developing a framework to represent the field's growing array of algorithms, and implementing that framework in software that supports a range of researchers and practitioners. To address these challenges, we have developed pyribs, a library built on a highly modular conceptual QD framework. By replacing components in the conceptual framework, and hence in pyribs, users can compose algorithms from across the QD literature; equally important, they can identify unexplored algorithm variations. Furthermore, pyribs makes this framework simple, flexible, and accessible, with a user-friendly API supported by extensive documentation and tutorials. This paper overviews the creation of pyribs, focusing on the conceptual framework that it implements and the design principles that have guided the library's development. Pyribs is available at https://pyribs.org}, + booktitle = {Proceedings of the Genetic and Evolutionary Computation Conference}, + pages = {220–229}, + numpages = {10}, + keywords = {software library, framework, quality diversity}, + location = {Lisbon, Portugal}, + series = {GECCO '23} +} +``` diff --git a/package/samplers/cmamae/__init__.py b/package/samplers/cmamae/__init__.py new file mode 100644 index 00000000..e724cd42 --- /dev/null +++ b/package/samplers/cmamae/__init__.py @@ -0,0 +1,4 @@ +from .sampler import CmaMaeSampler + + +__all__ = ["CmaMaeSampler"] diff --git a/package/samplers/cmamae/example.py b/package/samplers/cmamae/example.py new file mode 100644 index 00000000..5b68f7c0 --- /dev/null +++ b/package/samplers/cmamae/example.py @@ -0,0 +1,43 @@ +import optuna +from optuna.study import StudyDirection +import optunahub + + +module = optunahub.load_module("samplers/cmamae") +CmaMaeSampler = module.CmaMaeSampler + + +def objective(trial: optuna.trial.Trial) -> tuple[float, float, float]: + """Returns an objective followed by two measures.""" + x = trial.suggest_float("x", -10, 10) + y = trial.suggest_float("y", -10, 10) + return x**2 + y**2, x, y + + +if __name__ == "__main__": + sampler = CmaMaeSampler( + param_names=["x", "y"], + archive_dims=[20, 20], + archive_ranges=[(-1, 1), (-1, 1)], + archive_learning_rate=0.1, + archive_threshold_min=-10, + n_emitters=1, + emitter_x0={ + "x": 0, + "y": 0, + }, + emitter_sigma0=0.1, + emitter_batch_size=20, + ) + study = optuna.create_study( + sampler=sampler, + directions=[ + StudyDirection.MINIMIZE, + # The remaining directions are for the measures, which do not have + # an optimization direction. However, we set MINIMIZE as a + # placeholder direction. + StudyDirection.MINIMIZE, + StudyDirection.MINIMIZE, + ], + ) + study.optimize(objective, n_trials=10000) diff --git a/package/samplers/cmamae/requirements.txt b/package/samplers/cmamae/requirements.txt new file mode 100644 index 00000000..4572f714 --- /dev/null +++ b/package/samplers/cmamae/requirements.txt @@ -0,0 +1,3 @@ +optuna +optunahub +ribs diff --git a/package/samplers/cmamae/sampler.py b/package/samplers/cmamae/sampler.py new file mode 100644 index 00000000..15e0e4f2 --- /dev/null +++ b/package/samplers/cmamae/sampler.py @@ -0,0 +1,204 @@ +from __future__ import annotations + +from collections.abc import Sequence +from typing import Iterable + +import numpy as np +from optuna.distributions import BaseDistribution +from optuna.distributions import FloatDistribution +from optuna.study import Study +from optuna.study import StudyDirection +from optuna.trial import FrozenTrial +from optuna.trial import TrialState +import optunahub +from ribs.archives import GridArchive +from ribs.emitters import EvolutionStrategyEmitter +from ribs.schedulers import Scheduler + + +class CmaMaeSampler(optunahub.samplers.SimpleBaseSampler): + """A sampler using CMA-MAE as implemented in pyribs. + + `CMA-MAE `_ is a quality + diversity algorithm that has demonstrated state-of-the-art performance in a + variety of domains. `Pyribs `_ is a bare-bones Python + library for quality diversity optimization algorithms. For a primer on + CMA-MAE, quality diversity, and pyribs, we recommend referring to the series + of `pyribs tutorials `_. + + For simplicity, this implementation provides a default instantiation of + CMA-MAE with a `GridArchive + `_ and + `EvolutionStrategyEmitter + `_ + with improvement ranking, all wrapped up in a `Scheduler + `_. + However, it is possible to implement many variations of CMA-MAE and other + quality diversity algorithms using pyribs. + + Note that this sampler assumes the objective function will return a list of + values. The first value will be the objective, and the remaining values will + be the measures. + + Args: + param_names: List of names of parameters to optimize. + archive_dims: Number of archive cells in each dimension of the measure + space, e.g. ``[20, 30, 40]`` indicates there should be 3 dimensions + with 20, 30, and 40 cells. (The number of dimensions is implicitly + defined in the length of this argument). + archive_ranges: Upper and lower bound of each dimension of the measure + space for the archive, e.g. ``[(-1, 1), (-2, 2)]`` indicates the + first dimension should have bounds :math:`[-1,1]` (inclusive), and + the second dimension should have bounds :math:`[-2,2]` (inclusive). + ``ranges`` should be the same length as ``dims``. + archive_learning_rate: The learning rate for threshold updates in the + archive. + archive_threshold_min: The initial threshold value for all the cells in + the archive. + n_emitters: Number of emitters to use in CMA-MAE. + emitter_x0: Mapping from parameter names to their initial values. + emitter_sigma0: Initial step size / standard deviation of the + distribution from which solutions are sampled in the emitter. + emitter_batch_size: Number of solutions for each emitter to generate on + each iteration. + """ + + def __init__( + self, + *, + param_names: list[str], + archive_dims: list[int], + archive_ranges: list[tuple[float, float]], + archive_learning_rate: float, + archive_threshold_min: float, + n_emitters: int, + emitter_x0: dict[str, float], + emitter_sigma0: float, + emitter_batch_size: int, + ) -> None: + self._validate_params(param_names, emitter_x0) + self._param_names = param_names[:] + + # NOTE: SimpleBaseSampler must know Optuna search_space information. + search_space = {name: FloatDistribution(-1e9, 1e9) for name in self._param_names} + super().__init__(search_space=search_space) + + emitter_x0_np = self._convert_to_pyribs_params(emitter_x0) + + archive = GridArchive( + solution_dim=len(param_names), + dims=archive_dims, + ranges=archive_ranges, + learning_rate=archive_learning_rate, + threshold_min=archive_threshold_min, + ) + result_archive = GridArchive( + solution_dim=len(param_names), + dims=archive_dims, + ranges=archive_ranges, + ) + emitters = [ + EvolutionStrategyEmitter( + archive, + x0=emitter_x0_np, + sigma0=emitter_sigma0, + ranker="imp", + selection_rule="mu", + restart_rule="basic", + batch_size=emitter_batch_size, + ) + for _ in range(n_emitters) + ] + + # Number of solutions generated in each batch from pyribs. + self._batch_size = n_emitters * emitter_batch_size + + # Public to allow access for, e.g., visualization. + self.scheduler = Scheduler( + archive, + emitters, + result_archive=result_archive, + ) + + self._values_to_tell: list[list[float]] = [] + self._stored_trial_numbers: list[int] = [] + + def _validate_params(self, param_names: list[str], emitter_x0: dict[str, float]) -> None: + dim = len(param_names) + param_set = set(param_names) + if dim != len(param_set): + raise ValueError( + "Some elements in param_names are duplicated. Please make it a unique list." + ) + + if set(param_names) != emitter_x0.keys(): + raise ValueError( + "emitter_x0 does not contain the parameters listed in param_names. " + "Please provide an initial value for each parameter." + ) + + def _validate_param_names(self, given_param_names: Iterable[str]) -> None: + if set(self._param_names) != set(given_param_names): + raise ValueError( + "The given param names must match the param names " + "initially passed to this sampler." + ) + + def _convert_to_pyribs_params(self, params: dict[str, float]) -> np.ndarray: + np_params = np.empty(len(self._param_names), dtype=float) + for i, p in enumerate(self._param_names): + np_params[i] = params[p] + return np_params + + def _convert_to_optuna_params(self, params: np.ndarray) -> dict[str, float]: + dict_params = {} + for i, p in enumerate(self._param_names): + dict_params[p] = params[i] + return dict_params + + def sample_relative( + self, study: Study, trial: FrozenTrial, search_space: dict[str, BaseDistribution] + ) -> dict[str, float]: + self._validate_param_names(search_space.keys()) + + # Note: Batch optimization means we need to enqueue trials. + solutions = self.scheduler.ask() + next_params = self._convert_to_optuna_params(solutions[0]) + for solution in solutions[1:]: + params = self._convert_to_optuna_params(solution) + study.enqueue_trial(params) + + return next_params + + def after_trial( + self, + study: Study, + trial: FrozenTrial, + state: TrialState, + values: Sequence[float] | None, + ) -> None: + self._validate_param_names(trial.params.keys()) + + # Store the trial result. + direction0 = study.directions[0] + minimize_in_optuna = direction0 == StudyDirection.MINIMIZE + assert values is not None, "MyPy redefinition." + modified_values = list([float(v) for v in values]) + if minimize_in_optuna: + # The direction of the first objective (pyribs maximizes). + modified_values[0] = -values[0] + self._values_to_tell.append(modified_values) + self._stored_trial_numbers.append(trial.number) + + # If we have not retrieved the whole batch of solutions, then we should + # not tell() the results to the scheduler yet. + if len(self._values_to_tell) != self._batch_size: + return + + # Tell the batch results to external sampler once the batch is ready. + values_to_tell = np.asarray(self._values_to_tell)[np.argsort(self._stored_trial_numbers)] + self.scheduler.tell(objective=values_to_tell[:, 0], measures=values_to_tell[:, 1:]) + + # Empty the results. + self._values_to_tell = [] + self._stored_trial_numbers = []