Rework Envs

carl/context/context_space.py

@@ -0,0 +1,230 @@
+from __future__ import annotations
+
+from typing import Any, List, Union
+
+import gymnasium.spaces as spaces
+import numpy as np
+from ConfigSpace.configuration_space import Configuration, ConfigurationSpace
+from ConfigSpace.hyperparameters import (
+    Hyperparameter,
+    NormalFloatHyperparameter,
+    NumericalHyperparameter,
+    UniformFloatHyperparameter,
+)
+from omegaconf import DictConfig
+from search_space_encoding import search_space_to_config_space
+from typing_extensions import TypeAlias
+
+ContextFeature: TypeAlias = Hyperparameter
+
+
+class ContextSpace(ConfigurationSpace):
+    """
+    Baseclass for contexts of environments implemented in CARL. DO NOT USE DIRECTLY. USE SUBCLASSES!
+    """
+
+    def __init__(
+        self,
+        context_space: Union[List[Hyperparameter], str, DictConfig, ConfigurationSpace]
+        | None = None,
+        context_features: List[str] | None = None,
+        name: str = "context_space",
+        seed: int = None,
+    ):
+        super().__init__(name=name, seed=seed)
+
+        if isinstance(context_space, list):
+            self.add_hyperparameters(context_space)
+        else:
+            cs = search_space_to_config_space(context_space)
+            self.add_hyperparameters(cs.get_hyperparameters())
+
+        self._context_features: list[str] = context_features if context_features else []
+
+    def to_gym_space(
+        self, context_features: List[str] = None, as_dict: bool = False
+    ) -> spaces.Space:
+        if context_features is None:
+            context_features = [hp.name for hp in self.get_hyperparameters()]
+
+        if as_dict:
+            context_space = {}
+
+            for context_feature in context_features:
+                hp = self.get_hyperparameter(context_feature)
+                if isinstance(hp, NumericalHyperparameter):
+                    context_space[hp.name] = spaces.Box(low=hp.lower, high=hp.upper)
+                else:
+                    raise ValueError(
+                        f"Context features must be of type NumericalHyperparameter."
+                        f"Got {type(hp)}."
+                    )
+            return spaces.Dict(context_space)
+        else:
+            low = np.array(
+                [self.get_hyperparameter(cf).lower for cf in context_features]
+            )
+            high = np.array(
+                [self.get_hyperparameter(cf).upper for cf in context_features]
+            )
+
+            return spaces.Box(low=low, high=high, dtype=np.float32)
+
+    @property
+    def context_features(self) -> List[str]:
+        """
+        Context features that are sampled.
+
+        Returns
+        -------
+        List[str]
+            Context features.
+        """
+        return self._context_features
+
+    def sample_configuration(
+        self, size: int = 1
+    ) -> Union[Configuration, List[Configuration]]:
+        """
+        Samples values for all active context features. For all other values default is used.
+
+        Parameters
+        ----------
+        size : int = 1
+            Number of contexts to sample.
+
+        Returns
+        -------
+        Union[Configuration, List[Configuration]]
+            Sampled contexts.
+        """
+
+        def insert_defaults(cfg: Configuration) -> Configuration:
+            values = cfg.get_dictionary()
+            for feature in values.keys():
+                if feature not in self._context_features:
+                    values[feature] = self.get_hyperparameter(feature).default_value
+
+            return Configuration(self, values=values)
+
+        configurations = super().sample_configuration(size=size)
+
+        if size == 1:
+            configurations = insert_defaults(configurations)
+        else:
+            for i, configuration in enumerate(configurations):
+                configurations[i] = insert_defaults(configuration)
+
+        return configurations
+
+
+class CartPoleContextSpace(ContextSpace):
+    """
+    Context space for CARLCartPoleEnv.
+
+    Parameters
+    ----------
+    context_features : List[str]
+        Context features to vary, i.e. sample.
+
+    seed : int = None
+        Seed for PRNG.
+    """
+
+    def __init__(
+        self,
+        context_space: Union[List[Hyperparameter], str, DictConfig, ConfigurationSpace]
+        | None = None,
+        context_features: List[str] | None = None,
+        seed: int = None,
+    ):
+        if context_space is None:
+            context_space = CartPoleContextSpace.get_default()
+        super().__init__(
+            context_space, context_features, name=self.__class__.__name__, seed=seed
+        )
+
+    @staticmethod
+    def get_default() -> list[ContextFeature]:
+        """Get context features with default values
+
+        Returns
+        -------
+        list[ContextFeature]
+            List of context features
+        """
+        return [
+            UniformFloatHyperparameter(
+                "gravity", lower=0.1, upper=np.inf, default_value=9.8
+            ),
+            UniformFloatHyperparameter(
+                "masscart", lower=0.1, upper=10, default_value=1.0
+            ),
+            UniformFloatHyperparameter(
+                "masspole", lower=0.01, upper=1, default_value=0.1
+            ),
+            UniformFloatHyperparameter(
+                "pole_length", lower=0.05, upper=5, default_value=0.5
+            ),
+            UniformFloatHyperparameter(
+                "force_magnifier", lower=1, upper=100, default_value=10.0
+            ),
+            UniformFloatHyperparameter(
+                "update_interval", lower=0.002, upper=0.2, default_value=0.02
+            ),
+            UniformFloatHyperparameter(
+                "initial_state_lower", lower=-10, upper=10, default_value=-0.1
+            ),
+            UniformFloatHyperparameter(
+                "initial_state_upper", lower=-10, upper=10, default_value=0.1
+            ),  # TODO: We can't have inf. How to handle/annotate when transforming to obs space?
+        ]
+
+    def verify(self, configspace: ConfigurationSpace) -> bool:
+        # TODO Check whether HPs in of configspace in self.hyperparameters included
+        ...
+
+
+class ContextSampler(ConfigurationSpace):
+    def __init__(
+        self,
+        context_distributions: Union[
+            List[Hyperparameter], str, DictConfig, ConfigurationSpace
+        ],
+        seed: int,
+    ):
+        self.context_distributions = context_distributions
+        # TODO: pass seed
+        ...
+
+    def add_defaults(context_space: ContextSpace):
+        """Add defaults from env context space
+
+        Parameters
+        ----------
+        context_space : ContextSpace
+            _description_
+        """
+        ...
+
+    def sample_contexts(self, n_contexts: int) -> dict[int, dict[str, Any]]:
+        contexts = self.sample_configuration(size=n_contexts)
+        contexts = {i: C for i, C in enumerate(contexts)}
+        return contexts
+
+
+if __name__ == "__main__":
+    print("hello world")
+    context_space = CartPoleContextSpace()
+    print(context_space)
+    context = context_space.sample_configuration()
+    print(context)
+
+    context_sampler = ContextSampler(
+        [
+            NormalFloatHyperparameter("gravity", mu=9.8, sigma=1),
+        ],
+        seed=0,
+    )
+
+    contexts = context_sampler.sample_contexts(n_contexts=10)

