hunk.code_window:     covered_o = set()
       config["log_level"] = "ERROR"
       first_error = None
       for a_name, action_space in ACTION_SPACES_TO_TEST.items():
           for o_name, obs_space in OBSERVATION_SPACES_TO_TEST.items():
  
hunk.labels: keep keep add keep keep
hunk.after_edit:     torch = config.get("use_pytorch", False)
  
hunk.file_path: rllib/tests/test_supported_spaces.py
hunk.type: add
file: """BOHB (Bayesian Optimization with HyperBand)"""
  
  import copy
  import logging
  
  from ray.tune.suggest import SuggestionAlgorithm
  
  logger = logging.getLogger(__name__)
  
  
  class _BOHBJobWrapper():
      """Mock object for HpBandSter to process."""
  
      def __init__(self, loss, budget, config):
          self.result = {"loss": loss}
          self.kwargs = {"budget": budget, "config": config.copy()}
          self.exception = None
  
  
  class TuneBOHB(SuggestionAlgorithm):
      """BOHB suggestion component.
  
  
      Requires HpBandSter and ConfigSpace to be installed. You can install
      HpBandSter and ConfigSpace with: `pip install hpbandster ConfigSpace`.
  
      This should be used in conjunction with HyperBandForBOHB.
  
      Args:
          space (ConfigurationSpace): Continuous ConfigSpace search space.
              Parameters will be sampled from this space which will be used
              to run trials.
          bohb_config (dict): configuration for HpBandSter BOHB algorithm
          max_concurrent (int): Number of maximum concurrent trials. Defaults
              to 10.
          metric (str): The training result objective value attribute.
          mode (str): One of {min, max}. Determines whether objective is
              minimizing or maximizing the metric attribute.
  
      Example:
          >>> import ConfigSpace as CS
          >>> config_space = CS.ConfigurationSpace()
          >>> config_space.add_hyperparameter(
                  CS.UniformFloatHyperparameter('width', lower=0, upper=20))
          >>> config_space.add_hyperparameter(
                  CS.UniformFloatHyperparameter('height', lower=-100, upper=100))
          >>> config_space.add_hyperparameter(
                  CS.CategoricalHyperparameter(
                      name='activation', choices=['relu', 'tanh']))
          >>> algo = TuneBOHB(
                  config_space, max_concurrent=4, metric='mean_loss', mode='min')
          >>> bohb = HyperBandForBOHB(
                  time_attr='training_iteration',
                  metric='mean_loss',
                  mode='min',
                  max_t=100)
          >>> run(MyTrainableClass, scheduler=bohb, search_alg=algo)
  
      """
  
      def __init__(self,
                   space,
                   bohb_config=None,
                   max_concurrent=10,
                   metric="neg_mean_loss",
                   mode="max"):
          from hpbandster.optimizers.config_generators.bohb import BOHB
          assert BOHB is not None, "HpBandSter must be installed!"
          assert mode in ["min", "max"], "`mode` must be 'min' or 'max'!"
          self._max_concurrent = max_concurrent
          self.trial_to_params = {}
          self.running = set()
          self.paused = set()
          self.metric = metric
          if mode == "max":
              self._metric_op = -1.
          elif mode == "min":
              self._metric_op = 1.
          bohb_config = bohb_config or {}
          self.bohber = BOHB(space, **bohb_config)
          super(TuneBOHB, self).__init__()
  
      def _suggest(self, trial_id):
          if len(self.running) < self._max_concurrent:
              # This parameter is not used in hpbandster implementation.
              config, info = self.bohber.get_config(None)
              self.trial_to_params[trial_id] = copy.deepcopy(config)
              self.running.add(trial_id)
              return config
          return None
  
      def on_trial_result(self, trial_id, result):
          if trial_id not in self.paused:
              self.running.add(trial_id)
          if "hyperband_info" not in result:
              logger.warning("BOHB Info not detected in result. Are you using "
                             "HyperBandForBOHB as a scheduler?")
          elif "budget" in result.get("hyperband_info", {}):
              hbs_wrapper = self.to_wrapper(trial_id, result)
              self.bohber.new_result(hbs_wrapper)
  
      def on_trial_complete(self,
                            trial_id,
                            result=None,
                            error=False,
                            early_terminated=False):
          del self.trial_to_params[trial_id]
          if trial_id in self.paused:
              self.paused.remove(trial_id)
          if trial_id in self.running:
              self.running.remove(trial_id)
  
      def to_wrapper(self, trial_id, result):
          return _BOHBJobWrapper(self._metric_op * result[self.metric],
                                 result["hyperband_info"]["budget"],
                                 self.trial_to_params[trial_id])
  
      def on_pause(self, trial_id):
          self.paused.add(trial_id)
          self.running.remove(trial_id)
  
      def on_unpause(self, trial_id):
          self.paused.remove(trial_id)
          self.running.add(trial_id)
  