transition_criterion.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

# pyre-strict

from __future__ import annotations

from abc import abstractmethod
from typing import TYPE_CHECKING

from ax.core import MultiObjectiveOptimizationConfig
from ax.core.auxiliary import AuxiliaryExperiment, AuxiliaryExperimentPurpose
from ax.core.experiment import Experiment
from ax.core.trial_status import TrialStatus
from ax.core.utils import get_trial_indices_with_required_metrics
from ax.exceptions.core import DataRequiredError, UserInputError
from ax.exceptions.generation_strategy import (
    AxGenerationException,
    MaxParallelismReachedException,
)

if TYPE_CHECKING:
    from ax.generation_strategy.generation_node import GenerationNode

from ax.utils.common.base import SortableBase
from ax.utils.common.serialization import serialize_init_args


DATA_REQUIRED_MSG = (
    "All trials for current node {node_name} have been generated, "
    "but not enough data has been observed to proceed to the next "
    "Generation node. Try again when more data is available."
)


# ============================================================================
# Trial Counting Utility Functions
# ============================================================================


def get_trials_by_statuses(
    experiment: Experiment, statuses: list[TrialStatus]
) -> set[int]:
    """Get trial indices from the experiment with the specified statuses.

    Args:
        experiment: The experiment to query.
        statuses: The trial statuses to filter on.

    Returns:
        Set of trial indices with the specified statuses.
    """
    trials_with_statuses = set()
    for status in statuses:
        trials_with_statuses = trials_with_statuses.union(
            experiment.trial_indices_by_status[status]
        )
    return trials_with_statuses


def filter_trials_by_status(
    experiment: Experiment,
    only_in_statuses: list[TrialStatus] | None = None,
    not_in_statuses: list[TrialStatus] | None = None,
) -> set[int]:
    """Filter trial indices by status inclusion/exclusion.

    Args:
        experiment: The experiment to query.
        only_in_statuses: If provided, only include trials with these statuses.
        not_in_statuses: If provided, exclude trials with these statuses.

    Returns:
        Set of trial indices matching the filter criteria.
    """
    trials_to_check = set(experiment.trials.keys())
    if only_in_statuses is not None:
        trials_to_check = get_trials_by_statuses(
            experiment=experiment, statuses=only_in_statuses
        )
    if not_in_statuses is not None:
        trials_to_check -= get_trials_by_statuses(
            experiment=experiment, statuses=not_in_statuses
        )
    return trials_to_check


def count_trials_toward_threshold(
    experiment: Experiment,
    trials_from_node: set[int],
    only_in_statuses: list[TrialStatus] | None = None,
    not_in_statuses: list[TrialStatus] | None = None,
    use_all_trials_in_exp: bool = False,
    count_only_trials_with_data: bool = False,
    count_only_trials_without_data: bool = False,
) -> int:
    """Count trials contributing toward a threshold.

    Args:
        experiment: The experiment to query.
        trials_from_node: Set of trial indices generated by the current node.
        only_in_statuses: If provided, only count trials with these statuses.
        not_in_statuses: If provided, exclude trials with these statuses.
        use_all_trials_in_exp: If True, count all trials in the experiment.
            Otherwise, only count trials from the current node.
        count_only_trials_with_data: If True, only count trials with data.
        count_only_trials_without_data: If True, only count trials that are
            MISSING required metric data. Cannot be True if
            count_only_trials_with_data is also True.

    Returns:
        The number of trials contributing to the threshold.
    """
    if count_only_trials_with_data and count_only_trials_without_data:
        raise UserInputError(
            "Cannot set both count_only_trials_with_data and "
            "count_only_trials_without_data to True."
        )

    all_trials_to_check = filter_trials_by_status(
        experiment=experiment,
        only_in_statuses=only_in_statuses,
        not_in_statuses=not_in_statuses,
    )
    if count_only_trials_with_data:
        data_trial_indices = get_trial_indices_with_required_metrics(
            experiment=experiment,
            df=experiment.lookup_data().df,
            require_data_for_all_metrics=False,
        )
        all_trials_to_check = all_trials_to_check.intersection(data_trial_indices)
    elif count_only_trials_without_data:
        data_trial_indices = get_trial_indices_with_required_metrics(
            experiment=experiment,
            df=experiment.lookup_data().df,
            require_data_for_all_metrics=False,
        )
        all_trials_to_check = all_trials_to_check - data_trial_indices

    if use_all_trials_in_exp:
        return len(all_trials_to_check)

    return len(trials_from_node.intersection(all_trials_to_check))