carl/context/search_space_encoding.py

@@ -0,0 +1,142 @@
+from __future__ import annotations
+
+from typing import Any, Union
+
+import json
+
+from ConfigSpace import ConfigurationSpace  # type: ignore[import]
+from ConfigSpace.read_and_write import json as csjson  # type: ignore[import]
+from omegaconf import DictConfig, ListConfig
+
+
+class JSONCfgEncoder(json.JSONEncoder):
+    """Encode DictConfigs.
+
+    Convert DictConfigs to normal dicts.
+    """
+
+    def default(self, obj: Union[DictConfig, ListConfig, Any]) -> Any:
+        """Modify default of JSON encoder
+
+        Parameters
+        ----------
+        obj : DictConfig | ListConfig | Any
+            Object to encode
+
+        Returns
+        -------
+        Any
+            Encoded object
+        """
+        if isinstance(obj, DictConfig):
+            return dict(obj)
+        elif isinstance(obj, ListConfig):
+            parsed_list = []
+            for o in obj:
+                if type(o) == DictConfig:
+                    o = dict(o)
+                elif type(o) == ListConfig:
+                    o = list(o)
+                parsed_list.append(o)
+
+            return parsed_list  # [dict(o) for o in obj]
+        return json.JSONEncoder.default(self, obj)
+
+
+def search_space_to_config_space(
+    search_space: Union[str, DictConfig, ConfigurationSpace], seed: int = None
+) -> ConfigurationSpace:
+    """
+    Convert hydra search space to SMAC's configuration space.
+
+    See the [ConfigSpace docs](https://automl.github.io/ConfigSpace/master/API-Doc.html#) for information of how
+    to define a configuration (search) space.
+
+    In a yaml (hydra) config file, the smac.search space must take the form of:
+
+    search_space:
+        hyperparameters:
+            hyperparameter_name_0:
+                key1: value1
+                ...
+            hyperparameter_name_1:
+                key1: value1
+                key2: value2
+                ...
+
+
+    Parameters
+    ----------
+    search_space : Union[str, DictConfig, ConfigurationSpace]
+        The search space, either a DictConfig from a hydra yaml config file, or a path to a json configuration space
+        file in the format required of ConfigSpace.
+        If it already is a ConfigurationSpace, just optionally seed it.
+    seed : Optional[int]
+        Optional seed to seed configuration space.
+
+
+    Example of a json-serialized ConfigurationSpace file.
+    {
+      "hyperparameters": [
+        {
+          "name": "x0",
+          "type": "uniform_float",
+          "log": false,
+          "lower": -512.0,
+          "upper": 512.0,
+          "default": -3.0
+        },
+        {
+          "name": "x1",
+          "type": "uniform_float",
+          "log": false,
+          "lower": -512.0,
+          "upper": 512.0,
+          "default": -4.0
+        }
+      ],
+      "conditions": [],
+      "forbiddens": [],
+      "python_module_version": "0.4.17",
+      "json_format_version": 0.2
+    }
+
+
+    Returns
+    -------
+    ConfigurationSpace
+    """
+    if type(search_space) == str:
+        with open(search_space, "r") as f:
+            jason_string = f.read()
+        cs = csjson.read(jason_string)
+    elif type(search_space) == DictConfig:
+        # reorder hyperparameters as List[Dict]
+        hyperparameters = []
+        for name, cfg in search_space.hyperparameters.items():
+            cfg["name"] = name
+            if "default" not in cfg:
+                cfg["default"] = None
+            if "log" not in cfg:
+                cfg["log"] = False
+            hyperparameters.append(cfg)
+        search_space.hyperparameters = hyperparameters
+
+        if "conditions" not in search_space:
+            search_space["conditions"] = []
+
+        if "forbiddens" not in search_space:
+            search_space["forbiddens"] = []
+
+        jason_string = json.dumps(search_space, cls=JSONCfgEncoder)
+        cs = csjson.read(jason_string)
+    elif type(search_space) == ConfigurationSpace:
+        cs = search_space
+    else:
+        raise ValueError(
+            f"search_space must be of type str or DictConfig. Got {type(search_space)}."
+        )
+
+    if seed is not None:
+        cs.seed(seed=seed)
+    return cs

setup.py

@@ -90,6 +90,8 @@ def read_file(filepath: str) -> str:
         "PyYAML>=5.4.1",
         "tabulate>=0.8.9",
         "bs4>=0.0.1",
+        "ConfigSpace",
+        "omegaconf"
     ],
     extras_require=extras_require,
     test_suite="pytest",