Add generality-oriented BO features by StefanPSchmid · Pull Request #902 · emdgroup/baybe

StefanPSchmid · 2026-08-20T13:43:02Z

Generality-BO integration into BayBE

This PR creates a first draft for generality-oriented BO (done in this paper https://arxiv.org/abs/2502.18966) within BayBE.

Main features/changes:

Introduce GeneralityParameter, which triggers generality-oriented BO (inspired by TaskParameter)
Add recommend_generality on the Recommender level
Introduce GeneralityModel and GeneralityPosterior, which are models that perform the Aggregation described in the paper, over which the optimizer can then run

StefanPSchmid · 2026-08-20T13:49:11Z

+
+__all__ = [
+    "AggregationFunction",
+    "MeanAggregation",


These are the three relevant aggregation functions implemented in the paper;

Idea in those:

MeanAggregation: Optimize the average across all contexts

SigmoidAggregation: Optimize the number of contexts above a threshold

MinAggregation: Optimize the worst-case across all contexts

Not sure exactly how this is influenced by minimization (negation), and how to optimally bake this into the code (esp. Sigmoid and Min), but I tried accounting for that in the GeneralityModel

StefanPSchmid · 2026-08-20T13:54:48Z



+@define(frozen=True, slots=False)
+class GeneralityParameter(_DiscreteLabelLikeParameter):


GeneralityParameter, that should be used to trigger generality-oriented optimization within the recommenders; inspiration taken from the TaskParameter

Idea is to have a context (over which you want to be general), that can be encoded (not just OHE/INT, but any discrete parameter like Substance encoding)

StefanPSchmid · 2026-08-20T14:07:42Z

                numerical_measurements_must_be_within_tolerance=False,
            )

+        if searchspace._generality_parameter is not None:


Routing to generality-oriented optimization within the recommender

StefanPSchmid · 2026-08-20T14:32:47Z

+        assert len(params) == 1
+        return params[0]
+
+    def _split_by_generality(


maybe this can be combined with some other funcitonality in the searchspace, but essentially this removes the generalityparameter from the searchspace, since in the first step we want to optimize over all parameters except the generality one

StefanPSchmid · 2026-08-20T14:34:39Z

+        samples = self.rsample(torch.Size([self._n_variance_samples]))
+        return samples.var(dim=0)
+
+    def _aggregate(self, samples: Tensor) -> Tensor:


Perform the aggregation step (i.e. applying the generality metric), and then forward over it ... r is the number of contexts, m is the number of objectives

StefanPSchmid · 2026-08-20T14:35:45Z

+    from baybe.aggregation.base import AggregationFunction
+
+
+class _GeneralityPosterior(Posterior):


Posterior with many properties from the joint posterior (over design-space and context parameter), but that also performs the aggregation

StefanPSchmid · 2026-08-20T14:36:13Z

+        )
+
+    def _expand(self, X: Tensor) -> Tensor:
+        """Pair each candidate with all context values for the base GP.


done to assess generality

StefanPSchmid · 2026-08-20T14:36:44Z

+_register_generality_sampler()
+
+
+class _GeneralityModel(Model):


surrogate model for generality, that is based on the joint model but holds the GeneralityPosterior

StefanPSchmid · 2026-08-20T14:38:32Z

Claude generated since I am not experienced with tests and I assume the design will still change a bit; tried to keep it as close to the tests with MagicMock as I could

StefanPSchmid · 2026-08-20T20:12:39Z

+        objective=IdentityMCObjective(),
+    )
+
+    if not objective.is_multi_output:


Since generality-oriented BO is in a partial monitoring scenario, the "best_f" needed for some of the acquisition functions is not clear. What we do in the paper is to take the SurrogateModel and find the best predicted generality (PosteriorMean) and take that as the best_f

This is not possible for multi-objective, since multiple points on the generality-pareto-front could be "best", and there is no clear way of defining what best_f could be

StefanPSchmid · 2026-08-20T20:13:12Z

+        )
+        acqf_kwargs["best_f"] = best_f_scores.item()
+
+    user_attrs, _ = match_attributes(


this should give the X-acquisition function of choice the required arguments

StefanPSchmid · 2026-08-20T20:21:44Z

+        **{k: v for k, v in acqf_kwargs.items() if k in sig}
+    )
+
+    x_points, _ = recommender.optimizer(batch_size, botorch_acqf, x_subspace)


The optimization approach follows the "Sequential one-step lookahead" approach outlined in the paper, with a few differences. This approach was chosen because it performed well in our benchmarks, while being very fast (and Martin at one point mentioned to me that speed is important)

The workflow, with highlighted differences:

While the CurryBO implementation does batching by conditioning on points, here we first get a batch of X-points (which should all be distinct) at once, from the optimizer. Optimization is performed on the acquisition function on the GeneralityPosterior, which did the Aggregation over the different contexts already.

For each batch point, the recommender then fixes the search space, and picks a context (w) to optimize, by maximizing the standard deviation (picking the most uncertain / informative one). If I understood the code correctly, for discrete searchspaces, the (x,w) combinations that have already been measured should have been filtered out (if that flag is set to True) ... and for continuous/hybrid x-spaces, restricting (x,w) combinations makes no sense anyways

StefanPSchmid · 2026-08-20T20:23:20Z

This file implements the generality-oriented optimization logic; some nomenclature:

x is the "design space", i.e. the space you want to find your maximum in (which is also the search space, with the GeneralityParameter removed).
w, or context, is the variable over which you want to be general (and over which you aggregate).

Happy to discuss about nomenclature / better naming

StefanPSchmid · 2026-08-20T20:28:14Z

@@ -0,0 +1,138 @@
+"""Generality recommendation logic for BayesianRecommender (CurryBO algorithm)."""
+
+from __future__ import annotations


Few points I would also like to make you aware of / discuss:

How constraints could / should be handled. Especially for x,w combinations, currently not enabled. For the respective X space, constraints should work as implemented by the recommender, although the ones with subsets might need to be revisited (but as I understood from our previous discussions, that should get integrated into the optimizer?)

Due to the partial monitoring / generality-scenario, the best measurement is not automatically the optimum. In CurryBO, we thus not only predict the next point, but also the current optimum, based on the PosteriorMean of the generality function. I have no idea where such an additional part of the recommendation could be integrated, happy to hear your ideas! Thanks!

Copilot

Pull request overview

This PR adds a first integration of generality-oriented Bayesian optimization into BayBE by introducing a dedicated “context” parameter, context aggregation utilities, and a BoTorch model/posterior wrapper to optimize over designs while aggregating across contexts.

Changes:

Introduces GeneralityParameter and search space validation/splitting logic for separating design vs. generality context.
Adds a generality-aware recommendation path (recommend_generality) in the pure recommender stack.
Adds aggregation primitives and a BoTorch _GeneralityModel / _GeneralityPosterior, plus an initial test suite.

Reviewed changes

Copilot reviewed 13 out of 13 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
tests/test_generality.py	Adds unit/wiring tests for aggregation, search space splitting/fixing, model expansion, posterior aggregation, and recommender wiring.
baybe/aggregation/init.py	Exposes aggregation API for generality-oriented optimization.
baybe/aggregation/base.py	Defines the `AggregationFunction` base class for context aggregation.
baybe/aggregation/aggregations.py	Implements concrete aggregation strategies (mean/min/sigmoid).
baybe/parameters/categorical.py	Adds `GeneralityParameter` (context parameter wrapper) and aggregation configuration.
baybe/parameters/enum.py	Adds a new `_ParameterKind.GENERALITY` flag.
baybe/parameters/init.py	Re-exports `GeneralityParameter` from `baybe.parameters`.
baybe/searchspace/validation.py	Adds validation rules for `GeneralityParameter` count and Task/Generality incompatibility.
baybe/searchspace/core.py	Adds `_generality_parameter` access and `_split_by_generality()` helper for context-aware optimization.
baybe/surrogates/generality.py	Introduces `_GeneralityModel` and `_GeneralityPosterior` wrappers for context-aggregated posteriors.
baybe/recommenders/pure/bayesian/generality.py	Implements `recommend_generality()` (two-step design-then-context selection).
baybe/recommenders/pure/bayesian/core.py	Wires Bayesian recommender to dispatch into the generality recommendation path.
baybe/recommenders/pure/base.py	Adds top-level dispatch for search spaces containing a generality parameter.

Suppressed comments (1)

baybe/parameters/enum.py:38

_ParameterKind.from_parameter() doesn’t map GeneralityParameter to _ParameterKind.GENERALITY, so GeneralityParameter instances will still be classified as REGULAR via Parameter._kind. This makes the newly added GENERALITY flag effectively unused and can break any compatibility logic relying on _kind.

    @staticmethod
    def from_parameter(parameter: Parameter) -> _ParameterKind:
        """Determine the kind of a parameter from its type."""
        from baybe.parameters.categorical import TaskParameter

        if isinstance(parameter, TaskParameter):
            return _ParameterKind.TASK
        return _ParameterKind.REGULAR

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

        NotImplementedError: If more than one
            :class:`baybe.parameters.categorical.TaskParameter` is requested.


+        sample_row = self.discrete.exp_rep.head(1)
+        full_comp = self.transform(sample_row, allow_extra=True)
+        full_columns = list(full_comp.columns)
+


+@define(frozen=True, slots=False)
+class GeneralityParameter(_DiscreteLabelLikeParameter):
+    """Parameter marking a discrete parameter as the context dimension."""
+
+    context: DiscreteParameter = field()
+    """Parameter for the contexts."""
+
+    aggregation: AggregationFunction = field(factory=MeanAggregation)
+    """Aggregation mode over contexts."""
+
+    @property
+    def encoding(self):
+        """The encoding of the context parameter."""
+        return self.context.encoding
+
+    @override
+    @property
+    def values(self) -> tuple:
+        """The values of the context parameter."""
+        return self.context.values
+
+    @override
+    @cached_property
+    def comp_df(self) -> pd.DataFrame:
+        """The comp_df of the context parameter."""
+        return self.context.comp_df


StefanPSchmid added 12 commits August 18, 2026 13:14

Base class aggregation functions

0aa47b2

Aggregation functions

8e9dafe

Init aggregation functions

5d9e54c

Generality Parameter

fb9c99c

Generality Parameter ENUM and init

788b578

Add generality wiring to Pure Recommender

f857bd1

Add generality wiring to BayesianRecommender

dba3b46

Add generality splitting to SearchSpace

44690a8

Add GeneralityParameter to validation

a7247b3

Add generality recommendation logic

6f3e9bb

Add GeneralitySurrogate and GeneralityPosterior

8a58d5f

Add first draft tests

27b6258

StefanPSchmid commented Aug 20, 2026

View reviewed changes

StefanPSchmid marked this pull request as ready for review August 21, 2026 08:57

StefanPSchmid requested review from AdrianSosic and Scienfitz as code owners August 21, 2026 08:57

Copilot AI lite review requested due to automatic review settings August 21, 2026 08:57

StefanPSchmid requested a review from AVHopp as a code owner August 21, 2026 08:57

Copilot started reviewing on behalf of StefanPSchmid August 21, 2026 08:58 View session

Copilot AI reviewed Aug 21, 2026

View reviewed changes



		@define(frozen=True, slots=False)
		class GeneralityParameter(_DiscreteLabelLikeParameter):

		from baybe.aggregation.base import AggregationFunction


		class _GeneralityPosterior(Posterior):

		_register_generality_sampler()


		class _GeneralityModel(Model):

		@@ -0,0 +1,138 @@
		"""Generality recommendation logic for BayesianRecommender (CurryBO algorithm)."""

		from __future__ import annotations

		NotImplementedError: If more than one
		:class:`baybe.parameters.categorical.TaskParameter` is requested.

Uh oh!

Conversation

StefanPSchmid commented Aug 20, 2026

Generality-BO integration into BayBE

Uh oh!

StefanPSchmid Aug 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

StefanPSchmid Aug 20, 2026 •

edited

Loading