# ============================================================================
# PausingCriterion - for blocking generation without transitioning
# ============================================================================


class PausingCriterion(SortableBase):
    """A criterion that pauses generation from a GenerationNode without triggering
    a transition to another node.
    """

    @abstractmethod
    def is_met(
        self,
        experiment: Experiment,
        curr_node: GenerationNode,
    ) -> bool:
        """Returns True if this criterion's condition is met."""
        pass

    @abstractmethod
    def block_continued_generation_error(
        self,
        node_name: str,
        experiment: Experiment,
        trials_from_node: set[int],
    ) -> None:
        """Raises an appropriate error when generation is blocked."""
        pass

    @property
    def criterion_class(self) -> str:
        """Name of the class of this PausingCriterion."""
        return self.__class__.__name__

    def __repr__(self) -> str:
        return f"{self.criterion_class}({serialize_init_args(obj=self)})"

    @property
    def _unique_id(self) -> str:
        """Unique id for this PausingCriterion."""
        return str(self)


class TrialCountPausingCriterion(PausingCriterion):
    """Abstract base class for pausing criteria based on trial count thresholds.

    This class provides shared logic for pausing criteria that count trials toward
    a threshold. Subclasses only need to implement `block_continued_generation_error()`
    to define the specific error raised when generation is paused.

    Args:
        threshold: The maximum number of trials allowed before blocking generation.
        only_in_statuses: A list of trial statuses to filter on when checking the
            criterion threshold.
        not_in_statuses: A list of trial statuses to exclude when checking the
            criterion threshold.
        use_all_trials_in_exp: A flag to use all trials in the experiment, instead of
            only those generated by the current GenerationNode.
        count_only_trials_without_data: If True, only count trials that are MISSING
            required metric data. Defaults to False.
    """

    def __init__(
        self,
        threshold: int,
        only_in_statuses: list[TrialStatus] | None = None,
        not_in_statuses: list[TrialStatus] | None = None,
        use_all_trials_in_exp: bool = False,
        count_only_trials_without_data: bool = False,
    ) -> None:
        self.threshold = threshold
        self.only_in_statuses = only_in_statuses
        self.not_in_statuses = not_in_statuses
        self.use_all_trials_in_exp = use_all_trials_in_exp
        self.count_only_trials_without_data = count_only_trials_without_data

    def num_contributing_to_threshold(
        self, experiment: Experiment, trials_from_node: set[int]
    ) -> int:
        """Returns the number of trials contributing to the threshold."""
        return count_trials_toward_threshold(
            experiment=experiment,
            trials_from_node=trials_from_node,
            only_in_statuses=self.only_in_statuses,
            not_in_statuses=self.not_in_statuses,
            use_all_trials_in_exp=self.use_all_trials_in_exp,
            count_only_trials_without_data=self.count_only_trials_without_data,
        )

    def num_till_threshold(
        self, experiment: Experiment, trials_from_node: set[int]
    ) -> int:
        """Returns the number of trials available before hitting the threshold."""
        return self.threshold - self.num_contributing_to_threshold(
            experiment=experiment, trials_from_node=trials_from_node
        )

    def is_met(
        self,
        experiment: Experiment,
        curr_node: GenerationNode,
    ) -> bool:
        """Returns True if the trial count threshold has been reached."""
        return (
            self.num_contributing_to_threshold(
                experiment=experiment, trials_from_node=curr_node.trials_from_node
            )
            >= self.threshold
        )

    @abstractmethod
    def block_continued_generation_error(
        self,
        node_name: str,
        experiment: Experiment,
        trials_from_node: set[int],
    ) -> None:
        """Raises an appropriate error when generation is blocked.

        Subclasses must implement this to define the specific error behavior.
        """
        pass


class MaxGenerationParallelism(TrialCountPausingCriterion):
    """A PausingCriterion that pauses generation after a maximum number of trials
    have been generated for the current GenerationNode and are currently running.

    Args:
        threshold: The maximum number of trials allowed in the specified statuses.
        only_in_statuses: A list of trial statuses to filter on when checking the
            criterion threshold.
        not_in_statuses: A list of trial statuses to exclude when checking the
            criterion threshold.
        use_all_trials_in_exp: A flag to use all trials in the experiment, instead of
            only those generated by the current GenerationNode.
        count_only_trials_without_data: If True, only count trials that are MISSING
            required metric data. Defaults to False.
    """

    def __init__(
        self,
        threshold: int,
        only_in_statuses: list[TrialStatus] | None = None,
        not_in_statuses: list[TrialStatus] | None = None,
        use_all_trials_in_exp: bool = False,
        count_only_trials_without_data: bool = False,
    ) -> None:
        super().__init__(
            threshold=threshold,
            only_in_statuses=only_in_statuses,
            not_in_statuses=not_in_statuses,
            use_all_trials_in_exp=use_all_trials_in_exp,
            count_only_trials_without_data=count_only_trials_without_data,
        )

    def block_continued_generation_error(
        self,
        node_name: str,
        experiment: Experiment,
        trials_from_node: set[int],
    ) -> None:
        """Raises MaxParallelismReachedException."""
        raise MaxParallelismReachedException(
            node_name=node_name,
            num_running=self.num_contributing_to_threshold(
                experiment=experiment, trials_from_node=trials_from_node
            ),
        )


class MaxTrialsAwaitingData(TrialCountPausingCriterion):
    """A PausingCriterion that pauses generation after a maximum number of trials
    have been generated, waiting for data before allowing more generation.

    This criterion pauses generation from the associated GenerationNode when the
    threshold is met, but does NOT trigger a transition to another node. Use this
    when you want to enforce that a node generates at most a certain number of
    trials before requiring data.

    Args:
        threshold: The maximum number of trials allowed before blocking generation.
        only_in_statuses: A list of trial statuses to filter on when checking the
            criterion threshold.
        not_in_statuses: A list of trial statuses to exclude when checking the
            criterion threshold. Defaults to [FAILED, ABANDONED].
        use_all_trials_in_exp: A flag to use all trials in the experiment, instead of
            only those generated by the current GenerationNode. Defaults to False.
        count_only_trials_without_data: If True, only count trials that are MISSING
            required metric data. Use this to block generation until running trials
            have data. Defaults to False.
    """

    def __init__(
        self,
        threshold: int,
        only_in_statuses: list[TrialStatus] | None = None,
        not_in_statuses: list[TrialStatus] | None = None,
        use_all_trials_in_exp: bool = False,
        count_only_trials_without_data: bool = False,
    ) -> None:
        if not_in_statuses is None:
            not_in_statuses = [TrialStatus.FAILED, TrialStatus.ABANDONED]
        super().__init__(
            threshold=threshold,
            only_in_statuses=only_in_statuses,
            not_in_statuses=not_in_statuses,
            use_all_trials_in_exp=use_all_trials_in_exp,
            count_only_trials_without_data=count_only_trials_without_data,
        )

    def block_continued_generation_error(
        self,
        node_name: str,
        experiment: Experiment,
        trials_from_node: set[int],
    ) -> None:
        """Raises DataRequiredError when the trial threshold is reached."""
        raise DataRequiredError(DATA_REQUIRED_MSG.format(node_name=node_name))


class StagedTrialsPausingCriterion(TrialCountPausingCriterion):
    """Pauses generation when staged trials exist.

    Defaults to threshold=1, only_in_statuses=[TrialStatus.STAGED],
    use_all_trials_in_exp=True so that any staged trial in the experiment
    pauses generation.
    """

    def __init__(
        self,
        threshold: int = 1,
        use_all_trials_in_exp: bool = True,
    ) -> None:
        super().__init__(
            threshold=threshold,
            only_in_statuses=[TrialStatus.STAGED],
            use_all_trials_in_exp=use_all_trials_in_exp,
        )

    def block_continued_generation_error(
        self,
        node_name: str,
        experiment: Experiment,
        trials_from_node: set[int],
    ) -> None:
        staged = experiment.trial_indices_by_status[TrialStatus.STAGED]
        staged_str = ", ".join(str(i) for i in staged)
        raise AxGenerationException(
            f"You already have staged trials (indices: {staged_str}). "
            "Please either allocate and run the staged trials in the associated "
            "experimentation platform (QE, QRT, etc.) or abandon them to unblock "
            "generation of new candidate trials."
        )


class MissingOptimizationConfigPausingCriterion(PausingCriterion):
    """Pauses generation if no optimization_config exists."""

    def is_met(
        self,
        experiment: Experiment,
        curr_node: GenerationNode,
    ) -> bool:
        return experiment.optimization_config is None

    def block_continued_generation_error(
        self,
        node_name: str,
        experiment: Experiment,
        trials_from_node: set[int],
    ) -> None:
        raise AxGenerationException(
            "Your experiment does not currently have an optimization config set. "
            "In order to generate candidates with Bayesian Optimization, you must "
            "configure + save an optimization config on your experiment. This can "
            "be done via the UI or Client API."
        )


# ============================================================================
# TransitionCriterion - for node transitions
# ============================================================================


class TransitionCriterion(SortableBase):
    """
    Simple class to describe a condition which must be met for this GenerationNode to
    take an action such as generation, transition, etc.

    Args:
        transition_to: The name of the GenerationNode the GenerationStrategy should
            transition to when this criterion is met.
        continue_trial_generation: A flag to indicate that all generation for a given
            trial is not completed, and thus even after transition, the next node will
            continue to generate arms for the same trial. Example usage: in
            ``BatchTrial``s we may  enable generation of arms within a batch from
            different ``GenerationNodes`` by setting this flag to True.
    """

    _transition_to: str

    def __init__(
        self,
        transition_to: str,
        continue_trial_generation: bool | None = False,
    ) -> None:
        self._transition_to = transition_to
        self.continue_trial_generation = continue_trial_generation

    @property
    def transition_to(self) -> str:
        """The name of the next GenerationNode after this TransitionCriterion is
        completed.
        """
        return self._transition_to

    @abstractmethod
    def is_met(
        self,
        experiment: Experiment,
        curr_node: GenerationNode,
    ) -> bool:
        """If the criterion of this TransitionCriterion is met, returns True."""
        pass

    @property
    def criterion_class(self) -> str:
        """Name of the class of this TransitionCriterion."""
        return self.__class__.__name__

    def __repr__(self) -> str:
        return f"{self.criterion_class}({serialize_init_args(obj=self)})"

    @property
    def _unique_id(self) -> str:
        """Unique id for this TransitionCriterion."""
        return str(self)


class AutoTransitionAfterGen(TransitionCriterion):
    """A class to designate automatic transition from one GenerationNode to another.

    Args:
        transition_to: The name of the GenerationNode the GenerationStrategy should
            transition to next.
        continue_trial_generation: A flag to indicate that all generation for a given
            trial is not completed, and thus even after transition, the next node will
            continue to generate arms for the same trial. Example usage: in
            ``BatchTrial``s we may  enable generation of arms within a batch from
            different ``GenerationNodes`` by setting this flag to True.
    """

    def __init__(
        self,
        transition_to: str,
        continue_trial_generation: bool | None = True,
    ) -> None:
        super().__init__(
            transition_to=transition_to,
            continue_trial_generation=continue_trial_generation,
        )

    def is_met(
        self,
        experiment: Experiment,
        curr_node: GenerationNode,
    ) -> bool:
        """Return True as soon as any GeneratorRun is generated by this
        GenerationNode.
        """
        # Handle edge case where the InputConstructor for a GenerationNode
        # with this criterion requests no arms to be generated, therefore, indicating
        # that this GenerationNode should be skipped and so we can transition to the
        # next node as defined by this criterion.
        if curr_node._should_skip:
            return True
        last_gr_from_gs = curr_node.generation_strategy.last_generator_run
        return (
            last_gr_from_gs._generation_node_name == curr_node.name
            if last_gr_from_gs is not None
            else False
        )


class IsSingleObjective(TransitionCriterion):
    """A class to initiate transition based on whether the experiment is optimizing
    for a single objective or multiple objectives.

    Args:
        transition_to: The name of the GenerationNode the GenerationStrategy should
            transition to next.
        continue_trial_generation: A flag to indicate that all generation for a given
            trial is not completed, and thus even after transition, the next node will
            continue to generate arms for the same trial. Example usage: in
            ``BatchTrial``s we may  enable generation of arms within a batch from
            different ``GenerationNodes`` by setting this flag to True.
    """

    def __init__(
        self,
        transition_to: str,
        continue_trial_generation: bool | None = False,
    ) -> None:
        super().__init__(
            transition_to=transition_to,
            continue_trial_generation=continue_trial_generation,
        )

    def is_met(
        self,
        experiment: Experiment,
        curr_node: GenerationNode,
    ) -> bool:
        """Return True if the optimization config is not of type
        ``MultiObjectiveOptimizationConfig``."""
        return (
            not isinstance(
                experiment.optimization_config, MultiObjectiveOptimizationConfig
            )
            if experiment.optimization_config is not None
            else True
        )


class TrialBasedCriterion(TransitionCriterion):
    """Common class for transition criterion that are based on trial information.

    Args:
        threshold: The threshold as an integer for this criterion. Ex: If we want to
            generate at most 3 trials, then the threshold is 3.
        only_in_statuses: A list of trial statuses to filter on when checking the
            criterion threshold.
        not_in_statuses: A list of trial statuses to exclude when checking the
            criterion threshold.
        transition_to: The name of the GenerationNode the GenerationStrategy should
            transition to when this criterion is met, if it exists.
        use_all_trials_in_exp: A flag to use all trials in the experiment, instead of
            only those generated by the current GenerationNode.
        continue_trial_generation: A flag to indicate that all generation for a given
            trial is not completed, and thus even after transition, the next node will
            continue to generate arms for the same trial. Example usage: in
            ``BatchTrial``s we may  enable generation of arms within a batch from
            different ``GenerationNodes`` by setting this flag to True.
        count_only_trials_with_data: If set to True, only trials with data for all
            metrics in the opt config will be counted towards the ``threshold``.
            Defaults to False.
    """

    def __init__(
        self,
        threshold: int,
        transition_to: str,
        only_in_statuses: list[TrialStatus] | None = None,
        not_in_statuses: list[TrialStatus] | None = None,
        use_all_trials_in_exp: bool | None = False,
        continue_trial_generation: bool | None = False,
        count_only_trials_with_data: bool = False,
    ) -> None:
        self.threshold = threshold
        self.only_in_statuses = only_in_statuses
        self.not_in_statuses = not_in_statuses
        self.use_all_trials_in_exp = use_all_trials_in_exp
        self.count_only_trials_with_data = count_only_trials_with_data
        super().__init__(
            transition_to=transition_to,
            continue_trial_generation=continue_trial_generation,
        )

    def experiment_trials_by_status(
        self, experiment: Experiment, statuses: list[TrialStatus]
    ) -> set[int]:
        """Get the trial indices from the experiment with the desired statuses."""
        return get_trials_by_statuses(experiment=experiment, statuses=statuses)

    def all_trials_to_check(self, experiment: Experiment) -> set[int]:
        """All the trials to check that meet the provided status filters."""
        return filter_trials_by_status(
            experiment=experiment,
            only_in_statuses=self.only_in_statuses,
            not_in_statuses=self.not_in_statuses,
        )

    def num_contributing_to_threshold(
        self, experiment: Experiment, trials_from_node: set[int]
    ) -> int:
        """Returns the number of trials contributing to the threshold."""
        return count_trials_toward_threshold(
            experiment=experiment,
            trials_from_node=trials_from_node,
            only_in_statuses=self.only_in_statuses,
            not_in_statuses=self.not_in_statuses,
            use_all_trials_in_exp=bool(self.use_all_trials_in_exp),
            count_only_trials_with_data=self.count_only_trials_with_data,
        )

    def num_till_threshold(
        self, experiment: Experiment, trials_from_node: set[int]
    ) -> int:
        """Returns the number of trials needed to meet the threshold.

        Args:
            experiment: The experiment associated with this GenerationStrategy.
            trials_from_node: The set of trials generated by this GenerationNode.
        """
        return self.threshold - self.num_contributing_to_threshold(
            experiment=experiment, trials_from_node=trials_from_node
        )

    def is_met(
        self,
        experiment: Experiment,
        curr_node: GenerationNode,
    ) -> bool:
        """Returns if this criterion has been met given its constraints.
        Args:
            experiment: The experiment associated with this GenerationStrategy.
            trials_from_node: The set of trials generated by this GenerationNode.
            block_continued_generation: A flag to prevent continued generation from the
                associated GenerationNode if this criterion is met but other criterion
                remain unmet. Ex: ``MinTrials`` has not been met yet, but
                MinTrials has been reached. If this flag is set to true on MinTrials
                then we will raise an error, otherwise we will continue to generate
                trials until ``MinTrials`` is met (thus overriding MinTrials).
        """
        return (
            self.num_contributing_to_threshold(
                experiment=experiment, trials_from_node=curr_node.trials_from_node
            )
            >= self.threshold
        )


class MinTrials(TrialBasedCriterion):
    """
    Simple class to enforce a minimum threshold for the number of trials with the
    designated statuses being generated by a specific GenerationNode.

    Args:
        threshold: The threshold as an integer for this criterion. Ex: If we want to
            generate at most 3 trials, then the threshold is 3.
        only_in_statuses: A list of trial statuses to filter on when checking the
            criterion threshold.
        not_in_statuses: A list of trial statuses to exclude when checking the
            criterion threshold.
        transition_to: The name of the GenerationNode the GenerationStrategy should
            transition to when this criterion is met.
        use_all_trials_in_exp: A flag to use all trials in the experiment, instead of
            only those generated by the current GenerationNode.
        continue_trial_generation: A flag to indicate that all generation for a given
            trial is not completed, and thus even after transition, the next node will
            continue to generate arms for the same trial. Example usage: in
            ``BatchTrial``s we may  enable generation of arms within a batch from
            different ``GenerationNodes`` by setting this flag to True.
        count_only_trials_with_data: If set to True, only trials with data will be
            counted towards the ``threshold``. Defaults to False.
    """

    def __init__(
        self,
        threshold: int,
        transition_to: str,
        only_in_statuses: list[TrialStatus] | None = None,
        not_in_statuses: list[TrialStatus] | None = None,
        use_all_trials_in_exp: bool | None = False,
        continue_trial_generation: bool | None = False,
        count_only_trials_with_data: bool = False,
    ) -> None:
        super().__init__(
            threshold=threshold,
            transition_to=transition_to,
            only_in_statuses=only_in_statuses,
            not_in_statuses=not_in_statuses,
            use_all_trials_in_exp=use_all_trials_in_exp,
            continue_trial_generation=continue_trial_generation,
            count_only_trials_with_data=count_only_trials_with_data,
        )


class AuxiliaryExperimentCheck(TransitionCriterion):
    """A class to transition from one GenerationNode to another by checking if certain
    types of Auxiliary Experiment purposes exists.

    A common use case is to use auxiliary_experiment_purposes_to_include to transition
    to a node and auxiliary_experiment_purposes_to_exclude to transition away from it.

    Example usage: In Bayesian optimization with preference exploration (BOPE), we
    check if the preference exploration (PE) auxiliary experiment exists to indicate
    transition to the node that will generate candidates based on the learned
    objective. Since preference exploration is usually conducted after the exploratory
    batch is completed, we do not know at experiment creation time if the PE node
    should be used during the GenerationStrategy.

    Args:
        transition_to: The name of the GenerationNode the GenerationStrategy should
            transition to when this criterion is met, if it exists.
        auxiliary_experiment_purposes_to_include: Optional list of auxiliary experiment
            purposes we expect to have. This can be helpful when need to transition to
            a node based on AuxiliaryExperimentPurpose. Criterion is met when all
            inclusion and exclusion checks pass.
        auxiliary_experiment_purposes_to_exclude: Optional list of auxiliary experiment
            purpose we expect to not have. This can be helpful when need to transition
            out of a node based on AuxiliaryExperimentPurpose. Criterion is met when
            all inclusion and exclusion checks pass.
        continue_trial_generation: A flag to indicate that all generation for a given
            trial is not completed, and thus even after transition, the next node will
            continue to generate arms for the same trial. Example usage: in
            ``BatchTrial``s we may  enable generation of arms within a batch from
            different ``GenerationNodes`` by setting this flag to True.
    """

    def __init__(
        self,
        transition_to: str,
        auxiliary_experiment_purposes_to_include: (
            list[AuxiliaryExperimentPurpose] | None
        ) = None,
        auxiliary_experiment_purposes_to_exclude: (
            list[AuxiliaryExperimentPurpose] | None
        ) = None,
        continue_trial_generation: bool | None = False,
    ) -> None:
        super().__init__(
            transition_to=transition_to,
            continue_trial_generation=continue_trial_generation,
        )

        if (
            auxiliary_experiment_purposes_to_include is None
            and auxiliary_experiment_purposes_to_exclude is None
        ):
            raise UserInputError(
                f"{self.__class__} cannot have both "
                "`auxiliary_experiment_purposes_to_include` and "
                "`auxiliary_experiment_purposes_to_exclude` be None."
            )
        self.auxiliary_experiment_purposes_to_include = (
            auxiliary_experiment_purposes_to_include
        )
        self.auxiliary_experiment_purposes_to_exclude = (
            auxiliary_experiment_purposes_to_exclude
        )

    def check_aux_exp_purposes(
        self,
        aux_exp_by_purposes: dict[
            AuxiliaryExperimentPurpose, list[AuxiliaryExperiment]
        ],
        include: bool,
        expected_aux_exp_purposes: list[AuxiliaryExperimentPurpose] | None = None,
    ) -> bool:
        """Helper method to check if all elements in expected_aux_exp_purposes
        are in (or not in) aux_exp_purposes"""
        if expected_aux_exp_purposes is not None:
            for purpose in expected_aux_exp_purposes:
                purpose_present = (
                    purpose in aux_exp_by_purposes
                    and len(aux_exp_by_purposes[purpose]) > 0
                )
                if purpose_present != include:
                    return False
        return True

    def is_met(
        self,
        experiment: Experiment,
        curr_node: GenerationNode,
    ) -> bool:
        """Check if the experiment has auxiliary experiments for certain purpose."""
        inclusion_check = self.check_aux_exp_purposes(
            aux_exp_by_purposes=experiment.auxiliary_experiments_by_purpose,
            include=True,
            expected_aux_exp_purposes=self.auxiliary_experiment_purposes_to_include,
        )
        exclusion_check = self.check_aux_exp_purposes(
            aux_exp_by_purposes=experiment.auxiliary_experiments_by_purpose,
            include=False,
            expected_aux_exp_purposes=self.auxiliary_experiment_purposes_to_exclude,
        )
        return inclusion_check and exclusion_check